Fixed native-image universe factory.

tzezula · tzezula · commit a94c18b1a6f5 · 2025-10-31T13:09:40.000+01:00
diff --git a/compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/truffle/PartialEvaluator.java b/compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/truffle/PartialEvaluator.java
@@ -26,6 +26,7 @@
 
 import java.net.URI;
 import java.util.Map;
+import java.util.function.Function;
 import java.util.function.Supplier;
 
 import jdk.graal.compiler.annotation.EnumElement;
@@ -181,48 +182,85 @@ private static int actualStableDimensions(ResolvedJavaField field, int dimension
      * Gets an object describing how a read of {@code field} can be constant folded based on Truffle
      * annotations.
      *
+     * @param field the field for which to compute {@link ConstantFieldInfo}
+     * @param declaredAnnotationValues a map of method annotations keyed by their declaring
+     *            {@link ResolvedJavaType}
+     * @param types cached types known to the Truffle compiler
+     * @param unwrapType a function that unwraps a {@link ResolvedJavaType} obtained from
+     *            {@code types} to its underlying host {@link ResolvedJavaType}. In native-image,
+     *            known Truffle types are represented as {@code AnalysisType}, while the keys in
+     *            {@code declaredAnnotationValues} correspond to host JVMCI types. Therefore, each
+     *            {@code AnalysisType} must be unwrapped via
+     *            {@code OriginalClassProvider.getOriginalType(JavaType)} to correctly access the
+     *            corresponding annotation entries in {@code declaredAnnotationValues}.
+     *
      * @return {@code null} if there are no constant folding related Truffle annotations on
      *         {@code field}
      */
     public static ConstantFieldInfo computeConstantFieldInfo(ResolvedJavaField field,
                     Map<ResolvedJavaType, AnnotationValue> declaredAnnotationValues,
-                    KnownTruffleTypes types) {
-        if (declaredAnnotationValues.containsKey(types.Node_Child)) {
+                    KnownTruffleTypes types, Function<ResolvedJavaType, ResolvedJavaType> unwrapType) {
+        if (declaredAnnotationValues.containsKey(unwrapType.apply(types.Node_Child))) {
             return ConstantFieldInfo.CHILD;
         }
-        if (declaredAnnotationValues.containsKey(types.Node_Children)) {
+        if (declaredAnnotationValues.containsKey(unwrapType.apply(types.Node_Children))) {
             return ConstantFieldInfo.CHILDREN;
         }
-        AnnotationValue cf = declaredAnnotationValues.get(types.CompilerDirectives_CompilationFinal);
+        AnnotationValue cf = declaredAnnotationValues.get(unwrapType.apply(types.CompilerDirectives_CompilationFinal));
         if (cf != null) {
             int dimensions = actualStableDimensions(field, cf.get("dimensions", Integer.class));
             return ConstantFieldInfo.forDimensions(dimensions);
         }
         return null;
     }
 
+    /**
+     * Computes {@link PartialEvaluationMethodInfo} for the given {@code method} based on its
+     * declared annotations.
+     *
+     * <p>
+     * This method analyzes the annotations present on {@code method}, provided in
+     * {@code declaredAnnotationValues}, and determines partial evaluation behavior such as loop
+     * explosion, inlining, and specialization characteristics.
+     * </p>
+     *
+     * @param runtime the Truffle runtime used to determine whether
+     *            {@linkplain TruffleCompilerRuntime#isJavaInstrumentationActive() flight recorder
+     *            instrumentation is active}
+     * @param method the method for which to compute {@link PartialEvaluationMethodInfo}
+     * @param declaredAnnotationValues a map of method annotations keyed by their declaring
+     *            {@link ResolvedJavaType}
+     * @param types cached types known to the Truffle compiler
+     * @param unwrapType a function that unwraps a {@link ResolvedJavaType} obtained from
+     *            {@code types} to its underlying host {@link ResolvedJavaType}. In native-image,
+     *            known Truffle types are represented as {@code AnalysisType}, while the keys in
+     *            {@code declaredAnnotationValues} correspond to host JVMCI types. Therefore, each
+     *            {@code AnalysisType} must be unwrapped via
+     *            {@code OriginalClassProvider.getOriginalType(JavaType)} to correctly access the
+     *            corresponding annotation entries in {@code declaredAnnotationValues}.
+     */
     public static PartialEvaluationMethodInfo computePartialEvaluationMethodInfo(TruffleCompilerRuntime runtime,
                     ResolvedJavaMethod method,
                     Map<ResolvedJavaType, AnnotationValue> declaredAnnotationValues,
-                    KnownTruffleTypes types) {
-        AnnotationValue truffleBoundary = declaredAnnotationValues.get(types.CompilerDirectives_TruffleBoundary);
-        AnnotationValue truffleCallBoundary = declaredAnnotationValues.get(types.TruffleCallBoundary);
-        boolean specialization = declaredAnnotationValues.get(types.Specialization) != null;
-        AnnotationValue explodeLoop = declaredAnnotationValues.get(types.ExplodeLoop);
+                    KnownTruffleTypes types, Function<ResolvedJavaType, ResolvedJavaType> unwrapType) {
+        AnnotationValue truffleBoundary = declaredAnnotationValues.get(unwrapType.apply(types.CompilerDirectives_TruffleBoundary));
+        AnnotationValue truffleCallBoundary = declaredAnnotationValues.get(unwrapType.apply(types.TruffleCallBoundary));
+        boolean specialization = declaredAnnotationValues.get(unwrapType.apply(types.Specialization)) != null;
+        AnnotationValue explodeLoop = declaredAnnotationValues.get(unwrapType.apply(types.ExplodeLoop));
 
-        return new PartialEvaluationMethodInfo(getLoopExplosionKind(explodeLoop, types),
+        return new PartialEvaluationMethodInfo(getLoopExplosionKind(explodeLoop, unwrapType.apply(types.ExplodeLoop_LoopExplosionKind)),
                         getInlineKind(runtime, truffleBoundary, truffleCallBoundary, method, true, types),
                         getInlineKind(runtime, truffleBoundary, truffleCallBoundary, method, false, types),
                         method.canBeInlined(),
                         specialization);
     }
 
-    private static LoopExplosionKind getLoopExplosionKind(AnnotationValue value, KnownTruffleTypes types) {
+    private static LoopExplosionKind getLoopExplosionKind(AnnotationValue value, ResolvedJavaType expectedType) {
         if (value == null) {
             return LoopExplosionKind.NONE;
         }
         EnumElement enumElement = value.get("kind", EnumElement.class);
-        if (!types.ExplodeLoop_LoopExplosionKind.equals(enumElement.enumType)) {
+        if (!expectedType.equals(enumElement.enumType)) {
             throw new IllegalStateException("Incompatible ExplodeLoop change. ExplodeLoop.kind must be LoopExplosionKind.");
         }
         return Enum.valueOf(LoopExplosionKind.class, enumElement.name);
diff --git a/compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/truffle/hotspot/HotSpotPartialEvaluator.java b/compiler/src/jdk.graal.compiler/src/jdk/graal/compiler/truffle/hotspot/HotSpotPartialEvaluator.java
@@ -28,6 +28,7 @@
 import java.util.concurrent.ConcurrentHashMap;
 import java.util.concurrent.ConcurrentMap;
 import java.util.concurrent.atomic.AtomicReference;
+import java.util.function.Function;
 import java.util.function.Supplier;
 
 import org.graalvm.collections.EconomicMap;
@@ -181,7 +182,7 @@ protected Object createKey(ResolvedJavaMethod method) {
         @Override
         protected PartialEvaluationMethodInfo computeValue(ResolvedJavaMethod method) {
             Map<ResolvedJavaType, AnnotationValue> declaredAnnotationValues = AnnotationValueSupport.getDeclaredAnnotationValues(method);
-            PartialEvaluationMethodInfo methodInfo = computePartialEvaluationMethodInfo(config.runtime(), method, declaredAnnotationValues, getTypes());
+            PartialEvaluationMethodInfo methodInfo = computePartialEvaluationMethodInfo(config.runtime(), method, declaredAnnotationValues, getTypes(), Function.identity());
             /*
              * We can canonicalize the instances to reduce space required in the cache. There are
              * only a small number of possible instances of PartialEvaluationMethodInfo as it just
@@ -206,7 +207,7 @@ protected Object createKey(ResolvedJavaField field) {
         @Override
         protected ConstantFieldInfo computeValue(ResolvedJavaField field) {
             Map<ResolvedJavaType, AnnotationValue> declaredAnnotationValues = AnnotationValueSupport.getDeclaredAnnotationValues(field);
-            return computeConstantFieldInfo(field, declaredAnnotationValues, getTypes());
+            return computeConstantFieldInfo(field, declaredAnnotationValues, getTypes(), Function.identity());
         }
     }
 }
diff --git a/substratevm/src/com.oracle.svm.graal/src/com/oracle/svm/graal/hosted/runtimecompilation/RuntimeCompilationFeature.java b/substratevm/src/com.oracle.svm.graal/src/com/oracle/svm/graal/hosted/runtimecompilation/RuntimeCompilationFeature.java
@@ -31,6 +31,7 @@
 import static jdk.graal.compiler.java.BytecodeParserOptions.InlineDuringParsingMaxDepth;
 
 import java.lang.reflect.Executable;
+import java.util.ArrayList;
 import java.util.Arrays;
 import java.util.Collection;
 import java.util.List;
@@ -236,6 +237,7 @@ enum InlineDecision {
     private Function<ConstantFieldProvider, ConstantFieldProvider> constantFieldProviderWrapper = Function.identity();
     private Consumer<CallTreeInfo> blocklistChecker = _ -> {
     };
+    private List<Runnable> afterInstallRuntimeConfig;
 
     public HostedProviders getHostedProviders() {
         return hostedProviders;
@@ -249,6 +251,27 @@ public void setUniverseFactory(SubstrateUniverseFactory universeFactory) {
         this.universeFactory = universeFactory;
     }
 
+    /**
+     * Registers a {@code callback} to be invoked after
+     * {@link #beforeAnalysis(BeforeAnalysisAccess)} completes the
+     * {@link #installRuntimeConfig(BeforeAnalysisAccessImpl) runtime configuration setup}.
+     *
+     * <p>
+     * This hook is used by {@code TruffleFeature} to initialize {@code KnownTruffleTypes}, which
+     * are required by {@code SubstrateTruffleUniverseFactory}. The {@code KnownTruffleTypes} must
+     * be initialized after {@link #installRuntimeConfig(BeforeAnalysisAccessImpl)} has executed,
+     * but before the remaining logic in {@link #beforeAnalysis(BeforeAnalysisAccess)}, since that
+     * phase already relies on {@link SubstrateUniverseFactory} to create {@link SubstrateMethod}
+     * instances.
+     * </p>
+     */
+    public synchronized void addAfterInstallRuntimeConfigCallback(Runnable callback) {
+        if (afterInstallRuntimeConfig == null) {
+            afterInstallRuntimeConfig = new ArrayList<>();
+        }
+        afterInstallRuntimeConfig.add(callback);
+    }
+
     @SuppressWarnings("unused")
     private static boolean defaultAllowRuntimeCompilation(ResolvedJavaMethod method) {
         return false;
@@ -432,6 +455,18 @@ public void beforeAnalysis(BeforeAnalysisAccess c) {
         BeforeAnalysisAccessImpl config = (BeforeAnalysisAccessImpl) c;
         installRuntimeConfig(config);
 
+        List<Runnable> callbacks;
+        synchronized (this) {
+            if (afterInstallRuntimeConfig != null) {
+                callbacks = List.copyOf(afterInstallRuntimeConfig);
+            } else {
+                callbacks = List.of();
+            }
+        }
+        for (Runnable callback : callbacks) {
+            callback.run();
+        }
+
         SubstrateGraalRuntime graalRuntime = new SubstrateGraalRuntime();
         objectReplacer.setGraalRuntime(graalRuntime);
         objectReplacer.setAnalysisAccess(config);
diff --git a/substratevm/src/com.oracle.svm.truffle/src/com/oracle/svm/truffle/TruffleFeature.java b/substratevm/src/com.oracle.svm.truffle/src/com/oracle/svm/truffle/TruffleFeature.java
@@ -136,11 +136,13 @@
 import com.oracle.svm.hosted.FeatureImpl.DuringSetupAccessImpl;
 import com.oracle.svm.hosted.meta.HostedType;
 import com.oracle.svm.hosted.substitute.DeletedElementException;
+import com.oracle.svm.graal.TruffleRuntimeCompilationSupport;
 import com.oracle.svm.truffle.api.SubstrateThreadLocalHandshake;
 import com.oracle.svm.truffle.api.SubstrateThreadLocalHandshakeSnippets;
 import com.oracle.svm.truffle.api.SubstrateTruffleCompiler;
 import com.oracle.svm.truffle.api.SubstrateTruffleRuntime;
 import com.oracle.svm.truffle.api.SubstrateTruffleUniverseFactory;
+import com.oracle.svm.truffle.api.SubstrateKnownTruffleTypes;
 import com.oracle.svm.util.AnnotationUtil;
 import com.oracle.svm.util.LogUtils;
 import com.oracle.svm.util.ReflectionUtil;
@@ -233,6 +235,14 @@ public boolean getAsBoolean() {
     private final Set<ResolvedJavaMethod> warnMethods;
     private final Set<Pair<ResolvedJavaMethod, String>> neverPartOfCompilationViolations;
     Set<AnalysisMethod> runtimeCompiledMethods;
+    /**
+     * Holds the {@link KnownTruffleTypes} instance, which is initialized lazily after
+     * {@link RuntimeCompilationFeature#beforeAnalysis(BeforeAnalysisAccess)} completes runtime
+     * configuration setup.
+     *
+     * @see RuntimeCompilationFeature#addAfterInstallRuntimeConfigCallback(Runnable)
+     */
+    private volatile KnownTruffleTypes types;
 
     public TruffleFeature() {
         blocklistMethods = new HashSet<>();
@@ -241,6 +251,10 @@ public TruffleFeature() {
         neverPartOfCompilationViolations = ConcurrentHashMap.newKeySet();
     }
 
+    public KnownTruffleTypes getTypes() {
+        return types;
+    }
+
     @Override
     public void registerForeignCalls(SubstrateForeignCallsProvider foreignCalls) {
         foreignCalls.register(SubstrateThreadLocalHandshake.FOREIGN_POLL);
@@ -290,6 +304,12 @@ public void afterRegistration(AfterRegistrationAccess access) {
         UserError.guarantee(runtime instanceof SubstrateTruffleRuntime, "TruffleFeature requires SubstrateTruffleRuntime");
         SubstrateTruffleRuntime truffleRuntime = (SubstrateTruffleRuntime) Truffle.getRuntime();
         truffleRuntime.resetHosted();
+        RuntimeCompilationFeature runtimeCompilationFeature = RuntimeCompilationFeature.singleton();
+        runtimeCompilationFeature.addAfterInstallRuntimeConfigCallback(() -> {
+            Providers providers = TruffleRuntimeCompilationSupport.getRuntimeConfig().getProviders();
+            types = new SubstrateKnownTruffleTypes(truffleRuntime, providers.getMetaAccess(), providers.getConstantReflection());
+        });
+        runtimeCompilationFeature.setUniverseFactory(new SubstrateTruffleUniverseFactory(truffleRuntime, this));
     }
 
     @Override
@@ -382,10 +402,7 @@ public void beforeAnalysis(BeforeAnalysisAccess access) {
         truffleRuntime.initializeHostedKnownMethods(config.getUniverse().getOriginalMetaAccess());
 
         PartialEvaluator partialEvaluator = truffleCompiler.getPartialEvaluator();
-        KnownTruffleTypes types = partialEvaluator.getTypes();
         registerKnownTruffleFields(config, types);
-
-        runtimeCompilationFeature.setUniverseFactory(new SubstrateTruffleUniverseFactory(truffleRuntime, types));
         GraphBuilderConfiguration graphBuilderConfig = partialEvaluator.getGraphBuilderConfigPrototype();
 
         TruffleAllowInliningPredicate allowInliningPredicate = new TruffleAllowInliningPredicate(runtimeCompilationFeature.getHostedProviders().getReplacements(),
diff --git a/substratevm/src/com.oracle.svm.truffle/src/com/oracle/svm/truffle/TruffleSupport.java b/substratevm/src/com.oracle.svm.truffle/src/com/oracle/svm/truffle/TruffleSupport.java
@@ -28,6 +28,7 @@
 import java.util.function.Consumer;
 import java.util.function.Supplier;
 
+import jdk.graal.compiler.truffle.KnownTruffleTypes;
 import org.graalvm.nativeimage.ImageSingletons;
 
 import com.oracle.svm.core.SubstrateOptions;
@@ -36,7 +37,6 @@
 import com.oracle.svm.core.meta.SubstrateObjectConstant;
 import com.oracle.svm.graal.TruffleRuntimeCompilationSupport;
 import com.oracle.svm.graal.hosted.runtimecompilation.RuntimeCompilationFeature;
-import com.oracle.svm.truffle.api.SubstrateKnownTruffleTypes;
 import com.oracle.svm.truffle.api.SubstrateOptimizedCallTarget;
 import com.oracle.svm.truffle.api.SubstrateOptimizedCallTargetInstalledCode;
 import com.oracle.svm.truffle.api.SubstratePartialEvaluator;
@@ -57,7 +57,6 @@
 import com.oracle.truffle.runtime.OptimizedDirectCallNode;
 
 import jdk.graal.compiler.nodes.graphbuilderconf.GraphBuilderConfiguration;
-import jdk.graal.compiler.phases.util.Providers;
 import jdk.graal.compiler.truffle.EconomyPartialEvaluatorConfiguration;
 import jdk.graal.compiler.truffle.PartialEvaluatorConfiguration;
 import jdk.graal.compiler.truffle.TruffleCompilerConfiguration;
@@ -107,8 +106,7 @@ protected static TruffleCompilerConfiguration createSubstrateTruffleCompilerConf
         final GraphBuilderConfiguration.Plugins graphBuilderPlugins = runtimeCompilationFeature.getHostedProviders().getGraphBuilderPlugins();
         SubstrateBackend substrateBackend = TruffleRuntimeCompilationSupport.getRuntimeConfig().getBackendForNormalMethod();
         substrateBackend.setRuntimeToRuntimeInvokeMethod(optimizedCallTargetMethod);
-        Providers providers = TruffleRuntimeCompilationSupport.getRuntimeConfig().getProviders();
-        SubstrateKnownTruffleTypes types = new SubstrateKnownTruffleTypes(runtime, providers.getMetaAccess(), providers.getConstantReflection());
+        KnownTruffleTypes types = ImageSingletons.lookup(TruffleFeature.class).getTypes();
         final TruffleTierConfiguration firstTier = new TruffleTierConfiguration(new EconomyPartialEvaluatorConfiguration(), substrateBackend,
                         TruffleRuntimeCompilationSupport.getFirstTierProviders(), TruffleRuntimeCompilationSupport.getFirstTierSuites(), TruffleRuntimeCompilationSupport.getFirstTierLirSuites(),
                         types);
diff --git a/substratevm/src/com.oracle.svm.truffle/src/com/oracle/svm/truffle/api/SubstrateTruffleUniverseFactory.java b/substratevm/src/com.oracle.svm.truffle/src/com/oracle/svm/truffle/api/SubstrateTruffleUniverseFactory.java
diff --git a/truffle/src/com.oracle.truffle.runtime/src/com/oracle/truffle/runtime/OptimizedTruffleRuntime.java b/truffle/src/com.oracle.truffle.runtime/src/com/oracle/truffle/runtime/OptimizedTruffleRuntime.java
