Improve composition local performance

With this change the composer will cache the composition local
scope found by call to CompositionLocal.current and reuse it when
until it is invalidated by composition restart or the provider local
scope closes.

Includes additional synthetic benchmarks that show improvenent on
mulitple looks-ups in the same scope.

Test: ./gradlew :compose:r:r:tDUT
Change-Id: I21aa88323872ab932f343a917180adee1471e657
diff --git a/compose/runtime/runtime/compose-runtime-benchmark/src/androidTest/java/androidx/compose/runtime/benchmark/CompositionLocalBenchmark.kt b/compose/runtime/runtime/compose-runtime-benchmark/src/androidTest/java/androidx/compose/runtime/benchmark/CompositionLocalBenchmark.kt
new file mode 100644
index 0000000..e973ae4
--- /dev/null
+++ b/compose/runtime/runtime/compose-runtime-benchmark/src/androidTest/java/androidx/compose/runtime/benchmark/CompositionLocalBenchmark.kt
@@ -0,0 +1,208 @@
+/*
+ * Copyright 2022 The Android Open Source Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *      http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package androidx.compose.runtime.benchmark
+
+import androidx.compose.runtime.Composable
+import androidx.compose.runtime.CompositionLocalProvider
+import androidx.compose.runtime.compositionLocalOf
+import androidx.compose.runtime.getValue
+import androidx.compose.runtime.mutableStateOf
+import androidx.compose.runtime.setValue
+import androidx.test.annotation.UiThreadTest
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.LargeTest
+import kotlinx.coroutines.ExperimentalCoroutinesApi
+import org.junit.Ignore
+import org.junit.Test
+import org.junit.runner.RunWith
+
+val local = compositionLocalOf { 0 }
+
+@LargeTest
+@RunWith(AndroidJUnit4::class)
+@OptIn(ExperimentalCoroutinesApi::class)
+class CompositionLocalBenchmark : ComposeBenchmarkBase() {
+
+    @UiThreadTest
+    @Test
+    @Ignore // Only used for overhead comparison, not to be tracked.
+    fun compositionLocal_compose_depth_1_1() = runBlockingTestWithFrameClock {
+        measureCompose {
+            CompositionLocalProvider(local provides 100) {
+                DepthOf(1) {
+                    local.current
+                }
+            }
+        }
+    }
+
+    @UiThreadTest
+    @Test
+    @Ignore // Only used for overhead comparison, not to be tracked.
+    fun compositionLocal_compose_depth_1_10() = runBlockingTestWithFrameClock {
+        measureCompose {
+            CompositionLocalProvider(local provides 100) {
+                DepthOf(1) {
+                    repeat(10) { local.current }
+                }
+            }
+        }
+    }
+
+    @UiThreadTest
+    @Test
+    @Ignore // Only used for overhead comparison, not to be tracked.
+    fun compositionLocal_compose_depth_1_100() = runBlockingTestWithFrameClock {
+        measureCompose {
+            CompositionLocalProvider(local provides 100) {
+                DepthOf(1) {
+                    repeat(100) { local.current }
+                }
+            }
+        }
+    }
+
+    @UiThreadTest
+    @Test
+    @Ignore // Only used for overhead comparison, not to be tracked.
+    fun compositionLocal_compose_depth_100_1() = runBlockingTestWithFrameClock {
+        measureCompose {
+            CompositionLocalProvider(local provides 100) {
+                DepthOf(100) {
+                    local.current
+                }
+            }
+        }
+    }
+
+    @UiThreadTest
+    @Test
+    @Ignore // Only used for overhead comparison, not to be tracked.
+    fun compositionLocal_compose_depth_100_10() = runBlockingTestWithFrameClock {
+        measureCompose {
+            CompositionLocalProvider(local provides 100) {
+                DepthOf(100) {
+                    repeat(10) { local.current }
+                }
+            }
+        }
+    }
+
+    @UiThreadTest
+    @Test
+    @Ignore // Only used for overhead comparison, not to be tracked.
+    fun compositionLocal_compose_depth_100_100() = runBlockingTestWithFrameClock {
+        measureCompose {
+            CompositionLocalProvider(local provides 100) {
+                DepthOf(100) {
+                    repeat(100) { local.current }
+                }
+            }
+        }
+    }
+
+    @UiThreadTest
+    @Test
+    fun compositionLocal_compose_depth_10000_1() = runBlockingTestWithFrameClock {
+        measureCompose {
+            CompositionLocalProvider(local provides 100) {
+                DepthOf(10000) {
+                    local.current
+                }
+            }
+        }
+    }
+    @UiThreadTest
+    @Test
+    @Ignore // Only used for overhead comparison, not to be tracked.
+    fun compositionLocal_compose_depth_10000_10() = runBlockingTestWithFrameClock {
+        measureCompose {
+            CompositionLocalProvider(local provides 100) {
+                DepthOf(10000) {
+                    repeat(10) { local.current }
+                }
+            }
+        }
+    }
+
+    // This is the only one of the "compose" benchmarks that should be tracked.
+    @UiThreadTest
+    @Test
+    fun compositionLocal_compose_depth_10000_100() = runBlockingTestWithFrameClock {
+        measureCompose {
+            CompositionLocalProvider(local provides 100) {
+                DepthOf(10000) {
+                    repeat(100) { local.current }
+                }
+            }
+        }
+    }
+
+    @UiThreadTest
+    @Test
+    fun compositionLocal_recompose_depth_10000_1() = runBlockingTestWithFrameClock {
+        var data by mutableStateOf(0)
+        var sync: Int = 0
+
+        measureRecomposeSuspending {
+            compose {
+                DepthOf(10000) {
+                    // Force the read to occur in a way that is difficult for the compiler to figure
+                    // out that it is not used.
+                    sync = data
+                    repeat(1) { local.current }
+                }
+            }
+            update {
+                data++
+            }
+        }
+        if (sync > Int.MAX_VALUE / 2) {
+            println("This is just to fool the compiler into thinking sync is used")
+        }
+    }
+
+    @UiThreadTest
+    @Test
+    fun compositionLocal_recompose_depth_10000_100() = runBlockingTestWithFrameClock {
+        var data by mutableStateOf(0)
+        var sync: Int = 0
+
+        measureRecomposeSuspending {
+            compose {
+                DepthOf(10000) {
+                    // Force the read to occur in a way that is difficult for the compiler to figure
+                    // out that it is not used.
+                    sync = data
+                    repeat(100) { local.current }
+                }
+            }
+            update {
+                data++
+            }
+        }
+        if (sync > Int.MAX_VALUE / 2) {
+            println("This is just to fool the compiler into thinking sync is used")
+        }
+    }
+}
+
+@Composable
+fun DepthOf(count: Int, content: @Composable () -> Unit) {
+    if (count > 0) DepthOf(count - 1, content)
+    else content()
+}
\ No newline at end of file
diff --git a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/Composer.kt b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/Composer.kt
index 92c760d..a2454c1 100644
--- a/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/Composer.kt
+++ b/compose/runtime/runtime/src/commonMain/kotlin/androidx/compose/runtime/Composer.kt
@@ -1188,6 +1188,8 @@
 
     private var writer: SlotWriter = insertTable.openWriter().also { it.close() }
     private var writerHasAProvider = false
+    private var providerCache: CompositionLocalMap? = null
+
     private var insertAnchor: Anchor = insertTable.read { it.anchor(0) }
     private val insertFixups = mutableListOf<Change>()
 
@@ -1328,6 +1330,7 @@
         parentProvider = parentContext.getCompositionLocalScope()
         providersInvalidStack.push(providersInvalid.asInt())
         providersInvalid = changed(parentProvider)
+        providerCache = null
         if (!forceRecomposeScopes) {
             forceRecomposeScopes = parentContext.collectingParameterInformation
         }
@@ -1770,6 +1773,8 @@
      * Return the current [CompositionLocal] scope which was provided by a parent group.
      */
     private fun currentCompositionLocalScope(group: Int? = null): CompositionLocalMap {
+        if (group == null)
+            providerCache?.let { return it }
         if (inserting && writerHasAProvider) {
             var current = writer.parent
             while (current > 0) {
@@ -1777,7 +1782,9 @@
                     writer.groupObjectKey(current) == compositionLocalMap
                 ) {
                     @Suppress("UNCHECKED_CAST")
-                    return writer.groupAux(current) as CompositionLocalMap
+                    val providers = writer.groupAux(current) as CompositionLocalMap
+                    providerCache = providers
+                    return providers
                 }
                 current = writer.parent(current)
             }
@@ -1789,12 +1796,15 @@
                     reader.groupObjectKey(current) == compositionLocalMap
                 ) {
                     @Suppress("UNCHECKED_CAST")
-                    return providerUpdates[current]
+                    val providers = providerUpdates[current]
                         ?: reader.groupAux(current) as CompositionLocalMap
+                    providerCache = providers
+                    return providers
                 }
                 current = reader.parent(current)
             }
         }
+        providerCache = parentProvider
         return parentProvider
     }
 
@@ -1864,6 +1874,7 @@
         }
         providersInvalidStack.push(providersInvalid.asInt())
         providersInvalid = invalid
+        providerCache = providers
         start(compositionLocalMapKey, compositionLocalMap, false, providers)
     }
 
@@ -1872,6 +1883,7 @@
         endGroup()
         endGroup()
         providersInvalid = providersInvalidStack.pop().asBool()
+        providerCache = null
     }
 
     @InternalComposeApi
@@ -1929,6 +1941,7 @@
             // Append to the end of the table
             writer.skipToGroupEnd()
             writerHasAProvider = false
+            providerCache = null
         }
     }
 
@@ -2034,6 +2047,7 @@
                 // inserted into in the table.
                 reader.beginEmpty()
                 inserting = true
+                providerCache = null
                 ensureWriter()
                 writer.beginInsert()
                 val startIndex = writer.currentGroup
@@ -2294,6 +2308,10 @@
                     recomposeCompoundKey
                 )
 
+                // We have moved so the cached lookup of the provider is invalid
+                providerCache = null
+
+                // Invoke the scope's composition function
                 firstInRange.scope.compose(this)
 
                 // Restore the parent of the reader to the previous parent
@@ -2741,6 +2759,8 @@
         // needs to be created as a late change.
         if (inserting && !force) {
             writerHasAProvider = true
+            providerCache = null
+
             // Create an anchor to the movable group
             val anchor = writer.anchor(writer.parent(writer.parent))
             val reference = MovableContentStateReference(