Repository: cashapp/redwood
Issue: we need to remove widgets from our map!

we need to remove widgets from our map!

https://github.com/square/treehouse/blob/1ad7b7242bb7c68e31465a3e65ecfd7bded817ae/treehouse-widget/src/commonMain/kotlin/app/cash/treehouse/widget/WidgetDisplay.kt#L30

```text

package app.cash.treehouse.widget

import app.cash.treehouse.protocol.Diff
import app.cash.treehouse.protocol.WidgetDiff

class WidgetDisplay<T : Any>(
  private val root: Widget<T>,
  private val factory: Widget.Factory<T>,
  private val events: EventSink,
) {
  private val widgets = mutableMapOf(Diff.RootId to root)

  fun apply(diff: Diff) {
    for (widgetDiff in diff.widgetDiffs) {
      val widget = checkNotNull(widgets[widgetDiff.id]) {
        "Unknown widget ID ${widgetDiff.id}"
      }

      when (widgetDiff) {
        is WidgetDiff.Insert -> {
          val childWidget = factory.create(widget.value, widgetDiff.kind, widgetDiff.childId, events)
          widgets[widgetDiff.childId] = childWidget
          widget.children(widgetDiff.childrenIndex).insert(widgetDiff.index, childWidget.value)
        }
        is WidgetDiff.Move -> {
          widget.children(widgetDiff.childrenIndex).move(widgetDiff.fromIndex, widgetDiff.toIndex, widgetDiff.count)
        }
        is WidgetDiff.Remove -> {
          widget.children(widgetDiff.childrenIndex).remove(widgetDiff.index, widgetDiff.count)
          // TODO we need to remove widgets from our map!
        }
        WidgetDiff.Clear -> {
          widget.children(Diff.RootChildrenIndex).clear()
          widgets.clear()
          widgets[Diff.RootId] = root
        }
      }
    }

    for (propertyDiff in diff.propertyDiffs) {
      val widget = checkNotNull(widgets[propertyDiff.id]) {
        "Unknown widget iD ${propertyDiff.id}"
      }

      widget.apply(propertyDiff)
    }
  }
}

```

edf3f9a0d1576c577b8785b13ce44e6b2bcae7c4

Changed files:
treehouse/treehouse-protocol/src/commonTest/kotlin/app/cash/treehouse/protocol/ProtocolTest.kt
treehouse/treehouse-widget/src/commonMain/kotlin/app/cash/treehouse/widget/WidgetDisplay.kt
treehouse/treehouse-compose/src/commonMain/kotlin/app/cash/treehouse/compose/applier.kt
treehouse/treehouse-protocol/src/commonMain/kotlin/app/cash/treehouse/protocol/protocol.kt

diff --git a/treehouse/treehouse-compose/src/commonMain/kotlin/app/cash/treehouse/compose/applier.kt b/treehouse/treehouse-compose/src/commonMain/kotlin/app/cash/treehouse/compose/applier.kt
index 77520f3d4..4383be9a1 100644
--- a/treehouse/treehouse-compose/src/commonMain/kotlin/app/cash/treehouse/compose/applier.kt
+++ b/treehouse/treehouse-compose/src/commonMain/kotlin/app/cash/treehouse/compose/applier.kt
@@ -180,13 +180,19 @@ internal class ProtocolApplier(
   override fun remove(index: Int, count: Int) {
     // Children instances are never removed from their parents.
     val current = current as ProtocolChildrenNode
-
     val children = current.children
+
+    // TODO We should not have to track this and send it as part of the protocol.
+    //  Ideally this would be entirely encapsulated on the display-side with additional bookkeeping.
+    //  For now, we track it here and send it in the protocol as a simple solution.
+    val removedIds = ArrayList<Long>(count)
     for (i in index until index + count) {
-      nodes.remove(children[i].id)
+      removedIds.add(children[i].id)
     }
+
+    nodes.keys.removeAll(removedIds)
     children.remove(index, count)
-    childrenDiffs.add(ChildrenDiff.Remove(current.id, current.tag, index, count))
+    childrenDiffs.add(ChildrenDiff.Remove(current.id, current.tag, index, count, removedIds))
   }
 
   override fun move(from: Int, to: Int, count: Int) {
diff --git a/treehouse/treehouse-protocol/src/commonMain/kotlin/app/cash/treehouse/protocol/protocol.kt b/treehouse/treehouse-protocol/src/commonMain/kotlin/app/cash/treehouse/protocol/protocol.kt
index 42f316c5d..e16bb4a0f 100644
--- a/treehouse/treehouse-protocol/src/commonMain/kotlin/app/cash/treehouse/protocol/protocol.kt
+++ b/treehouse/treehouse-protocol/src/commonMain/kotlin/app/cash/treehouse/protocol/protocol.kt
@@ -84,7 +84,14 @@ public sealed class ChildrenDiff {
     override val tag: Int,
     val index: Int,
     val count: Int,
-  ) : ChildrenDiff()
+    val removedIds: List<Long>,
+  ) : ChildrenDiff() {
+    init {
+      require(count == removedIds.size) {
+        "Count $count != Removed ID list size ${removedIds.size}"
+      }
+    }
+  }
 
   public companion object {
     public const val RootId: Long = 0L
diff --git a/treehouse/treehouse-protocol/src/commonTest/kotlin/app/cash/treehouse/protocol/ProtocolTest.kt b/treehouse/treehouse-protocol/src/commonTest/kotlin/app/cash/treehouse/protocol/ProtocolTest.kt
index b6adaebf1..31a20c77a 100644
--- a/treehouse/treehouse-protocol/src/commonTest/kotlin/app/cash/treehouse/protocol/ProtocolTest.kt
+++ b/treehouse/treehouse-protocol/src/commonTest/kotlin/app/cash/treehouse/protocol/ProtocolTest.kt
@@ -22,6 +22,7 @@ import kotlinx.serialization.modules.SerializersModule
 import kotlinx.serialization.modules.polymorphic
 import kotlin.test.Test
 import kotlin.test.assertEquals
+import kotlin.test.assertFailsWith
 
 class ProtocolTest {
   private val format = Json {
@@ -51,7 +52,7 @@ class ProtocolTest {
         ChildrenDiff.Clear,
         ChildrenDiff.Insert(1, 2, 3, 4, 5),
         ChildrenDiff.Move(1, 2, 3, 4, 5),
-        ChildrenDiff.Remove(1, 2, 3, 4),
+        ChildrenDiff.Remove(1, 2, 3, 4, listOf(5, 6, 7, 8)),
       ),
       propertyDiffs = listOf(
         PropertyDiff(1, 2, "Hello"),
@@ -63,7 +64,7 @@ class ProtocolTest {
       """["clear",{}],""" +
       """["insert",{"id":1,"tag":2,"childId":3,"kind":4,"index":5}],""" +
       """["move",{"id":1,"tag":2,"fromIndex":3,"toIndex":4,"count":5}],""" +
-      """["remove",{"id":1,"tag":2,"index":3,"count":4}]""" +
+      """["remove",{"id":1,"tag":2,"index":3,"count":4,"removedIds":[5,6,7,8]}]""" +
       """],"propertyDiffs":[""" +
       """{"id":1,"tag":2,"value":["kotlin.String","Hello"]},""" +
       """{"id":1,"tag":2}""" +
@@ -80,6 +81,12 @@ class ProtocolTest {
     assertJsonRoundtrip(Diff.serializer(), model, json)
   }
 
+  @Test fun removeCountMustMatchListSize() {
+    assertFailsWith<IllegalArgumentException>("Count 4 != Removed ID list size 3") {
+      ChildrenDiff.Remove(1, 2, 3, 4, listOf(5, 6, 7))
+    }
+  }
+
   private fun <T> assertJsonRoundtrip(serializer: KSerializer<T>, model: T, json: String) {
     assertEquals(json, format.encodeToString(serializer, model))
     assertEquals(model, format.decodeFromString(serializer, json))
diff --git a/treehouse/treehouse-widget/src/commonMain/kotlin/app/cash/treehouse/widget/WidgetDisplay.kt b/treehouse/treehouse-widget/src/commonMain/kotlin/app/cash/treehouse/widget/WidgetDisplay.kt
index 1c04a3088..c3fdca8ae 100644
--- a/treehouse/treehouse-widget/src/commonMain/kotlin/app/cash/treehouse/widget/WidgetDisplay.kt
+++ b/treehouse/treehouse-widget/src/commonMain/kotlin/app/cash/treehouse/widget/WidgetDisplay.kt
@@ -55,7 +55,7 @@ public class WidgetDisplay<T : Any>(
         }
         is ChildrenDiff.Remove -> {
           children.remove(childrenDiff.index, childrenDiff.count)
-          // TODO we need to remove widgets from our map!
+          widgets.keys.removeAll(childrenDiff.removedIds)
         }
         ChildrenDiff.Clear -> {
           children.clear()