Replace VmObjectLike.iterateMemberValues() with VmObjectCursor (#1879)

Motivation:
- VmObjectLike.iterateMemberValues() is an internal iterator whose loop body is a Java lambda.
  It suffers from the following limitations:
  - `continue` and `break` cannot be used.
  - `return` becomes an (awkward) replacement for `continue` and `break`.
    It cannot be used to return from the enclosing method.
  - Mutating variables requires hoisting them (MutableBoolean, etc.).
  - Node.execute() calls are very slow because they are behind a Truffle boundary.
  - Callers of iterateMemberValues() do not specify what and how they want to iterate.
    This obscures the iteration code and makes it difficult (often impossible) to optimize.

Changes:
- Add interface VmObjectCursor, an external iterator that iterates over VmObject members.
  Its design is inspired by org.graalvm.collections.MapCursor.
- Add VmObjectCursor implementations specialized for different members, VmObject subclasses, and cursor options.
- Replace VmObjectLike.iterateMemberValues() with the following methods that return VmObjectCursor:
  - properties()
  - properties(CursorOption)
  - elements()
  - elements(CursorOption)
  - entries()
  - entries(CursorOption)
  - members
  - members(CursorOption)
- Replace all usages of iterateMemberValues() with one of the above methods.
  Before: object.iterateMembersValues((key, member, value) -> { if (member.isProp()) {...} });
  After : for (var cursor = object.properties(); cursor.advance(); ) {...}

Result:
- Iteration code is much simpler and clearer.
- Iteration code that calls Node.execute() is much faster because it is no longer behind a Truffle boundary.
- Other iteration code is potentially faster and can be better optimized in the future.

Co-authored-by: translatenix <119817707+translatenix@users.noreply.github.com>
Co-authored-by: Jen Basch <jbasch94@gmail.com>

Author: 3 people

pkl-core/src/main/java/org/pkl/core/EvaluatorImpl.java

@@ -15,6 +15,7 @@
  */
 package org.pkl.core;
 
+import com.oracle.truffle.api.nodes.IndirectCallNode;
 import java.io.IOException;
 import java.net.URI;
 import java.net.URISyntaxException;
@@ -195,7 +196,9 @@ public Map<String, FileOutput> evaluateOutputFiles(ModuleSource moduleSource) {
         (module) -> {
           var output = VmUtils.readModuleOutput(module);
           return VmUtils.readFilesProperty(
-              output, (fileOutput) -> new FileOutputImpl(this, fileOutput));
+              output,
+              (fileOutput) -> new FileOutputImpl(this, fileOutput),
+              IndirectCallNode.getUncached());
         });
   }
 

pkl-core/src/main/java/org/pkl/core/ast/expression/binary/EqualNode.java

@@ -19,6 +19,7 @@
 import com.oracle.truffle.api.dsl.Cached;
 import com.oracle.truffle.api.dsl.NodeChild;
 import com.oracle.truffle.api.dsl.Specialization;
+import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.NodeInfo;
 import com.oracle.truffle.api.source.SourceSection;
 import org.jspecify.annotations.Nullable;
@@ -34,6 +35,9 @@ protected EqualNode(SourceSection sourceSection) {
     super(sourceSection);
   }
 
+  /** When using this method, pass {@code null} in for {@code leftNode} and {@code rightNode}. */
+  public abstract boolean executeWith(VirtualFrame frame, Object left, Object right);
+
   @Specialization
   protected boolean eval(String left, String right) {
     return left.equals(right);

pkl-core/src/main/java/org/pkl/core/ast/expression/generator/GeneratorForNode.java

@@ -23,22 +23,31 @@
 import com.oracle.truffle.api.frame.FrameDescriptor;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.ExplodeLoop;
+import com.oracle.truffle.api.profiles.LoopConditionProfile;
 import com.oracle.truffle.api.source.SourceSection;
 import org.jspecify.annotations.Nullable;
 import org.pkl.core.ast.ExpressionNode;
+import org.pkl.core.ast.internal.ReadCursorValueNode;
+import org.pkl.core.ast.internal.ReadCursorValueNodeGen;
 import org.pkl.core.ast.type.TypeNode;
 import org.pkl.core.ast.type.UnresolvedTypeNode;
 import org.pkl.core.runtime.*;
+import org.pkl.core.runtime.Iterators.TruffleIterator;
 import org.pkl.core.util.ArrayUtils;
 
 public abstract class GeneratorForNode extends GeneratorMemberNode {
   private final FrameDescriptor generatorDescriptor;
+  private final LoopConditionProfile loopConditionProfile = LoopConditionProfile.create();
   @Child private ExpressionNode iterableNode;
   @Child private @Nullable UnresolvedTypeNode unresolvedKeyTypeNode;
   @Child private @Nullable UnresolvedTypeNode unresolvedValueTypeNode;
   @Children private final GeneratorMemberNode[] childNodes;
   @Child private @Nullable TypeNode keyTypeNode;
   @Child private @Nullable TypeNode valueTypeNode;
+
+  @Child
+  private ReadCursorValueNode readCursorValueNode = ReadCursorValueNodeGen.create(sourceSection);
+
   private final int keySlot;
   private final int valueSlot;
   private final int[] slotsToCopy;
@@ -101,40 +110,50 @@ protected void eval(VirtualFrame frame, Object parent, ObjectData data, VmDynami
 
   @Specialization
   protected void eval(VirtualFrame frame, Object parent, ObjectData data, VmList iterable) {
-    long idx = 0;
-    for (var element : iterable) {
-      executeIteration(frame, parent, data, idx++, element);
+    loopConditionProfile.profileCounted(iterable.getLength());
+    for (var idx = 0L; loopConditionProfile.inject(idx < iterable.getLength()); idx++) {
+      executeIteration(frame, parent, data, idx, iterable.get(idx));
     }
   }
 
   @Specialization
   protected void eval(VirtualFrame frame, Object parent, ObjectData data, VmMap iterable) {
-    for (var entry : iterable) {
+    loopConditionProfile.profileCounted(iterable.getLength());
+    var iterator = new TruffleIterator<>(iterable);
+    while (loopConditionProfile.inject(iterator.hasNext())) {
+      var entry = iterator.next();
       executeIteration(frame, parent, data, VmUtils.getKey(entry), VmUtils.getValue(entry));
     }
   }
 
   @Specialization
   protected void eval(VirtualFrame frame, Object parent, ObjectData data, VmSet iterable) {
-    long idx = 0;
-    for (var element : iterable) {
+    var idx = 0L;
+    loopConditionProfile.profileCounted(iterable.getLength());
+    var iterator = new TruffleIterator<>(iterable);
+    while (loopConditionProfile.inject(iterator.hasNext())) {
+      var element = iterator.next();
       executeIteration(frame, parent, data, idx++, element);
     }
   }
 
   @Specialization
   protected void eval(VirtualFrame frame, Object parent, ObjectData data, VmIntSeq iterable) {
     var length = iterable.getLength();
-    for (long key = 0, value = iterable.start; key < length; key++, value += iterable.step) {
+    loopConditionProfile.profileCounted(iterable.getLength());
+    for (long key = 0, value = iterable.start;
+        loopConditionProfile.inject(key < length);
+        key++, value += iterable.step) {
       executeIteration(frame, parent, data, key, value);
     }
   }
 
   @Specialization
   protected void eval(VirtualFrame frame, Object parent, ObjectData data, VmBytes iterable) {
-    long idx = 0;
-    for (var byt : iterable.getBytes()) {
-      executeIteration(frame, parent, data, idx++, (long) byt);
+    var bytes = iterable.getBytes();
+    loopConditionProfile.profileCounted(bytes.length);
+    for (var idx = 0; loopConditionProfile.inject(idx < bytes.length); idx++) {
+      executeIteration(frame, parent, data, (long) idx, (long) bytes[idx]);
     }
   }
 
@@ -150,16 +169,32 @@ protected void fallback(VirtualFrame frame, Object parent, ObjectData data, Obje
   }
 
   private void doEvalObject(VirtualFrame frame, VmObject iterable, Object parent, ObjectData data) {
-    var materializedFrame = frame.materialize();
-    iterable.forceAndIterateMemberValues(
-        (key, member, value) -> {
-          var convertedKey = member.isProp() ? key.toString() : key;
-          // TODO: Executing iteration behind a Truffle boundary is bad for performance.
-          // This and similar cases will be fixed in an upcoming PR that replaces method
-          // `(forceAnd)iterateMemberValues` with cursor-based external iterators.
-          executeIteration(materializedFrame, parent, data, convertedKey, value);
-          return true;
-        });
+    // TODO: skip types for module objects?
+    var cursor = iterable.members();
+    if (cursor.getLength() == -1) {
+      doEvalObjectProfiled(frame, parent, data, cursor);
+    } else {
+      doEvalObjectCounted(frame, parent, data, cursor);
+    }
+  }
+
+  private void doEvalObjectCounted(
+      VirtualFrame frame, Object parent, ObjectData data, VmObjectCursor cursor) {
+    loopConditionProfile.profileCounted(cursor.getLength());
+    while (loopConditionProfile.inject(cursor.advance())) {
+      var key = cursor.isProperty() ? cursor.keyToString() : cursor.key();
+      var value = readCursorValueNode.execute(frame, cursor);
+      executeIteration(frame, parent, data, key, value);
+    }
+  }
+
+  private void doEvalObjectProfiled(
+      VirtualFrame frame, Object parent, ObjectData data, VmObjectCursor cursor) {
+    while (loopConditionProfile.profile(cursor.advance())) {
+      var key = cursor.isProperty() ? cursor.keyToString() : cursor.key();
+      var value = readCursorValueNode.execute(frame, cursor);
+      executeIteration(frame, parent, data, key, value);
+    }
   }
 
   @ExplodeLoop

pkl-core/src/main/java/org/pkl/core/ast/expression/generator/GeneratorSpreadNode.java

@@ -23,8 +23,11 @@
 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.profiles.LoopConditionProfile;
 import com.oracle.truffle.api.source.SourceSection;
 import org.pkl.core.ast.ExpressionNode;
+import org.pkl.core.ast.internal.ReadCursorValueNode;
+import org.pkl.core.ast.internal.ReadCursorValueNodeGen;
 import org.pkl.core.ast.member.ObjectMember;
 import org.pkl.core.runtime.BaseModule;
 import org.pkl.core.runtime.Identifier;
@@ -41,12 +44,18 @@
 import org.pkl.core.runtime.VmMapping;
 import org.pkl.core.runtime.VmNull;
 import org.pkl.core.runtime.VmObject;
+import org.pkl.core.runtime.VmObjectCursor;
 import org.pkl.core.runtime.VmTyped;
 import org.pkl.core.runtime.VmUtils;
 
 @ImportStatic(BaseModule.class)
 public abstract class GeneratorSpreadNode extends GeneratorMemberNode {
+  private final LoopConditionProfile loopConditionProfile = LoopConditionProfile.create();
   @Child private ExpressionNode iterableNode;
+
+  @Child
+  private ReadCursorValueNode readCursorValueNode = ReadCursorValueNodeGen.create(sourceSection);
+
   private final boolean nullable;
 
   public GeneratorSpreadNode(
@@ -189,53 +198,73 @@ protected void fallback(VirtualFrame frame, Object parent, ObjectData data, Obje
   }
 
   protected void doEvalDynamic(VirtualFrame frame, ObjectData data, VmObject iterable) {
-    var materializedFrame = frame.materialize();
-    iterable.forceAndIterateMemberValues(
-        (key, member, value) -> {
-          if (member.isElement()) {
-            data.addElement(materializedFrame, createMember(member, value), this);
-          } else {
-            data.addMember(materializedFrame, key, createMember(member, value), this);
-          }
-          return true;
-        });
+    var cursor = iterable.members();
+    var counted = isCounted(cursor);
+    while (advance(cursor, counted)) {
+      var member = cursor.member();
+      var value = readCursorValueNode.execute(frame, cursor);
+      if (member.isElement()) {
+        data.addElement(frame, createMember(member, value), this);
+      } else {
+        data.addMember(frame, cursor.key(), createMember(member, value), this);
+      }
+    }
   }
 
   private void doEvalMapping(VirtualFrame frame, ObjectData data, VmObject iterable) {
-    var materializedFrame = frame.materialize();
-    iterable.forceAndIterateMemberValues(
-        (key, member, value) -> {
-          if (member.isElement() || member.isProp()) {
-            cannotHaveMember(BaseModule.getMappingClass(), member);
-          }
-          data.addMember(materializedFrame, key, createMember(member, value), this);
-          return true;
-        });
+    var cursor = iterable.members();
+    var counted = isCounted(cursor);
+    while (advance(cursor, counted)) {
+      var member = cursor.member();
+      if (member.isElement() || member.isProp()) {
+        cannotHaveMember(BaseModule.getMappingClass(), member);
+      }
+      data.addMember(
+          frame,
+          cursor.key(),
+          createMember(member, readCursorValueNode.execute(frame, cursor)),
+          this);
+    }
   }
 
   private void doEvalListing(VirtualFrame frame, ObjectData data, VmObject iterable) {
-    var materializedFrame = frame.materialize();
-    iterable.forceAndIterateMemberValues(
-        (key, member, value) -> {
-          if (member.isEntry() || member.isProp()) {
-            cannotHaveMember(getListingClass(), member);
-          }
-          data.addElement(materializedFrame, createMember(member, value), this);
-          return true;
-        });
+    var cursor = iterable.members();
+    var counted = isCounted(cursor);
+    while (advance(cursor, counted)) {
+      var member = cursor.member();
+      if (member.isEntry() || member.isProp()) {
+        cannotHaveMember(getListingClass(), member);
+      }
+      data.addElement(
+          frame, createMember(member, readCursorValueNode.execute(frame, cursor)), this);
+    }
   }
 
   private void doEvalTyped(VirtualFrame frame, VmClass clazz, ObjectData data, VmObject iterable) {
-    var materializedFrame = frame.materialize();
-    iterable.forceAndIterateMemberValues(
-        (key, member, value) -> {
-          if (member.isElement() || member.isEntry()) {
-            cannotHaveMember(clazz, member);
-          }
-          checkIsValidTypedProperty(clazz, member);
-          data.addProperty(materializedFrame, createMember(member, value), this);
-          return true;
-        });
+    var cursor = iterable.members();
+    var counted = isCounted(cursor);
+    while (advance(cursor, counted)) {
+      var member = cursor.member();
+      if (member.isElement() || member.isEntry()) {
+        cannotHaveMember(clazz, member);
+      }
+      checkIsValidTypedProperty(clazz, member);
+      data.addProperty(
+          frame, createMember(member, readCursorValueNode.execute(frame, cursor)), this);
+    }
+  }
+
+  private boolean isCounted(VmObjectCursor cursor) {
+    var length = cursor.getLength();
+    if (length == -1) return false;
+    loopConditionProfile.profileCounted(length);
+    return true;
+  }
+
+  private boolean advance(VmObjectCursor cursor, boolean counted) {
+    return counted
+        ? loopConditionProfile.inject(cursor.advance())
+        : loopConditionProfile.profile(cursor.advance());
   }
 
   // handles both `List` and `Set`
@@ -250,7 +279,7 @@ private void doEvalCollection(
           .withProgramValue("Value", iterable)
           .build();
     }
-    spreadIterable(frame, data, iterable);
+    spreadIterable(frame, data, iterable, iterable.getLength());
   }
 
   private void doEvalMap(VirtualFrame frame, VmClass parent, ObjectData data, VmMap iterable) {
@@ -278,7 +307,7 @@ private void doEvalIntSeq(
           .withProgramValue("Value", iterable)
           .build();
     }
-    spreadIterable(frame, data, iterable);
+    spreadIterable(frame, data, iterable, iterable.getLength());
   }
 
   private void doEvalBytes(VirtualFrame frame, VmClass parent, ObjectData data, VmBytes iterable) {
@@ -290,7 +319,7 @@ private void doEvalBytes(VirtualFrame frame, VmClass parent, ObjectData data, Vm
           .withProgramValue("Value", iterable)
           .build();
     }
-    spreadIterable(frame, data, iterable);
+    spreadIterable(frame, data, iterable, iterable.getLength());
   }
 
   private void cannotHaveMember(VmClass clazz, ObjectMember member) {
@@ -352,9 +381,11 @@ private ObjectMember createMember(ObjectMember prototype, Object value) {
     return result;
   }
 
-  private void spreadIterable(VirtualFrame frame, ObjectData data, Iterable<?> iterable) {
+  private void spreadIterable(
+      VirtualFrame frame, ObjectData data, Iterable<?> iterable, long length) {
+    loopConditionProfile.profileCounted(length);
     var iterator = new TruffleIterator<>(iterable);
-    while (iterator.hasNext()) {
+    while (loopConditionProfile.inject(iterator.hasNext())) {
       var elem = iterator.next();
       var member = VmUtils.createSyntheticObjectElement(String.valueOf(data.length()), elem);
       data.addElement(frame, member, this);

pkl-core/src/main/java/org/pkl/core/ast/expression/literal/ConstantEntriesLiteralNode.java

@@ -148,7 +148,7 @@ protected VmDynamic evalDynamicClass(
         "checkIsValidListingAmendment()",
         "checkMaxListingMemberIndex(0)"
       })
-  protected void evalListingClass(@SuppressWarnings("unused") VmClass parent) {
+  protected Object evalListingClass(@SuppressWarnings("unused") VmClass parent) {
     CompilerDirectives.transferToInterpreter();
     throw exceptionBuilder().unreachableCode().build();
   }

pkl-core/src/main/java/org/pkl/core/ast/expression/literal/ElementsEntriesLiteralNode.java

@@ -34,8 +34,9 @@
 import org.pkl.core.util.EconomicMaps;
 
 /**
- * Object literal that contains both elements and entries (and possibly properties). Example: `new
- * foo { "pigeon", [3] = "barn owl" }`
+ * Object literal that contains both elements and entries (and possibly properties).
+ *
+ * <p>Example: {@code new foo { "pigeon"; [3] = "barn owl" } }
  */
 @ImportStatic({BaseModule.class, VmUtils.class})
 public abstract class ElementsEntriesLiteralNode extends SpecializedObjectLiteralNode {

pkl-core/src/main/java/org/pkl/core/ast/expression/literal/ObjectLiteralNode.java

@@ -42,6 +42,7 @@ public abstract class ObjectLiteralNode extends ExpressionNode {
   protected final boolean isCustomThisScope;
   protected final @Nullable FrameDescriptor parametersDescriptor;
   @Children protected final UnresolvedTypeNode[] parameterTypes;
+  @Child protected @Nullable AmendFunctionNode amendFunctionNode;
 
   public ObjectLiteralNode(
       SourceSection sourceSection,