Optimize SystemFeaturesMetadata index lookup

Using a ~150 length string->int switch statement lookup turns out to
be both slower than an equivalent ArraySet lookup, and larger in terms
of code size (both dex + compiled). Switch to using a static ArraySet
(initialized in the zygote) for efficiency.

A follow-up change will add a microbenchmark for further validation.

Bug: 203143243
Test: m + presubmit + compare dex/odex/runtime perf
Flag: EXEMPT refactor
Change-Id: I1231149b928488cb716ec09a33ef1ec4ac7e2785

2 files changed, 90 insertions, 24 deletions

diff --git a/tools/systemfeatures/src/com/android/systemfeatures/SystemFeaturesMetadataProcessor.kt b/tools/systemfeatures/src/com/android/systemfeatures/SystemFeaturesMetadataProcessor.kt
index 4a6d4b1f49ef..c51c6d661314 100644
--- a/tools/systemfeatures/src/com/android/systemfeatures/SystemFeaturesMetadataProcessor.kt
+++ b/tools/systemfeatures/src/com/android/systemfeatures/SystemFeaturesMetadataProcessor.kt
@@ -18,9 +18,11 @@ package com.android.systemfeatures
 
 import android.annotation.SdkConstant
 import com.squareup.javapoet.ClassName
+import com.squareup.javapoet.CodeBlock
 import com.squareup.javapoet.FieldSpec
 import com.squareup.javapoet.JavaFile
 import com.squareup.javapoet.MethodSpec
+import com.squareup.javapoet.ParameterizedTypeName
 import com.squareup.javapoet.TypeSpec
 import java.io.IOException
 import javax.annotation.processing.AbstractProcessor
@@ -101,8 +103,8 @@ class SystemFeaturesMetadataProcessor : AbstractProcessor() {
             TypeSpec.classBuilder("SystemFeaturesMetadata")
                 .addModifiers(Modifier.PUBLIC, Modifier.FINAL)
                 .addJavadoc("@hide")
-                .addField(buildFeatureCount(featureVarNames))
-                .addMethod(buildFeatureIndexLookup(featureVarNames))
+                .addFeatureCount(featureVarNames)
+                .addFeatureIndexLookup(featureVarNames)
                 .build()
 
         try {
@@ -120,19 +122,55 @@ class SystemFeaturesMetadataProcessor : AbstractProcessor() {
         return true
     }
 
-    private fun buildFeatureCount(featureVarNames: Collection<String>): FieldSpec {
-        return FieldSpec.builder(Int::class.java, "SDK_FEATURE_COUNT")
-            .addModifiers(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL)
-            .addJavadoc(
-                "# of {@link android.annotation.SdkConstant}` features defined in PackageManager."
-            )
-            .addJavadoc("\n\n@hide")
-            .initializer("\$L", featureVarNames.size)
-            .build()
+    private fun TypeSpec.Builder.addFeatureCount(
+        featureVarNames: Collection<String>
+    ): TypeSpec.Builder {
+        return addField(
+            FieldSpec.builder(Int::class.java, "SDK_FEATURE_COUNT")
+                .addModifiers(Modifier.PUBLIC, Modifier.STATIC, Modifier.FINAL)
+                .addJavadoc(
+                    "# of {@link android.annotation.SdkConstant}` features in PackageManager."
+                )
+                .addJavadoc("\n\n@hide")
+                .initializer("\$L", featureVarNames.size)
+                .build()
+        )
     }
 
-    private fun buildFeatureIndexLookup(featureVarNames: Collection<String>): MethodSpec {
-        val methodBuilder =
+    private fun TypeSpec.Builder.addFeatureIndexLookup(
+        featureVarNames: Collection<String>
+    ): TypeSpec.Builder {
+        // NOTE: This was initially implemented in terms of a single, long switch() statement.
+        // However, this resulted in:
+        //   1) relatively large compiled code size for the lookup method (~20KB)
+        //   2) worse runtime lookup performance than a simple ArraySet
+        // The ArraySet approach adds just ~1KB to the code/image and is 2x faster at runtime.
+
+        // Provide the initial capacity of the ArraySet for efficiency.
+        addField(
+            FieldSpec.builder(
+                    ParameterizedTypeName.get(ARRAYSET_CLASS, ClassName.get(String::class.java)),
+                    "sFeatures",
+                )
+                .addModifiers(Modifier.PRIVATE, Modifier.STATIC, Modifier.FINAL)
+                .initializer("new ArraySet<>(\$L)", featureVarNames.size)
+                .build()
+        )
+
+        // Use a temp array + Collections.addAll() to minimizes the generated code size.
+        addStaticBlock(
+            CodeBlock.builder()
+                .add("final \$T[] features = {\n", String::class.java)
+                .indent()
+                .apply { featureVarNames.forEach { add("\$T.\$N,\n", PACKAGEMANAGER_CLASS, it) } }
+                .unindent()
+                .addStatement("}")
+                .addStatement("\$T.addAll(sFeatures, features)", COLLECTIONS_CLASS)
+                .build()
+        )
+
+        // Use ArraySet.indexOf to provide the implicit feature index mapping.
+        return addMethod(
             MethodSpec.methodBuilder("maybeGetSdkFeatureIndex")
                 .addModifiers(Modifier.PUBLIC, Modifier.STATIC)
                 .addJavadoc("@return an index in [0, SDK_FEATURE_COUNT) for features defined ")
@@ -140,21 +178,15 @@ class SystemFeaturesMetadataProcessor : AbstractProcessor() {
                 .addJavadoc("\n\n@hide")
                 .returns(Int::class.java)
                 .addParameter(String::class.java, "featureName")
-        methodBuilder.beginControlFlow("switch (featureName)")
-        featureVarNames.forEachIndexed { index, featureVarName ->
-            methodBuilder
-                .addCode("case \$T.\$N: ", PACKAGEMANAGER_CLASS, featureVarName)
-                .addStatement("return \$L", index)
-        }
-        methodBuilder
-            .addCode("default: ")
-            .addStatement("return -1")
-            .endControlFlow()
-        return methodBuilder.build()
+                .addStatement("return sFeatures.indexOf(featureName)")
+                .build()
+        )
     }
 
     companion object {
         private val SDK_CONSTANT_ANNOTATION_NAME = SdkConstant::class.qualifiedName
         private val PACKAGEMANAGER_CLASS = ClassName.get("android.content.pm", "PackageManager")
+        private val ARRAYSET_CLASS = ClassName.get("android.util", "ArraySet")
+        private val COLLECTIONS_CLASS = ClassName.get("java.util", "Collections")
     }
 }
diff --git a/tools/systemfeatures/tests/src/ArraySet.java b/tools/systemfeatures/tests/src/ArraySet.java
new file mode 100644
index 000000000000..0eb8f298bd89
--- /dev/null
+++ b/tools/systemfeatures/tests/src/ArraySet.java
@@ -0,0 +1,34 @@
+/*
+ * 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 android.util;
+
+import java.util.ArrayList;
+
+/** Stub for testing, we extend ArrayList to get indexOf() for free. */
+public final class ArraySet<K> extends ArrayList<K> {
+    public ArraySet(int capacity) {
+        super(capacity);
+    }
+
+    @Override
+    public boolean add(K k) {
+        if (!contains(k)) {
+            return super.add(k);
+        }
+        return false;
+    }
+}