Restore communal widgets after host is restored

This change adds an observer to the app widget host restored broadcast.
When this signal is received, a backed-up state is expected to have been
written on disk and read into memory. The widget ids are updated based
on the mapping provided by the broadcast, and database is wiped and
restored to the backed up state.

Test: atest CommunalWidgetDaoTest
Test: atest CommunalWidgetRepositoryImplTest
Test: atest CommunalBackupRestoreStartableTest
Test: manual; see instructions at go/glanceable-hub-br
Bug: 309809222
Flag: ACONFIG com.android.systemui.communal_hub TEAMFOOD
Change-Id: Ia67bb901e6976d1b7507672e58b53b6bf5b42c81

11 files changed, 588 insertions, 1 deletions

diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/CommunalBackupRestoreStartableTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/CommunalBackupRestoreStartableTest.kt
new file mode 100644
index 000000000000..722eb2b9b622
--- /dev/null
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/CommunalBackupRestoreStartableTest.kt
@@ -0,0 +1,152 @@
+/*
+ * Copyright (C) 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.systemui.communal
+
+import android.appwidget.AppWidgetManager
+import android.content.Context
+import android.content.Intent
+import android.content.mockedContext
+import androidx.test.ext.junit.runners.AndroidJUnit4
+import androidx.test.filters.SmallTest
+import com.android.systemui.SysuiTestCase
+import com.android.systemui.broadcast.FakeBroadcastDispatcher
+import com.android.systemui.broadcast.broadcastDispatcher
+import com.android.systemui.communal.domain.interactor.CommunalInteractor
+import com.android.systemui.communal.widgets.CommunalWidgetModule
+import com.android.systemui.kosmos.testScope
+import com.android.systemui.log.logcatLogBuffer
+import com.android.systemui.testKosmos
+import com.android.systemui.util.mockito.any
+import com.android.systemui.util.mockito.kotlinArgumentCaptor
+import com.google.common.truth.Truth.assertThat
+import kotlinx.coroutines.test.runTest
+import org.junit.Before
+import org.junit.Test
+import org.junit.runner.RunWith
+import org.mockito.Mock
+import org.mockito.Mockito.never
+import org.mockito.Mockito.verify
+import org.mockito.MockitoAnnotations
+
+@SmallTest
+@RunWith(AndroidJUnit4::class)
+class CommunalBackupRestoreStartableTest : SysuiTestCase() {
+    private val kosmos = testKosmos()
+    private val testScope = kosmos.testScope
+
+    @Mock private lateinit var communalInteractor: CommunalInteractor
+
+    private val mapCaptor = kotlinArgumentCaptor<Map<Int, Int>>()
+
+    private lateinit var context: Context
+    private lateinit var broadcastDispatcher: FakeBroadcastDispatcher
+    private lateinit var underTest: CommunalBackupRestoreStartable
+
+    @Before
+    fun setup() {
+        MockitoAnnotations.initMocks(this)
+
+        context = kosmos.mockedContext
+        broadcastDispatcher = kosmos.broadcastDispatcher
+
+        underTest =
+            CommunalBackupRestoreStartable(
+                broadcastDispatcher,
+                communalInteractor,
+                logcatLogBuffer("CommunalBackupRestoreStartable"),
+            )
+    }
+
+    @Test
+    fun testRestoreWidgetsUponHostRestored() =
+        testScope.runTest {
+            underTest.start()
+
+            // Verify restore widgets not called
+            verify(communalInteractor, never()).restoreWidgets(any())
+
+            // Trigger app widget host restored
+            val intent =
+                Intent().apply {
+                    action = AppWidgetManager.ACTION_APPWIDGET_HOST_RESTORED
+                    putExtra(
+                        AppWidgetManager.EXTRA_HOST_ID,
+                        CommunalWidgetModule.APP_WIDGET_HOST_ID
+                    )
+                    putExtra(AppWidgetManager.EXTRA_APPWIDGET_OLD_IDS, intArrayOf(1, 2, 3))
+                    putExtra(AppWidgetManager.EXTRA_APPWIDGET_IDS, intArrayOf(7, 8, 9))
+                }
+            broadcastDispatcher.sendIntentToMatchingReceiversOnly(context, intent)
+
+            // Verify restore widgets called
+            verify(communalInteractor).restoreWidgets(mapCaptor.capture())
+            val oldToNewWidgetIdMap = mapCaptor.value
+            assertThat(oldToNewWidgetIdMap)
+                .containsExactlyEntriesIn(
+                    mapOf(
+                        Pair(1, 7),
+                        Pair(2, 8),
+                        Pair(3, 9),
+                    )
+                )
+        }
+
+    @Test
+    fun testDoNotRestoreWidgetsIfNotForCommunalWidgetHost() =
+        testScope.runTest {
+            underTest.start()
+
+            // Trigger app widget host restored, but for another host
+            val hostId = CommunalWidgetModule.APP_WIDGET_HOST_ID + 1
+            val intent =
+                Intent().apply {
+                    action = AppWidgetManager.ACTION_APPWIDGET_HOST_RESTORED
+                    putExtra(AppWidgetManager.EXTRA_HOST_ID, hostId)
+                    putExtra(AppWidgetManager.EXTRA_APPWIDGET_OLD_IDS, intArrayOf(1, 2, 3))
+                    putExtra(AppWidgetManager.EXTRA_APPWIDGET_IDS, intArrayOf(7, 8, 9))
+                }
+            broadcastDispatcher.sendIntentToMatchingReceiversOnly(context, intent)
+
+            // Verify restore widgets not called
+            verify(communalInteractor, never()).restoreWidgets(any())
+        }
+
+    @Test
+    fun testAbortRestoreWidgetsIfOldToNewIdsMappingInvalid() =
+        testScope.runTest {
+            underTest.start()
+
+            // Trigger app widget host restored, but new ids list is one too many for old ids
+            val oldIds = intArrayOf(1, 2, 3)
+            val newIds = intArrayOf(6, 7, 8, 9)
+            val intent =
+                Intent().apply {
+                    action = AppWidgetManager.ACTION_APPWIDGET_HOST_RESTORED
+                    putExtra(
+                        AppWidgetManager.EXTRA_HOST_ID,
+                        CommunalWidgetModule.APP_WIDGET_HOST_ID
+                    )
+                    putExtra(AppWidgetManager.EXTRA_APPWIDGET_OLD_IDS, oldIds)
+                    putExtra(AppWidgetManager.EXTRA_APPWIDGET_IDS, newIds)
+                }
+            broadcastDispatcher.sendIntentToMatchingReceiversOnly(context, intent)
+
+            // Verify restore widgets aborted
+            verify(communalInteractor).abortRestoreWidgets()
+            verify(communalInteractor, never()).restoreWidgets(any())
+        }
+}
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryImplTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryImplTest.kt
index fb1471cc11e5..fe4d32d88612 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryImplTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/communal/data/repository/CommunalWidgetRepositoryImplTest.kt
@@ -22,13 +22,17 @@ import android.appwidget.AppWidgetProviderInfo
 import android.appwidget.AppWidgetProviderInfo.WIDGET_FEATURE_CONFIGURATION_OPTIONAL
 import android.appwidget.AppWidgetProviderInfo.WIDGET_FEATURE_RECONFIGURABLE
 import android.content.ComponentName
+import android.content.applicationContext
 import android.os.UserHandle
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.communal.data.backup.CommunalBackupUtils
 import com.android.systemui.communal.data.db.CommunalItemRank
 import com.android.systemui.communal.data.db.CommunalWidgetDao
 import com.android.systemui.communal.data.db.CommunalWidgetItem
+import com.android.systemui.communal.nano.CommunalHubState
+import com.android.systemui.communal.proto.toByteArray
 import com.android.systemui.communal.shared.model.CommunalWidgetContentModel
 import com.android.systemui.communal.widgets.CommunalAppWidgetHost
 import com.android.systemui.communal.widgets.CommunalWidgetHost
@@ -43,6 +47,7 @@ import com.android.systemui.res.R
 import com.android.systemui.testKosmos
 import com.android.systemui.util.mockito.any
 import com.android.systemui.util.mockito.whenever
+import com.android.systemui.util.mockito.withArgCaptor
 import com.google.common.truth.Truth.assertThat
 import java.util.Optional
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -71,6 +76,7 @@ class CommunalWidgetRepositoryImplTest : SysuiTestCase() {
     @Mock private lateinit var communalWidgetDao: CommunalWidgetDao
     @Mock private lateinit var backupManager: BackupManager
 
+    private lateinit var backupUtils: CommunalBackupUtils
     private lateinit var logBuffer: LogBuffer
     private lateinit var fakeWidgets: MutableStateFlow<Map<CommunalItemRank, CommunalWidgetItem>>
 
@@ -91,6 +97,7 @@ class CommunalWidgetRepositoryImplTest : SysuiTestCase() {
         MockitoAnnotations.initMocks(this)
         fakeWidgets = MutableStateFlow(emptyMap())
         logBuffer = logcatLogBuffer(name = "CommunalWidgetRepoImplTest")
+        backupUtils = CommunalBackupUtils(kosmos.applicationContext)
 
         setAppWidgetIds(emptyList())
 
@@ -109,6 +116,7 @@ class CommunalWidgetRepositoryImplTest : SysuiTestCase() {
                 communalWidgetDao,
                 logBuffer,
                 backupManager,
+                backupUtils,
             )
     }
 
@@ -288,6 +296,144 @@ class CommunalWidgetRepositoryImplTest : SysuiTestCase() {
             verify(backupManager).dataChanged()
         }
 
+    @Test
+    fun restoreWidgets_deleteStateFileIfRestoreFails() =
+        testScope.runTest {
+            // Write a state file that is invalid, and verify it is written
+            backupUtils.writeBytesToDisk(byteArrayOf(1, 2, 3, 4, 5, 6))
+            assertThat(backupUtils.fileExists()).isTrue()
+
+            // Try to restore widgets
+            underTest.restoreWidgets(emptyMap())
+            runCurrent()
+
+            // The restore should fail, and verify that the file is deleted
+            assertThat(backupUtils.fileExists()).isFalse()
+        }
+
+    @Test
+    fun restoreWidgets_deleteStateFileAfterWidgetsRestored() =
+        testScope.runTest {
+            // Write a state file, and verify it is written
+            backupUtils.writeBytesToDisk(fakeState.toByteArray())
+            assertThat(backupUtils.fileExists()).isTrue()
+
+            // Set up app widget host with widget ids
+            setAppWidgetIds(listOf(11, 12))
+
+            // Restore widgets
+            underTest.restoreWidgets(mapOf(Pair(1, 11), Pair(2, 12)))
+            runCurrent()
+
+            // Verify state restored
+            verify(communalWidgetDao).restoreCommunalHubState(any())
+
+            // Verify state file deleted
+            assertThat(backupUtils.fileExists()).isFalse()
+        }
+
+    @Test
+    fun restoreWidgets_restoredWidgetsNotRegisteredWithHostAreSkipped() =
+        testScope.runTest {
+            // Write fake state to file
+            backupUtils.writeBytesToDisk(fakeState.toByteArray())
+
+            // Set up app widget host with widget ids. Widget 12 (previously 2) is absent.
+            setAppWidgetIds(listOf(11))
+
+            // Restore widgets.
+            underTest.restoreWidgets(mapOf(Pair(1, 11), Pair(2, 12)))
+            runCurrent()
+
+            // Verify state restored, and widget 2 skipped
+            val restoredState =
+                withArgCaptor<CommunalHubState> {
+                    verify(communalWidgetDao).restoreCommunalHubState(capture())
+                }
+            val restoredWidgets = restoredState.widgets.toList()
+            assertThat(restoredWidgets).hasSize(1)
+
+            val restoredWidget = restoredWidgets.first()
+            val expectedWidget = fakeState.widgets.first()
+
+            // Verify widget id is updated, and the rest remain the same as expected
+            assertThat(restoredWidget.widgetId).isEqualTo(11)
+            assertThat(restoredWidget.componentName).isEqualTo(expectedWidget.componentName)
+            assertThat(restoredWidget.rank).isEqualTo(expectedWidget.rank)
+        }
+
+    @Test
+    fun restoreWidgets_registeredWidgetsNotRestoredAreRemoved() =
+        testScope.runTest {
+            // Write fake state to file
+            backupUtils.writeBytesToDisk(fakeState.toByteArray())
+
+            // Set up app widget host with widget ids. Widget 13 will not be restored.
+            setAppWidgetIds(listOf(11, 12, 13))
+
+            // Restore widgets.
+            underTest.restoreWidgets(mapOf(Pair(1, 11), Pair(2, 12)))
+            runCurrent()
+
+            // Verify widget 1 and 2 are restored, and are now 11 and 12.
+            val restoredState =
+                withArgCaptor<CommunalHubState> {
+                    verify(communalWidgetDao).restoreCommunalHubState(capture())
+                }
+            val restoredWidgets = restoredState.widgets.toList()
+            assertThat(restoredWidgets).hasSize(2)
+
+            val restoredWidget1 = restoredWidgets[0]
+            val expectedWidget1 = fakeState.widgets[0]
+            assertThat(restoredWidget1.widgetId).isEqualTo(11)
+            assertThat(restoredWidget1.componentName).isEqualTo(expectedWidget1.componentName)
+            assertThat(restoredWidget1.rank).isEqualTo(expectedWidget1.rank)
+
+            val restoredWidget2 = restoredWidgets[1]
+            val expectedWidget2 = fakeState.widgets[1]
+            assertThat(restoredWidget2.widgetId).isEqualTo(12)
+            assertThat(restoredWidget2.componentName).isEqualTo(expectedWidget2.componentName)
+            assertThat(restoredWidget2.rank).isEqualTo(expectedWidget2.rank)
+
+            // Verify widget 13 removed since it is not restored
+            verify(appWidgetHost).deleteAppWidgetId(13)
+        }
+
+    @Test
+    fun restoreWidgets_onlySomeWidgetsGotNewIds() =
+        testScope.runTest {
+            // Write fake state to file
+            backupUtils.writeBytesToDisk(fakeState.toByteArray())
+
+            // Set up app widget host with widget ids. Widget 2 gets a new id: 12, but widget 1 does
+            // not.
+            setAppWidgetIds(listOf(1, 12))
+
+            // Restore widgets.
+            underTest.restoreWidgets(mapOf(Pair(2, 12)))
+            runCurrent()
+
+            // Verify widget 1 and 2 are restored, and are now 1 and 12.
+            val restoredState =
+                withArgCaptor<CommunalHubState> {
+                    verify(communalWidgetDao).restoreCommunalHubState(capture())
+                }
+            val restoredWidgets = restoredState.widgets.toList()
+            assertThat(restoredWidgets).hasSize(2)
+
+            val restoredWidget1 = restoredWidgets[0]
+            val expectedWidget1 = fakeState.widgets[0]
+            assertThat(restoredWidget1.widgetId).isEqualTo(1)
+            assertThat(restoredWidget1.componentName).isEqualTo(expectedWidget1.componentName)
+            assertThat(restoredWidget1.rank).isEqualTo(expectedWidget1.rank)
+
+            val restoredWidget2 = restoredWidgets[1]
+            val expectedWidget2 = fakeState.widgets[1]
+            assertThat(restoredWidget2.widgetId).isEqualTo(12)
+            assertThat(restoredWidget2.componentName).isEqualTo(expectedWidget2.componentName)
+            assertThat(restoredWidget2.rank).isEqualTo(expectedWidget2.rank)
+        }
+
     private fun installedProviders(providers: List<AppWidgetProviderInfo>) {
         whenever(appWidgetManager.installedProviders).thenReturn(providers)
     }
@@ -305,5 +451,22 @@ class CommunalWidgetRepositoryImplTest : SysuiTestCase() {
                 widgetFeatures =
                     WIDGET_FEATURE_CONFIGURATION_OPTIONAL or WIDGET_FEATURE_RECONFIGURABLE
             }
+        val fakeState =
+            CommunalHubState().apply {
+                widgets =
+                    listOf(
+                            CommunalHubState.CommunalWidgetItem().apply {
+                                widgetId = 1
+                                componentName = "pk_name/fake_widget_1"
+                                rank = 1
+                            },
+                            CommunalHubState.CommunalWidgetItem().apply {
+                                widgetId = 2
+                                componentName = "pk_name/fake_widget_2"
+                                rank = 2
+                            },
+                        )
+                        .toTypedArray()
+            }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/CommunalBackupRestoreStartable.kt b/packages/SystemUI/src/com/android/systemui/communal/CommunalBackupRestoreStartable.kt
new file mode 100644
index 000000000000..cdeeb6ff0b23
--- /dev/null
+++ b/packages/SystemUI/src/com/android/systemui/communal/CommunalBackupRestoreStartable.kt
@@ -0,0 +1,83 @@
+/*
+ * Copyright (C) 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.systemui.communal
+
+import android.appwidget.AppWidgetManager
+import android.content.BroadcastReceiver
+import android.content.Context
+import android.content.Intent
+import android.content.IntentFilter
+import com.android.systemui.CoreStartable
+import com.android.systemui.broadcast.BroadcastDispatcher
+import com.android.systemui.communal.domain.interactor.CommunalInteractor
+import com.android.systemui.communal.widgets.CommunalWidgetModule
+import com.android.systemui.dagger.SysUISingleton
+import com.android.systemui.log.LogBuffer
+import com.android.systemui.log.core.Logger
+import com.android.systemui.log.dagger.CommunalLog
+import javax.inject.Inject
+
+@SysUISingleton
+class CommunalBackupRestoreStartable
+@Inject
+constructor(
+    private val broadcastDispatcher: BroadcastDispatcher,
+    private val communalInteractor: CommunalInteractor,
+    @CommunalLog logBuffer: LogBuffer,
+) : CoreStartable, BroadcastReceiver() {
+
+    private val logger = Logger(logBuffer, TAG)
+
+    override fun start() {
+        broadcastDispatcher.registerReceiver(
+            receiver = this,
+            filter = IntentFilter(AppWidgetManager.ACTION_APPWIDGET_HOST_RESTORED),
+        )
+    }
+
+    override fun onReceive(context: Context?, intent: Intent?) {
+        if (intent == null) {
+            logger.w("On app widget host restored, but intent is null")
+            return
+        }
+
+        if (intent.action != AppWidgetManager.ACTION_APPWIDGET_HOST_RESTORED) {
+            return
+        }
+
+        val hostId = intent.getIntExtra(AppWidgetManager.EXTRA_HOST_ID, 0)
+        if (hostId != CommunalWidgetModule.APP_WIDGET_HOST_ID) {
+            return
+        }
+
+        val oldIds = intent.getIntArrayExtra(AppWidgetManager.EXTRA_APPWIDGET_OLD_IDS)
+        val newIds = intent.getIntArrayExtra(AppWidgetManager.EXTRA_APPWIDGET_IDS)
+
+        if (oldIds == null || newIds == null || oldIds.size != newIds.size) {
+            logger.w("On app widget host restored, but old to new ids mapping is invalid")
+            communalInteractor.abortRestoreWidgets()
+            return
+        }
+
+        val oldToNewWidgetIdMap = oldIds.zip(newIds).toMap()
+        communalInteractor.restoreWidgets(oldToNewWidgetIdMap)
+    }
+
+    companion object {
+        const val TAG = "CommunalBackupRestoreStartable"
+    }
+}
diff --git a/packages/SystemUI/src/com/android/systemui/communal/dagger/CommunalModule.kt b/packages/SystemUI/src/com/android/systemui/communal/dagger/CommunalModule.kt
index 72dcb26b089a..3dadea5961e4 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/dagger/CommunalModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/dagger/CommunalModule.kt
@@ -16,6 +16,8 @@
 
 package com.android.systemui.communal.dagger
 
+import android.content.Context
+import com.android.systemui.communal.data.backup.CommunalBackupUtils
 import com.android.systemui.communal.data.db.CommunalDatabaseModule
 import com.android.systemui.communal.data.repository.CommunalMediaRepositoryModule
 import com.android.systemui.communal.data.repository.CommunalPrefsRepositoryModule
@@ -74,5 +76,13 @@ interface CommunalModule {
                 )
             return SceneDataSourceDelegator(applicationScope, config)
         }
+
+        @Provides
+        @SysUISingleton
+        fun providesCommunalBackupUtils(
+            @Application context: Context,
+        ): CommunalBackupUtils {
+            return CommunalBackupUtils(context)
+        }
     }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalWidgetDao.kt b/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalWidgetDao.kt
index dc53da80f654..d174fd1c97ea 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalWidgetDao.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/data/db/CommunalWidgetDao.kt
@@ -23,6 +23,7 @@ import androidx.room.Query
 import androidx.room.RoomDatabase
 import androidx.room.Transaction
 import androidx.sqlite.db.SupportSQLiteDatabase
+import com.android.systemui.communal.nano.CommunalHubState
 import com.android.systemui.communal.widgets.CommunalWidgetHost
 import com.android.systemui.communal.widgets.CommunalWidgetModule.Companion.DEFAULT_WIDGETS
 import com.android.systemui.dagger.qualifiers.Application
@@ -116,6 +117,10 @@ interface CommunalWidgetDao {
     @Query("UPDATE communal_item_rank_table SET rank = :order WHERE uid = :itemUid")
     fun updateItemRank(itemUid: Long, order: Int)
 
+    @Query("DELETE FROM communal_widget_table") fun clearCommunalWidgetsTable()
+
+    @Query("DELETE FROM communal_item_rank_table") fun clearCommunalItemRankTable()
+
     @Transaction
     fun updateWidgetOrder(widgetIdToPriorityMap: Map<Int, Int>) {
         widgetIdToPriorityMap.forEach { (id, priority) ->
@@ -154,4 +159,13 @@ interface CommunalWidgetDao {
         deleteWidgets(widget)
         return true
     }
+
+    /** Wipes current database and restores the snapshot represented by [state]. */
+    @Transaction
+    fun restoreCommunalHubState(state: CommunalHubState) {
+        clearCommunalWidgetsTable()
+        clearCommunalItemRankTable()
+
+        state.widgets.forEach { addWidget(it.widgetId, it.componentName, it.rank) }
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalWidgetRepository.kt b/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalWidgetRepository.kt
index 3b258dcde542..1f54e70fa21b 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalWidgetRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/data/repository/CommunalWidgetRepository.kt
@@ -21,9 +21,12 @@ import android.appwidget.AppWidgetManager
 import android.content.ComponentName
 import android.os.UserHandle
 import androidx.annotation.WorkerThread
+import com.android.systemui.communal.data.backup.CommunalBackupUtils
 import com.android.systemui.communal.data.db.CommunalItemRank
 import com.android.systemui.communal.data.db.CommunalWidgetDao
 import com.android.systemui.communal.data.db.CommunalWidgetItem
+import com.android.systemui.communal.nano.CommunalHubState
+import com.android.systemui.communal.proto.toCommunalHubState
 import com.android.systemui.communal.shared.model.CommunalWidgetContentModel
 import com.android.systemui.communal.widgets.CommunalAppWidgetHost
 import com.android.systemui.communal.widgets.CommunalWidgetHost
@@ -70,6 +73,15 @@ interface CommunalWidgetRepository {
      * @param widgetIdToPriorityMap mapping of the widget ids to the priority of the widget.
      */
     fun updateWidgetOrder(widgetIdToPriorityMap: Map<Int, Int>) {}
+
+    /**
+     * Restores the database by reading a state file from disk and updating the widget ids according
+     * to [oldToNewWidgetIdMap].
+     */
+    fun restoreWidgets(oldToNewWidgetIdMap: Map<Int, Int>)
+
+    /** Aborts the restore process and removes files from disk if necessary. */
+    fun abortRestoreWidgets()
 }
 
 @SysUISingleton
@@ -84,6 +96,7 @@ constructor(
     private val communalWidgetDao: CommunalWidgetDao,
     @CommunalLog logBuffer: LogBuffer,
     private val backupManager: BackupManager,
+    private val backupUtils: CommunalBackupUtils,
 ) : CommunalWidgetRepository {
     companion object {
         const val TAG = "CommunalWidgetRepository"
@@ -173,6 +186,75 @@ constructor(
         }
     }
 
+    override fun restoreWidgets(oldToNewWidgetIdMap: Map<Int, Int>) {
+        bgScope.launch {
+            // Read restored state file from disk
+            val state: CommunalHubState
+            try {
+                state = backupUtils.readBytesFromDisk().toCommunalHubState()
+            } catch (e: Exception) {
+                logger.e({ "Failed reading restore data from disk: $str1" }) {
+                    str1 = e.localizedMessage
+                }
+                abortRestoreWidgets()
+                return@launch
+            }
+
+            val widgetsWithHost = appWidgetHost.appWidgetIds.toList()
+            val widgetsToRemove = widgetsWithHost.toMutableList()
+
+            // Produce a new state to be restored, skipping invalid widgets
+            val newWidgets =
+                state.widgets.mapNotNull { restoredWidget ->
+                    val newWidgetId =
+                        oldToNewWidgetIdMap[restoredWidget.widgetId] ?: restoredWidget.widgetId
+
+                    // Skip if widget id is not registered with the host
+                    if (!widgetsWithHost.contains(newWidgetId)) {
+                        logger.d({
+                            "Skipped restoring widget (old:$int1 new:$int2) " +
+                                "because it is not registered with host"
+                        }) {
+                            int1 = restoredWidget.widgetId
+                            int2 = newWidgetId
+                        }
+                        return@mapNotNull null
+                    }
+
+                    widgetsToRemove.remove(newWidgetId)
+
+                    CommunalHubState.CommunalWidgetItem().apply {
+                        widgetId = newWidgetId
+                        componentName = restoredWidget.componentName
+                        rank = restoredWidget.rank
+                    }
+                }
+            val newState = CommunalHubState().apply { widgets = newWidgets.toTypedArray() }
+
+            // Restore database
+            logger.i("Restoring communal database $newState")
+            communalWidgetDao.restoreCommunalHubState(newState)
+
+            // Delete restored state file from disk
+            backupUtils.clear()
+
+            // Remove widgets from host that have not been restored
+            widgetsToRemove.forEach { widgetId ->
+                logger.i({ "Deleting widget $int1 from host since it has not been restored" }) {
+                    int1 = widgetId
+                }
+                appWidgetHost.deleteAppWidgetId(widgetId)
+            }
+        }
+    }
+
+    override fun abortRestoreWidgets() {
+        bgScope.launch {
+            logger.i("Restore widgets aborted")
+            backupUtils.clear()
+        }
+    }
+
     @WorkerThread
     private fun mapToContentModel(
         entry: Map.Entry<CommunalItemRank, CommunalWidgetItem>
diff --git a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt
index 373e1c9daa7b..f2275723c664 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/domain/interactor/CommunalInteractor.kt
@@ -171,6 +171,23 @@ constructor(
     }
 
     /**
+     * Repopulates the communal widgets database by first reading a backed-up state from disk and
+     * updating the widget ids indicated by [oldToNewWidgetIdMap]. The backed-up state is removed
+     * from disk afterwards.
+     */
+    fun restoreWidgets(oldToNewWidgetIdMap: Map<Int, Int>) {
+        widgetRepository.restoreWidgets(oldToNewWidgetIdMap)
+    }
+
+    /**
+     * Aborts the task of restoring widgets from a backup. The backed up state stored on disk is
+     * removed.
+     */
+    fun abortRestoreWidgets() {
+        widgetRepository.abortRestoreWidgets()
+    }
+
+    /**
      * Updates the transition state of the hub [SceneTransitionLayout].
      *
      * Note that you must call is with `null` when the UI is done or risk a memory leak.
diff --git a/packages/SystemUI/src/com/android/systemui/communal/widgets/CommunalWidgetModule.kt b/packages/SystemUI/src/com/android/systemui/communal/widgets/CommunalWidgetModule.kt
index 60fb8d4840cb..aa6516d54563 100644
--- a/packages/SystemUI/src/com/android/systemui/communal/widgets/CommunalWidgetModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/communal/widgets/CommunalWidgetModule.kt
@@ -38,7 +38,7 @@ import kotlinx.coroutines.CoroutineScope
 @Module
 interface CommunalWidgetModule {
     companion object {
-        private const val APP_WIDGET_HOST_ID = 116
+        const val APP_WIDGET_HOST_ID = 116
         const val DEFAULT_WIDGETS = "default_widgets"
 
         @SysUISingleton
diff --git a/packages/SystemUI/src/com/android/systemui/dagger/SystemUICoreStartableModule.kt b/packages/SystemUI/src/com/android/systemui/dagger/SystemUICoreStartableModule.kt
index 21ee5bd92328..98e65bd89aad 100644
--- a/packages/SystemUI/src/com/android/systemui/dagger/SystemUICoreStartableModule.kt
+++ b/packages/SystemUI/src/com/android/systemui/dagger/SystemUICoreStartableModule.kt
@@ -25,6 +25,7 @@ import com.android.systemui.back.domain.interactor.BackActionInteractor
 import com.android.systemui.biometrics.BiometricNotificationService
 import com.android.systemui.clipboardoverlay.ClipboardListener
 import com.android.systemui.communal.CommunalDreamStartable
+import com.android.systemui.communal.CommunalBackupRestoreStartable
 import com.android.systemui.communal.CommunalSceneStartable
 import com.android.systemui.communal.log.CommunalLoggerStartable
 import com.android.systemui.communal.widgets.CommunalAppWidgetHostStartable
@@ -341,6 +342,13 @@ abstract class SystemUICoreStartableModule {
 
     @Binds
     @IntoMap
+    @ClassKey(CommunalBackupRestoreStartable::class)
+    abstract fun bindCommunalBackupRestoreStartable(
+        impl: CommunalBackupRestoreStartable
+    ): CoreStartable
+
+    @Binds
+    @IntoMap
     @ClassKey(HomeControlsDreamStartable::class)
     abstract fun bindHomeControlsDreamStartable(impl: HomeControlsDreamStartable): CoreStartable
 }
diff --git a/packages/SystemUI/tests/src/com/android/systemui/communal/data/db/CommunalWidgetDaoTest.kt b/packages/SystemUI/tests/src/com/android/systemui/communal/data/db/CommunalWidgetDaoTest.kt
index 20dd913550d3..f77c7a672ae3 100644
--- a/packages/SystemUI/tests/src/com/android/systemui/communal/data/db/CommunalWidgetDaoTest.kt
+++ b/packages/SystemUI/tests/src/com/android/systemui/communal/data/db/CommunalWidgetDaoTest.kt
@@ -21,6 +21,7 @@ import androidx.room.Room
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
 import com.android.systemui.SysuiTestCase
+import com.android.systemui.communal.nano.CommunalHubState
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.lifecycle.InstantTaskExecutorRule
 import com.google.common.truth.Truth.assertThat
@@ -224,6 +225,42 @@ class CommunalWidgetDaoTest : SysuiTestCase() {
                 .inOrder()
         }
 
+    @Test
+    fun restoreCommunalHubState() =
+        testScope.runTest {
+            // Set up db
+            listOf(widgetInfo1, widgetInfo2, widgetInfo3).forEach { addWidget(it) }
+
+            // Restore db to fake state
+            communalWidgetDao.restoreCommunalHubState(fakeState)
+
+            // Verify db matches new state
+            val expected = mutableMapOf<CommunalItemRank, CommunalWidgetItem>()
+            fakeState.widgets.forEachIndexed { index, fakeWidget ->
+                // Auto-generated uid continues after the initial 3 widgets and starts at 4
+                val uid = index + 4L
+                val rank = CommunalItemRank(uid = uid, rank = fakeWidget.rank)
+                val widget =
+                    CommunalWidgetItem(
+                        uid = uid,
+                        widgetId = fakeWidget.widgetId,
+                        componentName = fakeWidget.componentName,
+                        itemId = rank.uid,
+                    )
+                expected[rank] = widget
+            }
+            val widgets by collectLastValue(communalWidgetDao.getWidgets())
+            assertThat(widgets).containsExactlyEntriesIn(expected)
+        }
+
+    private fun addWidget(metadata: FakeWidgetMetadata, priority: Int? = null) {
+        communalWidgetDao.addWidget(
+            widgetId = metadata.widgetId,
+            provider = metadata.provider,
+            priority = priority ?: metadata.priority,
+        )
+    }
+
     data class FakeWidgetMetadata(
         val widgetId: Int,
         val provider: ComponentName,
@@ -273,5 +310,22 @@ class CommunalWidgetDaoTest : SysuiTestCase() {
                 componentName = widgetInfo3.provider.flattenToString(),
                 itemId = communalItemRankEntry3.uid,
             )
+        val fakeState =
+            CommunalHubState().apply {
+                widgets =
+                    listOf(
+                            CommunalHubState.CommunalWidgetItem().apply {
+                                widgetId = 1
+                                componentName = "pk_name/fake_widget_1"
+                                rank = 1
+                            },
+                            CommunalHubState.CommunalWidgetItem().apply {
+                                widgetId = 2
+                                componentName = "pk_name/fake_widget_2"
+                                rank = 2
+                            },
+                        )
+                        .toTypedArray()
+            }
     }
 }
diff --git a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/FakeCommunalWidgetRepository.kt b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/FakeCommunalWidgetRepository.kt
index 4ed6fe27338a..329c0f1ab5b4 100644
--- a/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/FakeCommunalWidgetRepository.kt
+++ b/packages/SystemUI/tests/utils/src/com/android/systemui/communal/data/repository/FakeCommunalWidgetRepository.kt
@@ -46,6 +46,10 @@ class FakeCommunalWidgetRepository(private val coroutineScope: CoroutineScope) :
         _communalWidgets.value = _communalWidgets.value.filter { it.appWidgetId != widgetId }
     }
 
+    override fun restoreWidgets(oldToNewWidgetIdMap: Map<Int, Int>) {}
+
+    override fun abortRestoreWidgets() {}
+
     private fun onConfigured(id: Int, providerInfo: AppWidgetProviderInfo, priority: Int) {
         _communalWidgets.value += listOf(CommunalWidgetContentModel(id, providerInfo, priority))
     }