Do not store the initial intent's extra in the saved state.

CreationExtras's DEFAULT_ARGS_KEY vlaue gets saved in
SavedStateHandle for each view model the activity creates. Thus by
storing the ActivityModel in there we effectively duplicated the initial
intent's extra in the activity's saved state 4 times: DEFAULT_ARGS_KEY
contains the extras (put there by ComponentActivity) plus ActivityModel
per two view models we created.
This change makes Chooser and Resolver activities provide default
CreationExtras with empty DEFAULT_ARGS_KEY values and stores
ActivityModel in the new repository class (instead of the
SavedStateHandle instance).

Fix: 331897641
Test: manual testing with injected logging the values being put in the
 activity's saved state
Test: atest IntentResolver-tests-unit
Test: atest Intentresolver-tests-activity
Flag: EXEMPT bugfix

Change-Id: Ice2e51971476b2bb963f04275d7b180c85126288

10 files changed, 85 insertions, 57 deletions

diff --git a/java/src/com/android/intentresolver/ChooserActivity.java b/java/src/com/android/intentresolver/ChooserActivity.java
index 3db821c1..f7d81ca4 100644
--- a/java/src/com/android/intentresolver/ChooserActivity.java
+++ b/java/src/com/android/intentresolver/ChooserActivity.java
@@ -23,13 +23,12 @@ import static android.view.WindowManager.LayoutParams.SYSTEM_FLAG_HIDE_NON_SYSTE
 import static androidx.lifecycle.LifecycleKt.getCoroutineScope;
 
 import static com.android.intentresolver.ChooserActionFactory.EDIT_SOURCE;
-import static com.android.intentresolver.Flags.shareouselUpdateExcludeComponentsExtra;
 import static com.android.intentresolver.Flags.fixShortcutsFlashing;
+import static com.android.intentresolver.Flags.shareouselUpdateExcludeComponentsExtra;
 import static com.android.intentresolver.Flags.unselectFinalItem;
-import static com.android.intentresolver.ext.CreationExtrasExtKt.addDefaultArgs;
+import static com.android.intentresolver.ext.CreationExtrasExtKt.replaceDefaultArgs;
 import static com.android.intentresolver.profiles.MultiProfilePagerAdapter.PROFILE_PERSONAL;
 import static com.android.intentresolver.profiles.MultiProfilePagerAdapter.PROFILE_WORK;
-import static com.android.intentresolver.ui.model.ActivityModel.ACTIVITY_MODEL_KEY;
 import static com.android.internal.util.LatencyTracker.ACTION_LOAD_SHARE_SHEET;
 
 import static java.util.Objects.requireNonNull;
@@ -102,6 +101,7 @@ import com.android.intentresolver.chooser.TargetInfo;
 import com.android.intentresolver.contentpreview.ChooserContentPreviewUi;
 import com.android.intentresolver.contentpreview.HeadlineGeneratorImpl;
 import com.android.intentresolver.data.model.ChooserRequest;
+import com.android.intentresolver.data.repository.ActivityModelRepository;
 import com.android.intentresolver.data.repository.DevicePolicyResources;
 import com.android.intentresolver.domain.interactor.UserInteractor;
 import com.android.intentresolver.emptystate.CompositeEmptyStateProvider;
@@ -127,6 +127,7 @@ import com.android.intentresolver.profiles.MultiProfilePagerAdapter.ProfileType;
 import com.android.intentresolver.profiles.OnProfileSelectedListener;
 import com.android.intentresolver.profiles.OnSwitchOnWorkSelectedListener;
 import com.android.intentresolver.profiles.TabConfig;
+import com.android.intentresolver.shared.model.ActivityModel;
 import com.android.intentresolver.shared.model.Profile;
 import com.android.intentresolver.shortcuts.AppPredictorFactory;
 import com.android.intentresolver.shortcuts.ShortcutLoader;
@@ -134,7 +135,6 @@ import com.android.intentresolver.ui.ActionTitle;
 import com.android.intentresolver.ui.ProfilePagerResources;
 import com.android.intentresolver.ui.ShareResultSender;
 import com.android.intentresolver.ui.ShareResultSenderFactory;
-import com.android.intentresolver.ui.model.ActivityModel;
 import com.android.intentresolver.ui.viewmodel.ChooserViewModel;
 import com.android.intentresolver.widget.ActionRow;
 import com.android.intentresolver.widget.ImagePreviewView;
@@ -149,8 +149,6 @@ import com.google.common.collect.ImmutableList;
 
 import dagger.hilt.android.AndroidEntryPoint;
 
-import kotlin.Pair;
-
 import kotlinx.coroutines.CoroutineDispatcher;
 
 import java.util.ArrayList;
@@ -273,6 +271,7 @@ public class ChooserActivity extends Hilt_ChooserActivity implements
     @Inject public ClipboardManager mClipboardManager;
     @Inject public IntentForwarding mIntentForwarding;
     @Inject public ShareResultSenderFactory mShareResultSenderFactory;
+    @Inject public ActivityModelRepository mActivityModelRepository;
 
     private ActivityModel mActivityModel;
     private ChooserRequest mRequest;
@@ -331,15 +330,18 @@ public class ChooserActivity extends Hilt_ChooserActivity implements
     @NonNull
     @Override
     public CreationExtras getDefaultViewModelCreationExtras() {
-        return addDefaultArgs(
-                super.getDefaultViewModelCreationExtras(),
-                new Pair<>(ACTIVITY_MODEL_KEY, createActivityModel()));
+        // DEFAULT_ARGS_KEY extra is saved for each ViewModel we create. ComponentActivity puts the
+        // initial intent's extra into DEFAULT_ARGS_KEY thus we store these values 2 times (3 if we
+        // count the initial intent). We don't need those values to be saved as they don't capture
+        // the state.
+        return replaceDefaultArgs(super.getDefaultViewModelCreationExtras());
     }
 
     @Override
     protected void onCreate(Bundle savedInstanceState) {
         super.onCreate(savedInstanceState);
         Log.i(TAG, "onCreate");
+        mActivityModelRepository.initialize(this::createActivityModel);
 
         mTargetDataLoader = mChooserServiceFeatureFlags.chooserPayloadToggling()
                 ? mCachingTargetDataLoaderProvider.get()
diff --git a/java/src/com/android/intentresolver/ResolverActivity.java b/java/src/com/android/intentresolver/ResolverActivity.java
index a402fc72..2f220cf1 100644
--- a/java/src/com/android/intentresolver/ResolverActivity.java
+++ b/java/src/com/android/intentresolver/ResolverActivity.java
@@ -21,7 +21,7 @@ import static android.view.WindowManager.LayoutParams.SYSTEM_FLAG_HIDE_NON_SYSTE
 
 import static androidx.lifecycle.LifecycleKt.getCoroutineScope;
 
-import static com.android.intentresolver.ext.CreationExtrasExtKt.addDefaultArgs;
+import static com.android.intentresolver.ext.CreationExtrasExtKt.replaceDefaultArgs;
 import static com.android.internal.annotations.VisibleForTesting.Visibility.PROTECTED;
 
 import static java.util.Objects.requireNonNull;
@@ -85,6 +85,7 @@ import androidx.viewpager.widget.ViewPager;
 
 import com.android.intentresolver.chooser.DisplayResolveInfo;
 import com.android.intentresolver.chooser.TargetInfo;
+import com.android.intentresolver.data.repository.ActivityModelRepository;
 import com.android.intentresolver.data.repository.DevicePolicyResources;
 import com.android.intentresolver.domain.interactor.UserInteractor;
 import com.android.intentresolver.emptystate.CompositeEmptyStateProvider;
@@ -103,10 +104,10 @@ import com.android.intentresolver.profiles.OnProfileSelectedListener;
 import com.android.intentresolver.profiles.OnSwitchOnWorkSelectedListener;
 import com.android.intentresolver.profiles.ResolverMultiProfilePagerAdapter;
 import com.android.intentresolver.profiles.TabConfig;
+import com.android.intentresolver.shared.model.ActivityModel;
 import com.android.intentresolver.shared.model.Profile;
 import com.android.intentresolver.ui.ActionTitle;
 import com.android.intentresolver.ui.ProfilePagerResources;
-import com.android.intentresolver.ui.model.ActivityModel;
 import com.android.intentresolver.ui.model.ResolverRequest;
 import com.android.intentresolver.ui.viewmodel.ResolverViewModel;
 import com.android.intentresolver.widget.ResolverDrawerLayout;
@@ -119,8 +120,6 @@ import com.google.common.collect.ImmutableList;
 
 import dagger.hilt.android.AndroidEntryPoint;
 
-import kotlin.Pair;
-
 import kotlinx.coroutines.CoroutineDispatcher;
 
 import java.util.ArrayList;
@@ -150,6 +149,7 @@ public class ResolverActivity extends Hilt_ResolverActivity implements
     @Inject public ProfilePagerResources mProfilePagerResources;
     @Inject public IntentForwarding mIntentForwarding;
     @Inject public FeatureFlags mFeatureFlags;
+    @Inject public ActivityModelRepository mActivityModelRepository;
 
     private ResolverViewModel mViewModel;
     private ResolverRequest mRequest;
@@ -220,15 +220,14 @@ public class ResolverActivity extends Hilt_ResolverActivity implements
     @NonNull
     @Override
     public CreationExtras getDefaultViewModelCreationExtras() {
-        return addDefaultArgs(
-                super.getDefaultViewModelCreationExtras(),
-                new Pair<>(ActivityModel.ACTIVITY_MODEL_KEY, createActivityModel()));
+        return replaceDefaultArgs(super.getDefaultViewModelCreationExtras());
     }
 
     @Override
     protected void onCreate(Bundle savedInstanceState) {
         super.onCreate(savedInstanceState);
         Log.i(TAG, "onCreate");
+        mActivityModelRepository.initialize(this::createActivityModel);
         setTheme(R.style.Theme_DeviceDefault_Resolver);
         mResolverHelper.setInitializer(this::initialize);
     }
diff --git a/java/src/com/android/intentresolver/data/repository/ActivityModelRepository.kt b/java/src/com/android/intentresolver/data/repository/ActivityModelRepository.kt
new file mode 100644
index 00000000..7c3188d2
--- /dev/null
+++ b/java/src/com/android/intentresolver/data/repository/ActivityModelRepository.kt
@@ -0,0 +1,37 @@
+/*
+ * Copyright 2024 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.data.repository
+
+import com.android.intentresolver.shared.model.ActivityModel
+import dagger.hilt.android.scopes.ActivityRetainedScoped
+import javax.inject.Inject
+import kotlinx.atomicfu.atomic
+
+/** An [ActivityModel] repository that captures the first value. */
+@ActivityRetainedScoped
+class ActivityModelRepository @Inject constructor() {
+    private val _value = atomic<ActivityModel?>(null)
+
+    val value: ActivityModel
+        get() = requireNotNull(_value.value) { "Repository has not been initialized" }
+
+    fun initialize(block: () -> ActivityModel) {
+        if (_value.value == null) {
+            _value.compareAndSet(null, block())
+        }
+    }
+}
diff --git a/java/src/com/android/intentresolver/ext/CreationExtrasExt.kt b/java/src/com/android/intentresolver/ext/CreationExtrasExt.kt
index 2ba08c90..5635ec28 100644
--- a/java/src/com/android/intentresolver/ext/CreationExtrasExt.kt
+++ b/java/src/com/android/intentresolver/ext/CreationExtrasExt.kt
@@ -32,3 +32,9 @@ fun CreationExtras.addDefaultArgs(vararg values: Pair<String, Parcelable>): Crea
     defaultArgs.putAll(bundleOf(*values))
     return MutableCreationExtras(this).apply { set(DEFAULT_ARGS_KEY, defaultArgs) }
 }
+
+fun CreationExtras.replaceDefaultArgs(vararg values: Pair<String, Parcelable>): CreationExtras {
+    val mutableExtras = if (this is MutableCreationExtras) this else MutableCreationExtras(this)
+    mutableExtras[DEFAULT_ARGS_KEY] = bundleOf(*values)
+    return mutableExtras
+}
diff --git a/java/src/com/android/intentresolver/inject/ActivityModelModule.kt b/java/src/com/android/intentresolver/inject/ActivityModelModule.kt
index bbd25eb7..7201bd2b 100644
--- a/java/src/com/android/intentresolver/inject/ActivityModelModule.kt
+++ b/java/src/com/android/intentresolver/inject/ActivityModelModule.kt
@@ -19,9 +19,8 @@ package com.android.intentresolver.inject
 import android.content.Intent
 import android.net.Uri
 import android.service.chooser.ChooserAction
-import androidx.lifecycle.SavedStateHandle
 import com.android.intentresolver.data.model.ChooserRequest
-import com.android.intentresolver.ui.model.ActivityModel
+import com.android.intentresolver.data.repository.ActivityModelRepository
 import com.android.intentresolver.ui.viewmodel.readChooserRequest
 import com.android.intentresolver.util.ownedByCurrentUser
 import com.android.intentresolver.validation.Valid
@@ -37,26 +36,19 @@ import javax.inject.Qualifier
 @InstallIn(ViewModelComponent::class)
 object ActivityModelModule {
     @Provides
-    fun provideActivityModel(savedStateHandle: SavedStateHandle): ActivityModel =
-        requireNotNull(savedStateHandle[ActivityModel.ACTIVITY_MODEL_KEY]) {
-            "ActivityModel missing in SavedStateHandle! (${ActivityModel.ACTIVITY_MODEL_KEY})"
-        }
-
-    @Provides
     @ChooserIntent
-    fun chooserIntent(activityModel: ActivityModel): Intent = activityModel.intent
+    fun chooserIntent(activityModelRepo: ActivityModelRepository): Intent =
+        activityModelRepo.value.intent
 
     @Provides
     @ViewModelScoped
     fun provideInitialRequest(
-        activityModel: ActivityModel,
+        activityModelRepo: ActivityModelRepository,
         flags: ChooserServiceFlags,
-    ): ValidationResult<ChooserRequest> = readChooserRequest(activityModel, flags)
+    ): ValidationResult<ChooserRequest> = readChooserRequest(activityModelRepo.value, flags)
 
     @Provides
-    fun provideChooserRequest(
-        initialRequest: ValidationResult<ChooserRequest>,
-    ): ChooserRequest =
+    fun provideChooserRequest(initialRequest: ValidationResult<ChooserRequest>): ChooserRequest =
         requireNotNull((initialRequest as? Valid)?.value) {
             "initialRequest is Invalid, no chooser request available"
         }
diff --git a/java/src/com/android/intentresolver/ui/model/ActivityModel.kt b/java/src/com/android/intentresolver/shared/model/ActivityModel.kt
index 4bcdd69b..c5efdeba 100644
--- a/java/src/com/android/intentresolver/ui/model/ActivityModel.kt
+++ b/java/src/com/android/intentresolver/shared/model/ActivityModel.kt
@@ -1,5 +1,5 @@
 /*
- * Copyright (C) 2024 The Android Open Source Project
+ * Copyright 2024 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.
@@ -13,7 +13,7 @@
  * See the License for the specific language governing permissions and
  * limitations under the License.
  */
-package com.android.intentresolver.ui.model
+package com.android.intentresolver.shared.model
 
 import android.app.Activity
 import android.content.Intent
@@ -34,7 +34,7 @@ data class ActivityModel(
     /** The package of the sending app */
     val launchedFromPackage: String,
     /** The referrer as supplied to the activity. */
-    val referrer: Uri?
+    val referrer: Uri?,
 ) : Parcelable {
     constructor(
         source: Parcel
@@ -42,7 +42,7 @@ data class ActivityModel(
         intent = source.requireParcelable(),
         launchedFromUid = source.readInt(),
         launchedFromPackage = requireNotNull(source.readString()),
-        referrer = source.readParcelable()
+        referrer = source.readParcelable(),
     )
 
     /** A package name from referrer, if it is an android-app URI */
@@ -58,13 +58,12 @@ data class ActivityModel(
     }
 
     companion object {
-        const val ACTIVITY_MODEL_KEY = "com.android.intentresolver.ACTIVITY_MODEL"
-
         @JvmField
         @Suppress("unused")
         val CREATOR =
             object : Parcelable.Creator<ActivityModel> {
                 override fun newArray(size: Int) = arrayOfNulls<ActivityModel>(size)
+
                 override fun createFromParcel(source: Parcel) = ActivityModel(source)
             }
 
@@ -74,7 +73,7 @@ data class ActivityModel(
                 activity.intent,
                 activity.launchedFromUid,
                 Objects.requireNonNull<String>(activity.launchedFromPackage),
-                activity.referrer
+                activity.referrer,
             )
         }
     }
diff --git a/java/src/com/android/intentresolver/ui/viewmodel/ChooserRequestReader.kt b/java/src/com/android/intentresolver/ui/viewmodel/ChooserRequestReader.kt
index 4a194db9..13cadf37 100644
--- a/java/src/com/android/intentresolver/ui/viewmodel/ChooserRequestReader.kt
+++ b/java/src/com/android/intentresolver/ui/viewmodel/ChooserRequestReader.kt
@@ -49,7 +49,7 @@ import com.android.intentresolver.data.model.ChooserRequest
 import com.android.intentresolver.ext.hasSendAction
 import com.android.intentresolver.ext.ifMatch
 import com.android.intentresolver.inject.ChooserServiceFlags
-import com.android.intentresolver.ui.model.ActivityModel
+import com.android.intentresolver.shared.model.ActivityModel
 import com.android.intentresolver.util.hasValidIcon
 import com.android.intentresolver.validation.Validation
 import com.android.intentresolver.validation.ValidationResult
@@ -69,7 +69,7 @@ internal fun Intent.maybeAddSendActionFlags() =
 
 fun readChooserRequest(
     model: ActivityModel,
-    flags: ChooserServiceFlags
+    flags: ChooserServiceFlags,
 ): ValidationResult<ChooserRequest> {
     val extras = model.intent.extras ?: Bundle()
     @Suppress("DEPRECATION")
@@ -87,7 +87,7 @@ fun readChooserRequest(
                 ignored(
                     value<CharSequence>(EXTRA_TITLE),
                     "deprecated in P. You may wish to set a preview title by using EXTRA_TITLE " +
-                        "property of the wrapped EXTRA_INTENT."
+                        "property of the wrapped EXTRA_INTENT.",
                 )
                 null to R.string.chooseActivity
             } else {
diff --git a/java/src/com/android/intentresolver/ui/viewmodel/ChooserViewModel.kt b/java/src/com/android/intentresolver/ui/viewmodel/ChooserViewModel.kt
index 619e118a..e6f12750 100644
--- a/java/src/com/android/intentresolver/ui/viewmodel/ChooserViewModel.kt
+++ b/java/src/com/android/intentresolver/ui/viewmodel/ChooserViewModel.kt
@@ -17,7 +17,6 @@ package com.android.intentresolver.ui.viewmodel
 
 import android.content.ContentInterface
 import android.util.Log
-import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.ViewModel
 import androidx.lifecycle.viewModelScope
 import com.android.intentresolver.contentpreview.ImageLoader
@@ -26,11 +25,11 @@ import com.android.intentresolver.contentpreview.payloadtoggle.domain.interactor
 import com.android.intentresolver.contentpreview.payloadtoggle.domain.interactor.ProcessTargetIntentUpdatesInteractor
 import com.android.intentresolver.contentpreview.payloadtoggle.ui.viewmodel.ShareouselViewModel
 import com.android.intentresolver.data.model.ChooserRequest
+import com.android.intentresolver.data.repository.ActivityModelRepository
 import com.android.intentresolver.data.repository.ChooserRequestRepository
 import com.android.intentresolver.inject.Background
 import com.android.intentresolver.inject.ChooserServiceFlags
-import com.android.intentresolver.ui.model.ActivityModel
-import com.android.intentresolver.ui.model.ActivityModel.Companion.ACTIVITY_MODEL_KEY
+import com.android.intentresolver.shared.model.ActivityModel
 import com.android.intentresolver.validation.Invalid
 import com.android.intentresolver.validation.Valid
 import com.android.intentresolver.validation.ValidationResult
@@ -49,7 +48,7 @@ private const val TAG = "ChooserViewModel"
 class ChooserViewModel
 @Inject
 constructor(
-    args: SavedStateHandle,
+    activityModelRepository: ActivityModelRepository,
     private val shareouselViewModelProvider: Lazy<ShareouselViewModel>,
     private val processUpdatesInteractor: Lazy<ProcessTargetIntentUpdatesInteractor>,
     private val fetchPreviewsInteractor: Lazy<FetchPreviewsInteractor>,
@@ -67,10 +66,7 @@ constructor(
 ) : ViewModel() {
 
     /** Parcelable-only references provided from the creating Activity */
-    val activityModel: ActivityModel =
-        requireNotNull(args[ACTIVITY_MODEL_KEY]) {
-            "ActivityModel missing in SavedStateHandle! ($ACTIVITY_MODEL_KEY)"
-        }
+    val activityModel: ActivityModel = activityModelRepository.value
 
     val shareouselViewModel: ShareouselViewModel by lazy {
         // TODO: consolidate this logic, this would require a consolidated preview view model but
diff --git a/java/src/com/android/intentresolver/ui/viewmodel/ResolverRequestReader.kt b/java/src/com/android/intentresolver/ui/viewmodel/ResolverRequestReader.kt
index 856d9fdd..884be635 100644
--- a/java/src/com/android/intentresolver/ui/viewmodel/ResolverRequestReader.kt
+++ b/java/src/com/android/intentresolver/ui/viewmodel/ResolverRequestReader.kt
@@ -20,8 +20,8 @@ import android.os.Bundle
 import android.os.UserHandle
 import com.android.intentresolver.ResolverActivity.PROFILE_PERSONAL
 import com.android.intentresolver.ResolverActivity.PROFILE_WORK
+import com.android.intentresolver.shared.model.ActivityModel
 import com.android.intentresolver.shared.model.Profile
-import com.android.intentresolver.ui.model.ActivityModel
 import com.android.intentresolver.ui.model.ResolverRequest
 import com.android.intentresolver.validation.Validation
 import com.android.intentresolver.validation.ValidationResult
diff --git a/java/src/com/android/intentresolver/ui/viewmodel/ResolverViewModel.kt b/java/src/com/android/intentresolver/ui/viewmodel/ResolverViewModel.kt
index a3dc58a6..3511637b 100644
--- a/java/src/com/android/intentresolver/ui/viewmodel/ResolverViewModel.kt
+++ b/java/src/com/android/intentresolver/ui/viewmodel/ResolverViewModel.kt
@@ -17,10 +17,9 @@
 package com.android.intentresolver.ui.viewmodel
 
 import android.util.Log
-import androidx.lifecycle.SavedStateHandle
 import androidx.lifecycle.ViewModel
-import com.android.intentresolver.ui.model.ActivityModel
-import com.android.intentresolver.ui.model.ActivityModel.Companion.ACTIVITY_MODEL_KEY
+import com.android.intentresolver.data.repository.ActivityModelRepository
+import com.android.intentresolver.shared.model.ActivityModel
 import com.android.intentresolver.ui.model.ResolverRequest
 import com.android.intentresolver.validation.Invalid
 import com.android.intentresolver.validation.Valid
@@ -33,13 +32,11 @@ import kotlinx.coroutines.flow.asStateFlow
 private const val TAG = "ResolverViewModel"
 
 @HiltViewModel
-class ResolverViewModel @Inject constructor(args: SavedStateHandle) : ViewModel() {
+class ResolverViewModel @Inject constructor(activityModelrepo: ActivityModelRepository) :
+    ViewModel() {
 
     /** Parcelable-only references provided from the creating Activity */
-    val activityModel: ActivityModel =
-        requireNotNull(args[ACTIVITY_MODEL_KEY]) {
-            "ActivityModel missing in SavedStateHandle! ($ACTIVITY_MODEL_KEY)"
-        }
+    val activityModel: ActivityModel = activityModelrepo.value
 
     /**
      * Provided only for the express purpose of early exit in the event of an invalid request.