Fix some GraalVM PE issues (#1944)

* VmObject.getRootFirstObject() is a recursive call, which sometimes
  fails with "too deep inlining"
* Incorrect loop condition profile usage in some places
* Need to mark `checkedReturnTypeNode` final so that GraalVM knows
  that it is PE constant

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

@@ -44,7 +44,6 @@ import org.pkl.core.runtime.VmMap;
 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;
 
@@ -199,8 +198,7 @@ public abstract class GeneratorSpreadNode extends GeneratorMemberNode {
 
   protected void doEvalDynamic(VirtualFrame frame, ObjectData data, VmObject iterable) {
     var cursor = iterable.members();
-    var counted = isCounted(cursor);
-    while (advance(cursor, counted)) {
+    while (loopConditionProfile.profile(cursor.advance())) {
       var member = cursor.member();
       var value = readCursorValueNode.execute(frame, cursor);
       if (member.isElement()) {
@@ -213,8 +211,7 @@ public abstract class GeneratorSpreadNode extends GeneratorMemberNode {
 
   private void doEvalMapping(VirtualFrame frame, ObjectData data, VmObject iterable) {
     var cursor = iterable.members();
-    var counted = isCounted(cursor);
-    while (advance(cursor, counted)) {
+    while (loopConditionProfile.profile(cursor.advance())) {
       var member = cursor.member();
       if (member.isElement() || member.isProp()) {
         cannotHaveMember(BaseModule.getMappingClass(), member);
@@ -229,8 +226,8 @@ public abstract class GeneratorSpreadNode extends GeneratorMemberNode {
 
   private void doEvalListing(VirtualFrame frame, ObjectData data, VmObject iterable) {
     var cursor = iterable.members();
-    var counted = isCounted(cursor);
-    while (advance(cursor, counted)) {
+    loopConditionProfile.profileCounted(cursor.getLength());
+    while (loopConditionProfile.inject(cursor.advance())) {
       var member = cursor.member();
       if (member.isEntry() || member.isProp()) {
         cannotHaveMember(getListingClass(), member);
@@ -242,8 +239,7 @@ public abstract class GeneratorSpreadNode extends GeneratorMemberNode {
 
   private void doEvalTyped(VirtualFrame frame, VmClass clazz, ObjectData data, VmObject iterable) {
     var cursor = iterable.members();
-    var counted = isCounted(cursor);
-    while (advance(cursor, counted)) {
+    while (loopConditionProfile.profile(cursor.advance())) {
       var member = cursor.member();
       if (member.isElement() || member.isEntry()) {
         cannotHaveMember(clazz, member);
@@ -254,19 +250,6 @@ public abstract class GeneratorSpreadNode extends GeneratorMemberNode {
     }
   }
 
-  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`
   private void doEvalCollection(
       VirtualFrame frame, VmClass parent, ObjectData data, VmCollection iterable) {

pkl-core/src/main/java/org/pkl/core/ast/member/FunctionNode.java

@@ -54,7 +54,8 @@ public final class FunctionNode extends RegularMemberNode {
   @Children private TypeParameterNode[] typeParameterNodes;
 
   @Children private final TypeNode[] parameterTypeNodes;
-  private @Nullable TypeNode checkedReturnTypeNode;
+  // not marked `@Child` because it's the same as `returnTypeNode`.
+  private final @Nullable TypeNode checkedReturnTypeNode;
   @Child private @Nullable TypeNode returnTypeNode;
 
   @TruffleBoundary

pkl-core/src/main/java/org/pkl/core/ast/type/TypeNode.java

@@ -1257,6 +1257,8 @@ public abstract class TypeNode extends PklNode {
 
   public static final class CollectionTypeNode extends ObjectSlotTypeNode {
     @Child private TypeNode elementTypeNode;
+    private final LoopConditionProfile listLoopProfile = LoopConditionProfile.create();
+    private final LoopConditionProfile setLoopProfile = LoopConditionProfile.create();
 
     public CollectionTypeNode(SourceSection sourceSection, TypeNode elementTypeNode) {
       super(sourceSection);
@@ -1316,13 +1318,14 @@ public abstract class TypeNode extends PklNode {
     }
 
     @SuppressWarnings("DuplicatedCode")
-    @ExplodeLoop
     private Object evalList(VirtualFrame frame, VmList value) {
       var ret = value;
       var idx = 0;
+      listLoopProfile.profileCounted(value.getLength());
 
-      for (var elem : value) {
+      for (var i = 0; listLoopProfile.inject(i < value.getLength()); i++) {
         TruffleSafepoint.poll(this);
+        var elem = value.get(i);
         var result = elementTypeNode.executeLazily(frame, elem);
         if (result != elem) {
           ret = ret.replace(idx, result);
@@ -1330,27 +1333,31 @@ public abstract class TypeNode extends PklNode {
         idx++;
       }
 
-      LoopNode.reportLoopCount(this, idx);
+      LoopNode.reportLoopCount(this, value.getLength());
       return ret;
     }
 
     private Object evalListEagerly(VirtualFrame frame, VmList value) {
-      for (var elem : value) {
+      var length = value.getLength();
+      listLoopProfile.profileCounted(length);
+      for (var i = 0; listLoopProfile.inject(i < length); i++) {
         TruffleSafepoint.poll(this);
-        elementTypeNode.executeEagerly(frame, elem);
+        elementTypeNode.executeEagerly(frame, value.get(i));
       }
 
-      LoopNode.reportLoopCount(this, value.getLength());
+      LoopNode.reportLoopCount(this, length);
       return value;
     }
 
     private Object evalSet(VirtualFrame frame, VmSet value) {
-      for (var elem : value) {
+      var iterator = value.iterator();
+      var count = 0;
+      while (setLoopProfile.inject(iterator.hasNext())) {
         TruffleSafepoint.poll(this);
-        elementTypeNode.executeEagerly(frame, elem);
+        elementTypeNode.executeEagerly(frame, iterator.next());
+        count++;
       }
-
-      LoopNode.reportLoopCount(this, value.getLength());
+      LoopNode.reportLoopCount(this, count);
       return value;
     }
   }

pkl-core/src/main/java/org/pkl/core/runtime/VmObject.java

@@ -79,6 +79,8 @@ public abstract class VmObject extends VmObjectLike {
     return EconomicMaps.containsKey(members, key);
   }
 
+  // deeply recursive calls are bad for PE, which can happen for deeply nested objects.
+  @TruffleBoundary
   public final @Nullable ObjectMember getRootFirstMember(Object key) {
     if (parent != null) {
       var parentMember = parent.getRootFirstMember(key);