@@ -168,13 +168,13 @@ public void UnregisterAll(object recipient)
168168 // dictionary (a hashed collection) only costs O(1) in the best case, while
169169 // if we had tried to iterate the whole dictionary every time we would have
170170 // paid an O(n) minimum cost for each single remove operation.
171- this . typesMap . Remove ( mapping . TypeArguments ) ;
171+ this . typesMap . TryRemove ( mapping . TypeArguments , out _ ) ;
172172 }
173173 }
174174 }
175175
176176 // Remove the associated set in the recipients map
177- this . recipientsMap . Remove ( key ) ;
177+ this . recipientsMap . TryRemove ( key , out _ ) ;
178178 }
179179 }
180180
@@ -242,26 +242,24 @@ public void UnregisterAll<TToken>(object recipient, TToken token)
242242
243243 // Try to remove the registered handler for the input token,
244244 // for the current message type (unknown from here).
245- if ( holder . Remove ( token ) )
245+ if ( holder . TryRemove ( token , out _ ) &&
246+ holder . Count == 0 )
246247 {
247- if ( holder . Count == 0 )
248- {
249- // If the map is empty, remove the recipient entirely from its container
250- map . Remove ( key ) ;
248+ // If the map is empty, remove the recipient entirely from its container
249+ map . TryRemove ( key , out _ ) ;
251250
252- if ( map . Count == 0 )
251+ if ( map . Count == 0 )
252+ {
253+ // If no handlers are left at all for the recipient, across all
254+ // message types and token types, remove the set of mappings
255+ // entirely for the current recipient, and lost the strong
256+ // reference to it as well. This is the same situation that
257+ // would've been achieved by just calling UnregisterAll(recipient).
258+ set . Remove ( Unsafe . As < IMapping > ( map ) ) ;
259+
260+ if ( set . Count == 0 )
253261 {
254- // If no handlers are left at all for the recipient, across all
255- // message types and token types, remove the set of mappings
256- // entirely for the current recipient, and lost the strong
257- // reference to it as well. This is the same situation that
258- // would've been achieved by just calling UnregisterAll(recipient).
259- set . Remove ( Unsafe . As < IMapping > ( map ) ) ;
260-
261- if ( set . Count == 0 )
262- {
263- this . recipientsMap . Remove ( key ) ;
264- }
262+ this . recipientsMap . TryRemove ( key , out _ ) ;
265263 }
266264 }
267265 }
@@ -309,25 +307,23 @@ public void Unregister<TMessage, TToken>(object recipient, TToken token)
309307 }
310308
311309 // Remove the target handler
312- if ( dictionary ! . Remove ( token ) )
310+ if ( dictionary ! . TryRemove ( token , out _ ) &&
311+ dictionary . Count == 0 )
313312 {
314313 // If the map is empty, it means that the current recipient has no remaining
315314 // registered handlers for the current <TMessage, TToken> combination, regardless,
316315 // of the specific token value (ie. the channel used to receive messages of that type).
317316 // We can remove the map entirely from this container, and remove the link to the map itself
318317 // to the current mapping between existing registered recipients (or entire recipients too).
319- if ( dictionary . Count == 0 )
320- {
321- mapping . Remove ( key ) ;
318+ mapping . TryRemove ( key , out _ ) ;
322319
323- HashSet < IMapping > set = this . recipientsMap [ key ] ;
320+ HashSet < IMapping > set = this . recipientsMap [ key ] ;
324321
325- set . Remove ( mapping ) ;
322+ set . Remove ( mapping ) ;
326323
327- if ( set . Count == 0 )
328- {
329- this . recipientsMap . Remove ( key ) ;
330- }
324+ if ( set . Count == 0 )
325+ {
326+ this . recipientsMap . TryRemove ( key , out _ ) ;
331327 }
332328 }
333329 }
0 commit comments