Register all referencing contexts for API stubs

Bug: b/335803299
Change-Id: Iac464d7beaea858444777323d09ad7f1ee44d290
diff --git a/src/main/java/com/android/tools/r8/androidapi/ApiReferenceStubber.java b/src/main/java/com/android/tools/r8/androidapi/ApiReferenceStubber.java
index 07261a9..4e8c852 100644
--- a/src/main/java/com/android/tools/r8/androidapi/ApiReferenceStubber.java
+++ b/src/main/java/com/android/tools/r8/androidapi/ApiReferenceStubber.java
@@ -29,7 +29,6 @@
 import com.android.tools.r8.utils.AndroidApiLevel;
 import com.android.tools.r8.utils.ListUtils;
 import com.android.tools.r8.utils.ThreadUtils;
-import com.android.tools.r8.utils.WorkList;
 import com.google.common.collect.Sets;
 import java.util.Arrays;
 import java.util.Collection;
@@ -51,7 +50,6 @@
   private final Map<DexLibraryClass, Set<DexProgramClass>> referencingContexts =
       new ConcurrentHashMap<>();
   private final Set<DexLibraryClass> libraryClassesToMock = Sets.newConcurrentHashSet();
-  private final Set<DexType> seenTypes = Sets.newConcurrentHashSet();
   private final AndroidApiLevelCompute apiLevelCompute;
   private final ApiReferenceStubberEventConsumer eventConsumer;
 
@@ -150,25 +148,24 @@
     if (!type.isClassType() || isJavaType(type, appView.dexItemFactory())) {
       return;
     }
-    WorkList.newIdentityWorkList(type, seenTypes)
-        .process(
-            (classType, workList) -> {
-              DexClass clazz = appView.definitionFor(classType);
-              if (clazz == null || !clazz.isLibraryClass()) {
-                return;
-              }
-              ComputedApiLevel androidApiLevel =
-                  apiLevelCompute.computeApiLevelForLibraryReference(
-                      clazz.type, ComputedApiLevel.unknown());
-              if (androidApiLevel.isGreaterThan(appView.computedMinApiLevel())
-                  && androidApiLevel.isKnownApiLevel()) {
-                workList.addIfNotSeen(clazz.allImmediateSupertypes());
-                libraryClassesToMock.add(clazz.asLibraryClass());
-                referencingContexts
-                    .computeIfAbsent(clazz.asLibraryClass(), ignoreKey(Sets::newConcurrentHashSet))
-                    .add(context);
-              }
-            });
+    DexClass clazz = appView.definitionFor(type);
+    if (clazz == null || !clazz.isLibraryClass()) {
+      return;
+    }
+    DexLibraryClass libraryClass = clazz.asLibraryClass();
+    ComputedApiLevel androidApiLevel =
+        apiLevelCompute.computeApiLevelForLibraryReference(
+            libraryClass.type, ComputedApiLevel.unknown());
+    if (androidApiLevel.isGreaterThan(appView.computedMinApiLevel())
+        && androidApiLevel.isKnownApiLevel()) {
+      libraryClassesToMock.add(libraryClass);
+      referencingContexts
+          .computeIfAbsent(libraryClass, ignoreKey(Sets::newConcurrentHashSet))
+          .add(context);
+    }
+    for (DexType supertype : libraryClass.allImmediateSupertypes()) {
+      findReferencedLibraryClasses(supertype, context);
+    }
   }
 
   public static boolean isJavaType(DexType type, DexItemFactory factory) {
diff --git a/src/test/java/com/android/tools/r8/synthesis/globals/GlobalSyntheticStubContextRegressionTest.java b/src/test/java/com/android/tools/r8/synthesis/globals/GlobalSyntheticStubContextRegressionTest.java
new file mode 100644
index 0000000..54f5f0c
--- /dev/null
+++ b/src/test/java/com/android/tools/r8/synthesis/globals/GlobalSyntheticStubContextRegressionTest.java
@@ -0,0 +1,110 @@
+// Copyright (c) 2024, the R8 project authors. Please see the AUTHORS file
+// for details. All rights reserved. Use of this source code is governed by a
+// BSD-style license that can be found in the LICENSE file.
+
+package com.android.tools.r8.synthesis.globals;
+
+import static org.junit.Assert.assertEquals;
+import static org.junit.Assert.assertTrue;
+
+import com.android.tools.r8.OutputMode;
+import com.android.tools.r8.TestBase;
+import com.android.tools.r8.TestParameters;
+import com.android.tools.r8.TestParametersCollection;
+import com.android.tools.r8.ToolHelper;
+import com.android.tools.r8.transformers.ClassFileTransformer.MethodPredicate;
+import com.android.tools.r8.utils.AndroidApiLevel;
+import com.android.tools.r8.utils.StringUtils;
+import java.io.IOException;
+import java.nio.file.Path;
+import org.junit.Test;
+import org.junit.runner.RunWith;
+import org.junit.runners.Parameterized;
+
+/** Regression test for b/335803299. */
+@RunWith(Parameterized.class)
+public class GlobalSyntheticStubContextRegressionTest extends TestBase {
+
+  static final String EXPECTED =
+      StringUtils.lines("Hello", "all good...", "Hello again", "still good...");
+
+  private final TestParameters parameters;
+
+  @Parameterized.Parameters(name = "{0}")
+  public static TestParametersCollection data() {
+    return getTestParameters().withDexRuntimes().withApiLevel(AndroidApiLevel.N).build();
+  }
+
+  public GlobalSyntheticStubContextRegressionTest(TestParameters parameters) {
+    this.parameters = parameters;
+  }
+
+  @Test
+  public void testDexPerClassFileDir() throws Exception {
+    GlobalSyntheticsTestingConsumer globals = new GlobalSyntheticsTestingConsumer();
+    Path dexOut =
+        testForD8()
+            .apply(b -> b.getBuilder().setEnableExperimentalMissingLibraryApiModeling(true))
+            .addProgramClassFileData(
+                transformClass(TestClass1.class), transformClass(TestClass2.class))
+            .setMinApi(parameters)
+            .addLibraryFiles(ToolHelper.getAndroidJar(AndroidApiLevel.LATEST))
+            .apply(b -> b.getBuilder().setGlobalSyntheticsConsumer(globals))
+            .setOutputMode(OutputMode.DexFilePerClassFile)
+            .compile()
+            .writeToZip();
+
+    assertTrue(globals.hasGlobals());
+    assertEquals(2, globals.getProviders().size());
+
+    testForD8()
+        .addProgramFiles(dexOut)
+        .setMinApi(parameters)
+        .apply(b -> b.getBuilder().addGlobalSyntheticsResourceProviders(globals.getProviders()))
+        .run(parameters.getRuntime(), TestClass1.class)
+        .assertSuccessWithOutput(EXPECTED);
+  }
+
+  private byte[] transformClass(Class<?> clazz) throws IOException {
+    return transformer(clazz)
+        .transformTryCatchBlock(
+            MethodPredicate.all(),
+            (start, end, handler, type, visitor) -> {
+              assertEquals("java/lang/Exception", type);
+              visitor.visitTryCatchBlock(
+                  start, end, handler, "android/app/AuthenticationRequiredException");
+            })
+        .transform();
+  }
+
+  static class TestClass1 {
+
+    public static void main(String[] args) {
+      Runnable r =
+          () -> {
+            try {
+              System.out.println("Hello");
+            } catch (Exception e) {
+              System.out.println("fail...");
+              throw e;
+            }
+            System.out.println("all good...");
+          };
+      r.run();
+      TestClass2.main(args);
+    }
+  }
+
+  static class TestClass2 {
+
+    public static void main(String[] args) {
+      try {
+        System.out.println("Hello again");
+      } catch (Exception e) {
+        System.out.println("fail...");
+        throw e;
+      }
+      System.out.println("still good...");
+    }
+  }
+}