Snap for 8857176 from d7334228c4c4f9b0b9037fac934e023432509d8a to mainline-go-neuralnetworks-releaseaml_go_neu_330912000android13-mainline-go-neuralnetworks-release

Change-Id: Iddb286453648c9566dbf1e3c9fb947b7606b4c73

18 files changed, 523 insertions, 23 deletions

diff --git a/libraries/collectors-helper/adservices/Android.bp b/libraries/collectors-helper/adservices/Android.bp
new file mode 100644
index 000000000..35253277b
--- /dev/null
+++ b/libraries/collectors-helper/adservices/Android.bp
@@ -0,0 +1,34 @@
+// Copyright (C) 2022 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.
+
+// Used for collecting Lyric specific metrics..
+package {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+java_library {
+    name: "adservices-collector-helper",
+    defaults: ["tradefed_errorprone_defaults"],
+
+    srcs: [
+        "src/**/*.java",
+    ],
+    static_libs: [
+        "androidx.test.runner",
+        "collector-helper-utilities",
+        "guava",
+    ],
+
+    sdk_version: "current",
+}
diff --git a/libraries/collectors-helper/adservices/src/com/android/helpers/JSScriptEngineLatencyHelper.java b/libraries/collectors-helper/adservices/src/com/android/helpers/JSScriptEngineLatencyHelper.java
new file mode 100644
index 000000000..0aba67493
--- /dev/null
+++ b/libraries/collectors-helper/adservices/src/com/android/helpers/JSScriptEngineLatencyHelper.java
@@ -0,0 +1,161 @@
+/*
+ * Copyright (C) 2022 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.helpers;
+
+import android.util.Log;
+
+import com.google.common.annotations.VisibleForTesting;
+import com.google.common.collect.ImmutableMap;
+
+import java.io.BufferedReader;
+import java.io.IOException;
+import java.io.InputStream;
+import java.io.InputStreamReader;
+import java.time.Clock;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.time.format.DateTimeFormatter;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Collections;
+import java.util.List;
+import java.util.Map;
+import java.util.function.Supplier;
+import java.util.regex.Matcher;
+import java.util.regex.Pattern;
+
+public class JSScriptEngineLatencyHelper implements ICollectorHelper<Long> {
+    @VisibleForTesting public static final String SANDBOX_INIT_TIME_AVG = "SANDBOX_INIT_TIME_AVG";
+    public static final String ISOLATE_CREATE_TIME_AVG = "ISOLATE_CREATE_TIME_AVG";
+    public static final String WEBVIEW_EXECUTION_TIME_AVG = "WEBVIEW_EXECUTION_TIME_AVG";
+    public static final String JAVA_EXECUTION_TIME_AVG = "JAVA_EXECUTION_TIME_AVG";
+    public static final String NUM_ITERATIONS = "NUM_ITERATIONS";
+
+    private static final String TAG = JSScriptEngineLatencyHelper.class.getSimpleName();
+    private static final DateTimeFormatter LOG_TIME_FORMATTER =
+            DateTimeFormatter.ofPattern("MM-dd HH:mm:ss.SSS").withZone(ZoneId.systemDefault());
+
+    private static final Pattern sLatencyMetricPattern =
+            Pattern.compile("JSScriptEngine: \\((.*): (\\d+)\\)");
+    private static final String SANDBOX_INIT_TIME = "SANDBOX_INIT_TIME";
+    private static final String ISOLATE_CREATE_TIME = "ISOLATE_CREATE_TIME";
+    private static final String WEBVIEW_EXECUTION_TIME = "WEBVIEW_EXECUTION_TIME";
+    private static final String JAVA_EXECUTION_TIME = "JAVA_EXECUTION_TIME";
+
+    private Instant instant;
+    private final Clock mClock;
+    private final Supplier<MetricsEventStreamReader> metricsEventStreamSupplier;
+
+    public JSScriptEngineLatencyHelper() {
+        mClock = Clock.systemUTC();
+        metricsEventStreamSupplier = () -> new MetricsEventStreamReader();
+    }
+
+    @VisibleForTesting
+    public JSScriptEngineLatencyHelper(
+            Supplier<MetricsEventStreamReader> metricsEventStreamSupplier, Clock clock) {
+        this.metricsEventStreamSupplier = metricsEventStreamSupplier;
+        mClock = clock;
+    }
+
+    @Override
+    public boolean startCollecting() {
+        instant = mClock.instant();
+        return true;
+    }
+
+    @Override
+    public Map<String, Long> getMetrics() {
+        try {
+            return processOutput(metricsEventStreamSupplier.get().getMetricsEvents(instant));
+        } catch (Exception e) {
+            Log.e(TAG, "Failed to collect JSScriptEngine metrics.", e);
+        }
+
+        return Collections.emptyMap();
+    }
+
+    private Map<String, Long> processOutput(InputStream inputStream) throws IOException {
+        BufferedReader bufferedReader = new BufferedReader(new InputStreamReader(inputStream));
+
+        List<Long> sandboxInitTimes = new ArrayList<>();
+        List<Long> isolateCreateTimes = new ArrayList<>();
+        List<Long> javaProcessTimes = new ArrayList<>();
+        List<Long> webviewProcessTimes = new ArrayList<>();
+
+        String line = "";
+        while ((line = bufferedReader.readLine()) != null) {
+            Matcher matcher = sLatencyMetricPattern.matcher(line);
+            while (matcher.find()) {
+                /**
+                 * The lines from Logcat will look like: 06-13 18:09:24.058 20765 20781 D
+                 * JSScriptEngine: (JAVA_PROCESS_TIME: 43)
+                 */
+                String metric = matcher.group(1);
+                long latency = Long.parseLong(matcher.group(2));
+                if (SANDBOX_INIT_TIME.equals(metric)) {
+                    sandboxInitTimes.add(latency);
+                } else if (ISOLATE_CREATE_TIME.equals(metric)) {
+                    isolateCreateTimes.add(latency);
+                } else if (JAVA_EXECUTION_TIME.equals(metric)) {
+                    javaProcessTimes.add(latency);
+                } else if (WEBVIEW_EXECUTION_TIME.equals(metric)) {
+                    webviewProcessTimes.add(latency);
+                }
+            }
+        }
+
+        // Just getting average for now.
+        int defaultMetricVal = 0;
+        long sandboxInitTimeAvg = getAverage(sandboxInitTimes, defaultMetricVal);
+        long isolateCreateTimeAvg = getAverage(isolateCreateTimes, defaultMetricVal);
+        long webviewProcessTimeAvg = getAverage(webviewProcessTimes, defaultMetricVal);
+        long javaProcessTimeAvg = getAverage(javaProcessTimes, defaultMetricVal);
+
+        return ImmutableMap.of(
+                SANDBOX_INIT_TIME_AVG, sandboxInitTimeAvg,
+                ISOLATE_CREATE_TIME_AVG, isolateCreateTimeAvg,
+                WEBVIEW_EXECUTION_TIME_AVG, webviewProcessTimeAvg,
+                JAVA_EXECUTION_TIME_AVG, javaProcessTimeAvg,
+                NUM_ITERATIONS, (long) javaProcessTimes.size());
+    }
+
+    private Long getAverage(List<Long> list, long defaultValue) {
+        return (long) list.stream().mapToDouble(d -> d).average().orElse(defaultValue);
+    }
+
+    @Override
+    public boolean stopCollecting() {
+        return true;
+    }
+
+    @VisibleForTesting
+    public static class MetricsEventStreamReader {
+        /** Return JSScriptEngine logs that will be used to build the test metrics. */
+        public InputStream getMetricsEvents(Instant startTime) throws IOException {
+            ProcessBuilder pb =
+                    new ProcessBuilder(
+                            Arrays.asList(
+                                    "logcat",
+                                    "-s",
+                                    "JSScriptEngine:D",
+                                    "-t",
+                                    LOG_TIME_FORMATTER.format(startTime)));
+            return pb.start().getInputStream();
+        }
+    }
+}
diff --git a/libraries/collectors-helper/adservices/test/Android.bp b/libraries/collectors-helper/adservices/test/Android.bp
new file mode 100644
index 000000000..cc0e0e397
--- /dev/null
+++ b/libraries/collectors-helper/adservices/test/Android.bp
@@ -0,0 +1,35 @@
+// Copyright (C) 2022 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 {
+    default_applicable_licenses: ["Android-Apache-2.0"],
+}
+
+java_library {
+    name: "adservices-helper-test",
+    defaults: ["tradefed_errorprone_defaults"],
+
+    srcs: ["src/**/*.java"],
+
+    static_libs: [
+        "adservices-collector-helper",
+        "androidx.test.runner",
+        "app-collector-helper",
+        "junit",
+        "mockito-target",
+        "truth-prebuilt",
+    ],
+
+    sdk_version: "current",
+}
diff --git a/libraries/collectors-helper/adservices/test/src/com/android/helpers/JSScriptEngineLatencyHelperTest.java b/libraries/collectors-helper/adservices/test/src/com/android/helpers/JSScriptEngineLatencyHelperTest.java
new file mode 100644
index 000000000..bb952f3ea
--- /dev/null
+++ b/libraries/collectors-helper/adservices/test/src/com/android/helpers/JSScriptEngineLatencyHelperTest.java
@@ -0,0 +1,101 @@
+/*
+ * Copyright (C) 2022 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.helpers;
+
+import static com.google.common.truth.Truth.assertThat;
+
+import static org.mockito.Mockito.when;
+
+import org.junit.Before;
+import org.junit.Test;
+import org.mockito.Mock;
+import org.mockito.MockitoAnnotations;
+
+import java.io.ByteArrayInputStream;
+import java.io.IOException;
+import java.time.Clock;
+import java.time.Instant;
+import java.time.ZoneId;
+import java.util.Map;
+
+/**
+ * To run, use atest CollectorsHelperAospTest:com.android.helpers.JSScriptEngineLatencyHelperTest
+ */
+public class JSScriptEngineLatencyHelperTest {
+    @Mock private JSScriptEngineLatencyHelper.MetricsEventStreamReader mMetricsEventStreamReader;
+
+    @Mock private Clock mClock;
+
+    private JSScriptEngineLatencyHelper mJSScriptEngineLatencyHelper;
+    private Instant mInstant = Clock.systemUTC().instant();
+
+    @Before
+    public void setUp() {
+        MockitoAnnotations.initMocks(this);
+        when(mClock.getZone()).thenReturn(ZoneId.of("UTC")); // Used by DateTimeFormatter.
+        when(mClock.instant()).thenReturn(mInstant);
+        mJSScriptEngineLatencyHelper =
+                new JSScriptEngineLatencyHelper(() -> mMetricsEventStreamReader, mClock);
+        mJSScriptEngineLatencyHelper.startCollecting();
+    }
+
+    @Test
+    public void testInitTimeInLogcat() throws IOException {
+        String logcatOutput =
+                "06-13 18:09:24.022 20765 D JSScriptEngine: (SANDBOX_INIT_TIME: 102)\n"
+                    + "06-29 02:47:32.030 31075 D JSScriptEngine: (ISOLATE_CREATE_TIME: 1)\n"
+                    + "06-13 18:09:24.058 20765 D JSScriptEngine: (JAVA_EXECUTION_TIME: 43)\n"
+                    + "06-13 18:09:24.058 31075 D JSScriptEngine: (WEBVIEW_EXECUTION_TIME: 23)\n"
+                    + "06-13 18:09:24.129 20765 D JSScriptEngine: (SANDBOX_INIT_TIME: 45)\n"
+                    + "06-13 18:09:24.130 31075 D JSScriptEngine: (ISOLATE_CREATE_TIME: 1)\n"
+                    + "06-13 18:09:24.152 20765 D JSScriptEngine: (JAVA_EXECUTION_TIME: 19)\n"
+                    + "06-29 02:47:31.824 31075 D JSScriptEngine: (WEBVIEW_EXECUTION_TIME: 29)";
+
+        when(mMetricsEventStreamReader.getMetricsEvents(mInstant))
+                .thenReturn(new ByteArrayInputStream(logcatOutput.getBytes()));
+        Map<String, Long> actual = mJSScriptEngineLatencyHelper.getMetrics();
+
+        assertThat(actual.get(JSScriptEngineLatencyHelper.SANDBOX_INIT_TIME_AVG))
+                .isEqualTo((102 + 45) / 2);
+        assertThat(actual.get(JSScriptEngineLatencyHelper.ISOLATE_CREATE_TIME_AVG))
+                .isEqualTo((1 + 1) / 2);
+        assertThat(actual.get(JSScriptEngineLatencyHelper.JAVA_EXECUTION_TIME_AVG))
+                .isEqualTo((43 + 19) / 2);
+        assertThat(actual.get(JSScriptEngineLatencyHelper.WEBVIEW_EXECUTION_TIME_AVG))
+                .isEqualTo((23 + 29) / 2);
+        assertThat(actual.get(JSScriptEngineLatencyHelper.NUM_ITERATIONS)).isEqualTo(2);
+    }
+
+    @Test
+    public void testWithEmptyLogcat() throws IOException {
+        when(mMetricsEventStreamReader.getMetricsEvents(mInstant))
+                .thenReturn(new ByteArrayInputStream("".getBytes()));
+        Map<String, Long> actual = mJSScriptEngineLatencyHelper.getMetrics();
+        for (Long val : actual.values()) {
+            assertThat(val).isEqualTo(0);
+        }
+    }
+
+    @Test
+    public void testInputStreamThrowsException() throws IOException {
+        when(mMetricsEventStreamReader.getMetricsEvents(mInstant)).thenThrow(new IOException());
+        Map<String, Long> actual = mJSScriptEngineLatencyHelper.getMetrics();
+        for (Long val : actual.values()) {
+            assertThat(val).isEqualTo(0);
+        }
+    }
+}
diff --git a/libraries/collectors-helper/tests/Android.bp b/libraries/collectors-helper/tests/Android.bp
index 502e9fcf9..88412c193 100644
--- a/libraries/collectors-helper/tests/Android.bp
+++ b/libraries/collectors-helper/tests/Android.bp
@@ -21,10 +21,11 @@ android_test {
     defaults: ["tradefed_errorprone_defaults"],
 
     static_libs: [
-        "perfetto-helper-test",
+        "adservices-helper-test",
         "app-collector-helper-test",
         "jank-helper-test",
         "memory-helper-test",
+        "perfetto-helper-test",
         "system-helper-test",
     ],
 
diff --git a/libraries/device-collectors/src/main/Android.bp b/libraries/device-collectors/src/main/Android.bp
index 98f9ec560..68023fa62 100644
--- a/libraries/device-collectors/src/main/Android.bp
+++ b/libraries/device-collectors/src/main/Android.bp
@@ -23,6 +23,7 @@ java_library {
     srcs: ["java/**/*.java"],
 
     static_libs: [
+        "adservices-collector-helper",
         "androidx.annotation_annotation",
         "androidx.test.runner",
         "androidx.test.uiautomator",
@@ -32,10 +33,10 @@ java_library {
         "lyric-metric-helper",
         "memory-helper",
         "perfetto-helper",
+        "platform-test-annotations",
         "power-helper",
         "simpleperf-helper",
         "system-metric-helper",
-        "platform-test-annotations",
     ],
 
     sdk_version: "test_current",
diff --git a/libraries/device-collectors/src/main/java/android/device/collectors/JSScriptEngineListener.java b/libraries/device-collectors/src/main/java/android/device/collectors/JSScriptEngineListener.java
new file mode 100644
index 000000000..10f5b40c7
--- /dev/null
+++ b/libraries/device-collectors/src/main/java/android/device/collectors/JSScriptEngineListener.java
@@ -0,0 +1,28 @@
+/*
+ * Copyright (C) 2022 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.device.collectors;
+
+import android.device.collectors.annotations.OptionClass;
+
+import com.android.helpers.JSScriptEngineLatencyHelper;
+
+@OptionClass(alias = "jsscriptengine-latency-listener")
+public class JSScriptEngineListener extends BaseCollectionListener<Integer> {
+    public JSScriptEngineListener() {
+        createHelperInstance(new JSScriptEngineLatencyHelper());
+    }
+}
diff --git a/libraries/screenshot/src/androidTest/java/platform/test/screenshot/PixelPerfectMatcherTest.kt b/libraries/screenshot/src/androidTest/java/platform/test/screenshot/PixelPerfectMatcherTest.kt
index bbab1bb4e..498937d1b 100644
--- a/libraries/screenshot/src/androidTest/java/platform/test/screenshot/PixelPerfectMatcherTest.kt
+++ b/libraries/screenshot/src/androidTest/java/platform/test/screenshot/PixelPerfectMatcherTest.kt
@@ -20,10 +20,10 @@ import android.graphics.Rect
 import androidx.test.ext.junit.runners.AndroidJUnit4
 import androidx.test.filters.MediumTest
 import com.google.common.truth.Truth.assertThat
-import platform.test.screenshot.matchers.PixelPerfectMatcher
-import platform.test.screenshot.utils.loadBitmap
 import org.junit.Test
 import org.junit.runner.RunWith
+import platform.test.screenshot.matchers.PixelPerfectMatcher
+import platform.test.screenshot.utils.loadBitmap
 
 @RunWith(AndroidJUnit4::class)
 @MediumTest
@@ -66,9 +66,33 @@ class PixelPerfectMatcherTest {
         val result = matcher.compareBitmaps(
             first.toIntArray(), second.toIntArray(),
             first.width, first.height,
-            arrayOf(Rect(/* left= */1, /* top= */1, /* right= */4, /* bottom= */4))
+            listOf(Rect(/* left= */1, /* top= */1, /* right= */4, /* bottom= */4))
         )
 
         assertThat(result.matches).isTrue()
     }
+
+    @Test
+    fun performDiff_sameSize_partialCompare_checkDiffImage() {
+        val first = loadBitmap("qmc-folder1")
+        val second = loadBitmap("qmc-folder2")
+        val matcher = PixelPerfectMatcher()
+        val interestingRegion = Rect(/* left= */10, /* top= */15, /* right= */70, /* bottom= */50)
+        val result = matcher.compareBitmaps(
+            first.toIntArray(), second.toIntArray(),
+            first.width, first.height, listOf(interestingRegion)
+        )
+        val diffImage = result.diff!!.toIntArray()
+
+        assertThat(result.matches).isFalse()
+        for (i in 0..first.height - 1) {
+            for (j in 0..first.width - 1) {
+                val rowInRange = i >= interestingRegion.top && i <= interestingRegion.bottom
+                val colInRange = j >= interestingRegion.left && j <= interestingRegion.right
+                if (!(rowInRange && colInRange)) {
+                    assertThat(diffImage[i * first.width + j] == 0).isTrue()
+                }
+            }
+        }
+    }
 }
diff --git a/libraries/screenshot/src/androidTest/java/platform/test/screenshot/ScreenshotTestRuleTest.kt b/libraries/screenshot/src/androidTest/java/platform/test/screenshot/ScreenshotTestRuleTest.kt
index 0c99f817b..9aa9f8764 100644
--- a/libraries/screenshot/src/androidTest/java/platform/test/screenshot/ScreenshotTestRuleTest.kt
+++ b/libraries/screenshot/src/androidTest/java/platform/test/screenshot/ScreenshotTestRuleTest.kt
@@ -25,6 +25,7 @@ import com.google.common.truth.Truth.assertThat
 import java.io.File
 import java.io.FileInputStream
 import java.lang.AssertionError
+import java.util.ArrayList
 import org.junit.After
 import org.junit.Rule
 import org.junit.Test
@@ -107,10 +108,12 @@ class ScreenshotTestRuleTest {
     @Test
     fun performDiff_noPixelCompared() {
         val first = loadBitmap("round_rect_gray")
+        val regions = ArrayList<Rect>()
+        regions.add(Rect(/* left= */1, /* top= */1, /* right= */2, /* bottom=*/2))
 
         first.assertAgainstGolden(
             rule, "round_rect_green", matcher = MSSIMMatcher(),
-            regions = arrayOf(Rect(/* left= */1, /* top= */1, /* right= */2, /* bottom=*/2))
+            regions = regions
         )
 
         val resultProto = rule.getPathOnDeviceFor(RESULT_PROTO)
@@ -129,10 +132,12 @@ class ScreenshotTestRuleTest {
         val startWidth = 10 * first.width / 20
         val endWidth = 11 * first.width / 20
         val matcher = MSSIMMatcher()
+        val regions = ArrayList<Rect>()
+        regions.add(Rect(startWidth, startHeight, endWidth, endHeight))
+        regions.add(Rect())
 
         first.assertAgainstGolden(
-            rule, "qmc-folder2", matcher,
-            arrayOf(Rect(startWidth, startHeight, endWidth, endHeight))
+            rule, "qmc-folder2", matcher, regions
         )
 
         val resultProto = rule.getPathOnDeviceFor(RESULT_PROTO)
diff --git a/libraries/screenshot/src/main/java/platform/test/screenshot/ScreenshotTestRule.kt b/libraries/screenshot/src/main/java/platform/test/screenshot/ScreenshotTestRule.kt
index bc7295039..7532695fd 100644
--- a/libraries/screenshot/src/main/java/platform/test/screenshot/ScreenshotTestRule.kt
+++ b/libraries/screenshot/src/main/java/platform/test/screenshot/ScreenshotTestRule.kt
@@ -127,7 +127,7 @@ open class ScreenshotTestRule(
             actual = actual,
             goldenIdentifier = goldenIdentifier,
             matcher = matcher,
-            regions = null
+            regions = emptyList<Rect>()
         )
     }
 
@@ -154,7 +154,7 @@ open class ScreenshotTestRule(
         actual: Bitmap,
         goldenIdentifier: String,
         matcher: BitmapMatcher,
-        regions: Array<Rect>?
+        regions: List<Rect>
     ) {
         if (!goldenIdentifier.matches("^[A-Za-z0-9_-]+$".toRegex())) {
             throw IllegalArgumentException(
@@ -213,7 +213,7 @@ open class ScreenshotTestRule(
             actual = actual,
             comparisonStatistics = comparisonResult.comparisonStatistics,
             expected = highlightedBitmap(expected, regions),
-            diff = highlightedBitmap(comparisonResult.diff, regions)
+            diff = comparisonResult.diff
         )
 
         if (!comparisonResult.matches) {
@@ -349,8 +349,8 @@ open class ScreenshotTestRule(
         }
     }
 
-    private fun highlightedBitmap(original: Bitmap?, regions: Array<Rect>?): Bitmap? {
-        if (original == null || regions == null) {
+    private fun highlightedBitmap(original: Bitmap?, regions: List<Rect>): Bitmap? {
+        if (original == null || regions.isEmpty()) {
             return original
         }
         val bitmapArray = original.toIntArray()
@@ -428,7 +428,7 @@ fun Bitmap.assertAgainstGolden(
     rule: ScreenshotTestRule,
     goldenIdentifier: String,
     matcher: BitmapMatcher = MSSIMMatcher(),
-    regions: Array<Rect>? = null
+    regions: List<Rect> = emptyList<Rect>()
 ) {
     rule.assertBitmapAgainstGolden(this, goldenIdentifier, matcher = matcher, regions = regions)
 }
diff --git a/libraries/screenshot/src/main/java/platform/test/screenshot/matchers/BitmapMatcher.kt b/libraries/screenshot/src/main/java/platform/test/screenshot/matchers/BitmapMatcher.kt
index 8b4984e6c..077265ce9 100644
--- a/libraries/screenshot/src/main/java/platform/test/screenshot/matchers/BitmapMatcher.kt
+++ b/libraries/screenshot/src/main/java/platform/test/screenshot/matchers/BitmapMatcher.kt
@@ -18,6 +18,7 @@ package platform.test.screenshot.matchers
 
 import android.graphics.Bitmap
 import android.graphics.Rect
+import kotlin.collections.List
 import platform.test.screenshot.proto.ScreenshotResultProto.DiffResult.ComparisonStatistics
 
 /**
@@ -40,11 +41,11 @@ abstract class BitmapMatcher {
         given: IntArray,
         width: Int,
         height: Int,
-        regions: Array<Rect>? = null
+        regions: List<Rect> = emptyList<Rect>()
     ): MatchResult
 
-    protected fun getFilter(width: Int, height: Int, regions: Array<Rect>?): IntArray {
-        if (regions == null) { return IntArray(width * height) { 1 } }
+    protected fun getFilter(width: Int, height: Int, regions: List<Rect>): IntArray {
+        if (regions.isEmpty()) { return IntArray(width * height) { 1 } }
         val bitmapArray = IntArray(width * height) { 0 }
         for (region in regions) {
             for (i in region.top..region.bottom) {
diff --git a/libraries/screenshot/src/main/java/platform/test/screenshot/matchers/MSSIMMatcher.kt b/libraries/screenshot/src/main/java/platform/test/screenshot/matchers/MSSIMMatcher.kt
index a31f72388..c6103d398 100644
--- a/libraries/screenshot/src/main/java/platform/test/screenshot/matchers/MSSIMMatcher.kt
+++ b/libraries/screenshot/src/main/java/platform/test/screenshot/matchers/MSSIMMatcher.kt
@@ -19,6 +19,7 @@ package platform.test.screenshot.matchers
 import android.graphics.Color
 import android.graphics.Rect
 import androidx.annotation.FloatRange
+import kotlin.collections.List
 import kotlin.math.pow
 import platform.test.screenshot.proto.ScreenshotResultProto
 
@@ -46,7 +47,7 @@ class MSSIMMatcher(
         given: IntArray,
         width: Int,
         height: Int,
-        regions: Array<Rect>?
+        regions: List<Rect>
     ): MatchResult {
         val filter = getFilter(width, height, regions)
         val calSSIMResult = calculateSSIM(expected, given, width, height, filter)
@@ -71,7 +72,7 @@ class MSSIMMatcher(
 
         // Create diff
         val result = PixelPerfectMatcher()
-            .compareBitmaps(expected, given, width, height, null)
+            .compareBitmaps(expected, given, width, height, regions)
         return MatchResult(
             matches = false,
             diff = result.diff,
diff --git a/libraries/screenshot/src/main/java/platform/test/screenshot/matchers/PixelPerfectMatcher.kt b/libraries/screenshot/src/main/java/platform/test/screenshot/matchers/PixelPerfectMatcher.kt
index cd839a89f..1a82e4ba4 100644
--- a/libraries/screenshot/src/main/java/platform/test/screenshot/matchers/PixelPerfectMatcher.kt
+++ b/libraries/screenshot/src/main/java/platform/test/screenshot/matchers/PixelPerfectMatcher.kt
@@ -19,6 +19,7 @@ package platform.test.screenshot.matchers
 import android.graphics.Bitmap
 import android.graphics.Color
 import android.graphics.Rect
+import kotlin.collections.List
 import platform.test.screenshot.proto.ScreenshotResultProto
 
 /**
@@ -31,7 +32,7 @@ class PixelPerfectMatcher : BitmapMatcher() {
         given: IntArray,
         width: Int,
         height: Int,
-        regions: Array<Rect>?
+        regions: List<Rect>
     ): MatchResult {
         check(expected.size == given.size)
 
diff --git a/libraries/sts-common-util/device-side/src/com/android/sts/common/util/StsExtraBusinessLogicTestCase.java b/libraries/sts-common-util/device-side/src/com/android/sts/common/util/StsExtraBusinessLogicTestCase.java
index f9aa29462..83ca09662 100644
--- a/libraries/sts-common-util/device-side/src/com/android/sts/common/util/StsExtraBusinessLogicTestCase.java
+++ b/libraries/sts-common-util/device-side/src/com/android/sts/common/util/StsExtraBusinessLogicTestCase.java
@@ -75,6 +75,12 @@ public class StsExtraBusinessLogicTestCase extends ExtraBusinessLogicTestCase im
         return Boolean.parseBoolean(useKernelSplString);
     }
 
+    @Override
+    public boolean shouldSkipMainlineTests() {
+        return Boolean.parseBoolean(
+                InstrumentationRegistry.getArguments().getString("sts-should-skip-mainline"));
+    }
+
     /**
      * Specify the latest release bulletin. Control this from the command-line with the following:
      * --test-arg
diff --git a/libraries/sts-common-util/host-side/src/com/android/sts/common/CommandUtil.java b/libraries/sts-common-util/host-side/src/com/android/sts/common/CommandUtil.java
index bdd65e742..aaa9a90f6 100644
--- a/libraries/sts-common-util/host-side/src/com/android/sts/common/CommandUtil.java
+++ b/libraries/sts-common-util/host-side/src/com/android/sts/common/CommandUtil.java
@@ -61,7 +61,7 @@ public final class CommandUtil {
                 String.format(
                         "cmd failed: %s\ncode: %s\nstdout:\n%s\nstderr:\n%s",
                         cmd, res.getExitCode(), res.getStdout(), res.getStderr());
-        assertEquals(failMsg, res.getStatus(), CommandStatus.SUCCESS);
+        assertEquals(failMsg, CommandStatus.SUCCESS, res.getStatus());
         return res;
     }
 }
diff --git a/libraries/sts-common-util/host-side/src/com/android/sts/common/tradefed/testtype/StsExtraBusinessLogicHostTestBase.java b/libraries/sts-common-util/host-side/src/com/android/sts/common/tradefed/testtype/StsExtraBusinessLogicHostTestBase.java
index a0b92d0a4..80c5cc95b 100644
--- a/libraries/sts-common-util/host-side/src/com/android/sts/common/tradefed/testtype/StsExtraBusinessLogicHostTestBase.java
+++ b/libraries/sts-common-util/host-side/src/com/android/sts/common/tradefed/testtype/StsExtraBusinessLogicHostTestBase.java
@@ -84,6 +84,12 @@ public class StsExtraBusinessLogicHostTestBase extends ExtraBusinessLogicHostTes
         return Boolean.parseBoolean(useKernelSplString);
     }
 
+    @Override
+    public boolean shouldSkipMainlineTests() {
+        return Boolean.parseBoolean(
+                getBuild().getBuildAttributes().get("sts-should-skip-mainline"));
+    }
+
     /**
      * Specify the latest release bulletin. Control this from the command-line with the following
      * command line argument: --build-attribute "release-bulletin-spl=2021-06"
diff --git a/libraries/sts-common-util/util/src/com/android/sts/common/util/StsLogic.java b/libraries/sts-common-util/util/src/com/android/sts/common/util/StsLogic.java
index 99cea1f51..38fd3df90 100644
--- a/libraries/sts-common-util/util/src/com/android/sts/common/util/StsLogic.java
+++ b/libraries/sts-common-util/util/src/com/android/sts/common/util/StsLogic.java
@@ -45,8 +45,10 @@ public interface StsLogic {
                         "uploadSpl",
                         "uploadModificationTime",
                         "uploadKernelBugs",
+                        "uploadMainlineModules",
                         "declaredSpl",
                         "fridaAssetTemplate",
+                        "mainline",
                     });
     List<String> STS_EXTRA_BUSINESS_LOGIC_INCREMENTAL =
             Arrays.asList(
@@ -54,17 +56,25 @@ public interface StsLogic {
                         "uploadSpl",
                         "uploadModificationTime",
                         "uploadKernelBugs",
+                        "uploadMainlineModules",
                         "declaredSpl",
                         "incremental",
                         "fridaAssetTemplate",
+                        "mainline",
                     });
 
     // intentionally empty because declaredSpl and incremental skipping is not desired when
     // developing STS tests.
+    // Exceptions:
+    // * uploads
+    // * mainline skipping (this is behind an additional flag)
+    // * frida (necessary for frida tests)
     List<String> STS_EXTRA_BUSINESS_LOGIC_DEVELOP =
             Arrays.asList(
                     new String[] {
                         "fridaAssetTemplate",
+                        "uploadMainlineModules",
+                        "mainline",
                     });
 
     Description getTestDescription();
@@ -75,6 +85,8 @@ public interface StsLogic {
 
     boolean shouldUseKernelSpl();
 
+    boolean shouldSkipMainlineTests();
+
     LocalDate getReleaseBulletinSpl();
 
     static List<String> getExtraBusinessLogicForPlan(String stsDynamicPlan) {
@@ -273,6 +285,28 @@ public interface StsLogic {
         return minTestSpl.isAfter(deviceSpl);
     }
 
+    default boolean shouldSkipMainline() {
+        // check if the flag to skip mainline tests has been set to true
+        if (!shouldSkipMainlineTests()) {
+            return false;
+        }
+
+        long[] bugIds = getCveBugIds();
+        if (bugIds == null) {
+            // There were no @AsbSecurityTest annotations
+            logInfo(LOG_TAG, "not an ASB test");
+            return false;
+        }
+
+        Map<String, String> bugModulesMap = BusinessLogicMapStore.getMap("bugid_mainline_modules");
+        for (long bugId : bugIds) {
+            if (bugModulesMap.containsKey(Long.toString(bugId))) {
+                return true;
+            }
+        }
+        return false;
+    }
+
     default void skip(String message) {
         assumeTrue(message, false);
     }
diff --git a/libraries/sts-common-util/util/tests/src/com/android/sts/common/util/StsLogicTest.java b/libraries/sts-common-util/util/tests/src/com/android/sts/common/util/StsLogicTest.java
index 3530e6da6..af3d13d34 100644
--- a/libraries/sts-common-util/util/tests/src/com/android/sts/common/util/StsLogicTest.java
+++ b/libraries/sts-common-util/util/tests/src/com/android/sts/common/util/StsLogicTest.java
@@ -33,8 +33,9 @@ import java.util.Optional;
 @RunWith(DeviceJUnit4ClassRunner.class)
 public class StsLogicTest extends BaseHostJUnit4Test {
 
-    private static final long[] PLATFORM_BUG = {1_000_000_000};
-    private static final long[] KERNEL_BUG = {2_000_000_000};
+    private static final long[] PLATFORM_BUG = {1_000_000_000L};
+    private static final long[] KERNEL_BUG = {2_000_000_000L};
+    private static final long[] MAINLINE_BUG = {3_000_000_000L};
 
     static {
         new BusinessLogicSetStore()
@@ -43,6 +44,13 @@ public class StsLogicTest extends BaseHostJUnit4Test {
                         Arrays.stream(KERNEL_BUG).mapToObj(Long::toString).toArray(String[]::new));
         new BusinessLogicMapStore().putMap("security_bulletins", null);
         new BusinessLogicMapStore().putMap("sts_modification_times", null);
+        new BusinessLogicMapStore()
+                .putMap(
+                        "bugid_mainline_modules",
+                        "#",
+                        Arrays.stream(MAINLINE_BUG)
+                                .mapToObj((bugId) -> Long.toString(bugId) + "#module 1,module 2")
+                                .toArray(String[]::new));
     }
 
     @Test
@@ -116,6 +124,48 @@ public class StsLogicTest extends BaseHostJUnit4Test {
                 logic.getDeviceSpl());
     }
 
+    @Test
+    public final void testNoSkipMainlineNoFlag() throws Exception {
+        StsLogic logic =
+                new StsLogicMock()
+                        .setCveBugIds(MAINLINE_BUG);
+        assertFalse(
+                "shouldn't skip because the flag isn't set.",
+                logic.shouldSkipMainline());
+    }
+
+    @Test
+    public final void testSkipMainlineWithFlag() throws Exception {
+        StsLogic logic =
+                new StsLogicMock()
+                        .setCveBugIds(MAINLINE_BUG)
+                        .setShouldSkipMainlineTests(true);
+        assertTrue(
+                "should skip because the flag is set and this is a Mainline CVE.",
+                logic.shouldSkipMainline());
+    }
+
+    @Test
+    public final void testNoSkipMainlineNotCve() throws Exception {
+        StsLogic logic =
+                new StsLogicMock()
+                        .setCveBugIds(null)
+                        .setShouldSkipMainlineTests(true);
+        assertFalse(
+                "shouldn't Mainline skip because this test is not a CVE test.",
+                logic.shouldSkipMainline());
+    }
+
+    @Test
+    public final void testNoSkipMainlineNotMainlineCve() throws Exception {
+        StsLogic logic =
+                new StsLogicMock()
+                        .setShouldSkipMainlineTests(true);
+        assertFalse(
+                "shouldn't Mainline skip because this test is not a Mainline CVE.",
+                logic.shouldSkipMainline());
+    }
+
     private static class StsLogicMock implements StsLogic {
 
         private long[] cveBugIds = PLATFORM_BUG;
@@ -123,6 +173,7 @@ public class StsLogicTest extends BaseHostJUnit4Test {
         private LocalDate releaseBulletinSpl;
         private Optional<LocalDate> kernelBuildDate;
         private boolean shouldUseKernelSpl = false;
+        private boolean shouldSkipMainlineTests = false;
 
         {
             setPlatformSpl("2022-01-01");
@@ -159,6 +210,11 @@ public class StsLogicTest extends BaseHostJUnit4Test {
             return this;
         }
 
+        public StsLogicMock setShouldSkipMainlineTests(boolean shouldSkipMainlineTests) {
+            this.shouldSkipMainlineTests = shouldSkipMainlineTests;
+            return this;
+        }
+
         @Override
         public Description getTestDescription() {
             throw new UnsupportedOperationException(
@@ -187,6 +243,11 @@ public class StsLogicTest extends BaseHostJUnit4Test {
         }
 
         @Override
+        public boolean shouldSkipMainlineTests() {
+            return this.shouldSkipMainlineTests;
+        }
+
+        @Override
         public LocalDate getReleaseBulletinSpl() {
             return this.releaseBulletinSpl;
         }