Track preserved JNI elements

Author: DSouzaM

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jni/access/JNIAccessibleClass.java

@@ -39,19 +39,21 @@
 /**
  * Information on a class that can be looked up and accessed via JNI.
  */
-public final class JNIAccessibleClass {
+public final class JNIAccessibleClass extends JNIAccessibleElement {
     private final Class<?> classObject;
     private EconomicMap<JNIAccessibleMethodDescriptor, JNIAccessibleMethod> methods;
     private EconomicMap<CharSequence, JNIAccessibleField> fields;
 
     @Platforms(HOSTED_ONLY.class)
-    public JNIAccessibleClass(Class<?> clazz) {
+    public JNIAccessibleClass(Class<?> clazz, boolean preserved) {
+        super(preserved);
         assert clazz != null;
         this.classObject = clazz;
     }
 
     @Platforms(HOSTED_ONLY.class)
     JNIAccessibleClass() {
+        super(false);
         /* For negative queries */
         this.classObject = null;
     }
@@ -73,22 +75,28 @@ public JNIAccessibleField getField(CharSequence name) {
     }
 
     @Platforms(HOSTED_ONLY.class)
-    public void addFieldIfAbsent(String name, Function<String, JNIAccessibleField> mappingFunction) {
+    public void addOrUpdateField(String name, boolean updatedPreserved, Function<String, JNIAccessibleField> mappingFunction) {
         if (fields == null) {
             fields = ImageHeapMap.createNonLayeredMap(JNIReflectionDictionary.WRAPPED_CSTRING_EQUIVALENCE);
         }
-        if (!fields.containsKey(name)) {
+        JNIAccessibleField existing = fields.get(name);
+        if (existing == null) {
             fields.put(name, mappingFunction.apply(name));
+        } else {
+            existing.reportReregistered(updatedPreserved);
         }
     }
 
     @Platforms(HOSTED_ONLY.class)
-    public void addMethodIfAbsent(JNIAccessibleMethodDescriptor descriptor, Function<JNIAccessibleMethodDescriptor, JNIAccessibleMethod> mappingFunction) {
+    public void addOrUpdateMethod(JNIAccessibleMethodDescriptor descriptor, boolean updatedPreserved, Function<JNIAccessibleMethodDescriptor, JNIAccessibleMethod> mappingFunction) {
         if (methods == null) {
             methods = ImageHeapMap.createNonLayeredMap();
         }
-        if (!methods.containsKey(descriptor)) {
+        JNIAccessibleMethod existing = methods.get(descriptor);
+        if (existing == null) {
             methods.put(descriptor, mappingFunction.apply(descriptor));
+        } else {
+            existing.reportReregistered(updatedPreserved);
         }
     }
 

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jni/access/JNIAccessibleElement.java

@@ -0,0 +1,44 @@
+/*
+ * Copyright (c) 2025, 2025, Oracle and/or its affiliates. All rights reserved.
+ * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
+ *
+ * This code is free software; you can redistribute it and/or modify it
+ * under the terms of the GNU General Public License version 2 only, as
+ * published by the Free Software Foundation.  Oracle designates this
+ * particular file as subject to the "Classpath" exception as provided
+ * by Oracle in the LICENSE file that accompanied this code.
+ *
+ * This code is distributed in the hope that it will be useful, but WITHOUT
+ * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
+ * FITNESS FOR A PARTICULAR PURPOSE.  See the GNU General Public License
+ * version 2 for more details (a copy is included in the LICENSE file that
+ * accompanied this code).
+ *
+ * You should have received a copy of the GNU General Public License version
+ * 2 along with this work; if not, write to the Free Software Foundation,
+ * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
+ *
+ * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
+ * or visit www.oracle.com if you need additional information or have any
+ * questions.
+ */
+package com.oracle.svm.core.jni.access;
+
+public abstract class JNIAccessibleElement {
+    private volatile boolean preserved;
+
+    protected JNIAccessibleElement(boolean preserved) {
+        this.preserved = preserved;
+    }
+
+    public void reportReregistered(boolean reregisterPreserved) {
+        // State can only ever go from "preserved" to "not preserved".
+        if (!reregisterPreserved) {
+            this.preserved = false;
+        }
+    }
+
+    public boolean isPreserved() {
+        return preserved;
+    }
+}

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jni/access/JNIAccessibleField.java

@@ -61,7 +61,7 @@ public final class JNIAccessibleField extends JNIAccessibleMember {
     private static final UnsignedWord ID_OFFSET_MASK = ID_LAYER_NUMBER_MASK.subtract(1);
 
     public static JNIAccessibleField negativeFieldQuery(JNIAccessibleClass jniClass) {
-        return new JNIAccessibleField(jniClass, null, 0);
+        return new JNIAccessibleField(false, jniClass, null, 0);
     }
 
     /**
@@ -116,8 +116,8 @@ private static int getLayerNumberFromId(UnsignedWord id) {
     private UnsignedWord id = Word.zero();
 
     @Platforms(HOSTED_ONLY.class)
-    public JNIAccessibleField(JNIAccessibleClass declaringClass, JavaKind kind, int modifiers) {
-        super(declaringClass);
+    public JNIAccessibleField(boolean preserved, JNIAccessibleClass declaringClass, JavaKind kind, int modifiers) {
+        super(preserved, declaringClass);
 
         UnsignedWord bits = Modifier.isStatic(modifiers) ? ID_STATIC_FLAG : Word.zero();
         if (kind == null) {

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jni/access/JNIAccessibleMember.java

@@ -31,13 +31,14 @@
 import com.oracle.svm.core.BuildPhaseProvider.ReadyForCompilation;
 import com.oracle.svm.core.heap.UnknownObjectField;
 
-abstract class JNIAccessibleMember {
+abstract class JNIAccessibleMember extends JNIAccessibleElement {
     private final JNIAccessibleClass declaringClass;
 
     @UnknownObjectField(fullyQualifiedTypes = "org.graalvm.collections.EconomicMapImpl", canBeNull = true, availability = ReadyForCompilation.class) //
     private EconomicSet<Class<?>> hidingSubclasses;
 
-    JNIAccessibleMember(JNIAccessibleClass declaringClass) {
+    JNIAccessibleMember(boolean preserved, JNIAccessibleClass declaringClass) {
+        super(preserved);
         this.declaringClass = declaringClass;
     }
 

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jni/access/JNIAccessibleMethod.java

@@ -57,7 +57,7 @@ public final class JNIAccessibleMethod extends JNIAccessibleMember {
     public static final int NEW_OBJECT_TARGET_INVALID_FOR_ABSTRACT_TYPE = -1;
 
     public static JNIAccessibleMethod negativeMethodQuery(JNIAccessibleClass jniClass) {
-        return new JNIAccessibleMethod(jniClass, RuntimeMetadataDecoderImpl.NEGATIVE_FLAG_MASK);
+        return new JNIAccessibleMethod(false, jniClass, RuntimeMetadataDecoderImpl.NEGATIVE_FLAG_MASK);
     }
 
     @Platforms(HOSTED_ONLY.class)
@@ -104,8 +104,8 @@ public static ResolvedJavaField getCallVariantWrapperField(MetaAccessProvider me
     @SuppressWarnings("unused") private CodePointer valistNonvirtualWrapper;
 
     @Platforms(HOSTED_ONLY.class)
-    public JNIAccessibleMethod(JNIAccessibleClass declaringClass, int modifiers) {
-        super(declaringClass);
+    public JNIAccessibleMethod(boolean preserved, JNIAccessibleClass declaringClass, int modifiers) {
+        super(preserved, declaringClass);
         this.modifiers = modifiers;
     }
 

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jni/access/JNIAccessibleMethodDescriptor.java

@@ -69,7 +69,7 @@ public static JNIAccessibleMethodDescriptor of(String methodName, Class<?>[] par
         return of(methodName, parameterTypes, null);
     }
 
-    private static JNIAccessibleMethodDescriptor of(String methodName, Class<?>[] parameterTypes, Class<?> returnType) {
+    public static JNIAccessibleMethodDescriptor of(String methodName, Class<?>[] parameterTypes, Class<?> returnType) {
         StringBuilder sb = new StringBuilder("(");
         for (Class<?> type : parameterTypes) {
             sb.append(MetaUtil.toInternalName(type.getName()));

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/jni/access/JNIReflectionDictionary.java

@@ -155,14 +155,18 @@ private static void dump(boolean condition, String label) {
     }
 
     @Platforms(HOSTED_ONLY.class)
-    public JNIAccessibleClass addClassIfAbsent(Class<?> classObj, Function<Class<?>, JNIAccessibleClass> mappingFunction) {
-        if (!classesByClassObject.containsKey(classObj)) {
+    public JNIAccessibleClass addOrUpdateClass(Class<?> classObj, boolean updatedPreserved, Function<Class<?>, JNIAccessibleClass> mappingFunction) {
+        JNIAccessibleClass existing = classesByClassObject.get(classObj);
+        if (existing == null) {
             JNIAccessibleClass instance = mappingFunction.apply(classObj);
             classesByClassObject.put(classObj, instance);
             String name = instance.getJNIName();
             classesByName.put(name, instance);
+            return instance;
+        } else {
+            existing.reportReregistered(updatedPreserved);
+            return existing;
         }
-        return classesByClassObject.get(classObj);
     }
 
     @Platforms(HOSTED_ONLY.class)
@@ -181,6 +185,15 @@ public Iterable<JNIAccessibleClass> getClasses() {
     }
 
     public static Class<?> getClassObjectByName(CharSequence name) {
+        JNIAccessibleClass clazz = getJniAccessibleClass(name);
+        if (clazz != null) {
+            return clazz.getClassObject();
+        }
+        dump(true, "getClassObjectByName");
+        return null;
+    }
+
+    public static JNIAccessibleClass getJniAccessibleClass(CharSequence name) {
         for (var dictionary : layeredSingletons()) {
             JNIAccessibleClass clazz = dictionary.classesByName.get(name);
             if (clazz == null && !ClassNameSupport.isValidJNIName(name.toString())) {
@@ -189,12 +202,8 @@ public static Class<?> getClassObjectByName(CharSequence name) {
                 // trace if class exists (positive query) or name is valid (negative query)
                 MetadataTracer.singleton().traceJNIType(ClassNameSupport.jniNameToTypeName(name.toString()));
             }
-            clazz = checkClass(clazz, name.toString());
-            if (clazz != null) {
-                return clazz.getClassObject();
-            }
+            return checkClass(clazz, name.toString());
         }
-        dump(true, "getClassObjectByName");
         return null;
     }
 

substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/imagelayer/SVMImageLayerLoader.java

@@ -940,7 +940,7 @@ protected boolean delegateLoadMethod(PersistedAnalysisMethod.Reader methodData)
             if (wm.isReflectionExpandSignature()) {
                 ImageSingletons.lookup(ReflectionFeature.class).getOrCreateAccessor(member);
             } else if (wm.isJavaCallVariantWrapper()) {
-                JNIAccessFeature.singleton().addMethod(member, (FeatureImpl.DuringAnalysisAccessImpl) universe.getConcurrentAnalysisAccess());
+                JNIAccessFeature.singleton().addMethod(member, false, (FeatureImpl.DuringAnalysisAccessImpl) universe.getConcurrentAnalysisAccess());
             }
             return true;
         } else if (wrappedMethod.isPolymorphicSignature()) {