Correctly dispose in WithConstraints

Test: new test
Change-Id: Ie9e7b0694c89682e58cf3ead5b3252b6d0b13925

diff --git a/ui/ui-core/src/androidTest/java/androidx/ui/core/test/WithConstraintsTest.kt b/ui/ui-core/src/androidTest/java/androidx/ui/core/test/WithConstraintsTest.kt
index cc7b67c..242bec2 100644
--- a/ui/ui-core/src/androidTest/java/androidx/ui/core/test/WithConstraintsTest.kt
+++ b/ui/ui-core/src/androidTest/java/androidx/ui/core/test/WithConstraintsTest.kt

@@ -23,6 +23,8 @@
 import androidx.compose.Composable
 import androidx.compose.emptyContent
 import androidx.compose.getValue
+import androidx.compose.mutableStateOf
+import androidx.compose.onDispose
 import androidx.compose.remember
 import androidx.compose.setValue
 import androidx.compose.state
@@ -689,6 +691,32 @@
         assertTrue(childMeasureLatch.await(1, TimeUnit.SECONDS))
     }
 
+    @Test
+    fun onDisposeInsideWithConstraintsCalled() {
+        var emit by mutableStateOf(true)
+        val composedLatch = CountDownLatch(1)
+        val disposedLatch = CountDownLatch(1)
+        rule.runOnUiThreadIR {
+            activity.setContent {
+                if (emit) {
+                    WithConstraints { _, _ ->
+                        composedLatch.countDown()
+                        onDispose {
+                            disposedLatch.countDown()
+                        }
+                    }
+                }
+            }
+        }
+
+        assertTrue(composedLatch.await(1, TimeUnit.SECONDS))
+
+        rule.runOnUiThread {
+            emit = false
+        }
+        assertTrue(disposedLatch.await(1, TimeUnit.SECONDS))
+    }
+
     private fun takeScreenShot(size: Int): Bitmap {
         assertTrue(drawLatch.await(1, TimeUnit.SECONDS))
         val bitmap = rule.waitAndScreenShot()

diff --git a/ui/ui-core/src/main/java/androidx/ui/core/Layout.kt b/ui/ui-core/src/main/java/androidx/ui/core/Layout.kt
index 47b130a..030a960 100644
--- a/ui/ui-core/src/main/java/androidx/ui/core/Layout.kt
+++ b/ui/ui-core/src/main/java/androidx/ui/core/Layout.kt

@@ -18,12 +18,14 @@
 
 import android.content.Context
 import androidx.compose.Composable
+import androidx.compose.Composition
 import androidx.compose.CompositionReference
 import androidx.compose.FrameManager
 import androidx.compose.Recomposer
 import androidx.compose.Untracked
 import androidx.compose.compositionReference
 import androidx.compose.currentComposer
+import androidx.compose.onDispose
 import androidx.compose.remember
 import androidx.ui.unit.Density
 import androidx.ui.unit.IntPx
@@ -476,6 +478,9 @@
     if (!layoutNode.needsRemeasure && layoutNode.owner != null) {
         state.subcompose()
     }
+    onDispose {
+        state.composition?.dispose()
+    }
 }
 
 private class WithConstrainsState {
@@ -486,6 +491,7 @@
     var lastConstraints: Constraints? = null
     var children: @Composable() (Constraints, LayoutDirection) -> Unit = { _, _ -> }
     var forceRecompose = false
+    var composition: Composition? = null
 
     val measureBlocks = object : LayoutNode.NoIntrinsicsMeasureBlocks(
         error = "Intrinsic measurements are not supported by WithConstraints"
@@ -529,7 +535,7 @@
         val node = nodeRef.value!!
         val constraints = lastConstraints!!
         // TODO(b/150390669): Review use of @Untracked
-        subcomposeInto(context, node, recomposer, compositionRef) @Untracked {
+        composition = subcomposeInto(context, node, recomposer, compositionRef) @Untracked {
             children(constraints, node.measureScope.layoutDirection)
         }
         forceRecompose = false