[GR-64959] Support InitialLayerOnlyImageSingleton with different key

PullRequest: graal/20887

Author: Zeavee

substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/layeredimagesingleton/LayeredImageSingletonSupport.java

@@ -53,5 +53,7 @@ static LayeredImageSingletonSupport singleton() {
 
     void freezeLayeredImageSingletonMetadata();
 
+    boolean isInitialLayerOnlyImageSingleton(Class<?> key);
+
     JavaConstant getInitialLayerOnlyImageSingleton(Class<?> key);
 }

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

@@ -230,6 +230,7 @@ struct ImageSingletonKey {
   persistFlag @1 :Int32;
   objectId @2 :SingletonObjId;
   constantId @3 :ConstantId;
+  isInitialLayerOnly @4 :Bool;
 }
 
 struct ImageSingletonObject {

substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/ImageSingletonsSupportImpl.java

@@ -82,6 +82,12 @@ public void freezeLayeredImageSingletonMetadata() {
         HostedManagement.getAndAssertExists().freezeLayeredImageSingletonMetadata();
     }
 
+    @Override
+    public boolean isInitialLayerOnlyImageSingleton(Class<?> key) {
+        var loader = HostedImageLayerBuildingSupport.singleton().getSingletonLoader();
+        return loader.isInitialLayerOnlyImageSingleton(key);
+    }
+
     @Override
     public JavaConstant getInitialLayerOnlyImageSingleton(Class<?> key) {
         var loader = HostedImageLayerBuildingSupport.singleton().getSingletonLoader();

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

@@ -36,7 +36,6 @@
 import org.graalvm.collections.UnmodifiableEconomicMap;
 
 import com.oracle.svm.core.layeredimagesingleton.ImageSingletonLoader;
-import com.oracle.svm.core.layeredimagesingleton.InitialLayerOnlyImageSingleton;
 import com.oracle.svm.core.layeredimagesingleton.LayeredImageSingleton.PersistFlags;
 import com.oracle.svm.core.util.UserError;
 import com.oracle.svm.core.util.VMError;
@@ -107,11 +106,9 @@ public Map<Object, Set<Class<?>>> loadImageSingletons(Object forbiddenObject) {
                 Class<?> clazz = imageLayerBuildingSupport.lookupClass(false, className);
                 singletonInitializationMap.computeIfAbsent(forbiddenObject, (k) -> new HashSet<>());
                 singletonInitializationMap.get(forbiddenObject).add(clazz);
-                if (InitialLayerOnlyImageSingleton.class.isAssignableFrom(clazz)) {
+                if (entry.getIsInitialLayerOnly()) {
                     int constantId = entry.getConstantId();
-                    if (constantId != -1) {
-                        initialLayerKeyToIdMap.put(clazz, constantId);
-                    }
+                    initialLayerKeyToIdMap.put(clazz, constantId);
                 }
             } else {
                 assert persistInfo == PersistFlags.NOTHING : "Unexpected PersistFlags value: " + persistInfo;
@@ -124,6 +121,10 @@ public Map<Object, Set<Class<?>>> loadImageSingletons(Object forbiddenObject) {
         return singletonInitializationMap;
     }
 
+    public boolean isInitialLayerOnlyImageSingleton(Class<?> key) {
+        return initialLayerOnlySingletonConstantIds.containsKey(key);
+    }
+
     public JavaConstant loadInitialLayerOnlyImageSingleton(Class<?> key) {
         int constantId = initialLayerOnlySingletonConstantIds.getOrDefault(key, -1);
         if (constantId != -1) {

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

@@ -1117,6 +1117,7 @@ public void writeImageSingletonInfo(List<Map.Entry<Class<?>, Object>> layeredIma
                 constantId = initialLayerOnlySingletonMap.getOrDefault(initialLayerOnlyImageSingleton, -1);
             }
             sb.setConstantId(constantId);
+            sb.setIsInitialLayerOnly(singleton instanceof InitialLayerOnlyImageSingleton);
         }
 
         var sortedByIDs = singletonInfoMap.entrySet().stream()

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

@@ -3948,6 +3948,13 @@ public final void setConstantId(int value) {
         _setIntField(2, value);
       }
 
+      public final boolean getIsInitialLayerOnly() {
+        return _getBooleanField(96);
+      }
+      public final void setIsInitialLayerOnly(boolean value) {
+        _setBooleanField(96, value);
+      }
+
     }
 
     public static final class Reader extends com.oracle.svm.shaded.org.capnproto.StructReader {
@@ -3974,6 +3981,10 @@ public final int getConstantId() {
         return _getIntField(2);
       }
 
+      public final boolean getIsInitialLayerOnly() {
+        return _getBooleanField(96);
+      }
+
     }
 
   }

substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/snippets/SubstrateGraphBuilderPlugins.java

@@ -88,7 +88,6 @@
 import com.oracle.svm.core.imagelayer.LoadImageSingletonFactory;
 import com.oracle.svm.core.jdk.proxy.DynamicProxyRegistry;
 import com.oracle.svm.core.layeredimagesingleton.ApplicationLayerOnlyImageSingleton;
-import com.oracle.svm.core.layeredimagesingleton.InitialLayerOnlyImageSingleton;
 import com.oracle.svm.core.layeredimagesingleton.LayeredImageSingletonBuilderFlags;
 import com.oracle.svm.core.layeredimagesingleton.LayeredImageSingletonSupport;
 import com.oracle.svm.core.layeredimagesingleton.MultiLayeredImageSingleton;
@@ -1204,27 +1203,30 @@ public boolean apply(GraphBuilderContext b, ResolvedJavaMethod targetMethod, Rec
             public boolean apply(GraphBuilderContext b, ResolvedJavaMethod targetMethod, Receiver unused, ValueNode classNode) {
                 Class<?> key = constantObjectParameter(b, targetMethod, 0, Class.class, classNode);
 
-                if (ApplicationLayerOnlyImageSingleton.isAssignableFrom(key) &&
-                                ImageLayerBuildingSupport.buildingSharedLayer()) {
-                    /*
-                     * This singleton is only installed in the application layer heap. All other
-                     * layers looks refer to this singleton.
-                     */
-                    b.addPush(JavaKind.Object, LoadImageSingletonFactory.loadApplicationOnlyImageSingleton(key, b.getMetaAccess()));
-                    return true;
-                }
+                var layeredSingletonSupport = LayeredImageSingletonSupport.singleton();
+                if (ImageLayerBuildingSupport.buildingImageLayer()) {
+                    if (ApplicationLayerOnlyImageSingleton.isAssignableFrom(key) &&
+                                    ImageLayerBuildingSupport.buildingSharedLayer()) {
+                        /*
+                         * This singleton is only installed in the application layer heap. All other
+                         * layers looks refer to this singleton.
+                         */
+                        b.addPush(JavaKind.Object, LoadImageSingletonFactory.loadApplicationOnlyImageSingleton(key, b.getMetaAccess()));
+                        return true;
+                    }
 
-                if (InitialLayerOnlyImageSingleton.class.isAssignableFrom(key) && ImageLayerBuildingSupport.buildingExtensionLayer()) {
-                    /*
-                     * This singleton is only installed in the initial layer heap. When allowed, all
-                     * other layers lookups refer to this singleton.
-                     */
-                    JavaConstant initialSingleton = LayeredImageSingletonSupport.singleton().getInitialLayerOnlyImageSingleton(key);
-                    b.addPush(JavaKind.Object, ConstantNode.forConstant(initialSingleton, b.getMetaAccess(), b.getGraph()));
-                    return true;
+                    if (ImageLayerBuildingSupport.buildingExtensionLayer() && layeredSingletonSupport.isInitialLayerOnlyImageSingleton(key)) {
+                        /*
+                         * This singleton is only installed in the initial layer heap. When allowed,
+                         * all other layers lookups refer to this singleton.
+                         */
+                        JavaConstant initialSingleton = layeredSingletonSupport.getInitialLayerOnlyImageSingleton(key);
+                        b.addPush(JavaKind.Object, ConstantNode.forConstant(initialSingleton, b.getMetaAccess(), b.getGraph()));
+                        return true;
+                    }
                 }
 
-                Object singleton = LayeredImageSingletonSupport.singleton().lookup(key, true, true);
+                Object singleton = layeredSingletonSupport.lookup(key, true, true);
                 LayeredImageSingletonBuilderFlags.validateRuntimeLookup(singleton);
                 b.addPush(JavaKind.Object, ConstantNode.forConstant(b.getSnippetReflection().forObject(singleton), b.getMetaAccess()));
                 return true;