Port auto restore optimization from legacy v0 code to new v1 code

app/src/main/java/com/stevesoltys/seedvault/transport/restore/KVRestore.kt

@@ -34,8 +34,7 @@ private class KVRestoreState(
     /**
      * Optional [PackageInfo] for single package restore, optimizes restore of @pm@
      */
-    @Deprecated("TODO remove?")
+    val autoRestorePackageInfo: PackageInfo?
-    val pmPackageInfo: PackageInfo?
 )
 
 private val TAG = KVRestore::class.java.simpleName
@@ -70,16 +69,16 @@ internal class KVRestore(
      * It is possible that the system decides to not restore the package.
      * Then a new state will be initialized right away without calling other methods.
      *
-     * @param pmPackageInfo single optional [PackageInfo] to optimize restore of @pm@
+     * @param autoRestorePackageInfo single optional [PackageInfo] to optimize restore of @pm@
      */
     fun initializeState(
         version: Byte,
         token: Long,
         name: String,
         packageInfo: PackageInfo,
-        pmPackageInfo: PackageInfo? = null
+        autoRestorePackageInfo: PackageInfo? = null
     ) {
-        state = KVRestoreState(version, token, name, packageInfo, pmPackageInfo)
+        state = KVRestoreState(version, token, name, packageInfo, autoRestorePackageInfo)
     }
 
     /**
@@ -95,10 +94,20 @@ internal class KVRestore(
         // take legacy path for version 0
         if (state.version == 0x00.toByte()) return getRestoreDataV0(state, data)
 
+        val pmPackageName = state.autoRestorePackageInfo?.packageName
+        val isAutoRestore = state.packageInfo.packageName == MAGIC_PACKAGE_MANAGER &&
+            pmPackageName != null
         return try {
-            val db = getRestoreDb(state)
+            val db = if (isAutoRestore) getCachedRestoreDb(state) else downloadRestoreDb(state)
             val out = outputFactory.getBackupDataOutput(data)
-            db.getAll().sortedBy { it.first }.forEach { (key, value) ->
+            val records = if (isAutoRestore) {
+                val keys = listOf(ANCESTRAL_RECORD_KEY, GLOBAL_METADATA_KEY, pmPackageName)
+                Log.d(TAG, "Single package restore, restrict restore keys to $pmPackageName")
+                db.getAll().filter { it.first in keys }
+            } else {
+                db.getAll()
+            }
+            records.sortedBy { it.first }.forEach { (key, value) ->
                 val size = value.size
                 Log.v(TAG, "    ... key=$key size=$size")
                 out.writeEntityHeader(key, size)
@@ -125,7 +134,17 @@ internal class KVRestore(
     }
 
     @Throws(IOException::class, GeneralSecurityException::class, UnsupportedVersionException::class)
-    private suspend fun getRestoreDb(state: KVRestoreState): KVDb {
+    private suspend fun getCachedRestoreDb(state: KVRestoreState): KVDb {
+        val packageName = state.packageInfo.packageName
+        return if (dbManager.existsDb(packageName)) {
+            dbManager.getDb(packageName)
+        } else {
+            downloadRestoreDb(state)
+        }
+    }
+
+    @Throws(IOException::class, GeneralSecurityException::class, UnsupportedVersionException::class)
+    private suspend fun downloadRestoreDb(state: KVRestoreState): KVDb {
         val packageName = state.packageInfo.packageName
         plugin.getInputStream(state.token, state.name).use { inputStream ->
             headerReader.readVersion(inputStream, state.version)
@@ -198,18 +217,11 @@ internal class KVRestore(
         if (records.isEmpty()) return null
 
         // Decode the key filenames into keys then sort lexically by key
-        val contents = ArrayList<DecodedKey>()
+        val contents = ArrayList<DecodedKey>().apply {
-        for (recordKey in records) contents.add(DecodedKey(recordKey))
+            for (recordKey in records) add(DecodedKey(recordKey))
-        // remove keys that are not needed for single package @pm@ restore
+        }
-        val pmPackageName = state?.pmPackageInfo?.packageName
+        contents.sort()
-        val sortedKeys =
+        return contents
-            if (packageInfo.packageName == MAGIC_PACKAGE_MANAGER && pmPackageName != null) {
-                val keys = listOf(ANCESTRAL_RECORD_KEY, GLOBAL_METADATA_KEY, pmPackageName)
-                Log.d(TAG, "Single package restore, restrict restore keys to $pmPackageName")
-                contents.filterTo(ArrayList()) { it.key in keys }
-            } else contents
-        sortedKeys.sort()
-        return sortedKeys
     }
 
     /**

app/src/main/java/com/stevesoltys/seedvault/transport/restore/RestoreCoordinator.kt

@@ -20,8 +20,8 @@ import com.stevesoltys.seedvault.metadata.BackupType
 import com.stevesoltys.seedvault.metadata.DecryptionFailedException
 import com.stevesoltys.seedvault.metadata.MetadataManager
 import com.stevesoltys.seedvault.metadata.MetadataReader
-import com.stevesoltys.seedvault.settings.SettingsManager
 import com.stevesoltys.seedvault.plugins.StoragePlugin
+import com.stevesoltys.seedvault.settings.SettingsManager
 import com.stevesoltys.seedvault.ui.notification.BackupNotificationManager
 import java.io.IOException
 
@@ -31,7 +31,7 @@ private data class RestoreCoordinatorState(
     /**
      * Optional [PackageInfo] for single package restore, to reduce data needed to read for @pm@
      */
-    val pmPackageInfo: PackageInfo?,
+    val autoRestorePackageInfo: PackageInfo?,
     val backupMetadata: BackupMetadata
 ) {
     var currentPackage: String? = null
@@ -212,7 +212,7 @@ internal class RestoreCoordinator(
                             token = state.token,
                             name = name,
                             packageInfo = packageInfo,
-                            pmPackageInfo = state.pmPackageInfo
+                            autoRestorePackageInfo = state.autoRestorePackageInfo
                         )
                         state.currentPackage = packageName
                         TYPE_KEY_VALUE
@@ -249,7 +249,7 @@ internal class RestoreCoordinator(
                 // check key/value data first and if available, don't even check for full data
                 kv.hasDataForPackage(state.token, packageInfo) -> {
                     Log.i(TAG, "Found K/V data for $packageName.")
-                    kv.initializeState(0x00, state.token, "", packageInfo, state.pmPackageInfo)
+                    kv.initializeState(0x00, state.token, "", packageInfo, null)
                     state.currentPackage = packageName
                     TYPE_KEY_VALUE
                 }