1 /*
2 * Copyright (c) 2025, 2026, Oracle and/or its affiliates. All rights reserved.
3 * DO NOT ALTER OR REMOVE COPYRIGHT NOTICES OR THIS FILE HEADER.
4 *
5 * This code is free software; you can redistribute it and/or modify it
6 * under the terms of the GNU General Public License version 2 only, as
7 * published by the Free Software Foundation.
8 *
9 * This code is distributed in the hope that it will be useful, but WITHOUT
10 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
11 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
12 * version 2 for more details (a copy is included in the LICENSE file that
13 * accompanied this code).
14 *
15 * You should have received a copy of the GNU General Public License version
16 * 2 along with this work; if not, write to the Free Software Foundation,
17 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
18 *
19 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
20 * or visit www.oracle.com if you need additional information or have any
21 * questions.
22 */
23
24 import jdk.incubator.code.*;
25 import jdk.incubator.code.dialect.core.CoreOp;
26 import jdk.incubator.code.dialect.core.FunctionType;
27 import jdk.incubator.code.dialect.java.*;
28 import jdk.incubator.code.extern.ExternalizedOp;
29
30 import java.lang.invoke.MethodHandle;
31 import java.lang.invoke.MethodHandles;
32 import java.lang.invoke.MethodType;
33 import java.lang.invoke.VarHandle;
34 import java.lang.reflect.Array;
35 import java.util.*;
36 import java.util.function.Predicate;
37 import java.util.stream.Collectors;
38 import java.util.stream.Stream;
39
40 final class PartialEvaluator {
41 final Set<Value> constants;
42 final Predicate<Op> opConstant;
43
44 PartialEvaluator(Set<Value> constants, Predicate<Op> opConstant) {
45 this.constants = new LinkedHashSet<>(constants);
46 this.opConstant = opConstant;
47 }
48
49 public static
50 CoreOp.FuncOp evaluate(MethodHandles.Lookup l,
51 Predicate<Op> opConstant, Set<Value> constants,
52 CoreOp.FuncOp op) {
53 PartialEvaluator pe = new PartialEvaluator(constants, opConstant);
54 Body.Builder outBody = pe.evaluateBody(l, op.body());
55 return CoreOp.func(op.funcName(), outBody);
56 }
57
58
59 @SuppressWarnings("serial")
60 public static final class EvaluationException extends RuntimeException {
61 private EvaluationException(Throwable cause) {
62 super(cause);
63 }
64 }
65
66 static EvaluationException evaluationException(Throwable cause) {
67 return new EvaluationException(cause);
68 }
69
70 static final class BodyContext {
71 final BodyContext parent;
72
73 final Map<Block, List<Block>> evaluatedPredecessors;
74 final Map<Value, Object> evaluatedValues;
75
76 final Queue<Block> blockStack;
77 final BitSet visited;
78
79 BodyContext(Block entryBlock) {
80 this.parent = null;
81
82 this.evaluatedPredecessors = new HashMap<>();
83 this.evaluatedValues = new HashMap<>();
84 this.blockStack = new PriorityQueue<>(Comparator.comparingInt(Block::index));
85
86 this.visited = new BitSet();
87 }
88
89 Object getValue(Value v) {
90 Object rv = evaluatedValues.get(v);
91 if (rv != null) {
92 return rv;
93 }
94
95 throw evaluationException(new IllegalArgumentException("Undefined value: " + v));
96 }
97
98 void setValue(Value v, Object o) {
99 evaluatedValues.put(v, o);
100 }
101 }
102
103 Body.Builder evaluateBody(MethodHandles.Lookup l,
104 Body inBody) {
105 Block inEntryBlock = inBody.entryBlock();
106
107 Body.Builder outBody = Body.Builder.of(null, inBody.bodySignature());
108 Block.Builder outEntryBlock = outBody.entryBlock();
109
110 CodeContext cc = outEntryBlock.context();
111 cc.mapBlock(inEntryBlock, outEntryBlock);
112 cc.mapValues(inEntryBlock.parameters(), outEntryBlock.parameters());
113
114 evaluateEntryBlock(l, inEntryBlock, outEntryBlock, new BodyContext(inEntryBlock));
115
116 return outBody;
117 }
118
119 void evaluateEntryBlock(MethodHandles.Lookup l,
120 Block inEntryBlock,
121 Block.Builder outEntryBlock,
122 BodyContext bc) {
123 assert inEntryBlock.isEntryBlock();
124
125 Map<Block, LoopAnalyzer.Loop> loops = new HashMap<>();
126 Set<Block> loopNoPeeling = new HashSet<>();
127
128 // The first block cannot have any successors so the queue will have at least one entry
129 bc.blockStack.add(inEntryBlock);
130 while (!bc.blockStack.isEmpty()) {
131 final Block inBlock = bc.blockStack.poll();
132 if (bc.visited.get(inBlock.index())) {
133 continue;
134 }
135 bc.visited.set(inBlock.index());
136
137 final Block.Builder outBlock = outEntryBlock.context().getBlock(inBlock);
138
139 nopeel: if (inBlock.predecessors().size() > 1 && bc.evaluatedPredecessors.get(inBlock).size() == 1) {
140 // If we reached to this block through just one evaluated predecessor
141 Block inBlockPred = bc.evaluatedPredecessors.get(inBlock).getFirst();
142 Block.Reference inBlockRef = inBlockPred.terminatingOp().successors().stream()
143 .filter(r -> r.targetBlock() == inBlock)
144 .findFirst().get();
145 List<Value> args = inBlockRef.arguments();
146 List<Boolean> argConstant = args.stream().map(constants::contains).toList();
147
148 LoopAnalyzer.Loop loop = loops.computeIfAbsent(inBlock, b -> LoopAnalyzer.isLoop(inBlock).orElse(null));
149 if (loop != null && inBlockPred.isDominatedBy(loop.header())) {
150 // Entering loop header from latch
151 assert loop.latches().contains(inBlockPred);
152
153 // Linear constant path from each exiting block (or nearest evaluated present dominator) to loop header
154 boolean constantExits = true;
155 for (LoopAnalyzer.LoopExit loopExitPair : loop.exits()) {
156 Block loopExit = loopExitPair.exit();
157
158 // Find nearest evaluated dominator
159 List<Block> ePreds = bc.evaluatedPredecessors.get(loopExit);
160 while (ePreds == null) {
161 loopExit = loopExit.immediateDominator();
162 ePreds = bc.evaluatedPredecessors.get(loopExit);
163 }
164 assert loop.body().contains(loopExit);
165
166 if (ePreds.size() != 1 ||
167 !(loopExit.terminatingOp() instanceof CoreOp.ConditionalBranchOp cbr) ||
168 !constants.contains(cbr.result())) {
169 // If there are multiple encounters, or terminal op is not a constant conditional branch
170 constantExits = false;
171 break;
172 }
173 }
174
175 // Determine if constant args, before reset
176 boolean constantArgs = constants.containsAll(args);
177
178 // Reset state within loop body
179 for (Block block : loop.body()) {
180 // Reset visits, but not for loop header
181 if (block != loop.header()) {
182 bc.evaluatedPredecessors.remove(block);
183 bc.visited.set(block.index(), false);
184 }
185
186 // Reset constants
187 for (Op op : block.ops()) {
188 constants.remove(op.result());
189 }
190 constants.removeAll(block.parameters());
191
192 // Reset no peeling for any nested loops
193 loopNoPeeling.remove(block);
194 }
195
196 if (!constantExits || !constantArgs) {
197 // Finish peeling
198 // No constant exit and no constant args
199 loopNoPeeling.addAll(loop.latches());
200 break nopeel;
201 }
202 // Peel next iteration
203 }
204
205 // Propagate constant arguments
206 for (int i = 0; i < args.size(); i++) {
207 Value inArgument = args.get(i);
208 if (argConstant.get(i)) {
209 Block.Parameter inParameter = inBlock.parameters().get(i);
210
211 // Map input parameter to output argument
212 outBlock.context().mapValue(inParameter, outBlock.context().getValue(inArgument));
213 // Set parameter constant
214 constants.add(inParameter);
215 bc.setValue(inParameter, bc.getValue(inArgument));
216 }
217 }
218 }
219
220 // Process all but the terminating operation
221 int nops = inBlock.ops().size();
222 for (int i = 0; i < nops - 1; i++) {
223 Op op = inBlock.ops().get(i);
224
225 if (isConstant(op)) {
226 // Evaluate operation
227 // @@@ Handle exceptions
228 Object result = interpretOp(l, bc, op);
229 bc.setValue(op.result(), result);
230
231 if (op instanceof CoreOp.VarOp) {
232 // @@@ Do not turn into constant to avoid conflicts with the interpreter
233 // and its runtime representation of vars
234 outBlock.add(op);
235 } else {
236 // Result was evaluated, replace with constant operation
237 Op.Result constantResult = outBlock.add(CoreOp.constant(op.resultType(), result));
238 outBlock.context().mapValue(op.result(), constantResult);
239 }
240 } else {
241 // Copy unevaluated operation
242 Op.Result r = outBlock.add(op);
243 // Explicitly remap result, since the op can be copied more than once in pealed loops
244 // @@@ See comment Block.op code which implicitly limits this
245 outBlock.context().mapValue(op.result(), r);
246 }
247 }
248
249 // Process the terminating operation
250 Op to = inBlock.terminatingOp();
251 switch (to) {
252 case CoreOp.ConditionalBranchOp cb -> {
253 if (isConstant(to)) {
254 boolean p = switch (bc.getValue(cb.predicateOperand())) {
255 case Boolean bp -> bp;
256 case Integer ip ->
257 // @@@ This is required when lifting up from bytecode, since boolean values
258 // are erased to int values, abd the bytecode lifting implementation is not currently
259 // sophisticated enough to recover the type information
260 ip != 0;
261 default -> throw evaluationException(
262 new UnsupportedOperationException("Unsupported type input to operation: " + cb));
263 };
264
265 Block.Reference nextInBlockRef = p ? cb.trueBranch() : cb.falseBranch();
266 Block nextInBlock = nextInBlockRef.targetBlock();
267
268 // @@@ might be latch to loop
269 assert !inBlock.isDominatedBy(nextInBlock);
270
271 processBlock(bc, inBlock, nextInBlock, outBlock);
272
273 outBlock.add(CoreOp.branch(outBlock.context().getReferenceOrCreate(nextInBlockRef)));
274 } else {
275 // @@@ might be non-constant latch to loop
276 processBlock(bc, inBlock, cb.falseBranch().targetBlock(), outBlock);
277 processBlock(bc, inBlock, cb.trueBranch().targetBlock(), outBlock);
278
279 outBlock.add(to);
280 }
281 }
282 case CoreOp.BranchOp b -> {
283 Block.Reference nextInBlockRef = b.branch();
284 Block nextInBlock = nextInBlockRef.targetBlock();
285
286 if (inBlock.isDominatedBy(nextInBlock)) {
287 // latch to loop header
288 assert bc.visited.get(nextInBlock.index());
289 if (!loopNoPeeling.contains(inBlock) && constants.containsAll(nextInBlock.parameters())) {
290 // Reset loop body to peel off another iteration
291 bc.visited.set(nextInBlock.index(), false);
292 bc.evaluatedPredecessors.remove(nextInBlock);
293 }
294 }
295
296 processBlock(bc, inBlock, nextInBlock, outBlock);
297
298 outBlock.add(b);
299 }
300 case CoreOp.ReturnOp _ -> outBlock.add(to);
301 default -> throw evaluationException(
302 new UnsupportedOperationException("Unsupported terminating operation: " + to));
303 }
304 }
305 }
306
307 boolean isConstant(Op op) {
308 if (constants.contains(op.result())) {
309 return true;
310 } else if (constants.containsAll(op.operands()) && opConstant.test(op)) {
311 constants.add(op.result());
312 return true;
313 } else {
314 return false;
315 }
316 }
317
318 void processBlock(BodyContext bc, Block inBlock, Block nextInBlock, Block.Builder outBlock) {
319 bc.blockStack.add(nextInBlock);
320 if (!bc.evaluatedPredecessors.containsKey(nextInBlock)) {
321 // Copy block
322 Block.Builder nextOutBlock = outBlock.block(nextInBlock.parameterTypes());
323 outBlock.context().mapBlock(nextInBlock, nextOutBlock);
324 outBlock.context().mapValues(nextInBlock.parameters(), nextOutBlock.parameters());
325 }
326 bc.evaluatedPredecessors.computeIfAbsent(nextInBlock, _ -> new ArrayList<>()).add(inBlock);
327 }
328
329 @SuppressWarnings("unchecked")
330 public static <E extends Throwable> void eraseAndThrow(Throwable e) throws E {
331 throw (E) e;
332 }
333
334 // @@@ This could be shared with the interpreter if it was more extensible
335 Object interpretOp(MethodHandles.Lookup l, BodyContext bc, Op o) {
336 switch (o) {
337 case CoreOp.ConstantOp co -> {
338 if (co.resultType().equals(JavaType.J_L_CLASS)) {
339 return resolveToClass(l, (JavaType) co.value());
340 } else {
341 return co.value();
342 }
343 }
344 case JavaOp.InvokeOp co -> {
345 MethodType target = resolveToMethodType(l, o.opSignature());
346 MethodHandles.Lookup il = switch (co.invokeKind()) {
347 case STATIC, INSTANCE -> l;
348 case SUPER -> l.in(target.parameterType(0));
349 };
350 MethodHandle mh = resolveToMethodHandle(il, co.invokeReference(), co.invokeKind());
351
352 mh = mh.asType(target).asFixedArity();
353 Object[] values = o.operands().stream().map(bc::getValue).toArray();
354 return invoke(mh, values);
355 }
356 case JavaOp.NewOp no -> {
357 Object[] values = o.operands().stream().map(bc::getValue).toArray();
358 JavaType nType = (JavaType) no.resultType();
359 if (nType instanceof ArrayType at) {
360 if (values.length > at.dimensions()) {
361 throw evaluationException(new IllegalArgumentException("Bad constructor NewOp: " + no));
362 }
363 int[] lengths = Stream.of(values).mapToInt(v -> (int) v).toArray();
364 for (int length : lengths) {
365 nType = ((ArrayType) nType).componentType();
366 }
367 return Array.newInstance(resolveToClass(l, nType), lengths);
368 } else {
369 MethodHandle mh = constructorHandle(l, no.constructorReference().signature());
370 return invoke(mh, values);
371 }
372 }
373 case CoreOp.VarOp vo -> {
374 Object[] vbox = vo.isUninitialized()
375 ? new Object[] { null, false }
376 : new Object[] { bc.getValue(o.operands().get(0)) };
377 return vbox;
378 }
379 case CoreOp.VarAccessOp.VarLoadOp vlo -> {
380 // Cast to CoreOp.Var, since the instance may have originated as an external instance
381 // via a captured value map
382 Object[] vbox = (Object[]) bc.getValue(o.operands().get(0));
383 if (vbox.length == 2 && !((Boolean) vbox[1])) {
384 throw evaluationException(new IllegalStateException("Loading from uninitialized variable"));
385 }
386 return vbox[0];
387 }
388 case CoreOp.VarAccessOp.VarStoreOp vso -> {
389 Object[] vbox = (Object[]) bc.getValue(o.operands().get(0));
390 if (vbox.length == 2) {
391 vbox[1] = true;
392 }
393 vbox[0] = bc.getValue(o.operands().get(1));
394 return null;
395 }
396 case CoreOp.TupleOp to -> {
397 return o.operands().stream().map(bc::getValue).toList();
398 }
399 case CoreOp.TupleLoadOp tlo -> {
400 @SuppressWarnings("unchecked")
401 List<Object> tb = (List<Object>) bc.getValue(o.operands().get(0));
402 return tb.get(tlo.index());
403 }
404 case CoreOp.TupleWithOp two -> {
405 @SuppressWarnings("unchecked")
406 List<Object> tb = (List<Object>) bc.getValue(o.operands().get(0));
407 List<Object> copy = new ArrayList<>(tb);
408 copy.set(two.index(), bc.getValue(o.operands().get(1)));
409 return Collections.unmodifiableList(copy);
410 }
411 case JavaOp.FieldAccessOp.FieldLoadOp fo -> {
412 if (fo.operands().isEmpty()) {
413 VarHandle vh = fieldStaticHandle(l, fo.fieldReference());
414 return vh.get();
415 } else {
416 Object v = bc.getValue(o.operands().get(0));
417 VarHandle vh = fieldHandle(l, fo.fieldReference());
418 return vh.get(v);
419 }
420 }
421 case JavaOp.FieldAccessOp.FieldStoreOp fo -> {
422 if (fo.operands().size() == 1) {
423 Object v = bc.getValue(o.operands().get(0));
424 VarHandle vh = fieldStaticHandle(l, fo.fieldReference());
425 vh.set(v);
426 } else {
427 Object r = bc.getValue(o.operands().get(0));
428 Object v = bc.getValue(o.operands().get(1));
429 VarHandle vh = fieldHandle(l, fo.fieldReference());
430 vh.set(r, v);
431 }
432 return null;
433 }
434 case JavaOp.InstanceOfOp io -> {
435 Object v = bc.getValue(o.operands().get(0));
436 return isInstance(l, io.targetType(), v);
437 }
438 case JavaOp.CastOp co -> {
439 Object v = bc.getValue(o.operands().get(0));
440 return cast(l, co.targetType(), v);
441 }
442 case JavaOp.ArrayLengthOp arrayLengthOp -> {
443 Object a = bc.getValue(o.operands().get(0));
444 return Array.getLength(a);
445 }
446 case JavaOp.ArrayAccessOp.ArrayLoadOp arrayLoadOp -> {
447 Object a = bc.getValue(o.operands().get(0));
448 Object index = bc.getValue(o.operands().get(1));
449 return Array.get(a, (int) index);
450 }
451 case JavaOp.ArrayAccessOp.ArrayStoreOp arrayStoreOp -> {
452 Object a = bc.getValue(o.operands().get(0));
453 Object index = bc.getValue(o.operands().get(1));
454 Object v = bc.getValue(o.operands().get(2));
455 Array.set(a, (int) index, v);
456 return null;
457 }
458 case JavaOp.ArithmeticOperation arithmeticOperation -> {
459 MethodHandle mh = opHandle(l, externalizeOpName(o), o.opSignature());
460 Object[] values = o.operands().stream().map(bc::getValue).toArray();
461 return invoke(mh, values);
462 }
463 case JavaOp.ConvOp convOp -> {
464 MethodHandle mh = opHandle(l, externalizeOpName(o) + "_" + o.opSignature().returnType(), o.opSignature());
465 Object[] values = o.operands().stream().map(bc::getValue).toArray();
466 return invoke(mh, values);
467 }
468 case JavaOp.ConcatOp concatOp -> {
469 return o.operands().stream()
470 .map(bc::getValue)
471 .map(String::valueOf)
472 .collect(Collectors.joining());
473 }
474 // @@@
475 // case CoreOp.LambdaOp lambdaOp -> {
476 // interpretEntryBlock(l, lambdaOp.body().entryBlock(), oc, new HashMap<>());
477 // unevaluatedOperations.add(o);
478 // return null;
479 // }
480 // case CoreOp.FuncOp funcOp -> {
481 // interpretEntryBlock(l, funcOp.body().entryBlock(), oc, new HashMap<>());
482 // unevaluatedOperations.add(o);
483 // return null;
484 // }
485 case null, default -> throw evaluationException(
486 new UnsupportedOperationException("Unsupported operation: " + o));
487 }
488 }
489
490
491 static String externalizeOpName(Op op) {
492 return (op instanceof ExternalizedOp.Externalizable eop)
493 ? eop.externalizeOpName()
494 : op.getClass().getName();
495 }
496
497 static MethodHandle opHandle(MethodHandles.Lookup l, String opName, FunctionType ft) {
498 MethodType mt = resolveToMethodType(l, ft).erase();
499 try {
500 return MethodHandles.lookup().findStatic(InvokableLeafOps.class, opName, mt);
501 } catch (NoSuchMethodException | IllegalAccessException e) {
502 throw evaluationException(e);
503 }
504 }
505
506 static MethodHandle constructorHandle(MethodHandles.Lookup l, FunctionType ft) {
507 MethodType mt = resolveToMethodType(l, ft);
508
509 if (mt.returnType().isArray()) {
510 if (mt.parameterCount() != 1 || mt.parameterType(0) != int.class) {
511 throw evaluationException(new IllegalArgumentException("Bad constructor descriptor: " + ft));
512 }
513 return MethodHandles.arrayConstructor(mt.returnType());
514 } else {
515 try {
516 return l.findConstructor(mt.returnType(), mt.changeReturnType(void.class));
517 } catch (NoSuchMethodException | IllegalAccessException e) {
518 throw evaluationException(e);
519 }
520 }
521 }
522
523 static VarHandle fieldStaticHandle(MethodHandles.Lookup l, FieldRef d) {
524 return resolveToVarHandle(l, d);
525 }
526
527 static VarHandle fieldHandle(MethodHandles.Lookup l, FieldRef d) {
528 return resolveToVarHandle(l, d);
529 }
530
531 static Object isInstance(MethodHandles.Lookup l, CodeType d, Object v) {
532 Class<?> c = resolveToClass(l, d);
533 return c.isInstance(v);
534 }
535
536 static Object cast(MethodHandles.Lookup l, CodeType d, Object v) {
537 Class<?> c = resolveToClass(l, d);
538 return c.cast(v);
539 }
540
541 static MethodHandle resolveToMethodHandle(MethodHandles.Lookup l, MethodRef d, JavaOp.InvokeOp.InvokeKind kind) {
542 try {
543 return d.resolveToHandle(l, kind);
544 } catch (ReflectiveOperationException e) {
545 throw evaluationException(e);
546 }
547 }
548
549 static VarHandle resolveToVarHandle(MethodHandles.Lookup l, FieldRef d) {
550 try {
551 return d.resolveToHandle(l);
552 } catch (ReflectiveOperationException e) {
553 throw evaluationException(e);
554 }
555 }
556
557 public static MethodType resolveToMethodType(MethodHandles.Lookup l, FunctionType ft) {
558 try {
559 return MethodRef.toNominalDescriptor(ft).resolveConstantDesc(l);
560 } catch (ReflectiveOperationException e) {
561 throw evaluationException(e);
562 }
563 }
564
565 public static Class<?> resolveToClass(MethodHandles.Lookup l, CodeType d) {
566 try {
567 if (d instanceof JavaType jt) {
568 return (Class<?>) jt.erasure().resolve(l);
569 } else {
570 throw new ReflectiveOperationException();
571 }
572 } catch (ReflectiveOperationException e) {
573 throw evaluationException(e);
574 }
575 }
576
577 static Object invoke(MethodHandle m, Object... args) {
578 try {
579 return m.invokeWithArguments(args);
580 } catch (RuntimeException | Error e) {
581 throw e;
582 } catch (Throwable e) {
583 eraseAndThrow(e);
584 throw new InternalError("should not reach here");
585 }
586 }
587 }