Don't create a derived State for Element.alpha

This CL changes the logic in Modifier.element() so that we don't create
a derived State object for each element. This paves the way towards
using the Modifier Node API for Modifier.element().

Bug: 291566282
Test: atest PlatformComposeSceneTransitionLayoutTests
Flag: NA
Change-Id: Ia778559659a44e4cb7785882cf56e9bb67e2d803
diff --git a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt
index d853172..6153e19 100644
--- a/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt
+++ b/packages/SystemUI/compose/scene/src/com/android/compose/animation/scene/Element.kt
@@ -18,7 +18,6 @@
 
 import androidx.compose.runtime.Composable
 import androidx.compose.runtime.DisposableEffect
-import androidx.compose.runtime.SideEffect
 import androidx.compose.runtime.derivedStateOf
 import androidx.compose.runtime.getValue
 import androidx.compose.runtime.movableContentOf
@@ -36,7 +35,6 @@
 import androidx.compose.ui.geometry.isUnspecified
 import androidx.compose.ui.geometry.lerp
 import androidx.compose.ui.graphics.drawscope.scale
-import androidx.compose.ui.graphics.graphicsLayer
 import androidx.compose.ui.layout.IntermediateMeasureScope
 import androidx.compose.ui.layout.Measurable
 import androidx.compose.ui.layout.Placeable
@@ -47,7 +45,6 @@
 import androidx.compose.ui.unit.round
 import com.android.compose.animation.scene.transformation.PropertyTransformation
 import com.android.compose.animation.scene.transformation.SharedElementTransformation
-import com.android.compose.modifiers.thenIf
 import com.android.compose.ui.util.lerp
 
 /** An element on screen, that can be composed in one or more scenes. */
@@ -146,8 +143,6 @@
 
             element
         }
-    val lastSharedValues = element.lastSharedValues
-    val lastSceneValues = sceneValues.lastValues
 
     DisposableEffect(scene, sceneValues, element) {
         onDispose {
@@ -160,18 +155,6 @@
         }
     }
 
-    val alpha =
-        remember(layoutImpl, element, scene, sceneValues) {
-            derivedStateOf { elementAlpha(layoutImpl, element, scene, sceneValues) }
-        }
-    val isOpaque by remember(alpha) { derivedStateOf { alpha.value == 1f } }
-    SideEffect {
-        if (isOpaque) {
-            lastSharedValues.alpha = 1f
-            lastSceneValues.alpha = 1f
-        }
-    }
-
     val drawScale by
         remember(layoutImpl, element, scene, sceneValues) {
             derivedStateOf { getDrawScale(layoutImpl, element, scene, sceneValues) }
@@ -200,14 +183,6 @@
                 place(layoutImpl, scene, element, sceneValues, placeable, placementScope = this)
             }
         }
-        .thenIf(!isOpaque) {
-            Modifier.graphicsLayer {
-                val alpha = alpha.value
-                this.alpha = alpha
-                lastSharedValues.alpha = alpha
-                lastSceneValues.alpha = alpha
-            }
-        }
         .testTag(key.testTag)
 }
 
@@ -325,6 +300,61 @@
     }
 }
 
+/**
+ * Whether the element is opaque or not.
+ *
+ * Important: The logic here should closely match the logic in [elementAlpha]. Note that we don't
+ * reuse [elementAlpha] and simply check if alpha == 1f because [isElementOpaque] is checked during
+ * placement and we don't want to read the transition progress in that phase.
+ */
+private fun isElementOpaque(
+    layoutImpl: SceneTransitionLayoutImpl,
+    element: Element,
+    scene: Scene,
+    sceneValues: Element.TargetValues,
+): Boolean {
+    val state = layoutImpl.state.transitionState
+
+    if (state !is TransitionState.Transition || state.fromScene == state.toScene) {
+        return true
+    }
+
+    if (!layoutImpl.isTransitionReady(state)) {
+        val lastValue =
+            sceneValues.lastValues.alpha.takeIf { it != Element.AlphaUnspecified }
+                ?: element.lastSharedValues.alpha.takeIf { it != Element.AlphaUnspecified } ?: 1f
+
+        return lastValue == 1f
+    }
+
+    val fromScene = state.fromScene
+    val toScene = state.toScene
+    val fromValues = element.sceneValues[fromScene]
+    val toValues = element.sceneValues[toScene]
+
+    if (fromValues == null && toValues == null) {
+        error("This should not happen, element $element is neither in $fromScene or $toScene")
+    }
+
+    val isSharedElement = fromValues != null && toValues != null
+    if (isSharedElement && isSharedElementEnabled(layoutImpl, state, element.key)) {
+        return true
+    }
+
+    return layoutImpl.transitions
+        .transitionSpec(fromScene, toScene)
+        .transformations(element.key, scene.key)
+        .alpha == null
+}
+
+/**
+ * Whether the element is opaque or not.
+ *
+ * Important: The logic here should closely match the logic in [isElementOpaque]. Note that we don't
+ * reuse [elementAlpha] in [isElementOpaque] and simply check if alpha == 1f because
+ * [isElementOpaque] is checked during placement and we don't want to read the transition progress
+ * in that phase.
+ */
 private fun elementAlpha(
     layoutImpl: SceneTransitionLayoutImpl,
     element: Element,
@@ -446,6 +476,8 @@
         }
 
         val currentOffset = lookaheadScopeCoordinates.localPositionOf(coords, Offset.Zero)
+        val lastSharedValues = element.lastSharedValues
+        val lastValues = sceneValues.lastValues
         val targetOffset =
             computeValue(
                 layoutImpl,
@@ -456,19 +488,31 @@
                 idleValue = targetOffsetInScene,
                 currentValue = { currentOffset },
                 lastValue = {
-                    sceneValues.lastValues.offset.takeIf { it.isSpecified }
-                        ?: element.lastSharedValues.offset.takeIf { it.isSpecified }
-                            ?: currentOffset
+                    lastValues.offset.takeIf { it.isSpecified }
+                        ?: lastSharedValues.offset.takeIf { it.isSpecified } ?: currentOffset
                 },
                 ::lerp,
             )
 
-        element.lastSharedValues.offset = targetOffset
-        sceneValues.lastValues.offset = targetOffset
+        lastSharedValues.offset = targetOffset
+        lastValues.offset = targetOffset
 
-        // TODO(b/291071158): Call placeWithLayer() if offset != IntOffset.Zero and size is not
-        // animated once b/305195729 is fixed. Test that drawing is not invalidated in that case.
-        placeable.place((targetOffset - currentOffset).round())
+        val offset = (targetOffset - currentOffset).round()
+        if (isElementOpaque(layoutImpl, element, scene, sceneValues)) {
+            // TODO(b/291071158): Call placeWithLayer() if offset != IntOffset.Zero and size is not
+            // animated once b/305195729 is fixed. Test that drawing is not invalidated in that
+            // case.
+            placeable.place(offset)
+            lastSharedValues.alpha = 1f
+            lastValues.alpha = 1f
+        } else {
+            placeable.placeWithLayer(offset) {
+                val alpha = elementAlpha(layoutImpl, element, scene, sceneValues)
+                this.alpha = alpha
+                lastSharedValues.alpha = alpha
+                lastValues.alpha = alpha
+            }
+        }
     }
 }