Overview
Truffle Node Rewriting is a fundamental optimization technique in the GraalVM Truffle framework that allows dynamic specialization of AST nodes based on runtime information. This enables high-performance language implementation through adaptive optimization.
Key Concepts
- Node Specialization: Transforming nodes to more specific versions
- Polymorphic Inline Caches: Caching multiple execution paths
- AST Transformation: Dynamic modification of abstract syntax trees
- Profile-Guided Optimization: Using runtime feedback for optimization
Basic Node Structure
1. Foundation Node Classes
import com.oracle.truffle.api.nodes.Node; import com.oracle.truffle.api.nodes.NodeInterface; import com.oracle.truffle.api.dsl.Specialization; // Base node interface public interface ExpressionNode extends NodeInterface { Object execute(VirtualFrame frame); } // Abstract base node public abstract class BaseNode extends Node implements ExpressionNode { // Node rewriting hook protected abstract ExpressionNode copy(); // Report optimization for debugging protected final void reportRewrite(String from, String to) { System.out.printf("Rewriting %s -> %s at %s%n", from, to, getSourceSection()); } } 2. Simple Arithmetic Nodes
public class ArithmeticNodes { // Generic add node that specializes based on types public abstract static class AddNode extends BaseNode { @Specialization(rewriteOn = ArithmeticException.class) protected int addInt(int left, int right) { return Math.addExact(left, right); } @Specialization(rewriteOn = ArithmeticException.class) protected long addLong(long left, long right) { return Math.addExact(left, right); } @Specialization protected double addDouble(double left, double right) { return left + right; } @Specialization protected String addString(String left, String right) { return left + right; } @Specialization(replaces = {"addInt", "addLong", "addDouble", "addString"}) protected Object addGeneric(Object left, Object right) { // Fallback for mixed types if (left instanceof Number && right instanceof Number) { return ((Number) left).doubleValue() + ((Number) right).doubleValue(); } return left.toString() + right.toString(); } } // Specialized node for integer addition only public abstract static class IntAddNode extends BaseNode { @Child private ExpressionNode leftNode; @Child private ExpressionNode rightNode; public IntAddNode(ExpressionNode leftNode, ExpressionNode rightNode) { this.leftNode = leftNode; this.rightNode = rightNode; } @Specialization protected int addInt(VirtualFrame frame) { int left = (int) leftNode.execute(frame); int right = (int) rightNode.execute(frame); return Math.addExact(left, right); } @Override protected ExpressionNode copy() { return new IntAddNode((ExpressionNode) leftNode.copy(), (ExpressionNode) rightNode.copy()); } } } Node Rewriting Patterns
1. Specialization-based Rewriting
public class SpecializationRewriting { // Variable access node that specializes based on usage patterns public abstract static class VariableNode extends BaseNode { protected final String name; protected final FrameSlot slot; public VariableNode(String name, FrameSlot slot) { this.name = name; this.slot = slot; } // Initially uninitialized specialization @Specialization(guards = "isUninitialized(frame)") protected Object readUninitialized(VirtualFrame frame) { throw new UninitializedVariableException(name); } // Integer specialization @Specialization(guards = "isInt(frame)") protected int readInt(VirtualFrame frame) { return frame.getInt(slot); } // Double specialization @Specialization(guards = "isDouble(frame)") protected double readDouble(VirtualFrame frame) { return frame.getDouble(slot); } // String specialization @Specialization(guards = "isString(frame)") protected String readString(VirtualFrame frame) { return (String) frame.getObject(slot); } // Guard methods protected boolean isUninitialized(VirtualFrame frame) { return !frame.isObject(slot) && !frame.isInt(slot) && !frame.isDouble(slot); } protected boolean isInt(VirtualFrame frame) { return frame.isInt(slot); } protected boolean isDouble(VirtualFrame frame) { return frame.isDouble(slot); } protected boolean isString(VirtualFrame frame) { Object value = frame.getObject(slot); return value instanceof String; } @Override protected ExpressionNode copy() { return new VariableNode(name, slot) {}; } } public static class UninitializedVariableException extends RuntimeException { public UninitializedVariableException(String name) { super("Variable '" + name + "' is uninitialized"); } } } 2. Manual Node Replacement
public class ManualNodeReplacement { // Generic comparison node that replaces itself with specialized versions public static class ComparisonNode extends BaseNode { @Child private ExpressionNode left; @Child private ExpressionNode right; private int intComparisons = 0; private int doubleComparisons = 0; private final int specializationThreshold = 10; public ComparisonNode(ExpressionNode left, ExpressionNode right) { this.left = left; this.right = right; } @Override public Object execute(VirtualFrame frame) { Object leftValue = left.execute(frame); Object rightValue = right.execute(frame); // Profile the types if (leftValue instanceof Integer && rightValue instanceof Integer) { intComparisons++; } else if (leftValue instanceof Double && rightValue instanceof Double) { doubleComparisons++; } // Check if we should specialize if (intComparisons >= specializationThreshold) { return replaceWithIntComparison().execute(frame); } else if (doubleComparisons >= specializationThreshold) { return replaceWithDoubleComparison().execute(frame); } // Generic comparison return compareGeneric(leftValue, rightValue); } private ExpressionNode replaceWithIntComparison() { reportRewrite("GenericComparison", "IntComparison"); IntComparisonNode specialized = new IntComparisonNode(left, right); return replace(specialized); } private ExpressionNode replaceWithDoubleComparison() { reportRewrite("GenericComparison", "DoubleComparison"); DoubleComparisonNode specialized = new DoubleComparisonNode(left, right); return replace(specialized); } private boolean compareGeneric(Object left, Object right) { if (left instanceof Number && right instanceof Number) { return ((Number) left).doubleValue() == ((Number) right).doubleValue(); } return left.equals(right); } @Override protected ExpressionNode copy() { return new ComparisonNode((ExpressionNode) left.copy(), (ExpressionNode) right.copy()); } } // Specialized integer comparison public static class IntComparisonNode extends BaseNode { @Child private ExpressionNode left; @Child private ExpressionNode right; public IntComparisonNode(ExpressionNode left, ExpressionNode right) { this.left = left; this.right = right; } @Override public Object execute(VirtualFrame frame) { int leftVal = (int) left.execute(frame); int rightVal = (int) right.execute(frame); return leftVal == rightVal; } @Override protected ExpressionNode copy() { return new IntComparisonNode((ExpressionNode) left.copy(), (ExpressionNode) right.copy()); } } // Specialized double comparison public static class DoubleComparisonNode extends BaseNode { @Child private ExpressionNode left; @Child private ExpressionNode right; public DoubleComparisonNode(ExpressionNode left, ExpressionNode right) { this.left = left; this.right = right; } @Override public Object execute(VirtualFrame frame) { double leftVal = (double) left.execute(frame); double rightVal = (double) right.execute(frame); return Math.abs(leftVal - rightVal) < 1e-10; } @Override protected ExpressionNode copy() { return new DoubleComparisonNode((ExpressionNode) left.copy(), (ExpressionNode) right.copy()); } } } Advanced Rewriting Techniques
1. Polymorphic Inline Caches
public class PolymorphicInlineCache { // Method call node with polymorphic inline cache public static class MethodCallNode extends BaseNode { private final String methodName; @Children private ExpressionNode[] arguments; // Cache for monomorphic, bimorphic, and megamorphic cases private static final int MAX_INLINE_CACHE_SIZE = 3; private final MethodCacheEntry[] cache = new MethodCacheEntry[MAX_INLINE_CACHE_SIZE]; private int cacheSize = 0; public MethodCallNode(String methodName, ExpressionNode[] arguments) { this.methodName = methodName; this.arguments = arguments; } @Override public Object execute(VirtualFrame frame) { Object receiver = arguments[0].execute(frame); // Try cache lookup first MethodCacheEntry cached = findInCache(receiver.getClass()); if (cached != null) { return cached.execute(receiver, frame, arguments); } // Cache miss - resolve method Method method = resolveMethod(receiver, methodName); MethodCacheEntry newEntry = new MethodCacheEntry(receiver.getClass(), method); // Update cache if (cacheSize < MAX_INLINE_CACHE_SIZE) { cache[cacheSize++] = newEntry; reportRewrite("CacheMiss", "Cached-" + receiver.getClass().getSimpleName()); } else { // Megamorphic case - replace with generic dispatch return replaceWithGenericDispatch().execute(frame); } return newEntry.execute(receiver, frame, arguments); } private MethodCacheEntry findInCache(Class<?> receiverClass) { for (int i = 0; i < cacheSize; i++) { if (cache[i].receiverClass == receiverClass) { return cache[i]; } } return null; } private ExpressionNode replaceWithGenericDispatch() { reportRewrite("PolymorphicCall", "GenericDispatch"); GenericMethodCallNode generic = new GenericMethodCallNode(methodName, arguments); return replace(generic); } private Method resolveMethod(Object receiver, String methodName) { // Simplified method resolution try { return receiver.getClass().getMethod(methodName); } catch (NoSuchMethodException e) { throw new RuntimeException("Method not found: " + methodName); } } @Override protected ExpressionNode copy() { ExpressionNode[] argsCopy = new ExpressionNode[arguments.length]; for (int i = 0; i < arguments.length; i++) { argsCopy[i] = (ExpressionNode) arguments[i].copy(); } return new MethodCallNode(methodName, argsCopy); } } // Cache entry for method dispatch private static class MethodCacheEntry { final Class<?> receiverClass; final Method method; MethodCacheEntry(Class<?> receiverClass, Method method) { this.receiverClass = receiverClass; this.method = method; } Object execute(Object receiver, VirtualFrame frame, ExpressionNode[] arguments) { try { Object[] args = new Object[arguments.length - 1]; for (int i = 1; i < arguments.length; i++) { args[i - 1] = arguments[i].execute(frame); } return method.invoke(receiver, args); } catch (Exception e) { throw new RuntimeException("Method invocation failed", e); } } } // Generic dispatch for megamorphic cases public static class GenericMethodCallNode extends BaseNode { private final String methodName; @Children private ExpressionNode[] arguments; public GenericMethodCallNode(String methodName, ExpressionNode[] arguments) { this.methodName = methodName; this.arguments = arguments; } @Override public Object execute(VirtualFrame frame) { Object receiver = arguments[0].execute(frame); Object[] args = new Object[arguments.length - 1]; for (int i = 1; i < arguments.length; i++) { args[i - 1] = arguments[i].execute(frame); } // Use reflection for generic dispatch try { Method method = receiver.getClass().getMethod(methodName); return method.invoke(receiver, args); } catch (Exception e) { throw new RuntimeException("Method invocation failed", e); } } @Override protected ExpressionNode copy() { ExpressionNode[] argsCopy = new ExpressionNode[arguments.length]; for (int i = 0; i < arguments.length; i++) { argsCopy[i] = (ExpressionNode) arguments[i].copy(); } return new GenericMethodCallNode(methodName, argsCopy); } } } 2. Loop Optimization through Rewriting
public class LoopRewriting { // Loop node that specializes based on iteration patterns public static class WhileLoopNode extends BaseNode { @Child private ExpressionNode condition; @Child private ExpressionNode body; private int iterationCount = 0; private static final int UNROLL_THRESHOLD = 100; private static final int UNROLL_FACTOR = 4; public WhileLoopNode(ExpressionNode condition, ExpressionNode body) { this.condition = condition; this.body = body; } @Override public Object execute(VirtualFrame frame) { iterationCount = 0; Object result = null; while ((Boolean) condition.execute(frame)) { result = body.execute(frame); iterationCount++; // Check if we should unroll the loop if (iterationCount >= UNROLL_THRESHOLD) { return replaceWithUnrolledLoop().execute(frame); } } return result; } private ExpressionNode replaceWithUnrolledLoop() { reportRewrite("WhileLoop", "UnrolledLoop"); UnrolledLoopNode unrolled = new UnrolledLoopNode(condition, body, UNROLL_FACTOR); return replace(unrolled); } @Override protected ExpressionNode copy() { return new WhileLoopNode((ExpressionNode) condition.copy(), (ExpressionNode) body.copy()); } } // Unrolled loop implementation public static class UnrolledLoopNode extends BaseNode { @Child private ExpressionNode condition; @Child private ExpressionNode body; private final int unrollFactor; public UnrolledLoopNode(ExpressionNode condition, ExpressionNode body, int unrollFactor) { this.condition = condition; this.body = body; this.unrollFactor = unrollFactor; } @Override public Object execute(VirtualFrame frame) { Object result = null; while (true) { // Check condition if (!(Boolean) condition.execute(frame)) { break; } // Execute unrolled iterations for (int i = 0; i < unrollFactor; i++) { result = body.execute(frame); // Check condition after each iteration if (!(Boolean) condition.execute(frame)) { return result; } } } return result; } @Override protected ExpressionNode copy() { return new UnrolledLoopNode((ExpressionNode) condition.copy(), (ExpressionNode) body.copy(), unrollFactor); } } } DSL-Based Node Generation
1. Using Truffle DSL for Automatic Specialization
import com.oracle.truffle.api.dsl.Fallback; import com.oracle.truffle.api.dsl.NodeChild; import com.oracle.truffle.api.dsl.NodeChildren; import com.oracle.truffle.api.dsl.Specialization; @NodeChildren({ @NodeChild("left"), @NodeChild("right") }) public abstract class DSLAddNode extends BaseNode { @Specialization(rewriteOn = ArithmeticException.class) protected int add(int left, int right) { return Math.addExact(left, right); } @Specialization(rewriteOn = ArithmeticException.class) protected long add(long left, long right) { return Math.addExact(left, right); } @Specialization protected double add(double left, double right) { return left + right; } @Specialization protected String add(String left, String right) { return left + right; } @Fallback protected Object add(Object left, Object right) { // Generic fallback if (left instanceof Number && right instanceof Number) { return ((Number) left).doubleValue() + ((Number) right).doubleValue(); } return left.toString() + right.toString(); } // DSL automatically generates execute method and node rewriting logic @Override protected ExpressionNode copy() { // DSL nodes handle copying automatically in most cases return this; } } 2. Custom DSL Node with Guards
@NodeChildren({ @NodeChild("value") }) public abstract class TypeCheckNode extends BaseNode { protected static boolean isInt(Object value) { return value instanceof Integer; } protected static boolean isDouble(Object value) { return value instanceof Double; } protected static boolean isString(Object value) { return value instanceof String; } @Specialization(guards = "isInt(value)") protected String checkInt(Object value) { return "Integer: " + value; } @Specialization(guards = "isDouble(value)") protected String checkDouble(Object value) { return "Double: " + value; } @Specialization(guards = "isString(value)") protected String checkString(Object value) { return "String: " + value; } @Fallback protected String checkOther(Object value) { return "Other: " + value.getClass().getSimpleName(); } } Practical Implementation Examples
Example 1: Simple Language Interpreter with Node Rewriting
public class SimpleLanguageInterpreter { // AST node for literal values public static class LiteralNode extends BaseNode { private final Object value; public LiteralNode(Object value) { this.value = value; } @Override public Object execute(VirtualFrame frame) { return value; } @Override protected ExpressionNode copy() { return new LiteralNode(value); } } // AST node for variable assignment public static class AssignmentNode extends BaseNode { private final String variableName; @Child private ExpressionNode valueNode; private FrameSlot slot; public AssignmentNode(String variableName, ExpressionNode valueNode, FrameDescriptor frameDescriptor) { this.variableName = variableName; this.valueNode = valueNode; this.slot = frameDescriptor.findOrAddFrameSlot(variableName); } @Override public Object execute(VirtualFrame frame) { Object value = valueNode.execute(frame); frame.setObject(slot, value); return value; } @Override protected ExpressionNode copy() { return new AssignmentNode(variableName, (ExpressionNode) valueNode.copy(), slot.getFrameDescriptor()); } } // If statement node with potential rewriting public static class IfNode extends BaseNode { @Child private ExpressionNode condition; @Child private ExpressionNode thenBranch; @Child private ExpressionNode elseBranch; private int trueCount = 0; private int falseCount = 0; private static final int BIAS_THRESHOLD = 100; public IfNode(ExpressionNode condition, ExpressionNode thenBranch, ExpressionNode elseBranch) { this.condition = condition; this.thenBranch = thenBranch; this.elseBranch = elseBranch; } @Override public Object execute(VirtualFrame frame) { boolean cond = (Boolean) condition.execute(frame); // Profile branch distribution if (cond) { trueCount++; if (trueCount > BIAS_THRESHOLD && falseCount < BIAS_THRESHOLD / 10) { return replaceWithThenOnly().execute(frame); } } else { falseCount++; if (falseCount > BIAS_THRESHOLD && trueCount < BIAS_THRESHOLD / 10) { return replaceWithElseOnly().execute(frame); } } return cond ? thenBranch.execute(frame) : elseBranch.execute(frame); } private ExpressionNode replaceWithThenOnly() { reportRewrite("IfNode", "ThenOnlyNode"); ThenOnlyNode specialized = new ThenOnlyNode(condition, thenBranch); return replace(specialized); } private ExpressionNode replaceWithElseOnly() { reportRewrite("IfNode", "ElseOnlyNode"); ElseOnlyNode specialized = new ElseOnlyNode(condition, elseBranch); return replace(specialized); } @Override protected ExpressionNode copy() { return new IfNode((ExpressionNode) condition.copy(), (ExpressionNode) thenBranch.copy(), (ExpressionNode) elseBranch.copy()); } } // Specialized node for mostly-true conditions public static class ThenOnlyNode extends BaseNode { @Child private ExpressionNode condition; @Child private ExpressionNode thenBranch; public ThenOnlyNode(ExpressionNode condition, ExpressionNode thenBranch) { this.condition = condition; this.thenBranch = thenBranch; } @Override public Object execute(VirtualFrame frame) { // Assume condition is true (deoptimize if not) if (!(Boolean) condition.execute(frame)) { throw new DeoptimizeException("Unexpected false condition"); } return thenBranch.execute(frame); } @Override protected ExpressionNode copy() { return new ThenOnlyNode((ExpressionNode) condition.copy(), (ExpressionNode) thenBranch.copy()); } } // Specialized node for mostly-false conditions public static class ElseOnlyNode extends BaseNode { @Child private ExpressionNode condition; @Child private ExpressionNode elseBranch; public ElseOnlyNode(ExpressionNode condition, ExpressionNode elseBranch) { this.condition = condition; this.elseBranch = elseBranch; } @Override public Object execute(VirtualFrame frame) { // Assume condition is false (deoptimize if not) if ((Boolean) condition.execute(frame)) { throw new DeoptimizeException("Unexpected true condition"); } return elseBranch.execute(frame); } @Override protected ExpressionNode copy() { return new ElseOnlyNode((ExpressionNode) condition.copy(), (ExpressionNode) elseBranch.copy()); } } public static class DeoptimizeException extends RuntimeException { public DeoptimizeException(String message) { super(message); } } } Example 2: Array Access Optimization
public class ArrayOptimization { // Generic array access node public static class ArrayAccessNode extends BaseNode { @Child private ExpressionNode arrayNode; @Child private ExpressionNode indexNode; private int boundsCheckSuccess = 0; private static final int BOUNDS_CHECK_THRESHOLD = 1000; public ArrayAccessNode(ExpressionNode arrayNode, ExpressionNode indexNode) { this.arrayNode = arrayNode; this.indexNode = indexNode; } @Override public Object execute(VirtualFrame frame) { Object[] array = (Object[]) arrayNode.execute(frame); int index = (int) indexNode.execute(frame); // Bounds check if (index < 0 || index >= array.length) { throw new ArrayIndexOutOfBoundsException(index); } boundsCheckSuccess++; // Optimize if bounds checks consistently succeed if (boundsCheckSuccess > BOUNDS_CHECK_THRESHOLD) { return replaceWithUncheckedAccess().execute(frame); } return array[index]; } private ExpressionNode replaceWithUncheckedAccess() { reportRewrite("CheckedArrayAccess", "UncheckedArrayAccess"); UncheckedArrayAccessNode optimized = new UncheckedArrayAccessNode(arrayNode, indexNode); return replace(optimized); } @Override protected ExpressionNode copy() { return new ArrayAccessNode((ExpressionNode) arrayNode.copy(), (ExpressionNode) indexNode.copy()); } } // Optimized array access without bounds checking public static class UncheckedArrayAccessNode extends BaseNode { @Child private ExpressionNode arrayNode; @Child private ExpressionNode indexNode; public UncheckedArrayAccessNode(ExpressionNode arrayNode, ExpressionNode indexNode) { this.arrayNode = arrayNode; this.indexNode = indexNode; } @Override public Object execute(VirtualFrame frame) { Object[] array = (Object[]) arrayNode.execute(frame); int index = (int) indexNode.execute(frame); // No bounds check - rely on hardware exception try { return array[index]; } catch (ArrayIndexOutOfBoundsException e) { // Deoptimize back to checked version throw new DeoptimizeException("Bounds check failed"); } } @Override protected ExpressionNode copy() { return new UncheckedArrayAccessNode((ExpressionNode) arrayNode.copy(), (ExpressionNode) indexNode.copy()); } } } Best Practices for Node Rewriting
1. Profiling and Thresholds
public class RewritingBestPractices { public abstract static class OptimizableNode extends BaseNode { protected static final int WARMUP_THRESHOLD = 100; protected static final int OPTIMIZATION_THRESHOLD = 1000; protected static final int DEOPTIMIZATION_THRESHOLD = 10; protected int executionCount = 0; protected int optimizationFailures = 0; protected boolean shouldOptimize() { return executionCount > OPTIMIZATION_THRESHOLD && optimizationFailures < DEOPTIMIZATION_THRESHOLD; } protected boolean shouldDeoptimize() { return optimizationFailures >= DEOPTIMIZATION_THRESHOLD; } protected void recordOptimizationFailure() { optimizationFailures++; } protected void recordExecution() { executionCount++; } } } 2. Incremental Specialization
public class IncrementalSpecialization { public static class IncrementalAddNode extends BaseNode { @Child private ExpressionNode left; @Child private ExpressionNode right; private int intExecutions = 0; private int doubleExecutions = 0; private int stringExecutions = 0; public IncrementalAddNode(ExpressionNode left, ExpressionNode right) { this.left = left; this.right = right; } @Override public Object execute(VirtualFrame frame) { Object leftVal = left.execute(frame); Object rightVal = right.execute(frame); // Profile types if (leftVal instanceof Integer && rightVal instanceof Integer) { intExecutions++; if (intExecutions > 100) { return replaceWithIntAdd().execute(frame); } return (int) leftVal + (int) rightVal; } else if (leftVal instanceof Double && rightVal instanceof Double) { doubleExecutions++; if (doubleExecutions > 100) { return replaceWithDoubleAdd().execute(frame); } return (double) leftVal + (double) rightVal; } else if (leftVal instanceof String && rightVal instanceof String) { stringExecutions++; if (stringExecutions > 100) { return replaceWithStringConcat().execute(frame); } return (String) leftVal + (String) rightVal; } // Generic fallback return leftVal.toString() + rightVal.toString(); } private ExpressionNode replaceWithIntAdd() { reportRewrite("GenericAdd", "IntAdd"); return replace(new IntAddNode(left, right)); } private ExpressionNode replaceWithDoubleAdd() { reportRewrite("GenericAdd", "DoubleAdd"); return replace(new DoubleAddNode(left, right)); } private ExpressionNode replaceWithStringConcat() { reportRewrite("GenericAdd", "StringConcat"); return replace(new StringConcatNode(left, right)); } @Override protected ExpressionNode copy() { return new IncrementalAddNode((ExpressionNode) left.copy(), (ExpressionNode) right.copy()); } } } Truffle Node Rewriting is a powerful technique that enables high-performance language implementations through dynamic specialization. By understanding and applying these patterns, you can create adaptive ASTs that optimize themselves based on runtime behavior.