Make SerializationSupport.constructorAccessors use DynamicHub id instead of Class Object as key

Author: Zeavee

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/hub/DynamicHub.java

@@ -85,7 +85,6 @@
 import java.util.function.BiFunction;
 import java.util.function.IntFunction;
 
-import com.oracle.svm.shared.singletons.MultiLayeredImageSingleton;
 import org.graalvm.collections.EconomicSet;
 import org.graalvm.nativeimage.AnnotationAccess;
 import org.graalvm.nativeimage.ImageSingletons;
@@ -142,6 +141,7 @@
 import com.oracle.svm.core.util.BasedOnJDKFile;
 import com.oracle.svm.core.util.LazyFinalReference;
 import com.oracle.svm.guest.staging.Uninterruptible;
+import com.oracle.svm.shared.singletons.MultiLayeredImageSingleton;
 import com.oracle.svm.shared.util.VMError;
 import com.oracle.svm.util.AnnotationUtil;
 import com.oracle.svm.util.GuestAccess;
@@ -2545,7 +2545,7 @@ public FieldAccessor newFieldAccessor(Field field0, boolean override) {
 
     @Substitute
     private Constructor<?> generateConstructor(Class<?> cl, Constructor<?> constructorToCall) {
-        ConstructorAccessor acc = (ConstructorAccessor) SerializationSupport.getSerializationConstructorAccessor(cl, constructorToCall.getDeclaringClass());
+        ConstructorAccessor acc = (ConstructorAccessor) SerializationSupport.getRuntimeSerializationConstructorAccessor(cl, constructorToCall.getDeclaringClass());
         /*
          * Unlike other root constructors, this constructor is not copied for mutation but directly
          * mutated, as it is not cached. To cache this constructor, setAccessible call must be done

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/reflect/serialize/SerializationRegistry.java

@@ -27,9 +27,10 @@
 
 import java.io.Serializable;
 import java.lang.invoke.SerializedLambda;
-import java.lang.reflect.Constructor;
 import java.lang.reflect.Modifier;
-import java.util.Objects;
+import java.util.function.Consumer;
+import java.util.function.Function;
+import java.util.function.Predicate;
 
 import org.graalvm.collections.EconomicMap;
 import org.graalvm.collections.MapCursor;
@@ -44,20 +45,19 @@
 import com.oracle.svm.core.heap.UnknownObjectField;
 import com.oracle.svm.core.hub.DynamicHub;
 import com.oracle.svm.core.layeredimagesingleton.LayeredImageSingletonSupport;
-import com.oracle.svm.shared.singletons.MultiLayeredImageSingleton;
 import com.oracle.svm.core.metadata.MetadataTracer;
 import com.oracle.svm.core.reflect.SubstrateConstructorAccessor;
+import com.oracle.svm.shared.singletons.MultiLayeredImageSingleton;
 import com.oracle.svm.shared.singletons.traits.BuiltinTraits.AllAccess;
 import com.oracle.svm.shared.singletons.traits.BuiltinTraits.NoLayeredCallbacks;
 import com.oracle.svm.shared.singletons.traits.SingletonLayeredInstallationKind.MultiLayer;
 import com.oracle.svm.shared.singletons.traits.SingletonTraits;
-import com.oracle.svm.core.util.ImageHeapMap;
 import com.oracle.svm.shared.util.VMError;
 
 import jdk.graal.compiler.java.LambdaUtils;
 
 @SingletonTraits(access = AllAccess.class, layeredCallbacks = NoLayeredCallbacks.class, layeredInstallationKind = MultiLayer.class)
-public class SerializationSupport implements SerializationRegistry {
+public class SerializationSupport {
 
     @Platforms(Platform.HOSTED_ONLY.class)
     public static SerializationSupport currentLayer() {
@@ -108,60 +108,52 @@ private StubForAbstractClass() {
         }
     }
 
-    private Constructor<?> stubConstructor;
+    private DynamicHub stubConstructorClass;
+    private DynamicHub serializedLambdaClass;
 
-    public static final class SerializationLookupKey {
-        private final Class<?> declaringClass;
-        private final Class<?> targetConstructorClass;
-
-        private SerializationLookupKey(Class<?> declaringClass, Class<?> targetConstructorClass) {
-            assert declaringClass != null && targetConstructorClass != null;
-            this.declaringClass = declaringClass;
-            this.targetConstructorClass = targetConstructorClass;
-        }
+    /**
+     * The ids are either a {@link DynamicHubKey} or {@code DynamicHub.typeID}.
+     */
+    public record SerializationLookupKey(Object declaringClassId, Object targetConstructorClassId) {
+    }
 
-        public Class<?> getDeclaringClass() {
-            return declaringClass;
-        }
+    @UnknownObjectField(fullyQualifiedTypes = "org.graalvm.collections.EconomicMapImpl", availability = AfterCompilation.class) //
+    private final EconomicMap<SerializationLookupKey, Object> constructorAccessors;
 
-        public Class<?> getTargetConstructorClass() {
-            return targetConstructorClass;
-        }
+    /**
+     * The constructor accessors need to be rescanned manually because the
+     * {@link SerializationSupport#constructorAccessors} map is only available after compilation.
+     */
+    @Platforms(Platform.HOSTED_ONLY.class)
+    private Consumer<Object> scanObject;
 
-        @Override
-        public boolean equals(Object o) {
-            if (this == o) {
-                return true;
-            }
-            if (o == null || getClass() != o.getClass()) {
-                return false;
-            }
-            SerializationLookupKey that = (SerializationLookupKey) o;
-            return declaringClass.equals(that.declaringClass) && targetConstructorClass.equals(that.targetConstructorClass);
-        }
+    public SerializationSupport() {
+        constructorAccessors = EconomicMap.create();
+    }
 
-        @Override
-        public int hashCode() {
-            return Objects.hash(declaringClass, targetConstructorClass);
-        }
+    public void setStubConstructor(DynamicHub stubConstructorClass) {
+        VMError.guarantee(this.stubConstructorClass == null, "Cannot reset stubConstructor");
+        this.stubConstructorClass = stubConstructorClass;
     }
 
-    private final EconomicMap<SerializationLookupKey, Object> constructorAccessors;
+    public void setSerializedLambdaClass(DynamicHub serializedLambdaClass) {
+        this.serializedLambdaClass = serializedLambdaClass;
+    }
 
-    @Platforms(Platform.HOSTED_ONLY.class)
-    public SerializationSupport() {
-        constructorAccessors = ImageHeapMap.create("constructorAccessors");
+    private DynamicHub getSerializedLambdaClass() {
+        return serializedLambdaClass;
     }
 
-    public void setStubConstructor(Constructor<?> stubConstructor) {
-        VMError.guarantee(this.stubConstructor == null, "Cannot reset stubConstructor");
-        this.stubConstructor = stubConstructor;
+    public void setScanObject(Consumer<Object> scanObject) {
+        this.scanObject = scanObject;
     }
 
     @Platforms(Platform.HOSTED_ONLY.class)
-    public Object addConstructorAccessor(Class<?> declaringClass, Class<?> targetConstructorClass, Object constructorAccessor) {
+    public Object addConstructorAccessor(DynamicHub declaringClass, DynamicHub targetConstructorClass, Object constructorAccessor) {
         VMError.guarantee(constructorAccessor instanceof SubstrateConstructorAccessor, "Not a SubstrateConstructorAccessor: %s", constructorAccessor);
-        SerializationLookupKey key = new SerializationLookupKey(declaringClass, targetConstructorClass);
+        SerializationLookupKey key = BuildPhaseProvider.isHostedUniverseBuilt() ? new SerializationLookupKey(declaringClass.getTypeID(), targetConstructorClass.getTypeID())
+                        : new SerializationLookupKey(new DynamicHubKey(declaringClass), new DynamicHubKey(targetConstructorClass));
+        scanObject.accept(constructorAccessor);
         return constructorAccessors.putIfAbsent(key, constructorAccessor);
     }
 
@@ -193,7 +185,7 @@ public String getClassLoaderSerializationLookupKey(ClassLoader classLoader) {
         while (constructorAccessorsCursor.advance()) {
             if (constructorAccessorsCursor.getValue().getClass().getClassLoader() == classLoader) {
                 var key = constructorAccessorsCursor.getKey();
-                return key.declaringClass.getName() + key.targetConstructorClass.getName();
+                return key.declaringClassId + " " + key.targetConstructorClassId;
             }
         }
         throw VMError.shouldNotReachHere("No constructor accessor uses the class loader %s", classLoader);
@@ -234,16 +226,33 @@ public void registerSerializationTargetClass(AccessCondition cnd, DynamicHub hub
     }
 
     public void replaceHubKeyWithTypeID() {
-        EconomicMap<Integer, RuntimeDynamicAccessMetadata> newEntries = EconomicMap.create();
-        var cursor = classes.getEntries();
+        replaceHubKeyWithTypeID(classes, key -> key instanceof DynamicHubKey, SerializationSupport::getTypeID);
+        replaceHubKeyWithTypeID(constructorAccessors, SerializationSupport::shouldReplaceSerializationLookupKey, SerializationSupport::replaceSerializationLookupKey);
+    }
+
+    private static <T, U> void replaceHubKeyWithTypeID(EconomicMap<T, U> map, Predicate<T> condition, Function<T, T> converter) {
+        EconomicMap<T, U> newEntries = EconomicMap.create();
+        var cursor = map.getEntries();
         while (cursor.advance()) {
-            Object key = cursor.getKey();
-            if (key instanceof DynamicHubKey hubKey) {
-                newEntries.put(hubKey.getTypeID(), cursor.getValue());
+            T key = cursor.getKey();
+            if (condition.test(key)) {
+                newEntries.put(converter.apply(key), cursor.getValue());
                 cursor.remove();
             }
         }
-        classes.putAll(newEntries);
+        map.putAll(newEntries);
+    }
+
+    private static boolean shouldReplaceSerializationLookupKey(SerializationLookupKey key) {
+        return key.declaringClassId() instanceof DynamicHubKey && key.targetConstructorClassId() instanceof DynamicHubKey;
+    }
+
+    private static SerializationLookupKey replaceSerializationLookupKey(SerializationLookupKey key) {
+        return new SerializationLookupKey(getTypeID(key.declaringClassId()), getTypeID(key.targetConstructorClassId()));
+    }
+
+    private static int getTypeID(Object classId) {
+        return ((DynamicHubKey) classId).getTypeID();
     }
 
     @Platforms(Platform.HOSTED_ONLY.class)
@@ -261,28 +270,24 @@ public boolean isLambdaCapturingClassRegistered(String lambdaCapturingClass) {
         return lambdaCapturingClasses.containsKey(lambdaCapturingClass);
     }
 
-    public static Object getSerializationConstructorAccessor(Class<?> serializationTargetClass, Class<?> targetConstructorClass) {
+    public static Object getRuntimeSerializationConstructorAccessor(Class<?> serializationTargetClass, Class<?> targetConstructorClass) {
+        SubstrateUtil.guaranteeRuntimeOnly();
         Class<?> declaringClass = serializationTargetClass;
 
         if (LambdaUtils.isLambdaClass(declaringClass)) {
             declaringClass = SerializedLambda.class;
         }
 
-        if (SubstrateUtil.HOSTED) {
-            Object constructorAccessor = currentLayer().getSerializationConstructorAccessor0(declaringClass, targetConstructorClass);
+        if (MetadataTracer.enabled()) {
+            MetadataTracer.singleton().traceSerializationType(declaringClass);
+        }
+        for (var singleton : layeredSingletons()) {
+            DynamicHub declaringHub = SubstrateUtil.cast(declaringClass, DynamicHub.class);
+            DynamicHub targetConstructorHub = SubstrateUtil.cast(targetConstructorClass, DynamicHub.class);
+            Object constructorAccessor = singleton.getSerializationConstructorAccessor0(declaringHub, targetConstructorHub, declaringClass.getModifiers());
             if (constructorAccessor != null) {
                 return constructorAccessor;
             }
-        } else {
-            if (MetadataTracer.enabled()) {
-                MetadataTracer.singleton().traceSerializationType(declaringClass);
-            }
-            for (var singleton : layeredSingletons()) {
-                Object constructorAccessor = singleton.getSerializationConstructorAccessor0(declaringClass, targetConstructorClass);
-                if (constructorAccessor != null) {
-                    return constructorAccessor;
-                }
-            }
         }
 
         String targetConstructorClassName = targetConstructorClass.getName();
@@ -291,23 +296,42 @@ public static Object getSerializationConstructorAccessor(Class<?> serializationT
         return null;
     }
 
-    @Override
-    public Object getSerializationConstructorAccessor0(Class<?> declaringClass, Class<?> rawTargetConstructorClass) {
-        VMError.guarantee(stubConstructor != null, "Called too early, no stub constructor yet.");
-        Class<?> targetConstructorClass = Modifier.isAbstract(declaringClass.getModifiers()) ? stubConstructor.getDeclaringClass() : rawTargetConstructorClass;
-        return constructorAccessors.get(new SerializationLookupKey(declaringClass, targetConstructorClass));
+    @Platforms(Platform.HOSTED_ONLY.class)
+    public static Object getHostedSerializationConstructorAccessor(DynamicHub serializationTargetClass, DynamicHub targetConstructorClass) {
+        SerializationSupport serializationSupport = currentLayer();
+        DynamicHub declaringClass = serializationTargetClass;
+
+        if (LambdaUtils.isLambdaClass(declaringClass.getHostedJavaClass())) {
+            declaringClass = serializationSupport.getSerializedLambdaClass();
+        }
+
+        VMError.guarantee(BuildPhaseProvider.isHostedUniverseBuilt(), "Called too early, hosted universe was not built yet.");
+        Object constructorAccessor = serializationSupport.getSerializationConstructorAccessor0(declaringClass, targetConstructorClass, declaringClass.getModifiers());
+        if (constructorAccessor != null) {
+            return constructorAccessor;
+        }
+
+        String targetConstructorClassName = targetConstructorClass.getName();
+        MissingSerializationRegistrationUtils.reportSerialization(declaringClass.getHostedJavaClass(),
+                        "type '" + declaringClass.getTypeName() + "' with target constructor class '" + targetConstructorClassName + "'");
+        return null;
+    }
+
+    public Object getSerializationConstructorAccessor0(DynamicHub declaringHub, DynamicHub rawTargetConstructorHub, int modifiers) {
+        VMError.guarantee(stubConstructorClass != null, "Called too early, no stub constructor yet.");
+        DynamicHub targetConstructorHub = Modifier.isAbstract(modifiers) ? stubConstructorClass : rawTargetConstructorHub;
+        return constructorAccessors.get(new SerializationLookupKey(declaringHub.getTypeID(), targetConstructorHub.getTypeID()));
     }
 
     public static boolean isRegisteredForSerialization(DynamicHub hub) {
-        for (SerializationRegistry singleton : SerializationSupport.layeredSingletons()) {
+        for (SerializationSupport singleton : SerializationSupport.layeredSingletons()) {
             if (singleton.isRegisteredForSerialization0(hub)) {
                 return true;
             }
         }
         return false;
     }
 
-    @Override
     public boolean isRegisteredForSerialization0(DynamicHub dynamicHub) {
         var conditionSet = classes.get(dynamicHub.getTypeID());
         return conditionSet != null && conditionSet.satisfied();

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/reflect/serialize/SerializationSupport.java

@@ -55,8 +55,8 @@ struct PersistedAnalysisType {
   wrappedType :union {
     none @31 :Void; # default
     serializationGenerated :group {
-      rawDeclaringClass @32 :Text;
-      rawTargetConstructor @33 :Text;
+      rawDeclaringClassId @32 :TypeId;
+      rawTargetConstructorId @33 :TypeId;
     }
     lambda :group {
       capturingClass @34 :Text;

substratevm/src/com.oracle.svm.hosted/resources/SharedLayerSnapshotCapnProtoSchema.capnp

@@ -92,7 +92,6 @@
 import com.oracle.svm.core.meta.MethodRef;
 import com.oracle.svm.core.reflect.proxy.DynamicProxySupport;
 import com.oracle.svm.core.reflect.serialize.SerializationSupport;
-import com.oracle.svm.shared.util.VMError;
 import com.oracle.svm.hosted.FeatureImpl;
 import com.oracle.svm.hosted.SVMHost;
 import com.oracle.svm.hosted.classinitialization.ClassInitializationSupport;
@@ -128,6 +127,7 @@
 import com.oracle.svm.shaded.org.capnproto.PrimitiveList;
 import com.oracle.svm.shaded.org.capnproto.StructList;
 import com.oracle.svm.shaded.org.capnproto.Text;
+import com.oracle.svm.shared.util.VMError;
 import com.oracle.svm.util.AnnotationUtil;
 import com.oracle.svm.util.GuestAccess;
 import com.oracle.svm.util.JVMCIReflectionUtil;
@@ -446,14 +446,18 @@ protected boolean delegateLoadType(PersistedAnalysisType.Reader typeData) {
         }
         if (wrappedType.isSerializationGenerated()) {
             SerializationGenerated.Reader sg = wrappedType.getSerializationGenerated();
-            String rawDeclaringClassName = sg.getRawDeclaringClass().toString();
-            String rawTargetConstructorClassName = sg.getRawTargetConstructor().toString();
-            Class<?> rawDeclaringClass = imageLayerBuildingSupport.lookupClass(false, rawDeclaringClassName);
-            Class<?> rawTargetConstructorClass = imageLayerBuildingSupport.lookupClass(false, rawTargetConstructorClassName);
+            int rawDeclaringClassId = sg.getRawDeclaringClassId();
+            int rawTargetConstructorClassId = sg.getRawTargetConstructorId();
+            AnalysisType rawDeclaringType = getAnalysisTypeForBaseLayerId(rawDeclaringClassId);
+            AnalysisType rawTargetConstructorType = getAnalysisTypeForBaseLayerId(rawTargetConstructorClassId);
+            Class<?> rawDeclaringClass = rawDeclaringType.getJavaClass();
+            Class<?> rawTargetConstructorClass = rawTargetConstructorType.getJavaClass();
             Constructor<?> rawTargetConstructor = ReflectionUtil.lookupConstructor(rawTargetConstructorClass);
             Constructor<?> constructor = ReflectionFactory.getReflectionFactory().newConstructorForSerialization(rawDeclaringClass, rawTargetConstructor);
-            SerializationSupport.currentLayer().addConstructorAccessor(rawDeclaringClass, rawTargetConstructorClass, SerializationFeature.getConstructorAccessor(constructor));
-            Class<?> constructorAccessor = SerializationSupport.getSerializationConstructorAccessor(rawDeclaringClass, rawTargetConstructorClass).getClass();
+            DynamicHub rawDeclaringHub = typeToHub(rawDeclaringType);
+            DynamicHub rawTargetConstructorHub = typeToHub(rawTargetConstructorType);
+            SerializationSupport.currentLayer().addConstructorAccessor(rawDeclaringHub, rawTargetConstructorHub, SerializationFeature.getConstructorAccessor(constructor));
+            Class<?> constructorAccessor = SerializationSupport.getHostedSerializationConstructorAccessor(rawDeclaringHub, rawTargetConstructorHub).getClass();
             metaAccess.lookupJavaType(constructorAccessor);
             return true;
         } else if (wrappedType.isLambda()) {