From 2c2f50113dd205004a37389c210f3600adebf97b Mon Sep 17 00:00:00 2001 From: Daniel Chao Date: Wed, 2 Sep 2026 10:01:25 -0700 Subject: [PATCH] Make enclosing frame lookup more PE friendly (#1842) The truffle compiler currently _always_ bails out of compilation for `VmUtils.getFrame()` calls, because it determines that some branch of this code can possibly return a materialized frame. This adjusts the code so that the hot path (levelsUp == 0) does not touch `VmUtils.getFrame()`, and avoids unnecessarily materializing the frame. --- .../org/pkl/core/ast/builder/AstBuilder.java | 7 +--- .../AbstractInvokeLexicalMethodNode.java | 10 +++-- .../core/ast/frame/GetEnclosingFrameNode.java | 38 ------------------- .../pkl/core/ast/frame/ReadFrameSlotNode.java | 37 ++++++++++++++---- .../java/org/pkl/core/runtime/VmUtils.java | 25 ++++++++++-- 5 files changed, 58 insertions(+), 59 deletions(-) delete mode 100644 pkl-core/src/main/java/org/pkl/core/ast/frame/GetEnclosingFrameNode.java diff --git a/pkl-core/src/main/java/org/pkl/core/ast/builder/AstBuilder.java b/pkl-core/src/main/java/org/pkl/core/ast/builder/AstBuilder.java index 36e1b83ef..6ea0ecb9c 100644 --- a/pkl-core/src/main/java/org/pkl/core/ast/builder/AstBuilder.java +++ b/pkl-core/src/main/java/org/pkl/core/ast/builder/AstBuilder.java @@ -151,7 +151,6 @@ import org.pkl.core.ast.expression.unary.ReadOrNullNodeGen; import org.pkl.core.ast.expression.unary.ThrowNodeGen; import org.pkl.core.ast.expression.unary.TraceNode; import org.pkl.core.ast.expression.unary.UnaryMinusNodeGen; -import org.pkl.core.ast.frame.GetEnclosingFrameNode; import org.pkl.core.ast.frame.ReadExactFrameSlotNodeGen; import org.pkl.core.ast.frame.ReadFrameSlotNodeGen; import org.pkl.core.ast.internal.GetBaseModuleClassNode; @@ -844,11 +843,9 @@ public class AstBuilder extends AbstractAstBuilder { // } return p.levelsUp() == 0 && !p.needsFrameSkip() ? ReadExactFrameSlotNodeGen.create(sourceSection, p.slot()) - : ReadFrameSlotNodeGen.create( - sourceSection, p.slot(), new GetEnclosingFrameNode(p.levelsUp())); + : ReadFrameSlotNodeGen.create(sourceSection, p.slot(), p.levelsUp()); } else if (resolution instanceof Parameter p) { - return ReadFrameSlotNodeGen.create( - sourceSection, p.slot(), new GetEnclosingFrameNode(p.levelsUp())); + return ReadFrameSlotNodeGen.create(sourceSection, p.slot(), p.levelsUp()); } else if (resolution instanceof ImplicitBaseProperty) { return ReadPropertyNodeGen.create( sourceSection, diff --git a/pkl-core/src/main/java/org/pkl/core/ast/expression/member/AbstractInvokeLexicalMethodNode.java b/pkl-core/src/main/java/org/pkl/core/ast/expression/member/AbstractInvokeLexicalMethodNode.java index bb019be8e..c27306e88 100644 --- a/pkl-core/src/main/java/org/pkl/core/ast/expression/member/AbstractInvokeLexicalMethodNode.java +++ b/pkl-core/src/main/java/org/pkl/core/ast/expression/member/AbstractInvokeLexicalMethodNode.java @@ -37,9 +37,11 @@ public abstract sealed class AbstractInvokeLexicalMethodNode extends AbstractInv @Override public final Object executeGeneric(VirtualFrame frame) { - var capturedFrame = VmUtils.getFrame(frame, levelsUp); - var owner = VmUtils.getOwner(capturedFrame); - var receiver = VmUtils.getReceiver(capturedFrame); - return invoke(frame, owner, receiver); + var owner = VmUtils.getOwner(frame); + if (levelsUp == 0 && !owner.isParseTimeInvisibleScope()) { + return invoke(frame, owner, VmUtils.getReceiver(frame)); + } + var enclosingFrame = VmUtils.getEnclosingFrame(owner, levelsUp); + return invoke(frame, VmUtils.getOwner(enclosingFrame), VmUtils.getReceiver(enclosingFrame)); } } diff --git a/pkl-core/src/main/java/org/pkl/core/ast/frame/GetEnclosingFrameNode.java b/pkl-core/src/main/java/org/pkl/core/ast/frame/GetEnclosingFrameNode.java deleted file mode 100644 index c9a0e4731..000000000 --- a/pkl-core/src/main/java/org/pkl/core/ast/frame/GetEnclosingFrameNode.java +++ /dev/null @@ -1,38 +0,0 @@ -/* - * Copyright © 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. - * You may obtain a copy of the License at - * - * https://www.apache.org/licenses/LICENSE-2.0 - * - * Unless required by applicable law or agreed to in writing, software - * distributed under the License is distributed on an "AS IS" BASIS, - * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. - * See the License for the specific language governing permissions and - * limitations under the License. - */ -package org.pkl.core.ast.frame; - -import com.oracle.truffle.api.frame.VirtualFrame; -import org.pkl.core.ast.ExpressionNode; -import org.pkl.core.runtime.VmUtils; - -public final class GetEnclosingFrameNode extends ExpressionNode { - private final int levelsUp; - - public GetEnclosingFrameNode(int levelsUp) { - this.levelsUp = levelsUp; - } - - @Override - public VirtualFrame executeGeneric(VirtualFrame frame) { - return VmUtils.getFrame(frame, levelsUp); - } - - @Override - public boolean isInstrumentable() { - return false; - } -} diff --git a/pkl-core/src/main/java/org/pkl/core/ast/frame/ReadFrameSlotNode.java b/pkl-core/src/main/java/org/pkl/core/ast/frame/ReadFrameSlotNode.java index 8c3083a31..089b303bc 100644 --- a/pkl-core/src/main/java/org/pkl/core/ast/frame/ReadFrameSlotNode.java +++ b/pkl-core/src/main/java/org/pkl/core/ast/frame/ReadFrameSlotNode.java @@ -15,44 +15,65 @@ */ package org.pkl.core.ast.frame; -import com.oracle.truffle.api.dsl.NodeChild; import com.oracle.truffle.api.dsl.Specialization; import com.oracle.truffle.api.frame.FrameSlotTypeException; import com.oracle.truffle.api.frame.VirtualFrame; import com.oracle.truffle.api.source.SourceSection; import org.pkl.core.ast.ExpressionNode; +import org.pkl.core.runtime.VmUtils; -@NodeChild(value = "enclosingFrame", type = GetEnclosingFrameNode.class) public abstract class ReadFrameSlotNode extends ExpressionNode { private final int slot; + private final int levelsUp; - protected ReadFrameSlotNode(SourceSection sourceSection, int slot) { + protected ReadFrameSlotNode(SourceSection sourceSection, int slot, int levelsUp) { super(sourceSection); this.slot = slot; + this.levelsUp = levelsUp; } @Specialization(rewriteOn = FrameSlotTypeException.class) protected long evalInt(VirtualFrame frame) throws FrameSlotTypeException { - return frame.getLong(slot); + var owner = VmUtils.getOwner(frame); + if (levelsUp == 0 && !owner.isParseTimeInvisibleScope()) { + return frame.getLong(slot); + } + return VmUtils.getEnclosingFrame(owner, levelsUp).getLong(slot); } @Specialization(rewriteOn = FrameSlotTypeException.class) protected double evalFloat(VirtualFrame frame) throws FrameSlotTypeException { - return frame.getDouble(slot); + var owner = VmUtils.getOwner(frame); + if (levelsUp == 0 && !owner.isParseTimeInvisibleScope()) { + return frame.getDouble(slot); + } + return VmUtils.getEnclosingFrame(owner, levelsUp).getDouble(slot); } @Specialization(rewriteOn = FrameSlotTypeException.class) protected boolean evalBoolean(VirtualFrame frame) throws FrameSlotTypeException { - return frame.getBoolean(slot); + var owner = VmUtils.getOwner(frame); + if (levelsUp == 0 && !owner.isParseTimeInvisibleScope()) { + return frame.getBoolean(slot); + } + return VmUtils.getEnclosingFrame(owner, levelsUp).getBoolean(slot); } @Specialization(rewriteOn = FrameSlotTypeException.class) protected Object evalObject(VirtualFrame frame) throws FrameSlotTypeException { - return frame.getObject(slot); + var owner = VmUtils.getOwner(frame); + if (levelsUp == 0 && !owner.isParseTimeInvisibleScope()) { + return frame.getObject(slot); + } + return VmUtils.getEnclosingFrame(owner, levelsUp).getObject(slot); } @Specialization(replaces = {"evalInt", "evalFloat", "evalBoolean", "evalObject"}) protected Object evalGeneric(VirtualFrame frame) { - return frame.getValue(slot); + var owner = VmUtils.getOwner(frame); + if (levelsUp == 0 && !owner.isParseTimeInvisibleScope()) { + return frame.getValue(slot); + } + return VmUtils.getEnclosingFrame(owner, levelsUp).getValue(slot); } } diff --git a/pkl-core/src/main/java/org/pkl/core/runtime/VmUtils.java b/pkl-core/src/main/java/org/pkl/core/runtime/VmUtils.java index 3c722a6cf..17891b1e3 100644 --- a/pkl-core/src/main/java/org/pkl/core/runtime/VmUtils.java +++ b/pkl-core/src/main/java/org/pkl/core/runtime/VmUtils.java @@ -171,14 +171,31 @@ public final class VmUtils { } public static VmObjectLike getOwner(VirtualFrame frame, int levelsUp) { - return getOwner(getFrame(frame, levelsUp)); + var owner = getOwner(frame); + if (levelsUp == 0 && !owner.isParseTimeInvisibleScope()) { + return owner; + } + return getOwner(getEnclosingFrame(owner, levelsUp)); } public static Object getReceiver(VirtualFrame frame, int levelsUp) { - return getReceiver(getFrame(frame, levelsUp)); + var owner = getOwner(frame); + if (levelsUp == 0 && !owner.isParseTimeInvisibleScope()) { + return getReceiver(frame); + } + return getReceiver(getEnclosingFrame(owner, levelsUp)); } - public static VirtualFrame getFrame(VirtualFrame frame, int levelsUp) { + public static MaterializedFrame getEnclosingFrame(VmObjectLike owner, int levelsUp) { + assert !(levelsUp == 0 && !owner.isParseTimeInvisibleScope()) + : "Must check for levelsUp == 0 && owner.isParseTimeInvisibleScope() before calling this method"; + var isInvisibleScope = owner.isParseTimeInvisibleScope(); + var enclosingFrame = owner.getEnclosingFrame(); + var remainingLevels = isInvisibleScope ? levelsUp : levelsUp - 1; + return doGetFrame(enclosingFrame, remainingLevels); + } + + private static MaterializedFrame doGetFrame(MaterializedFrame frame, int levelsUp) { frame = skipInvisibleScopes(frame); if (levelsUp == 0) { return frame; @@ -194,7 +211,7 @@ public final class VmUtils { return frame; } - private static VirtualFrame skipInvisibleScopes(VirtualFrame frame) { + private static MaterializedFrame skipInvisibleScopes(MaterializedFrame frame) { var owner = getOwner(frame); while (owner.isParseTimeInvisibleScope()) { frame = owner.getEnclosingFrame();