Merge "[Custom Key Glyph] Retrieve key glyph in Shortcut Helper" into main
diff --git a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperCategoriesRepositoryTest.kt b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperCategoriesRepositoryTest.kt
index 9e20e7d..620b8b6 100644
--- a/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperCategoriesRepositoryTest.kt
+++ b/packages/SystemUI/multivalentTests/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperCategoriesRepositoryTest.kt
@@ -16,7 +16,11 @@
 
 package com.android.systemui.keyboard.shortcut.data.repository
 
+import android.graphics.drawable.Drawable
+import android.hardware.input.KeyGlyphMap
 import android.hardware.input.fakeInputManager
+import android.platform.test.annotations.DisableFlags
+import android.platform.test.annotations.EnableFlags
 import android.view.KeyEvent.KEYCODE_1
 import android.view.KeyEvent.KEYCODE_A
 import android.view.KeyEvent.KEYCODE_B
@@ -26,10 +30,12 @@
 import android.view.KeyEvent.KEYCODE_F
 import android.view.KeyEvent.KEYCODE_G
 import android.view.KeyEvent.META_FUNCTION_ON
+import android.view.KeyEvent.META_META_ON
 import android.view.KeyboardShortcutGroup
 import android.view.KeyboardShortcutInfo
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.SmallTest
+import com.android.systemui.Flags.FLAG_SHORTCUT_HELPER_KEY_GLYPH
 import com.android.systemui.SysuiTestCase
 import com.android.systemui.coroutines.collectLastValue
 import com.android.systemui.keyboard.shortcut.data.source.FakeKeyboardShortcutGroupsSource
@@ -49,6 +55,7 @@
 import com.android.systemui.keyboard.shortcut.shortcutHelperTestHelper
 import com.android.systemui.kosmos.testDispatcher
 import com.android.systemui.kosmos.testScope
+import com.android.systemui.res.R
 import com.android.systemui.testKosmos
 import com.google.common.truth.Truth.assertThat
 import kotlinx.coroutines.ExperimentalCoroutinesApi
@@ -57,6 +64,9 @@
 import org.junit.Before
 import org.junit.Test
 import org.junit.runner.RunWith
+import org.mockito.ArgumentMatchers.anyInt
+import org.mockito.Mockito.mock
+import org.mockito.kotlin.whenever
 
 @OptIn(ExperimentalCoroutinesApi::class)
 @SmallTest
@@ -187,6 +197,79 @@
                 )
         }
 
+    @EnableFlags(FLAG_SHORTCUT_HELPER_KEY_GLYPH)
+    @Test
+    fun modifierMappedToCustomDrawableWhenKeyGlyphMapExists() =
+        testScope.runTest {
+            val metaDrawable = mock(Drawable::class.java)
+            val keyGlyph = mock(KeyGlyphMap::class.java)
+            whenever(keyGlyph.getDrawableForModifierState(context, META_META_ON))
+                .thenReturn(metaDrawable)
+            whenever(kosmos.fakeInputManager.inputManager.getKeyGlyphMap(anyInt()))
+                .thenReturn(keyGlyph)
+            fakeSystemSource.setGroups(simpleGroup(simpleShortcutInfo(KEYCODE_1, META_META_ON)))
+            helper.toggle(deviceId = 123)
+
+            val categories by collectLastValue(repo.categories)
+            val systemCategory = categories?.firstOrNull { it.type == ShortcutCategoryType.System }
+
+            val expectedCategory =
+                ShortcutCategory(
+                    type = ShortcutCategoryType.System,
+                    simpleSubCategory(
+                        simpleDrawableModifierShortcut("1", modifierDrawable = metaDrawable)
+                    ),
+                )
+
+            assertThat(systemCategory).isEqualTo(expectedCategory)
+        }
+
+    @EnableFlags(FLAG_SHORTCUT_HELPER_KEY_GLYPH)
+    @Test
+    fun modifierMappedToDefaultDrawableWhenNoKeyGlyphMapExists() =
+        testScope.runTest {
+            fakeSystemSource.setGroups(simpleGroup(simpleShortcutInfo(KEYCODE_1, META_META_ON)))
+            helper.toggle(deviceId = 123)
+
+            val categories by collectLastValue(repo.categories)
+            val systemCategory = categories?.firstOrNull { it.type == ShortcutCategoryType.System }
+
+            val expectedCategory =
+                ShortcutCategory(
+                    type = ShortcutCategoryType.System,
+                    simpleSubCategory(
+                        simpleResIdModifierShortcut("1", modifierResId = R.drawable.ic_ksh_key_meta)
+                    ),
+                )
+            assertThat(systemCategory).isEqualTo(expectedCategory)
+        }
+
+    @DisableFlags(FLAG_SHORTCUT_HELPER_KEY_GLYPH)
+    @Test
+    fun modifierMappedToDefaultDrawableWhenKeyGlyphDisabled() =
+        testScope.runTest {
+            val metaDrawable = mock(Drawable::class.java)
+            val keyGlyph = mock(KeyGlyphMap::class.java)
+            whenever(keyGlyph.getDrawableForModifierState(context, META_META_ON))
+                .thenReturn(metaDrawable)
+            whenever(kosmos.fakeInputManager.inputManager.getKeyGlyphMap(anyInt()))
+                .thenReturn(keyGlyph)
+            fakeSystemSource.setGroups(simpleGroup(simpleShortcutInfo(KEYCODE_1, META_META_ON)))
+            helper.toggle(deviceId = 123)
+
+            val categories by collectLastValue(repo.categories)
+            val systemCategory = categories?.firstOrNull { it.type == ShortcutCategoryType.System }
+
+            val expectedCategory =
+                ShortcutCategory(
+                    type = ShortcutCategoryType.System,
+                    simpleSubCategory(
+                        simpleResIdModifierShortcut("1", modifierResId = R.drawable.ic_ksh_key_meta)
+                    ),
+                )
+            assertThat(systemCategory).isEqualTo(expectedCategory)
+        }
+
     private fun simpleSubCategory(vararg shortcuts: Shortcut) =
         ShortcutSubCategory(simpleGroupLabel, shortcuts.asList())
 
@@ -196,6 +279,37 @@
             commands = listOf(ShortcutCommand(keys.map { ShortcutKey.Text(it) })),
         )
 
+    private fun simpleDrawableModifierShortcut(
+        vararg keys: String,
+        modifierDrawable: Drawable,
+    ): Shortcut {
+        val keyShortcuts = keys.map { ShortcutKey.Text(it) }
+        return Shortcut(
+            label = simpleShortcutLabel,
+            commands =
+                listOf(
+                    ShortcutCommand(
+                        listOf(ShortcutKey.Icon.DrawableIcon(drawable = modifierDrawable)) +
+                            keyShortcuts
+                    )
+                ),
+        )
+    }
+
+    private fun simpleResIdModifierShortcut(vararg keys: String, modifierResId: Int): Shortcut {
+        val keyShortcuts = keys.map { ShortcutKey.Text(it) }
+        return Shortcut(
+            label = simpleShortcutLabel,
+            commands =
+                listOf(
+                    ShortcutCommand(
+                        listOf(ShortcutKey.Icon.ResIdIcon(drawableResId = modifierResId)) +
+                            keyShortcuts
+                    )
+                ),
+        )
+    }
+
     private fun simpleGroup(vararg shortcuts: KeyboardShortcutInfo) =
         KeyboardShortcutGroup(simpleGroupLabel, shortcuts.asList())
 
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperCategoriesRepository.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperCategoriesRepository.kt
index a085887..12dd581 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperCategoriesRepository.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/data/repository/ShortcutHelperCategoriesRepository.kt
@@ -19,12 +19,14 @@
 import android.content.Context
 import android.graphics.drawable.Icon
 import android.hardware.input.InputManager
+import android.hardware.input.KeyGlyphMap
 import android.util.Log
 import android.view.InputDevice
 import android.view.KeyCharacterMap
 import android.view.KeyEvent
 import android.view.KeyboardShortcutGroup
 import android.view.KeyboardShortcutInfo
+import com.android.systemui.Flags.shortcutHelperKeyGlyph
 import com.android.systemui.dagger.SysUISingleton
 import com.android.systemui.dagger.qualifiers.Background
 import com.android.systemui.keyboard.shortcut.data.source.KeyboardShortcutGroupsSource
@@ -142,7 +144,10 @@
         return if (type == null) {
             null
         } else {
+            val keyGlyphMap =
+                if (shortcutHelperKeyGlyph()) inputManager.getKeyGlyphMap(inputDevice.id) else null
             toShortcutCategory(
+                keyGlyphMap,
                 inputDevice.keyCharacterMap,
                 type,
                 groups,
@@ -163,6 +168,7 @@
     }
 
     private fun toShortcutCategory(
+        keyGlyphMap: KeyGlyphMap?,
         keyCharacterMap: KeyCharacterMap,
         type: ShortcutCategoryType,
         shortcutGroups: List<KeyboardShortcutGroup>,
@@ -175,6 +181,7 @@
                     ShortcutSubCategory(
                         shortcutGroup.label.toString(),
                         toShortcuts(
+                            keyGlyphMap,
                             keyCharacterMap,
                             shortcutGroup.items,
                             keepIcons,
@@ -192,6 +199,7 @@
     }
 
     private fun toShortcuts(
+        keyGlyphMap: KeyGlyphMap?,
         keyCharacterMap: KeyCharacterMap,
         infoList: List<KeyboardShortcutInfo>,
         keepIcons: Boolean,
@@ -203,14 +211,16 @@
                 // keycode, or they could have a baseCharacter instead of a keycode.
                 it.keycode == KeyEvent.KEYCODE_UNKNOWN || supportedKeyCodes.contains(it.keycode)
             }
-            .mapNotNull { toShortcut(keyCharacterMap, it, keepIcons) }
+            .mapNotNull { toShortcut(keyGlyphMap, keyCharacterMap, it, keepIcons) }
 
     private fun toShortcut(
+        keyGlyphMap: KeyGlyphMap?,
         keyCharacterMap: KeyCharacterMap,
         shortcutInfo: KeyboardShortcutInfo,
         keepIcon: Boolean,
     ): Shortcut? {
-        val shortcutCommand = toShortcutCommand(keyCharacterMap, shortcutInfo) ?: return null
+        val shortcutCommand =
+            toShortcutCommand(keyGlyphMap, keyCharacterMap, shortcutInfo) ?: return null
         return Shortcut(
             label = shortcutInfo.label!!.toString(),
             icon = toShortcutIcon(keepIcon, shortcutInfo),
@@ -235,6 +245,7 @@
     }
 
     private fun toShortcutCommand(
+        keyGlyphMap: KeyGlyphMap?,
         keyCharacterMap: KeyCharacterMap,
         info: KeyboardShortcutInfo,
     ): ShortcutCommand? {
@@ -242,7 +253,7 @@
         var remainingModifiers = info.modifiers
         SUPPORTED_MODIFIERS.forEach { supportedModifier ->
             if ((supportedModifier and remainingModifiers) != 0) {
-                keys += toShortcutModifierKey(supportedModifier) ?: return null
+                keys += toShortcutModifierKey(keyGlyphMap, supportedModifier) ?: return null
                 // "Remove" the modifier from the remaining modifiers
                 remainingModifiers = remainingModifiers and supportedModifier.inv()
             }
@@ -253,7 +264,9 @@
             return null
         }
         if (info.keycode != 0 || info.baseCharacter > Char.MIN_VALUE) {
-            keys += toShortcutKey(keyCharacterMap, info.keycode, info.baseCharacter) ?: return null
+            keys +=
+                toShortcutKey(keyGlyphMap, keyCharacterMap, info.keycode, info.baseCharacter)
+                    ?: return null
         }
         if (keys.isEmpty()) {
             Log.wtf(TAG, "No keys for $info")
@@ -262,10 +275,15 @@
         return ShortcutCommand(keys)
     }
 
-    private fun toShortcutModifierKey(modifierMask: Int): ShortcutKey? {
+    private fun toShortcutModifierKey(keyGlyphMap: KeyGlyphMap?, modifierMask: Int): ShortcutKey? {
+        val modifierDrawable = keyGlyphMap?.getDrawableForModifierState(context, modifierMask)
+        if (modifierDrawable != null) {
+            return ShortcutKey.Icon.DrawableIcon(drawable = modifierDrawable)
+        }
+
         val iconResId = ShortcutHelperKeys.keyIcons[modifierMask]
         if (iconResId != null) {
-            return ShortcutKey.Icon(iconResId)
+            return ShortcutKey.Icon.ResIdIcon(iconResId)
         }
 
         val modifierLabel = ShortcutHelperKeys.modifierLabels[modifierMask]
@@ -277,13 +295,19 @@
     }
 
     private fun toShortcutKey(
+        keyGlyphMap: KeyGlyphMap?,
         keyCharacterMap: KeyCharacterMap,
         keyCode: Int,
         baseCharacter: Char = Char.MIN_VALUE,
     ): ShortcutKey? {
+        val keycodeDrawable = keyGlyphMap?.getDrawableForKeycode(context, keyCode)
+        if (keycodeDrawable != null) {
+            return ShortcutKey.Icon.DrawableIcon(drawable = keycodeDrawable)
+        }
+
         val iconResId = ShortcutHelperKeys.keyIcons[keyCode]
         if (iconResId != null) {
-            return ShortcutKey.Icon(iconResId)
+            return ShortcutKey.Icon.ResIdIcon(iconResId)
         }
         if (baseCharacter > Char.MIN_VALUE) {
             return ShortcutKey.Text(baseCharacter.uppercase())
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/shared/model/ShortcutCommand.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/shared/model/ShortcutCommand.kt
index e5b8096..28451ae 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/shared/model/ShortcutCommand.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/shared/model/ShortcutCommand.kt
@@ -28,7 +28,7 @@
     }
 
     fun key(@DrawableRes drawableResId: Int) {
-        keys += ShortcutKey.Icon(drawableResId)
+        keys += ShortcutKey.Icon.ResIdIcon(drawableResId)
     }
 
     fun build() = ShortcutCommand(keys)
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/shared/model/ShortcutKey.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/shared/model/ShortcutKey.kt
index 1abb78c..1a609ea 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/shared/model/ShortcutKey.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/shared/model/ShortcutKey.kt
@@ -16,10 +16,15 @@
 
 package com.android.systemui.keyboard.shortcut.shared.model
 
+import android.graphics.drawable.Drawable
 import androidx.annotation.DrawableRes
 
 sealed interface ShortcutKey {
     data class Text(val value: String) : ShortcutKey
 
-    data class Icon(@DrawableRes val drawableResId: Int) : ShortcutKey
+    sealed interface Icon : ShortcutKey {
+        data class ResIdIcon(@DrawableRes val drawableResId: Int) : Icon
+
+        data class DrawableIcon(val drawable: Drawable) : Icon
+    }
 }
diff --git a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt
index d537056..abddc70 100644
--- a/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt
+++ b/packages/SystemUI/src/com/android/systemui/keyboard/shortcut/ui/composable/ShortcutHelper.kt
@@ -647,7 +647,11 @@
 @Composable
 private fun BoxScope.ShortcutIconKey(key: ShortcutKey.Icon) {
     Icon(
-        painter = painterResource(key.drawableResId),
+        painter =
+            when (key) {
+                is ShortcutKey.Icon.ResIdIcon -> painterResource(key.drawableResId)
+                is ShortcutKey.Icon.DrawableIcon -> rememberDrawablePainter(drawable = key.drawable)
+            },
         contentDescription = null,
         modifier = Modifier.align(Alignment.Center).padding(6.dp),
     )