Optimize super property and super subscript nodes (#1853)

This adds cacheing to ReadSuper* nodes so that the same
node being executed can possibly re-use a previously cached value.

pkl-core/src/main/java/org/pkl/core/ast/builder/AstBuilder.java

@@ -125,8 +125,8 @@ import org.pkl.core.ast.expression.member.ReadAmbiguousLocalityPropertyNode;
 import org.pkl.core.ast.expression.member.ReadLexicalLocalPropertyNode;
 import org.pkl.core.ast.expression.member.ReadPropertyNodeGen;
 import org.pkl.core.ast.expression.member.ReadQualifiedLocalPropertyNode;
-import org.pkl.core.ast.expression.member.ReadSuperEntryNode;
-import org.pkl.core.ast.expression.member.ReadSuperPropertyNode;
+import org.pkl.core.ast.expression.member.ReadSuperEntryNodeGen;
+import org.pkl.core.ast.expression.member.ReadSuperPropertyNodeGen;
 import org.pkl.core.ast.expression.primary.ExecuteCustomThisWithRootNode;
 import org.pkl.core.ast.expression.primary.GetEnclosingReceiverNode;
 import org.pkl.core.ast.expression.primary.GetMemberKeyNode;
@@ -1164,12 +1164,12 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
     }
 
     // superproperty call
-    return new ReadSuperPropertyNode(createSourceSection(expr), memberName, needsConst);
+    return ReadSuperPropertyNodeGen.create(createSourceSection(expr), memberName, needsConst);
   }
 
   @Override
   public ExpressionNode visitSuperSubscriptExpr(SuperSubscriptExpr expr) {
-    return new ReadSuperEntryNode(createSourceSection(expr), visitExpr(expr.getArg()));
+    return ReadSuperEntryNodeGen.create(createSourceSection(expr), visitExpr(expr.getArg()));
   }
 
   @Override
@@ -2078,7 +2078,7 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
             bodyNode =
                 doVisitObjectBody(
                     objectBodies,
-                    new ReadSuperPropertyNode(
+                    ReadSuperPropertyNodeGen.create(
                         unavailableSourceSection(),
                         scope.getName(),
                         scope.getConstLevel() == ConstLevel.ALL));
@@ -2689,7 +2689,7 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
             bodyNode =
                 doVisitObjectBody(
                     body,
-                    new ReadSuperPropertyNode(
+                    ReadSuperPropertyNodeGen.create(
                         unavailableSourceSection(),
                         scope.getName(),
                         // Never need a const check for amends declarations. In `foo { ... }`:
@@ -3099,7 +3099,8 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
             var objectBody =
                 doVisitObjectBody(
                     objectBodyCtxs,
-                    new ReadSuperEntryNode(unavailableSourceSection(), new GetMemberKeyNode()));
+                    ReadSuperEntryNodeGen.create(
+                        unavailableSourceSection(), new GetMemberKeyNode()));
             if (isForGeneratorScope) {
               objectBody =
                   new RestoreForBindingsNode(

pkl-core/src/main/java/org/pkl/core/ast/expression/member/ReadSuperEntryNode.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.
@@ -15,21 +15,29 @@
  */
 package org.pkl.core.ast.expression.member;
 
+import com.oracle.truffle.api.dsl.Cached;
+import com.oracle.truffle.api.dsl.Executed;
+import com.oracle.truffle.api.dsl.ImportStatic;
+import com.oracle.truffle.api.dsl.Specialization;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.IndirectCallNode;
 import com.oracle.truffle.api.source.SourceSection;
+import org.jspecify.annotations.Nullable;
 import org.pkl.core.ast.ExpressionNode;
 import org.pkl.core.ast.lambda.ApplyVmFunction1Node;
 import org.pkl.core.ast.lambda.ApplyVmFunction1NodeGen;
+import org.pkl.core.ast.member.ObjectMember;
 import org.pkl.core.runtime.*;
+import org.pkl.core.util.Pair;
 
 /**
- * An expression of the form `super[key]`.
+ * An expression of the form {@code super[key]}.
  *
- * <p>Note: Reading an entry (`object[key]`) is the subscript operator (SubscriptNode).
+ * <p>Note: Reading an entry ({@code object[key]}) is the subscript operator (SubscriptNode).
  */
-public class ReadSuperEntryNode extends ExpressionNode {
-  @Child private ExpressionNode keyNode;
+@ImportStatic(VmUtils.class)
+public abstract class ReadSuperEntryNode extends ExpressionNode {
+  @Child @Executed protected ExpressionNode keyNode;
   @Child private IndirectCallNode callNode = IndirectCallNode.create();
   @Child private ApplyVmFunction1Node applyLambdaNode = ApplyVmFunction1NodeGen.create();
 
@@ -38,13 +46,62 @@ public class ReadSuperEntryNode extends ExpressionNode {
     this.keyNode = keyNode;
   }
 
-  @Override
-  public Object executeGeneric(VirtualFrame frame) {
-    var key = keyNode.executeGeneric(frame);
+  @SuppressWarnings("unused")
+  @Specialization(
+      guards = {
+        // don't use `Object.equals` for these because it will deep-force the owner/receiver
+        "getObjectReceiver(frame) == receiver",
+        "getOwner(frame) == initialOwner",
+        // fine to use `Object.equals` here because an object key is deep-forced upon lookup anyway
+        "key.equals(cachedKey)"
+      })
+  protected Object evalCached(
+      VirtualFrame frame,
+      Object key,
+      @Cached("key") Object cachedKey,
+      @Cached("getObjectReceiver(frame)") VmObjectLike receiver,
+      @Cached("getOwner(frame)") VmObjectLike initialOwner,
+      @Cached("getOwnerAndMember(initialOwner, key)")
+          @Nullable Pair<VmObjectLike, ObjectMember> ownerAndMember,
+      @Cached("doEval(key, receiver, ownerAndMember)") Object result) {
+    return result;
+  }
 
+  @Specialization(replaces = "evalCached")
+  protected Object evalUncached(VirtualFrame frame, Object key) {
     var receiver = VmUtils.getObjectReceiver(frame);
-
     var initialOwner = VmUtils.getOwner(frame);
+    var ownerAndMember = getOwnerAndMember(initialOwner, key);
+    return doEval(key, receiver, ownerAndMember);
+  }
+
+  protected final Object doEval(
+      Object key,
+      VmObjectLike receiver,
+      @Nullable Pair<VmObjectLike, ObjectMember> ownerAndMember) {
+    if (ownerAndMember == null) {
+      // not found -> apply lambda contained in `default` property
+      var defaultFunction =
+          (VmFunction) VmUtils.readMemberOrNull(receiver, Identifier.DEFAULT, callNode);
+      assert defaultFunction != null;
+      return applyLambdaNode.execute(defaultFunction, key);
+    }
+    var owner = ownerAndMember.getFirst();
+    var member = ownerAndMember.getSecond();
+    var constantValue = member.getConstantValue();
+    if (constantValue != null) return constantValue; // TODO: type check
+
+    return callNode.call(
+        member.getCallTarget(),
+        // TODO: should the marker only turn off constraint checking, not overall type checking?
+        receiver,
+        owner,
+        key,
+        VmUtils.SKIP_TYPECHECK_MARKER);
+  }
+
+  protected @Nullable Pair<VmObjectLike, ObjectMember> getOwnerAndMember(
+      VmObjectLike initialOwner, Object key) {
     while (initialOwner instanceof VmFunction) {
       initialOwner = initialOwner.getEnclosingOwner();
     }
@@ -52,26 +109,10 @@ public class ReadSuperEntryNode extends ExpressionNode {
     initialOwner = initialOwner.getParent();
 
     for (var owner = initialOwner; owner != null; owner = owner.getParent()) {
-      var member = owner.getMember(key);
-      if (member == null) continue;
-
-      var constantValue = member.getConstantValue();
-      if (constantValue != null) return constantValue; // TODO: type check
-
-      // caching the result of a super call is tricky (function of both receiver and owner)
-      return callNode.call(
-          member.getCallTarget(),
-          // TODO: should the marker only turn off constraint checking, not overall type checking?
-          receiver,
-          owner,
-          key,
-          VmUtils.SKIP_TYPECHECK_MARKER);
+      var property = owner.getMember(key);
+      if (property == null) continue;
+      return Pair.of(owner, property);
     }
-
-    // not found -> apply lambda contained in `default` property
-    var defaultFunction =
-        (VmFunction) VmUtils.readMemberOrNull(receiver, Identifier.DEFAULT, callNode);
-    assert defaultFunction != null;
-    return applyLambdaNode.execute(defaultFunction, key);
+    return null;
   }
 }

pkl-core/src/main/java/org/pkl/core/ast/expression/member/ReadSuperPropertyNode.java

@@ -1,5 +1,5 @@
 /*
- * Copyright © 2024 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.
@@ -16,13 +16,20 @@
 package org.pkl.core.ast.expression.member;
 
 import com.oracle.truffle.api.CompilerDirectives;
+import com.oracle.truffle.api.dsl.Cached;
+import com.oracle.truffle.api.dsl.ImportStatic;
+import com.oracle.truffle.api.dsl.Specialization;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.IndirectCallNode;
 import com.oracle.truffle.api.source.SourceSection;
+import org.jspecify.annotations.Nullable;
 import org.pkl.core.ast.ExpressionNode;
+import org.pkl.core.ast.member.ObjectMember;
 import org.pkl.core.runtime.*;
+import org.pkl.core.util.Pair;
 
-public final class ReadSuperPropertyNode extends ExpressionNode {
+@ImportStatic(VmUtils.class)
+public abstract class ReadSuperPropertyNode extends ExpressionNode {
 
   private final Identifier propertyName;
   private final boolean needsConst;
@@ -36,16 +43,46 @@ public final class ReadSuperPropertyNode extends ExpressionNode {
     this.needsConst = needsConst;
   }
 
-  // TODO: how can this be optimized?
-  // (result not cached and expensive lookups on every execution)
-  public Object executeGeneric(VirtualFrame frame) {
-    var receiver = VmUtils.getObjectReceiver(frame);
+  @SuppressWarnings("unused")
+  @Specialization(
+      guards = {"getObjectReceiver(frame) == receiver", "getOwner(frame) == initialOwner"})
+  public Object evalCached(
+      VirtualFrame frame,
+      @Cached("getObjectReceiver(frame)") VmObjectLike receiver,
+      @Cached("getOwner(frame)") VmObjectLike initialOwner,
+      @Cached("getPropertyAndOwner(initialOwner)")
+          @Nullable Pair<VmObjectLike, ObjectMember> ownerAndProperty,
+      @Cached("doEval(receiver, ownerAndProperty)") Object result) {
+    return result;
+  }
 
-    // start from the parent of the owner of the `super.<propertyName>` expression
-    // skip any function object owners (same as when resolving `this`)
-    // `receiver` must be passed on unchanged to make sure that overridden properties still take
-    // effect
-    var initialOwner = VmUtils.getOwner(frame);
+  @Specialization(replaces = "evalCached")
+  public Object evalUncached(VirtualFrame frame) {
+    return doEval(VmUtils.getObjectReceiver(frame), getPropertyAndOwner(VmUtils.getOwner(frame)));
+  }
+
+  protected Object doEval(
+      VmObjectLike receiver, @Nullable Pair<VmObjectLike, ObjectMember> ownerAndProperty) {
+    if (ownerAndProperty == null) {
+      // TODO: refine when to return VmDynamic.empty() and when to fail
+      return VmDynamic.empty();
+    }
+    var owner = ownerAndProperty.getFirst();
+    var property = ownerAndProperty.getSecond();
+    var constantValue = property.getConstantValue();
+    if (constantValue != null) return constantValue; // TODO: type check
+
+    return callNode.call(
+        property.getCallTarget(),
+        // TODO: should the marker only turn off constraint checking, not overall type checking?
+        receiver,
+        owner,
+        propertyName,
+        VmUtils.SKIP_TYPECHECK_MARKER);
+  }
+
+  protected @Nullable Pair<VmObjectLike, ObjectMember> getPropertyAndOwner(
+      VmObjectLike initialOwner) {
     while (initialOwner instanceof VmFunction) {
       initialOwner = initialOwner.getEnclosingOwner();
     }
@@ -59,21 +96,8 @@ public final class ReadSuperPropertyNode extends ExpressionNode {
         CompilerDirectives.transferToInterpreter();
         throw exceptionBuilder().evalError("propertyMustBeConst", propertyName.toString()).build();
       }
-
-      var constantValue = property.getConstantValue();
-      if (constantValue != null) return constantValue; // TODO: type check
-
-      // caching the result of a super call is tricky (function of both receiver and owner)
-      return callNode.call(
-          property.getCallTarget(),
-          // TODO: should the marker only turn off constraint checking, not overall type checking?
-          receiver,
-          owner,
-          propertyName,
-          VmUtils.SKIP_TYPECHECK_MARKER);
+      return Pair.of(owner, property);
     }
-
-    // TODO: refine when to return VmDynamic.empty() and when to fail
-    return VmDynamic.empty();
+    return null;
   }
 }

pkl-core/src/test/files/LanguageSnippetTests/input/objects/super6.pkl

@@ -0,0 +1,27 @@
+class Bird {
+  name: String
+  tag: String
+}
+
+hidden birdTypes: Mapping<String, Bird> = new {
+  ["pigeon"] {
+    name = "Pigeon"
+    tag = "Feathers"
+  }
+  ["cormorant"] {
+    name = "Cormorant"
+    tag = "Feathers"
+  }
+}
+
+res1 =
+  birdTypes.fold(
+    new Mapping {},
+    (acc, name, birds) -> (acc) {
+      // on different reads, `super[birds.tag]` gives a different owner
+      // however, the receiver is the same (receiver is `res1`)
+      [birds.tag] = (super[birds.tag]) {
+        [name] = birds
+      }
+    },
+  )

pkl-core/src/test/files/LanguageSnippetTests/input/objects/super7.pkl

@@ -0,0 +1,29 @@
+class Bird {
+  name: String
+}
+
+hidden birdTypes: Mapping<String, Bird> = new {
+  ["pigeon"] {
+    name = "Pigeon"
+  }
+  ["cormorant"] {
+    name = "Cormorant"
+  }
+}
+
+class Category {
+  tag: Mapping<String, Bird>
+}
+
+res1 =
+  birdTypes.fold(
+    new Category {},
+    (acc, name, birds) -> (acc) {
+      // on different reads, `super.tag` should give a different owner.
+      // however, the receiver is the same (receiver is `res1`)
+      tag = (super.tag) {
+        [name] = birds
+      }
+    },
+  )
+

pkl-core/src/test/files/LanguageSnippetTests/output/objects/super6.pcf

@@ -0,0 +1,12 @@
+res1 {
+  ["Feathers"] {
+    ["pigeon"] {
+      name = "Pigeon"
+      tag = "Feathers"
+    }
+    ["cormorant"] {
+      name = "Cormorant"
+      tag = "Feathers"
+    }
+  }
+}

pkl-core/src/test/files/LanguageSnippetTests/output/objects/super7.pcf

@@ -0,0 +1,10 @@
+res1 {
+  tag {
+    ["pigeon"] {
+      name = "Pigeon"
+    }
+    ["cormorant"] {
+      name = "Cormorant"
+    }
+  }
+}