prototype full support for InitialLayerOnly image singletons when the key differs from the implementation

teshull · Zeavee · commit ac1398cb8150 · 2025-06-03T20:50:35.000+02:00
diff --git a/substratevm/src/com.oracle.svm.core/src/com/oracle/svm/core/layeredimagesingleton/LayeredImageSingletonSupport.java b/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);
 }
diff --git a/substratevm/src/com.oracle.svm.hosted/resources/SharedLayerSnapshotCapnProtoSchema.capnp b/substratevm/src/com.oracle.svm.hosted/resources/SharedLayerSnapshotCapnProtoSchema.capnp
@@ -228,6 +228,7 @@ struct ImageSingletonKey {
   persistFlag @1 :Int32;
   objectId @2 :SingletonObjId;
   constantId @3 :ConstantId;
+  isInitialLayerOnly @4 :Bool;
 }
 
 struct ImageSingletonObject {
diff --git a/substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/ImageSingletonsSupportImpl.java b/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();
diff --git a/substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/imagelayer/SVMImageLayerSingletonLoader.java b/substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/imagelayer/SVMImageLayerSingletonLoader.java
@@ -36,18 +36,17 @@
 import java.util.stream.Stream;
 import java.util.stream.StreamSupport;
 
-import com.oracle.svm.shaded.org.capnproto.Text;
 import org.graalvm.collections.EconomicMap;
 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;
 import com.oracle.svm.hosted.imagelayer.SharedLayerSnapshotCapnProtoSchemaHolder.ImageSingletonKey;
 import com.oracle.svm.hosted.imagelayer.SharedLayerSnapshotCapnProtoSchemaHolder.ImageSingletonObject;
 import com.oracle.svm.hosted.imagelayer.SharedLayerSnapshotCapnProtoSchemaHolder.KeyStoreEntry;
+import com.oracle.svm.shaded.org.capnproto.Text;
 import com.oracle.svm.util.ReflectionUtil;
 
 import jdk.vm.ci.meta.JavaConstant;
@@ -112,11 +111,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;
@@ -129,6 +126,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) {
diff --git a/substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/imagelayer/SVMImageLayerWriter.java b/substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/imagelayer/SVMImageLayerWriter.java
@@ -1131,6 +1131,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()
diff --git a/substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/imagelayer/SharedLayerSnapshotCapnProtoSchemaHolder.java b/substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/imagelayer/SharedLayerSnapshotCapnProtoSchemaHolder.java
@@ -3937,6 +3937,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 {
@@ -3963,6 +3970,10 @@ public final int getConstantId() {
         return _getIntField(2);
       }
 
+      public final boolean getIsInitialLayerOnly() {
+        return _getBooleanField(96);
+      }
+
     }
 
   }
diff --git a/substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/snippets/SubstrateGraphBuilderPlugins.java b/substratevm/src/com.oracle.svm.hosted/src/com/oracle/svm/hosted/snippets/SubstrateGraphBuilderPlugins.java
@@ -91,7 +91,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.LayeredImageSingleton;
 import com.oracle.svm.core.layeredimagesingleton.LayeredImageSingletonBuilderFlags;
 import com.oracle.svm.core.layeredimagesingleton.LayeredImageSingletonSupport;
@@ -1202,27 +1201,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);
                 if (singleton instanceof LayeredImageSingleton layeredSingleton) {
                     if (!layeredSingleton.getImageBuilderFlags().contains(LayeredImageSingletonBuilderFlags.RUNTIME_ACCESS)) {
                         /*

Original file line number	Diff line number	Diff line change
`@@ -53,5 +53,7 @@ static LayeredImageSingletonSupport singleton() {`
`53`	`53`
`54`	`54`	`void freezeLayeredImageSingletonMetadata();`
`55`	`55`
	`56`	`+ boolean isInitialLayerOnlyImageSingleton(Class<?> key);`
	`57`	`+`
`56`	`58`	`JavaConstant getInitialLayerOnlyImageSingleton(Class<?> key);`
`57`	`59`	`}`
Original file line number	Diff line number	Diff line change
`@@ -228,6 +228,7 @@ struct ImageSingletonKey {`
`228`	`228`	`persistFlag @1 :Int32;`
`229`	`229`	`objectId @2 :SingletonObjId;`
`230`	`230`	`constantId @3 :ConstantId;`
	`231`	`+ isInitialLayerOnly @4 :Bool;`
`231`	`232`	`}`
`232`	`233`
`233`	`234`	`struct ImageSingletonObject {`
Original file line number	Diff line number	Diff line change
`@@ -1131,6 +1131,7 @@ public void writeImageSingletonInfo(List<Map.Entry<Class<?>, Object>> layeredIma`
`1131`	`1131`	`constantId = initialLayerOnlySingletonMap.getOrDefault(initialLayerOnlyImageSingleton, -1);`
`1132`	`1132`	`}`
`1133`	`1133`	`sb.setConstantId(constantId);`
	`1134`	`+ sb.setIsInitialLayerOnly(singleton instanceof InitialLayerOnlyImageSingleton);`
`1134`	`1135`	`}`
`1135`	`1136`
`1136`	`1137`	`var sortedByIDs = singletonInfoMap.entrySet().stream()`
Original file line number	Diff line number	Diff line change
`@@ -3937,6 +3937,13 @@ public final void setConstantId(int value) {`
`3937`	`3937`	`_setIntField(2, value);`
`3938`	`3938`	`}`
`3939`	`3939`
	`3940`	`+ public final boolean getIsInitialLayerOnly() {`
	`3941`	`+ return _getBooleanField(96);`
	`3942`	`+ }`
	`3943`	`+ public final void setIsInitialLayerOnly(boolean value) {`
	`3944`	`+ _setBooleanField(96, value);`
	`3945`	`+ }`
	`3946`	`+`
`3940`	`3947`	`}`
`3941`	`3948`
`3942`	`3949`	`public static final class Reader extends com.oracle.svm.shaded.org.capnproto.StructReader {`
`@@ -3963,6 +3970,10 @@ public final int getConstantId() {`
`3963`	`3970`	`return _getIntField(2);`
`3964`	`3971`	`}`
`3965`	`3972`
	`3973`	`+ public final boolean getIsInitialLayerOnly() {`
	`3974`	`+ return _getBooleanField(96);`
	`3975`	`+ }`
	`3976`	`+`
`3966`	`3977`	`}`
`3967`	`3978`
`3968`	`3979`	`}`