1 package jdk.incubator.code.analysis;
2
3 import jdk.incubator.code.*;
4 import jdk.incubator.code.dialect.core.CoreOp;
5
6 import java.util.HashMap;
7 import java.util.List;
8 import java.util.Map;
9 import java.util.function.BiConsumer;
10
11 import static jdk.incubator.code.dialect.core.CoreOp.branch;
12 import static jdk.incubator.code.dialect.core.CoreOp.return_;
13
14 /**
15 * Functionality for inlining code.
16 */
17 public final class Inliner {
18
19 private Inliner() {}
20
21 /**
22 * An inline consumer that inserts a return operation with a value, if non-null.
23 */
24 public static final BiConsumer<Block.Builder, Value> INLINE_RETURN = (block, value) -> {
25 block.op(value != null ? return_(value) : CoreOp.return_());
26 };
27
28 /**
29 * Inlines the invokable operation into the given block builder and returns the block builder from which to
30 * continue building.
31 * <p>
32 * This method {@link Block.Builder#transformBody(Body, List, CopyContext, OpTransformer) transforms} the
33 * body of the invokable operation with the given arguments, a new context, and an operation transformer that
34 * replaces return operations by applying the given consumer to a block builder and a return value.
35 * <p>
36 * The operation transformer copies all operations except return operations whose nearest invokable operation
37 * ancestor is the given the invokable operation. When such a return operation is encountered, then on
38 * first encounter of its grandparent body a return block builder is computed and used for this return operation
39 * and encounters of subsequent return operations with the same grandparent body.
40 * <p>
41 * If the grandparent body has only one block then operation transformer's block builder is the return
42 * block builder. Otherwise, if the grandparent body has one or more blocks then the return block builder is
43 * created from the operation transformer's block builder. The created return block builder will have a block
44 * parameter whose type corresponds to the return type, or will have no parameter for void return.
45 * The computation finishes by applying the return block builder and a return value to the inlining consumer.
46 * If the grandparent body has only one block then the return value is the value mapped from the return
47 * operation's operand, or is null for void return. Otherwise, if the grandparent body has one or more blocks
48 * then the value is the block parameter of the created return block builder, or is null for void return.
49 * <p>
50 * For every encounter of a return operation the associated return block builder is compared against the
51 * operation transformer's block builder. If they are not equal then a branch operation is added to the
52 * operation transformer's block builder whose successor is the return block builder with a block argument
53 * that is the value mapped from the return operation's operand, or with no block argument for void return.
54 * @apiNote
55 * It is easier to inline an invokable op if its body is in lowered form (there are no operations in the blocks
56 * of the body that are lowerable). This ensures a single exit point can be created (paired with the single
57 * entry point). If there are one or more nested return operations, then there is unlikely to be a single exit.
58 * Transforming the model to create a single exit point while preserving nested structure is in general
59 * non-trivial and outside the scope of this method. In such cases the invokable operation can be transformed
60 * with a lowering transformation after which it can then be inlined.
61 *
62 * @param _this the block builder
63 * @param invokableOp the invokable operation
64 * @param args the arguments to map to the invokable operation's parameters
65 * @param inlineConsumer the consumer applied to process the return from the invokable operation.
66 * This is called once for each grandparent body of a return operation, with a block to
67 * build replacement operations and the return value, or null for void return.
68 * @return the block builder to continue building from
69 * @param <O> The invokable type
70 */
71 public static <O extends Op & Op.Invokable>
72 Block.Builder inline(Block.Builder _this, O invokableOp, List<? extends Value> args,
73 BiConsumer<Block.Builder, Value> inlineConsumer) {
74 Map<Body, Block.Builder> returnBlocks = new HashMap<>();
75 // Create new context, ensuring inlining is isolated
76 _this.transformBody(invokableOp.body(), args, CopyContext.create(), (block, op) -> {
77 // If the return operation is associated with the invokable operation
78 if (op instanceof CoreOp.ReturnOp rop && getNearestInvokeableAncestorOp(op) == invokableOp) {
79 // Compute the return block
80 Block.Builder returnBlock = returnBlocks.computeIfAbsent(rop.ancestorBody(), _body -> {
81 Block.Builder rb;
82 // If the body has one block we know there is just one return op declared, otherwise there may
83 // one or more. If so, create a new block that joins all the returns.
84 // Note: we could count all return op in a body to avoid creating a new block for a body
85 // with two or more blocks with only one returnOp is declared.
86 Value r;
87 if (rop.ancestorBody().blocks().size() != 1) {
88 List<TypeElement> param = rop.returnValue() != null
89 ? List.of(invokableOp.invokableType().returnType())
90 : List.of();
91 rb = block.block(param);
92 r = !param.isEmpty()
93 ? rb.parameters().get(0)
94 : null;
95 } else {
96 r = rop.returnValue() != null
97 ? block.context().getValue(rop.returnValue())
98 : null;
99 rb = block;
100 }
101
102 // Inline the return
103 inlineConsumer.accept(rb, r);
104
105 return rb;
106 });
107
108 // Replace the return op with a branch to the return block, if needed
109 if (!returnBlock.equals(block)) {
110 // Replace return op with branch to return block, with given return value
111 List<Value> arg = rop.returnValue() != null
112 ? List.of(block.context().getValue(rop.returnValue()))
113 : List.of();
114 block.op(branch(returnBlock.successor(arg)));
115 }
116
117 return block;
118 }
119
120 block.op(op);
121 return block;
122 });
123
124
125 Block.Builder builder = returnBlocks.get(invokableOp.body());
126 return builder != null ? builder : _this;
127 }
128
129 private static Op getNearestInvokeableAncestorOp(Op op) {
130 do {
131 op = op.ancestorOp();
132 } while (!(op instanceof Op.Invokable));
133 return op;
134 }
135 }