Add experimental snapshot unsafeEnter/Leave

Add methods to enter and leave a snapshot without any scoping
guardrails. Useful for writing abstractions like SnapshotContextElement
that need to enter and leave the snapshot in response to callbacks that
cannot be bridged into the inline function scope of enter {}.

Test: SnapshotTests.kt
Relnote: Added experimental Snapshot.unsafeEnter/unsafeLeave
Change-Id: I108f3f633fdeed300c1260d62effd2749e38bbb3
diff --git a/compose/runtime/runtime/api/public_plus_experimental_current.txt b/compose/runtime/runtime/api/public_plus_experimental_current.txt
index 6e27e40..4a2b67c 100644
--- a/compose/runtime/runtime/api/public_plus_experimental_current.txt
+++ b/compose/runtime/runtime/api/public_plus_experimental_current.txt
@@ -845,6 +845,8 @@
     method public abstract androidx.compose.runtime.snapshots.Snapshot getRoot();
     method public abstract boolean hasPendingChanges();
     method public abstract androidx.compose.runtime.snapshots.Snapshot takeNestedSnapshot(optional kotlin.jvm.functions.Function1<java.lang.Object,kotlin.Unit>? readObserver);
+    method @androidx.compose.runtime.ExperimentalComposeApi public final androidx.compose.runtime.snapshots.Snapshot? unsafeEnter();
+    method @androidx.compose.runtime.ExperimentalComposeApi public final void unsafeLeave(androidx.compose.runtime.snapshots.Snapshot? oldSnapshot);
     property public int id;
     property public abstract boolean readOnly;
     property public abstract androidx.compose.runtime.snapshots.Snapshot root;
diff --git a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/snapshots/Snapshot.kt b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/snapshots/Snapshot.kt
index f8ffc53..fcb4815 100644
--- a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/snapshots/Snapshot.kt
+++ b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/snapshots/Snapshot.kt
@@ -21,6 +21,7 @@
 import androidx.compose.runtime.AtomicReference
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.DisallowComposableCalls
+import androidx.compose.runtime.ExperimentalComposeApi
 import androidx.compose.runtime.InternalComposeApi
 import androidx.compose.runtime.SnapshotThreadLocal
 import androidx.compose.runtime.synchronized
@@ -104,7 +105,7 @@
      * or a nested call to [enter] is called. When [block] returns, the previous current snapshot
      * is restored if there was one.
      *
-     * All changes to state object inside [block] are isolated to this snapshot and are not
+     * All changes to state objects inside [block] are isolated to this snapshot and are not
      * visible to other snapshot or as global state. If this is a [readOnly] snapshot, any
      * changes to state objects will throw an [IllegalStateException].
      *
@@ -137,6 +138,40 @@
         threadSnapshot.set(snapshot)
     }
 
+    /**
+     * Enter the snapshot, returning the previous [Snapshot] for leaving this snapshot later
+     * using [unsafeLeave]. Prefer [enter] or [asContextElement] instead of using [unsafeEnter]
+     * directly to prevent mismatched [unsafeEnter]/[unsafeLeave] calls.
+     *
+     * After returning all state objects have the value associated with this snapshot.
+     * The value of [currentSnapshot] will be this snapshot until [unsafeLeave] is called
+     * with the returned [Snapshot] or another call to [unsafeEnter] or [enter]
+     * is made.
+     *
+     * All changes to state objects until another snapshot is entered or this snapshot is left
+     * are isolated to this snapshot and are not visible to other snapshot or as global state.
+     * If this is a [readOnly] snapshot, any changes to state objects will throw an
+     * [IllegalStateException].
+     *
+     * For a [MutableSnapshot], changes made to a snapshot can be applied
+     * atomically to the global state (or to its parent snapshot if it is a nested snapshot) by
+     * calling [MutableSnapshot.apply].
+     */
+    @ExperimentalComposeApi
+    fun unsafeEnter(): Snapshot? = makeCurrent()
+
+    /**
+     * Leave the snapshot, restoring the [oldSnapshot] before returning.
+     * See [unsafeEnter].
+     */
+    @ExperimentalComposeApi
+    fun unsafeLeave(oldSnapshot: Snapshot?) {
+        check(threadSnapshot.get() === this) {
+            "Cannot leave snapshot; $this is not the current snapshot"
+        }
+        restoreCurrent(oldSnapshot)
+    }
+
     internal var disposed = false
 
     /*
diff --git a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/snapshots/SnapshotContextElement.kt b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/snapshots/SnapshotContextElement.kt
index 1ecd64b..b4a7f86 100644
--- a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/snapshots/SnapshotContextElement.kt
+++ b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/snapshots/SnapshotContextElement.kt
@@ -52,9 +52,9 @@
         get() = SnapshotContextElement
 
     override fun updateThreadContext(context: CoroutineContext): Snapshot? =
-        snapshot.makeCurrent()
+        snapshot.unsafeEnter()
 
     override fun restoreThreadContext(context: CoroutineContext, oldState: Snapshot?) {
-        snapshot.restoreCurrent(oldState)
+        snapshot.unsafeLeave(oldState)
     }
 }
diff --git a/compose/runtime/runtime/src/commonTest/kotlin/androidx/compose/runtime/snapshots/SnapshotTests.kt b/compose/runtime/runtime/src/commonTest/kotlin/androidx/compose/runtime/snapshots/SnapshotTests.kt
index b51a20a..8d279be 100644
--- a/compose/runtime/runtime/src/commonTest/kotlin/androidx/compose/runtime/snapshots/SnapshotTests.kt
+++ b/compose/runtime/runtime/src/commonTest/kotlin/androidx/compose/runtime/snapshots/SnapshotTests.kt
@@ -18,6 +18,7 @@
 
 package androidx.compose.runtime.snapshots
 
+import androidx.compose.runtime.ExperimentalComposeApi
 import androidx.compose.runtime.InternalComposeApi
 import androidx.compose.runtime.MutableState
 import androidx.compose.runtime.State
@@ -27,6 +28,7 @@
 import androidx.compose.runtime.neverEqualPolicy
 import androidx.compose.runtime.referentialEqualityPolicy
 import androidx.compose.runtime.setValue
+import androidx.compose.runtime.snapshots.Snapshot.Companion.current
 import androidx.compose.runtime.snapshots.Snapshot.Companion.openSnapshotCount
 import androidx.compose.runtime.snapshots.Snapshot.Companion.takeMutableSnapshot
 import androidx.compose.runtime.snapshots.Snapshot.Companion.takeSnapshot
@@ -37,6 +39,8 @@
 import kotlin.test.assertEquals
 import kotlin.test.assertFailsWith
 import kotlin.test.assertFalse
+import kotlin.test.assertNotSame
+import kotlin.test.assertSame
 import kotlin.test.assertTrue
 import kotlin.test.fail
 
@@ -899,6 +903,39 @@
         mutable1.dispose()
     }
 
+    @OptIn(ExperimentalComposeApi::class)
+    @Test
+    fun testUnsafeSnapshotEnterAndLeave() {
+        val snapshot = takeSnapshot()
+        try {
+            val oldSnapshot = snapshot.unsafeEnter()
+            try {
+                assertSame(snapshot, current, "expected taken snapshot to be current")
+            } finally {
+                snapshot.unsafeLeave(oldSnapshot)
+            }
+            assertNotSame(snapshot, current, "expected taken snapshot not to be current")
+        } finally {
+            snapshot.dispose()
+        }
+    }
+
+    @OptIn(ExperimentalComposeApi::class)
+    @Test
+    fun testUnsafeSnapshotLeaveThrowsIfNotCurrent() {
+        val snapshot = takeSnapshot()
+        try {
+            try {
+                snapshot.unsafeLeave(null)
+                fail("unsafeLeave should have thrown")
+            } catch (ise: IllegalStateException) {
+                // expected
+            }
+        } finally {
+            snapshot.dispose()
+        }
+    }
+
     private var count = 0
 
     @BeforeTest