From 0aaee88325ca6362764f37797a42867b62163fdf Mon Sep 17 00:00:00 2001
From: Gregory Conrad <gregorysconrad@gmail.com>
Date: Sat, 17 Aug 2024 11:15:28 -0400
Subject: [PATCH] fix: add workaround for flutter hot reload bug

Fixes #209
---
 .../lib/src/widgets/consumer.dart             | 36 +++++++++----------
 1 file changed, 18 insertions(+), 18 deletions(-)

diff --git a/packages/flutter_rearch/lib/src/widgets/consumer.dart b/packages/flutter_rearch/lib/src/widgets/consumer.dart
index 72ae72f..e4f5b72 100644
--- a/packages/flutter_rearch/lib/src/widgets/consumer.dart
+++ b/packages/flutter_rearch/lib/src/widgets/consumer.dart
@@ -74,24 +74,24 @@ class _RearchElement extends ComponentElement {
       // We need to do some dependency management here to ensure that we
       // get all the capsule updates we need, but nothing more.
 
-      // First, let's remove any no-longer used dependencies.
-      capsuleToRemoveDependency.entries
-          .where((e) => !capsulesUsedInCurrBuild!.contains(e.key))
-          .toList()
-          .forEach((e) {
-        e.value(); // dispose() from onNextUpdate
-        capsuleToRemoveDependency.remove(e.key);
-      });
-
-      // Next, for each capsule used in the current build,
-      // let's make sure we depend upon it.
-      capsulesUsedInCurrBuild!
-          .where((cap) => !capsuleToRemoveDependency.containsKey(cap))
-          .forEach((cap) {
-        capsuleToRemoveDependency[cap] = container.onNextUpdate(cap, () {
-          markNeedsBuild();
-          capsuleToRemoveDependency.remove(cap);
-        });
+      // First, ensure we depend upon every capsule used in the current build.
+      for (final capsule in capsulesUsedInCurrBuild!) {
+        capsuleToRemoveDependency.putIfAbsent(
+          capsule,
+          () => container.onNextUpdate(capsule, () {
+            markNeedsBuild();
+            capsuleToRemoveDependency.remove(capsule);
+          }),
+        );
+      }
+
+      // Next, remove any no-longer used dependencies.
+      capsuleToRemoveDependency.keys
+          .toSet()
+          .difference(capsulesUsedInCurrBuild!)
+          .forEach((capsule) {
+        capsuleToRemoveDependency[capsule]!.call();
+        capsuleToRemoveDependency.remove(capsule);
       });
 
       // Finally, let's reset everything for the next build.