Fix some direct boot related crashes and bugs

java/src/org/futo/inputmethod/latin/LatinIME.kt

@@ -80,6 +80,7 @@ import org.futo.inputmethod.latin.uix.setSetting
 import org.futo.inputmethod.latin.uix.theme.ThemeOption
 import org.futo.inputmethod.latin.uix.theme.ThemeOptions
 import org.futo.inputmethod.latin.uix.theme.applyWindowColors
+import org.futo.inputmethod.latin.uix.theme.orDefault
 import org.futo.inputmethod.latin.uix.theme.presets.DefaultDarkScheme
 import org.futo.inputmethod.latin.xlm.LanguageModelFacilitator
 import org.futo.inputmethod.updates.scheduleUpdateCheckingJob
@@ -89,6 +90,7 @@ import org.futo.inputmethod.v2keyboard.KeyboardSettings
 import org.futo.inputmethod.v2keyboard.KeyboardSizeSettingKind
 import org.futo.inputmethod.v2keyboard.KeyboardSizeStateProvider
 import org.futo.inputmethod.v2keyboard.KeyboardSizingCalculator
+import org.futo.inputmethod.v2keyboard.LayoutManager
 
 private class UnlockedBroadcastReceiver(val onDeviceUnlocked: () -> Unit) : BroadcastReceiver() {
     override fun onReceive(context: Context?, intent: Intent?) {
@@ -336,7 +338,9 @@ class LatinIME : InputMethodServiceCompose(), LatinIMELegacy.SuggestionStripCont
     override fun onCreate() {
         super.onCreate()
 
-        DataStoreHelper.init(this, lifecycleScope)
+        LayoutManager.init(this)
+
+        DataStoreHelper.init(this)
 
         val filter = IntentFilter(Intent.ACTION_USER_UNLOCKED)
         registerReceiver(unlockReceiver, filter)
@@ -356,12 +360,7 @@ class LatinIME : InputMethodServiceCompose(), LatinIMELegacy.SuggestionStripCont
         )
 
         getSettingBlocking(THEME_KEY).let {
-            val themeOptionFromSettings = ThemeOptions[it]
-            val themeOption = when {
-                themeOptionFromSettings == null -> DefaultDarkScheme
-                !themeOptionFromSettings.available(this@LatinIME) -> DefaultDarkScheme
-                else -> themeOptionFromSettings
-            }
+            val themeOption = ThemeOptions[it].orDefault(this@LatinIME)
 
             activeThemeOption = themeOption
             activeColorScheme = themeOption.obtainColors(this@LatinIME)
@@ -859,6 +858,8 @@ class LatinIME : InputMethodServiceCompose(), LatinIMELegacy.SuggestionStripCont
 
         uixManager.onPersistentStatesUnlocked()
 
+        updateTheme(ThemeOptions[getSettingBlocking(THEME_KEY)].orDefault(this))
+
         // TODO: Spell checker service
     }
 

java/src/org/futo/inputmethod/latin/uix/ImportResourceActivity.kt

@@ -51,8 +51,7 @@ import org.futo.inputmethod.latin.uix.theme.StatusBarColorSetter
 import org.futo.inputmethod.latin.uix.theme.ThemeOption
 import org.futo.inputmethod.latin.uix.theme.ThemeOptions
 import org.futo.inputmethod.latin.uix.theme.UixThemeWrapper
-import org.futo.inputmethod.latin.uix.theme.presets.DefaultDarkScheme
-import org.futo.inputmethod.latin.uix.theme.presets.VoiceInputTheme
+import org.futo.inputmethod.latin.uix.theme.orDefault
 import org.futo.inputmethod.latin.utils.SubtypeLocaleUtils
 import org.futo.inputmethod.latin.xlm.ModelPaths
 import org.futo.voiceinput.shared.BUILTIN_ENGLISH_MODEL
@@ -498,14 +497,7 @@ class ImportResourceActivity : ComponentActivity() {
         }
 
         deferGetSetting(THEME_KEY) {
-            val themeOptionFromSettings = ThemeOptions[it]
-            val themeOption = when {
-                themeOptionFromSettings == null -> DefaultDarkScheme
-                !themeOptionFromSettings.available(this) -> DefaultDarkScheme
-                else -> themeOptionFromSettings
-            }
-
-            this.themeOption.value = themeOption
+            this.themeOption.value = ThemeOptions[it].orDefault(this)
         }
     }
 }
\ No newline at end of file

java/src/org/futo/inputmethod/latin/uix/Settings.kt

@@ -23,18 +23,16 @@ import kotlinx.coroutines.Job
 import kotlinx.coroutines.SupervisorJob
 import kotlinx.coroutines.flow.Flow
 import kotlinx.coroutines.flow.MutableSharedFlow
-import kotlinx.coroutines.flow.first
 import kotlinx.coroutines.flow.map
-import kotlinx.coroutines.flow.take
 import kotlinx.coroutines.launch
 import kotlinx.coroutines.runBlocking
+import kotlinx.coroutines.sync.Mutex
 import kotlinx.coroutines.withContext
 import org.futo.inputmethod.latin.ActiveSubtype
 import org.futo.inputmethod.latin.Subtypes
 import org.futo.inputmethod.latin.SubtypesSetting
 import org.futo.inputmethod.latin.uix.theme.presets.ClassicMaterialDark
 import org.futo.inputmethod.latin.uix.theme.presets.DynamicSystemTheme
-import org.futo.inputmethod.v2keyboard.LayoutManager
 
 // Used before first unlock (direct boot)
 private object DefaultDataStore : DataStore<Preferences> {
@@ -84,40 +82,45 @@ private object DefaultDataStore : DataStore<Preferences> {
 // Set and used after first unlock (direct boot)
 private var unlockedDataStore: DataStore<Preferences>? = null
 
+// To prevent two threads trying to create a datastore at once
+private val dataStoreCreationMutex = Mutex()
+
 // Initializes unlockedDataStore, or uses DefaultDataStore if device is still locked (direct boot)
 @OptIn(DelicateCoroutinesApi::class)
 val Context.dataStore: DataStore<Preferences>
     get() {
         return unlockedDataStore ?: run {
             val userManager = getSystemService(Context.USER_SERVICE) as UserManager
-            return if (userManager.isUserUnlocked) {
-                // The device has been unlocked
-                val newDataStore = PreferenceDataStoreFactory.create(
-                    corruptionHandler = null,
-                    migrations = listOf(),
-                    scope = CoroutineScope(Dispatchers.IO + SupervisorJob())
-                ) {
-                    applicationContext.preferencesDataStoreFile("settings")
-                }
+            return if (userManager.isUserUnlocked && dataStoreCreationMutex.tryLock()) {
+                try {
+                    // The device has been unlocked
+                    val newDataStore = PreferenceDataStoreFactory.create(
+                        corruptionHandler = null,
+                        migrations = listOf(),
+                        scope = CoroutineScope(Dispatchers.IO + SupervisorJob())
+                    ) {
+                        applicationContext.preferencesDataStoreFile("settings")
+                    }
 
-                unlockedDataStore = newDataStore
+                    unlockedDataStore = newDataStore
 
-                // Send new values to the DefaultDataStore for any listeners
-                GlobalScope.launch {
-                    newDataStore.data.collect { value ->
-                        DefaultDataStore.sharedData.emit(value)
+                    // Send new values to the DefaultDataStore for any listeners
+                    GlobalScope.launch {
+                        newDataStore.data.collect { value ->
+                            DefaultDataStore.sharedData.emit(value)
+                        }
                     }
-                }
 
-                newDataStore
+                    newDataStore
+                } finally {
+                    dataStoreCreationMutex.unlock()
+                }
             } else {
-                // The device is still locked, return default data store
+                // The device is still locked (or an unlocked datastore is in the process of being
+                // created), return default data store
                 if (!DefaultDataStore.subtypesInitialized) {
                     DefaultDataStore.subtypesInitialized = true
 
-                    // Necessary to init LayoutManager in advance to avoid race condition
-                    LayoutManager.init(this@dataStore)
-
                     GlobalScope.launch {
                         DefaultDataStore.updateSubtypes(Subtypes.getDirectBootInitialLayouts(this@dataStore))
                     }
@@ -147,16 +150,17 @@ val Context.isDirectBootUnlocked: Boolean
     }
 
 class DataStoreHelper {
+    @OptIn(DelicateCoroutinesApi::class)
     companion object {
         private var initialized: Boolean = false
         private var currentPreferences: Preferences = preferencesOf()
 
         @JvmStatic
-        fun init(context: Context, scope: CoroutineScope) {
+        fun init(context: Context) {
             if(initialized) return
             initialized = true
 
-            scope.launch {
+            GlobalScope.launch {
                 context.dataStore.data.collect {
                     currentPreferences = it
                 }

java/src/org/futo/inputmethod/latin/uix/settings/SettingsActivity.kt

@@ -46,8 +46,7 @@ import org.futo.inputmethod.latin.uix.getSettingFlow
 import org.futo.inputmethod.latin.uix.theme.ThemeOption
 import org.futo.inputmethod.latin.uix.theme.ThemeOptions
 import org.futo.inputmethod.latin.uix.theme.UixThemeAuto
-import org.futo.inputmethod.latin.uix.theme.presets.DefaultDarkScheme
-import org.futo.inputmethod.latin.uix.theme.presets.DefaultLightScheme
+import org.futo.inputmethod.latin.uix.theme.orDefault
 import org.futo.inputmethod.latin.xlm.ModelPaths
 import org.futo.inputmethod.updates.checkForUpdateAndSaveToPreferences
 import java.io.File
@@ -199,12 +198,7 @@ class SettingsActivity : ComponentActivity(), DynamicThemeProviderOwner {
 
         lifecycleScope.launch {
             getSettingFlow(THEME_KEY).collect {
-                val themeOptionFromSettings = ThemeOptions[it]
-                val themeOption = when {
-                    themeOptionFromSettings == null -> DefaultDarkScheme
-                    !themeOptionFromSettings.available(this@SettingsActivity) -> DefaultDarkScheme
-                    else -> themeOptionFromSettings
-                }
+                val themeOption = ThemeOptions[it].orDefault(this@SettingsActivity)
 
                 this@SettingsActivity.themeOption.value = themeOption
                 this@SettingsActivity.themeProvider = BasicThemeProvider(
@@ -217,12 +211,7 @@ class SettingsActivity : ComponentActivity(), DynamicThemeProviderOwner {
         }
 
         getSettingBlocking(THEME_KEY).let {
-            val themeOptionFromSettings = ThemeOptions[it]
-            val themeOption = when {
-                themeOptionFromSettings == null -> DefaultDarkScheme
-                !themeOptionFromSettings.available(this) -> DefaultLightScheme
-                else -> themeOptionFromSettings
-            }
+            val themeOption = ThemeOptions[it].orDefault(this@SettingsActivity)
 
             this.themeOption.value = themeOption
             this.themeProvider = BasicThemeProvider(

java/src/org/futo/inputmethod/latin/uix/theme/ThemeOptions.kt

@@ -56,4 +56,15 @@ val ThemeOptions = mapOf(
     HotDog.key to HotDog
 )
 
-val ThemeOptionKeys = ThemeOptions.keys
\ No newline at end of file
+val ThemeOptionKeys = ThemeOptions.keys
+
+fun ThemeOption?.orDefault(context: Context): ThemeOption {
+    val themeOptionFromSettings = this
+    val themeOption = when {
+        themeOptionFromSettings == null -> DefaultDarkScheme
+        !themeOptionFromSettings.available(context) -> DefaultDarkScheme
+        else -> themeOptionFromSettings
+    }
+
+    return themeOption
+}
\ No newline at end of file

java/src/org/futo/inputmethod/latin/uix/voiceinput/downloader/DownloadActivity.kt

@@ -56,8 +56,7 @@ import org.futo.inputmethod.latin.uix.settings.useDataStore
 import org.futo.inputmethod.latin.uix.theme.ThemeOption
 import org.futo.inputmethod.latin.uix.theme.ThemeOptions
 import org.futo.inputmethod.latin.uix.theme.UixThemeWrapper
-import org.futo.inputmethod.latin.uix.theme.presets.DefaultDarkScheme
-import org.futo.inputmethod.latin.uix.theme.presets.VoiceInputTheme
+import org.futo.inputmethod.latin.uix.theme.orDefault
 import org.futo.voiceinput.shared.ui.theme.Typography
 import java.io.File
 import java.io.IOException
@@ -418,14 +417,7 @@ class DownloadActivity : ComponentActivity() {
         isDownloading = false
 
         deferGetSetting(THEME_KEY) {
-            val themeOptionFromSettings = ThemeOptions[it]
-            val themeOption = when {
-                themeOptionFromSettings == null -> DefaultDarkScheme
-                !themeOptionFromSettings.available(this) -> DefaultDarkScheme
-                else -> themeOptionFromSettings
-            }
-
-            this.themeOption.value = themeOption
+            this.themeOption.value = ThemeOptions[it].orDefault(this)
         }
 
         updateContent()

java/src/org/futo/inputmethod/v2keyboard/LayoutManager.kt

@@ -70,21 +70,21 @@ object LayoutManager {
         }
     }
 
+    private fun ensureInitialized() {
+        if(!initialized) throw IllegalStateException("LayoutManager method called without being initialized")
+    }
     fun getLayout(context: Context, name: String): Keyboard {
-        init(context)
-
+        ensureInitialized()
         return layoutsById?.get(name) ?: throw IllegalArgumentException("Failed to find keyboard layout $name. Available layouts: ${layoutsById?.keys}")
     }
 
     fun getLayoutMapping(context: Context): Map<Locale, List<String>> {
-        init(context)
-
+        ensureInitialized()
         return localeToLayoutsMappings!!
     }
 
     fun getAllLayoutNames(context: Context): List<String> {
-        init(context)
-
+        ensureInitialized()
         return getAllLayoutPaths(context.assets).map {
             it.split("/").last().split(".yaml").first()
         }