Generalize shared elements transition logic

Update shared elements transition logic in a way that allows an
ImagePreviewView implementation to specify multiple transition elements.

Flag: IntentResolver package entirely behind the CHOOSER_UNBUNDLED which
is in teamfood

Bug: 262280076
Test: manual shcreenshot animation test
Test atest IntentResolverUnitTests

Change-Id: Ia7cf5634bb2d907c5cdb56a22f838447a158dd25

8 files changed, 117 insertions, 85 deletions

diff --git a/java/res/layout/image_preview_view.xml b/java/res/layout/image_preview_view.xml
index d2f94690..8730fc30 100644
--- a/java/res/layout/image_preview_view.xml
+++ b/java/res/layout/image_preview_view.xml
@@ -25,6 +25,7 @@
 
     <com.android.intentresolver.widget.RoundedRectImageView
         android:id="@androidprv:id/content_preview_image_1_large"
+        android:transitionName="screenshot_preview_image"
         android:layout_width="120dp"
         android:layout_height="104dp"
         android:layout_alignParentTop="true"
diff --git a/java/src/com/android/intentresolver/ChooserActivity.java b/java/src/com/android/intentresolver/ChooserActivity.java
index dce8bde1..ca1f3f9a 100644
--- a/java/src/com/android/intentresolver/ChooserActivity.java
+++ b/java/src/com/android/intentresolver/ChooserActivity.java
@@ -756,16 +756,13 @@ public class ChooserActivity extends ResolverActivity implements
                         : R.layout.chooser_action_row,
                 parent,
                 imageLoader,
-                mEnterTransitionAnimationDelegate::markImagePreviewReady,
+                mEnterTransitionAnimationDelegate,
                 getContentResolver(),
                 this::isImageType);
 
         if (layout != null) {
             adjustPreviewWidth(getResources().getConfiguration().orientation, layout);
         }
-        if (previewType != ChooserContentPreviewUi.CONTENT_PREVIEW_IMAGE) {
-            mEnterTransitionAnimationDelegate.markImagePreviewReady(false);
-        }
 
         return layout;
     }
diff --git a/java/src/com/android/intentresolver/ChooserContentPreviewUi.java b/java/src/com/android/intentresolver/ChooserContentPreviewUi.java
index 27645da4..67cfbff2 100644
--- a/java/src/com/android/intentresolver/ChooserContentPreviewUi.java
+++ b/java/src/com/android/intentresolver/ChooserContentPreviewUi.java
@@ -47,6 +47,7 @@ import androidx.annotation.Nullable;
 
 import com.android.intentresolver.widget.ActionRow;
 import com.android.intentresolver.widget.ImagePreviewView;
+import com.android.intentresolver.widget.ImagePreviewView.TransitionElementStatusCallback;
 import com.android.intentresolver.widget.RoundedRectImageView;
 import com.android.internal.annotations.VisibleForTesting;
 
@@ -55,7 +56,7 @@ import java.util.ArrayList;
 import java.util.HashMap;
 import java.util.List;
 import java.util.Map;
-import java.util.function.Consumer;
+import java.util.stream.Collectors;
 
 /**
  * Collection of helpers for building the content preview UI displayed in {@link ChooserActivity}.
@@ -170,11 +171,14 @@ public final class ChooserContentPreviewUi {
             @LayoutRes int actionRowLayout,
             ViewGroup parent,
             ImageLoader previewImageLoader,
-            Consumer<Boolean> onTransitionTargetReady,
+            TransitionElementStatusCallback transitionElementStatusCallback,
             ContentResolver contentResolver,
             ImageMimeTypeClassifier imageClassifier) {
         ViewGroup layout = null;
 
+        if (previewType != CONTENT_PREVIEW_IMAGE) {
+            transitionElementStatusCallback.onAllTransitionElementsReady();
+        }
         List<ActionRow.Action> customActions = actionFactory.createCustomActions();
         switch (previewType) {
             case CONTENT_PREVIEW_TEXT:
@@ -197,7 +201,7 @@ public final class ChooserContentPreviewUi {
                                 customActions),
                         parent,
                         previewImageLoader,
-                        onTransitionTargetReady,
+                        transitionElementStatusCallback,
                         contentResolver,
                         imageClassifier,
                         actionRowLayout);
@@ -326,7 +330,7 @@ public final class ChooserContentPreviewUi {
             List<ActionRow.Action> actions,
             ViewGroup parent,
             ImageLoader imageLoader,
-            Consumer<Boolean> onTransitionTargetReady,
+            TransitionElementStatusCallback transitionElementStatusCallback,
             ContentResolver contentResolver,
             ImageMimeTypeClassifier imageClassifier,
             @LayoutRes int actionRowLayout) {
@@ -340,32 +344,26 @@ public final class ChooserContentPreviewUi {
             actionRow.setActions(actions);
         }
 
-        final ArrayList<Uri> imageUris = new ArrayList<>();
         String action = targetIntent.getAction();
-        if (Intent.ACTION_SEND.equals(action)) {
-            // TODO: why don't we use image classifier in this case as well?
-            Uri uri = targetIntent.getParcelableExtra(Intent.EXTRA_STREAM);
-            imageUris.add(uri);
-        } else {
-            List<Uri> uris = targetIntent.getParcelableArrayListExtra(Intent.EXTRA_STREAM);
-            for (Uri uri : uris) {
-                if (imageClassifier.isImageType(contentResolver.getType(uri))) {
-                    imageUris.add(uri);
-                }
-            }
-        }
+        // TODO: why don't we use image classifier for single-element ACTION_SEND?
+        final List<Uri> imageUris = Intent.ACTION_SEND.equals(action)
+                ? extractContentUris(targetIntent)
+                : extractContentUris(targetIntent)
+                        .stream()
+                        .filter(uri ->
+                                imageClassifier.isImageType(contentResolver.getType(uri))
+                        )
+                        .collect(Collectors.toList());
 
         if (imageUris.size() == 0) {
             Log.i(TAG, "Attempted to display image preview area with zero"
                     + " available images detected in EXTRA_STREAM list");
             ((View) imagePreview).setVisibility(View.GONE);
-            onTransitionTargetReady.accept(false);
+            transitionElementStatusCallback.onAllTransitionElementsReady();
             return contentPreviewLayout;
         }
 
-        imagePreview.setSharedElementTransitionTarget(
-                ChooserActivity.FIRST_IMAGE_PREVIEW_TRANSITION_NAME,
-                onTransitionTargetReady);
+        imagePreview.setTransitionElementStatusCallback(transitionElementStatusCallback);
         imagePreview.setImages(imageUris, imageLoader);
 
         return contentPreviewLayout;
@@ -398,7 +396,7 @@ public final class ChooserContentPreviewUi {
         ViewGroup contentPreviewLayout = (ViewGroup) layoutInflater.inflate(
                 R.layout.chooser_grid_preview_file, parent, false);
 
-        List<Uri> uris = extractFileUris(targetIntent);
+        List<Uri> uris = extractContentUris(targetIntent);
         final int uriCount = uris.size();
 
         if (uriCount == 0) {
@@ -442,7 +440,7 @@ public final class ChooserContentPreviewUi {
         return contentPreviewLayout;
     }
 
-    private static List<Uri> extractFileUris(Intent targetIntent) {
+    private static List<Uri> extractContentUris(Intent targetIntent) {
         List<Uri> uris = new ArrayList<>();
         if (Intent.ACTION_SEND.equals(targetIntent.getAction())) {
             Uri uri = targetIntent.getParcelableExtra(Intent.EXTRA_STREAM);
diff --git a/java/src/com/android/intentresolver/EnterTransitionAnimationDelegate.kt b/java/src/com/android/intentresolver/EnterTransitionAnimationDelegate.kt
index 31aeea44..b1178aa5 100644
--- a/java/src/com/android/intentresolver/EnterTransitionAnimationDelegate.kt
+++ b/java/src/com/android/intentresolver/EnterTransitionAnimationDelegate.kt
@@ -19,6 +19,7 @@ import android.app.SharedElementCallback
 import android.view.View
 import androidx.activity.ComponentActivity
 import androidx.lifecycle.lifecycleScope
+import com.android.intentresolver.widget.ImagePreviewView.TransitionElementStatusCallback
 import com.android.internal.annotations.VisibleForTesting
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.delay
@@ -29,13 +30,13 @@ import java.util.function.Supplier
  * A helper class to track app's readiness for the scene transition animation.
  * The app is ready when both the image is laid out and the drawer offset is calculated.
  */
-@VisibleForTesting
+@VisibleForTesting(visibility = VisibleForTesting.Visibility.PACKAGE)
 class EnterTransitionAnimationDelegate(
     private val activity: ComponentActivity,
     private val transitionTargetSupplier: Supplier<View?>,
-) : View.OnLayoutChangeListener {
+) : View.OnLayoutChangeListener, TransitionElementStatusCallback {
 
-    private var removeSharedElements = false
+    private val transitionElements = HashSet<String>()
     private var previewReady = false
     private var offsetCalculated = false
     private var timeoutJob: Job? = null
@@ -65,14 +66,15 @@ class EnterTransitionAnimationDelegate(
         // We only mark the preview readiness and not the offset readiness
         // (see [#markOffsetCalculated()]) as this is what legacy logic, effectively, did. We might
         // want to review that aspect separately.
-        markImagePreviewReady(runTransitionAnimation = false)
+        onAllTransitionElementsReady()
     }
 
-    fun markImagePreviewReady(runTransitionAnimation: Boolean) {
+    override fun onTransitionElementReady(name: String) {
+        transitionElements.add(name)
+    }
+
+    override fun onAllTransitionElementsReady() {
         timeoutJob?.cancel()
-        if (!runTransitionAnimation) {
-            removeSharedElements = true
-        }
         if (!previewReady) {
             previewReady = true
             maybeStartListenForLayout()
@@ -90,11 +92,8 @@ class EnterTransitionAnimationDelegate(
         names: MutableList<String>,
         sharedElements: MutableMap<String, View>
     ) {
-        if (removeSharedElements) {
-            names.remove(ChooserActivity.FIRST_IMAGE_PREVIEW_TRANSITION_NAME)
-            sharedElements.remove(ChooserActivity.FIRST_IMAGE_PREVIEW_TRANSITION_NAME)
-        }
-        removeSharedElements = false
+        names.removeAll { !transitionElements.contains(it) }
+        sharedElements.entries.removeAll { !transitionElements.contains(it.key) }
     }
 
     private fun maybeStartListenForLayout() {
@@ -119,7 +118,7 @@ class EnterTransitionAnimationDelegate(
     }
 
     private fun startPostponedEnterTransition() {
-        if (!removeSharedElements && activity.isActivityTransitionRunning) {
+        if (transitionElements.isNotEmpty() && activity.isActivityTransitionRunning) {
             // Disable the window animations as it interferes with the transition animation.
             activity.window.setWindowAnimations(0)
         }
diff --git a/java/src/com/android/intentresolver/widget/ChooserImagePreviewView.kt b/java/src/com/android/intentresolver/widget/ChooserImagePreviewView.kt
index dd1dd286..bf10bfaa 100644
--- a/java/src/com/android/intentresolver/widget/ChooserImagePreviewView.kt
+++ b/java/src/com/android/intentresolver/widget/ChooserImagePreviewView.kt
@@ -21,17 +21,15 @@ import android.content.Context
 import android.net.Uri
 import android.util.AttributeSet
 import android.view.LayoutInflater
-import android.view.View
-import android.view.ViewTreeObserver
 import android.view.animation.DecelerateInterpolator
 import android.widget.RelativeLayout
 import androidx.core.view.isVisible
 import com.android.intentresolver.R
+import com.android.intentresolver.widget.ImagePreviewView.TransitionElementStatusCallback
 import kotlinx.coroutines.Job
 import kotlinx.coroutines.MainScope
 import kotlinx.coroutines.coroutineScope
 import kotlinx.coroutines.launch
-import java.util.function.Consumer
 import com.android.internal.R as IntR
 
 private const val IMAGE_FADE_IN_MILLIS = 150L
@@ -56,7 +54,7 @@ class ChooserImagePreviewView : RelativeLayout, ImagePreviewView {
     private lateinit var thirdImage: RoundedRectImageView
 
     private var loadImageJob: Job? = null
-    private var onTransitionViewReadyCallback: Consumer<Boolean>? = null
+    private var transitionStatusElementCallback: TransitionElementStatusCallback? = null
 
     override fun onFinishInflate() {
         LayoutInflater.from(context).inflate(R.layout.image_preview_view, this, true)
@@ -67,17 +65,12 @@ class ChooserImagePreviewView : RelativeLayout, ImagePreviewView {
     }
 
     /**
-     * Specifies a transition animation target name and a readiness callback. The callback will be
-     * invoked once when the view preparation is done i.e. either when an image is loaded into it
-     * and it is laid out (and it is ready to be draw) or image loading has failed.
+     * Specifies a transition animation target readiness callback. The callback will be
+     * invoked once when views preparation is done.
      * Should be called before [setImages].
-     * @param name, transition name
-     * @param onViewReady, a callback that will be invoked with `true` if the view is ready to
-     * receive transition animation (the image was loaded successfully) and with `false` otherwise.
      */
-    override fun setSharedElementTransitionTarget(name: String, onViewReady: Consumer<Boolean>) {
-        mainImage.transitionName = name
-        onTransitionViewReadyCallback = onViewReady
+    override fun setTransitionElementStatusCallback(callback: TransitionElementStatusCallback?) {
+        transitionStatusElementCallback = callback
     }
 
     override fun setImages(uris: List<Uri>, imageLoader: ImageLoader) {
@@ -97,7 +90,7 @@ class ChooserImagePreviewView : RelativeLayout, ImagePreviewView {
         secondLargeImage.isVisible = false
         secondSmallImage.isVisible = false
         thirdImage.isVisible = false
-        invokeTransitionViewReadyCallback(runTransitionAnimation = false)
+        invokeTransitionViewReadyCallback()
     }
 
     private suspend fun showOneImage(uris: List<Uri>, imageLoader: ImageLoader) {
@@ -138,7 +131,7 @@ class ChooserImagePreviewView : RelativeLayout, ImagePreviewView {
         if (bitmap == null) {
             view.isVisible = false
             if (view === mainImage) {
-                invokeTransitionViewReadyCallback(runTransitionAnimation = false)
+                invokeTransitionViewReadyCallback()
             }
         } else {
             view.isVisible = true
@@ -150,26 +143,20 @@ class ChooserImagePreviewView : RelativeLayout, ImagePreviewView {
                 duration = IMAGE_FADE_IN_MILLIS
                 start()
             }
-            if (view === mainImage && onTransitionViewReadyCallback != null) {
-                setupPreDrawListener(mainImage)
+            if (view === mainImage && transitionStatusElementCallback != null) {
+                view.waitForPreDraw()
+                invokeTransitionViewReadyCallback()
             }
         }
     }
 
-    private fun setupPreDrawListener(view: View) {
-        view.viewTreeObserver.addOnPreDrawListener(
-            object : ViewTreeObserver.OnPreDrawListener {
-                override fun onPreDraw(): Boolean {
-                    view.viewTreeObserver.removeOnPreDrawListener(this)
-                    invokeTransitionViewReadyCallback(runTransitionAnimation = true)
-                    return true
-                }
+    private fun invokeTransitionViewReadyCallback() {
+        transitionStatusElementCallback?.apply {
+            if (mainImage.isVisible && mainImage.drawable != null) {
+                mainImage.transitionName?.let { onTransitionElementReady(it) }
             }
-        )
-    }
-
-    private fun invokeTransitionViewReadyCallback(runTransitionAnimation: Boolean) {
-        onTransitionViewReadyCallback?.accept(runTransitionAnimation)
-        onTransitionViewReadyCallback = null
+            onAllTransitionElementsReady()
+        }
+        transitionStatusElementCallback = null
     }
 }
diff --git a/java/src/com/android/intentresolver/widget/ImagePreviewView.kt b/java/src/com/android/intentresolver/widget/ImagePreviewView.kt
index a5756054..a166ef27 100644
--- a/java/src/com/android/intentresolver/widget/ImagePreviewView.kt
+++ b/java/src/com/android/intentresolver/widget/ImagePreviewView.kt
@@ -18,22 +18,32 @@ package com.android.intentresolver.widget
 
 import android.graphics.Bitmap
 import android.net.Uri
-import java.util.function.Consumer
 
 internal typealias ImageLoader = suspend (Uri) -> Bitmap?
 
 interface ImagePreviewView {
+    fun setTransitionElementStatusCallback(callback: TransitionElementStatusCallback?)
+    fun setImages(uris: List<Uri>, imageLoader: ImageLoader)
 
     /**
-     * Specifies a transition animation target name and a readiness callback. The callback will be
-     * invoked once when the view preparation is done i.e. either when an image is loaded into it
-     * and it is laid out (and it is ready to be draw) or image loading has failed.
-     * Should be called before [setImages].
-     * @param name, transition name
-     * @param onViewReady, a callback that will be invoked with `true` if the view is ready to
-     * receive transition animation (the image was loaded successfully) and with `false` otherwise.
+     * [ImagePreviewView] progressively prepares views for shared element transition and reports
+     * each successful preparation with [onTransitionElementReady] call followed by
+     * closing [onAllTransitionElementsReady] invocation. Thus the overall invocation pattern is
+     * zero or more [onTransitionElementReady] calls followed by the final
+     * [onAllTransitionElementsReady] call.
      */
-    fun setSharedElementTransitionTarget(name: String, onViewReady: Consumer<Boolean>)
+    interface TransitionElementStatusCallback {
+        /**
+         * Invoked when a view for a shared transition animation element is ready i.e. the image
+         * is loaded and the view is laid out.
+         * @param name shared element name.
+         */
+        fun onTransitionElementReady(name: String)
 
-    fun setImages(uris: List<Uri>, imageLoader: ImageLoader)
+        /**
+         * Indicates that all supported transition elements have been reported with
+         * [onTransitionElementReady].
+         */
+        fun onAllTransitionElementsReady()
+    }
 }
diff --git a/java/src/com/android/intentresolver/widget/ViewExtensions.kt b/java/src/com/android/intentresolver/widget/ViewExtensions.kt
new file mode 100644
index 00000000..11b7c146
--- /dev/null
+++ b/java/src/com/android/intentresolver/widget/ViewExtensions.kt
@@ -0,0 +1,39 @@
+/*
+ * Copyright (C) 2023 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 com.android.intentresolver.widget
+
+import android.util.Log
+import android.view.View
+import androidx.core.view.OneShotPreDrawListener
+import kotlinx.coroutines.suspendCancellableCoroutine
+import java.util.concurrent.atomic.AtomicBoolean
+
+internal suspend fun View.waitForPreDraw(): Unit = suspendCancellableCoroutine { continuation ->
+    val isResumed = AtomicBoolean(false)
+    val callback = OneShotPreDrawListener.add(
+        this,
+        Runnable {
+            if (isResumed.compareAndSet(false, true)) {
+                continuation.resumeWith(Result.success(Unit))
+            } else {
+                // it's not really expected but in some unknown corner-case let's not crash
+                Log.e("waitForPreDraw", "An attempt to resume a completed coroutine", Exception())
+            }
+        }
+    )
+    continuation.invokeOnCancellation { callback.removeListener() }
+}
diff --git a/java/tests/src/com/android/intentresolver/EnterTransitionAnimationDelegateTest.kt b/java/tests/src/com/android/intentresolver/EnterTransitionAnimationDelegateTest.kt
index ffe89400..9ea9dfa7 100644
--- a/java/tests/src/com/android/intentresolver/EnterTransitionAnimationDelegateTest.kt
+++ b/java/tests/src/com/android/intentresolver/EnterTransitionAnimationDelegateTest.kt
@@ -39,6 +39,7 @@ private const val TIMEOUT_MS = 200
 
 @OptIn(ExperimentalCoroutinesApi::class)
 class EnterTransitionAnimationDelegateTest {
+    private val elementName = "shared-element"
     private val scheduler = TestCoroutineScheduler()
     private val dispatcher = StandardTestDispatcher(scheduler)
     private val lifecycleOwner = TestLifecycleOwner()
@@ -89,9 +90,9 @@ class EnterTransitionAnimationDelegateTest {
     fun test_postponeTransition_animation_resumes_only_once() {
         testSubject.postponeTransition()
         testSubject.markOffsetCalculated()
-        testSubject.markImagePreviewReady(true)
+        testSubject.onTransitionElementReady(elementName)
         testSubject.markOffsetCalculated()
-        testSubject.markImagePreviewReady(true)
+        testSubject.onTransitionElementReady(elementName)
 
         scheduler.advanceTimeBy(TIMEOUT_MS + 1L)
         verify(activity, times(1)).startPostponedEnterTransition()
@@ -105,7 +106,7 @@ class EnterTransitionAnimationDelegateTest {
         testSubject.markOffsetCalculated()
         verify(activity, never()).startPostponedEnterTransition()
 
-        testSubject.markImagePreviewReady(true)
+        testSubject.onAllTransitionElementsReady()
         verify(activity, times(1)).startPostponedEnterTransition()
     }
 }