From a6db476c70e049666bd6a406d520e53689625c69 Mon Sep 17 00:00:00 2001
From: Jen Basch <jbasch@apple.com>
Date: Mon, 23 Mar 2026 07:42:40 -0700
Subject: [PATCH] Fix module reflection when instrumentation is active (#1464)

---
 .../main/java/org/pkl/core/runtime/VmTyped.java   | 12 ++++++++----
 .../LanguageSnippetTests/input/api/reflect5.pkl   | 15 +++++++++++++++
 .../LanguageSnippetTests/output/api/reflect5.pcf  |  1 +
 3 files changed, 24 insertions(+), 4 deletions(-)
 create mode 100644 pkl-core/src/test/files/LanguageSnippetTests/input/api/reflect5.pkl
 create mode 100644 pkl-core/src/test/files/LanguageSnippetTests/output/api/reflect5.pcf

diff --git a/pkl-core/src/main/java/org/pkl/core/runtime/VmTyped.java b/pkl-core/src/main/java/org/pkl/core/runtime/VmTyped.java
index da768bb6..a9e70938 100644
--- a/pkl-core/src/main/java/org/pkl/core/runtime/VmTyped.java
+++ b/pkl-core/src/main/java/org/pkl/core/runtime/VmTyped.java
@@ -1,5 +1,5 @@
 /*
- * Copyright © 2024-2025 Apple Inc. and the Pkl project authors. All rights reserved.
+ * Copyright © 2024-2026 Apple Inc. and the Pkl project authors. All rights reserved.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -18,11 +18,13 @@ package org.pkl.core.runtime;
 import com.oracle.truffle.api.CompilerDirectives.CompilationFinal;
 import com.oracle.truffle.api.CompilerDirectives.TruffleBoundary;
 import com.oracle.truffle.api.frame.MaterializedFrame;
+import com.oracle.truffle.api.instrumentation.InstrumentableNode.WrapperNode;
 import org.graalvm.collections.EconomicMap;
 import org.graalvm.collections.UnmodifiableEconomicMap;
 import org.pkl.core.Composite;
 import org.pkl.core.PModule;
 import org.pkl.core.PObject;
+import org.pkl.core.ast.ExpressionNode;
 import org.pkl.core.ast.expression.unary.ImportNode;
 import org.pkl.core.ast.member.ObjectMember;
 import org.pkl.core.util.EconomicMaps;
@@ -103,9 +105,11 @@ public final class VmTyped extends VmObject {
       if (member.isImport()) {
         var memberNode = member.getMemberNode();
         assert memberNode != null; // import is never a constant
-        builder.add(
-            member.getName().toString(),
-            ((ImportNode) memberNode.getBodyNode()).getImportUri().toString());
+        var bodyNode = memberNode.getBodyNode();
+        if (bodyNode instanceof WrapperNode wrapper) {
+          bodyNode = (ExpressionNode) wrapper.getDelegateNode();
+        }
+        builder.add(member.getName().toString(), ((ImportNode) bodyNode).getImportUri().toString());
       }
     }
     return builder.build();
diff --git a/pkl-core/src/test/files/LanguageSnippetTests/input/api/reflect5.pkl b/pkl-core/src/test/files/LanguageSnippetTests/input/api/reflect5.pkl
new file mode 100644
index 00000000..fab288af
--- /dev/null
+++ b/pkl-core/src/test/files/LanguageSnippetTests/input/api/reflect5.pkl
@@ -0,0 +1,15 @@
+import "pkl:reflect"
+import "pkl:test"
+
+class Foo {
+  bar: String(startsWith("a"))
+}
+
+// regression test for VmTyped.getImports()
+// which walks module members that are imports and casts their bodies to ImportNode
+// the catch activates instrumentation by failing a constraint check
+// so that the reflected import member bodies are wrapper nodes instead of ImportNode directly
+// the fix for this unwraps the wrapper before casting to ImportNode
+foo =
+  let (_ = test.catchOrNull(() -> new Foo { bar = "baz" }.bar))
+    reflect.Module(module).imports
diff --git a/pkl-core/src/test/files/LanguageSnippetTests/output/api/reflect5.pcf b/pkl-core/src/test/files/LanguageSnippetTests/output/api/reflect5.pcf
new file mode 100644
index 00000000..be1e4a46
--- /dev/null
+++ b/pkl-core/src/test/files/LanguageSnippetTests/output/api/reflect5.pcf
@@ -0,0 +1 @@
+foo = Map("reflect", "pkl:reflect", "test", "pkl:test")