Fix dlsym lookup trampoline for @CriticalNative.

For @CriticalNative we do not have the JniEnv* argument to
retrieve the Thread*, we must use the thread register.

The test 178-app-image-native-method was supposed to cover
this but the profile was missing the necessary references.
We add those references and add $noinline$opt$ tags to make
sure the test actually runs the way it was supposed to and
add extra calls to the critical methods as they can take
different paths on subsequent execution.

Also add an extra test for class initialization checks that
ensures correct behavior for planned direct calls from
compiled managed code to registered @CriticalNative code.

Test: testrunner.py --target --optimizing -t 178
Bug: 112189621
Change-Id: I6a63d990bc08236ebaac7dacb0f5979d835ee321
diff --git a/dex2oat/driver/compiler_driver.cc b/dex2oat/driver/compiler_driver.cc
index 40b999b..1b29e3b 100644
--- a/dex2oat/driver/compiler_driver.cc
+++ b/dex2oat/driver/compiler_driver.cc
@@ -301,7 +301,9 @@
 
 std::unique_ptr<const std::vector<uint8_t>>
 CompilerDriver::CreateJniDlsymLookupCriticalTrampoline() const {
-  CREATE_TRAMPOLINE(JNI, kJniAbi, pDlsymLookupCritical)
+  // @CriticalNative calls do not have the `JNIEnv*` parameter, so this trampoline uses the
+  // architecture-dependent access to `Thread*` using the managed code ABI, i.e. `kQuickAbi`.
+  CREATE_TRAMPOLINE(JNI, kQuickAbi, pDlsymLookupCritical)
 }
 
 std::unique_ptr<const std::vector<uint8_t>> CompilerDriver::CreateQuickGenericJniTrampoline()
diff --git a/test/178-app-image-native-method/expected.txt b/test/178-app-image-native-method/expected.txt
index b4e5ece..c6014b6 100644
--- a/test/178-app-image-native-method/expected.txt
+++ b/test/178-app-image-native-method/expected.txt
@@ -6,6 +6,16 @@
 testMissingFast
 testMissingCritical
 testCriticalSignatures
+test
+testFast
+testCritical
+testMissing
+testMissingFast
+testMissingCritical
+testCriticalSignatures
+testCriticalClinitCheck passed
+initializingCriticalClinitCheck finished
+testCriticalClinitCheck passed
 JNI_OnLoad called
 test
 testFast
@@ -14,3 +24,13 @@
 testMissingFast
 testMissingCritical
 testCriticalSignatures
+test
+testFast
+testCritical
+testMissing
+testMissingFast
+testMissingCritical
+testCriticalSignatures
+testCriticalClinitCheck passed
+initializingCriticalClinitCheck finished
+testCriticalClinitCheck passed
diff --git a/test/178-app-image-native-method/native_methods.cc b/test/178-app-image-native-method/native_methods.cc
index becda81..0669deb 100644
--- a/test/178-app-image-native-method/native_methods.cc
+++ b/test/178-app-image-native-method/native_methods.cc
@@ -609,4 +609,33 @@
   return 42;
 }
 
+extern "C" JNIEXPORT jint JNICALL Java_CriticalClinitCheck_nativeMethodVoid() {
+  return 42;
+}
+
+extern "C" JNIEXPORT jint JNICALL Java_CriticalClinitCheck_nativeMethod(jint i) {
+  return i;
+}
+
+extern "C" JNIEXPORT jint JNICALL Java_CriticalClinitCheck_nativeMethodWithManyParameters(
+    jint i1, jlong l1, jfloat f1, jdouble d1,
+    jint i2, jlong l2, jfloat f2, jdouble d2,
+    jint i3, jlong l3, jfloat f3, jdouble d3,
+    jint i4, jlong l4, jfloat f4, jdouble d4,
+    jint i5, jlong l5, jfloat f5, jdouble d5,
+    jint i6, jlong l6, jfloat f6, jdouble d6,
+    jint i7, jlong l7, jfloat f7, jdouble d7,
+    jint i8, jlong l8, jfloat f8, jdouble d8) {
+  bool ok = VerifyManyParameters(
+      i1, l1, f1, d1,
+      i2, l2, f2, d2,
+      i3, l3, f3, d3,
+      i4, l4, f4, d4,
+      i5, l5, f5, d5,
+      i6, l6, f6, d6,
+      i7, l7, f7, d7,
+      i8, l8, f8, d8);
+  return ok ? 42 : -1;
+}
+
 }  // namespace art
diff --git a/test/178-app-image-native-method/profile b/test/178-app-image-native-method/profile
index 597dde1..1c25dee 100644
--- a/test/178-app-image-native-method/profile
+++ b/test/178-app-image-native-method/profile
@@ -1,8 +1,17 @@
 LMain;
 LTest;
-HSPLMain;->test()V
-HSPLMain;->testFast()V
-HSPLMain;->testCritical()V
-HSPLMain;->testMissing()V
-HSPLMain;->testMissingFast()V
-HSPLMain;->testMissingCritical()V
+LTestFast;
+LTestCritical;
+LTestMissing;
+LTestMissingFast;
+LTestMissingCritical;
+LCriticalSignatures;
+LCriticalClinitCheck;
+HSPLMain;->$noinline$opt$test()V
+HSPLMain;->$noinline$opt$testFast()V
+HSPLMain;->$noinline$opt$testCritical()V
+HSPLMain;->$noinline$opt$testMissing()V
+HSPLMain;->$noinline$opt$testMissingFast()V
+HSPLMain;->$noinline$opt$testMissingCritical()V
+HSPLMain;->$noinline$opt$testCriticalSignatures()V
+HSPLMain;->$noinline$opt$testCriticalClinitCheck()V
diff --git a/test/178-app-image-native-method/run b/test/178-app-image-native-method/run
index f4b07f0..7cd0d57 100644
--- a/test/178-app-image-native-method/run
+++ b/test/178-app-image-native-method/run
@@ -14,8 +14,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-# Use a profile to put specific classes in the app image.
-${RUN} $@ --profile -Xcompiler-option --compiler-filter=speed-profile
+# Use a profile to put specific classes in the app image. Increase the large
+# method threshold to compile Main.$noinline$opt$testCriticalSignatures().
+${RUN} $@ --profile -Xcompiler-option --compiler-filter=speed-profile \
+    -Xcompiler-option --large-method-max=2000
 return_status1=$?
 
 # Also run with the verify filter to avoid compiling JNI stubs.
diff --git a/test/178-app-image-native-method/src/Main.java b/test/178-app-image-native-method/src/Main.java
index 9043081..a9c6f0a 100644
--- a/test/178-app-image-native-method/src/Main.java
+++ b/test/178-app-image-native-method/src/Main.java
@@ -32,17 +32,30 @@
     new CriticalSignatures();
     makeVisiblyInitialized();  // Make sure they are visibly initialized.
 
-    test();
-    testFast();
-    testCritical();
-    testMissing();
-    testMissingFast();
-    testMissingCritical();
+    $noinline$opt$test();
+    $noinline$opt$testFast();
+    $noinline$opt$testCritical();
+    $noinline$opt$testMissing();
+    $noinline$opt$testMissingFast();
+    $noinline$opt$testMissingCritical();
+    $noinline$opt$testCriticalSignatures();
 
-    testCriticalSignatures();
+    // For calls from AOT-compiled code, the first call shall use the resolution method
+    // retrieved from the .bss and the second call should find the .bss entry populated
+    // with the target method. Re-run these tests to exercise the second path.
+    $noinline$opt$test();
+    $noinline$opt$testFast();
+    $noinline$opt$testCritical();
+    $noinline$opt$testMissing();
+    $noinline$opt$testMissingFast();
+    $noinline$opt$testMissingCritical();
+    $noinline$opt$testCriticalSignatures();
+
+    new CriticalClinitCheck();
+    sTestCriticalClinitCheckOtherThread.join();
   }
 
-  static void test() {
+  static void $noinline$opt$test() {
     System.out.println("test");
     assertEquals(42, Test.nativeMethodVoid());
     assertEquals(42, Test.nativeMethod(42));
@@ -57,7 +70,7 @@
         81, 82L, 83.0f, 84.0d));
   }
 
-  static void testFast() {
+  static void $noinline$opt$testFast() {
     System.out.println("testFast");
     assertEquals(42, TestFast.nativeMethodVoid());
     assertEquals(42, TestFast.nativeMethod(42));
@@ -72,7 +85,7 @@
         81, 82L, 83.0f, 84.0d));
   }
 
-  static void testCritical() {
+  static void $noinline$opt$testCritical() {
     System.out.println("testCritical");
     assertEquals(42, TestCritical.nativeMethodVoid());
     assertEquals(42, TestCritical.nativeMethod(42));
@@ -87,7 +100,7 @@
         81, 82L, 83.0f, 84.0d));
   }
 
-  static void testMissing() {
+  static void $noinline$opt$testMissing() {
     System.out.println("testMissing");
 
     try {
@@ -114,7 +127,7 @@
     } catch (LinkageError expected) {}
   }
 
-  static void testMissingFast() {
+  static void $noinline$opt$testMissingFast() {
     System.out.println("testMissingFast");
 
     try {
@@ -141,7 +154,7 @@
     } catch (LinkageError expected) {}
   }
 
-  static void testMissingCritical() {
+  static void $noinline$opt$testMissingCritical() {
     System.out.println("testMissingCritical");
 
     try {
@@ -168,7 +181,7 @@
     } catch (LinkageError expected) {}
   }
 
-  static void testCriticalSignatures() {
+  static void $noinline$opt$testCriticalSignatures() {
     System.out.println("testCriticalSignatures");
     long l = 0xf00000000L;
     assertEquals(42, CriticalSignatures.nativeILFFFFD(1, l + 2L, 3.0f, 4.0f, 5.0f, 6.0f, 7.0));
@@ -356,6 +369,44 @@
         254));
   }
 
+  static void initializingCriticalClinitCheck() {
+    // Called from CriticalClinitCheck.<clinit>().
+    // Test @CriticalNative calls on the initializing thread.
+    $noinline$opt$testCriticalClinitCheck();
+    sTestCriticalClinitCheckOtherThread = new Thread() {
+      public void run() {
+        $noinline$opt$testCriticalClinitCheck();
+      }
+    };
+    sTestCriticalClinitCheckOtherThread.start();
+    // Sleep for a second to give the other thread an opportunity to run.
+    // We're testing that it performs a clinit check and blocks until we
+    // exit the class initializer after writing the output below.
+    try {
+      Thread.sleep(1000);
+    } catch (InterruptedException ie) {
+      throw new Error(ie);
+    }
+    System.out.println("initializingCriticalClinitCheck finished");
+  }
+
+  static void $noinline$opt$testCriticalClinitCheck() {
+    assertEquals(42, CriticalClinitCheck.nativeMethodVoid());
+    assertEquals(42, CriticalClinitCheck.nativeMethod(42));
+    assertEquals(42, CriticalClinitCheck.nativeMethodWithManyParameters(
+        11, 12L, 13.0f, 14.0d,
+        21, 22L, 23.0f, 24.0d,
+        31, 32L, 33.0f, 34.0d,
+        41, 42L, 43.0f, 44.0d,
+        51, 52L, 53.0f, 54.0d,
+        61, 62L, 63.0f, 64.0d,
+        71, 72L, 73.0f, 74.0d,
+        81, 82L, 83.0f, 84.0d));
+    System.out.println("testCriticalClinitCheck passed");
+  }
+
+  static Thread sTestCriticalClinitCheckOtherThread = null;
+
   static void assertEquals(int expected, int actual) {
     if (expected != actual) {
       throw new AssertionError("Expected " + expected + " got " + actual);
@@ -749,3 +800,26 @@
       long l252,
       int i254);
 }
+
+class CriticalClinitCheck {
+  @CriticalNative
+  public static native int nativeMethodVoid();
+
+  @CriticalNative
+  public static native int nativeMethod(int i);
+
+  @CriticalNative
+  public static native int nativeMethodWithManyParameters(
+      int i1, long l1, float f1, double d1,
+      int i2, long l2, float f2, double d2,
+      int i3, long l3, float f3, double d3,
+      int i4, long l4, float f4, double d4,
+      int i5, long l5, float f5, double d5,
+      int i6, long l6, float f6, double d6,
+      int i7, long l7, float f7, double d7,
+      int i8, long l8, float f8, double d8);
+
+  static {
+    Main.initializingCriticalClinitCheck();
+  }
+}