1 /*
2 * Copyright (c) 2024, 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. Oracle designates this
8 * particular file as subject to the "Classpath" exception as provided
9 * by Oracle in the LICENSE file that accompanied this code.
10 *
11 * This code is distributed in the hope that it will be useful, but WITHOUT
12 * ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or
13 * FITNESS FOR A PARTICULAR PURPOSE. See the GNU General Public License
14 * version 2 for more details (a copy is included in the LICENSE file that
15 * accompanied this code).
16 *
17 * You should have received a copy of the GNU General Public License version
18 * 2 along with this work; if not, write to the Free Software Foundation,
19 * Inc., 51 Franklin St, Fifth Floor, Boston, MA 02110-1301 USA.
20 *
21 * Please contact Oracle, 500 Oracle Parkway, Redwood Shores, CA 94065 USA
22 * or visit www.oracle.com if you need additional information or have any
23 * questions.
24 */
25
26 package jdk.incubator.code;
27
28 import jdk.incubator.code.dialect.core.CoreOp;
29 import jdk.incubator.code.dialect.core.CoreType;
30 import jdk.incubator.code.dialect.core.FunctionType;
31 import jdk.incubator.code.dialect.core.VarType;
32
33 import java.util.*;
34 import java.util.function.Consumer;
35
36 /**
37 * An operation and a mapping from the operation's {@link Op#operands() operands} and
38 * {@link Op#capturedValues() captured values} to runtime values. This pair is referred
39 * to as the quoted form of an operation, and represents an operation in a runtime context
40 * independent of its surrounding code model.
41 *
42 * @param <T> the type of operation that is quoted
43 */
44 public final class Quoted<T extends Op> {
45 private final T op;
46 private final SequencedMap<Value, Object> operandsAndCapturedValues;
47
48 /**
49 * Constructs the quoted form of a given operation.
50 *
51 * @param op the operation.
52 * @param operandsAndCapturedValues map containing the operation's operands and captured values to runtime values
53 * @throws IllegalArgumentException if the map's key set does not contain set consisting of operation's operands
54 * and the captured values
55 * @see Op#capturedValues()
56 * @see Op#operands()
57 */
58 public Quoted(T op, Map<Value, Object> operandsAndCapturedValues) {
59 SequencedMap<Value, Object> m = new LinkedHashMap<>();
60 for (Value value : op.operands()) {
61 m.put(value, runtimeValue(operandsAndCapturedValues, value));
62 }
63 for (Value value : op.capturedValues()) {
64 m.put(value, runtimeValue(operandsAndCapturedValues, value));
65 }
66
67 this.op = op;
68 this.operandsAndCapturedValues = Collections.unmodifiableSequencedMap(m);
69 }
70
71 static Object runtimeValue(Map<Value, Object> operandsAndCapturedValues, Value value) {
72 if (!operandsAndCapturedValues.containsKey(value)) {
73 throw new IllegalArgumentException("Value is not present as a key in the map of values");
74 }
75 return operandsAndCapturedValues.get(value);
76 }
77
78 /**
79 * Returns the operation.
80 *
81 * @return the operation.
82 */
83 public T op() {
84 return op;
85 }
86
87 /**
88 * Returns the map of captured values to runtime values.
89 * <p>
90 * The captured values key set has the same elements and same encounter order as
91 * operation's captured values, specifically the following expression evaluates to true:
92 * {@snippet lang=java :
93 * op().capturedValues().equals(new ArrayList<>(capturedValues().keySet()));
94 * }
95 *
96 * @return the mapping of captured values.
97 */
98 public SequencedMap<Value, Object> capturedValues() {
99 SequencedMap<Value, Object> m = new LinkedHashMap<>();
100 for (Value cv : op.capturedValues()) {
101 m.put(cv, operandsAndCapturedValues.get(cv));
102 }
103 return m;
104 }
105
106 /**
107 * Returns the map of operands to runtime values.
108 * <p>
109 * The key set has the same elements and same encounter order as the sequenced set of operation's operands,
110 * specifically the following expression evaluates to true:
111 * {@snippet lang = java:
112 * op.operands().equals(new ArrayList<>(operands().keySet()))
113 *}
114 *
115 * @return the map of operands.
116 */
117 public SequencedMap<Value, Object> operands() {
118 SequencedMap<Value, Object> m = new LinkedHashMap<>();
119 for (Value operand : op.operands()) {
120 // putIfAbsent is used because a value may be used as operand more than once
121 m.putIfAbsent(operand, operandsAndCapturedValues.get(operand));
122 }
123 return m;
124 }
125
126 /**
127 * Returns the map of operands and captured values to runtime values.
128 * <p>
129 * The map's key set is equal to the sequenced set of the quoted operation's
130 * operands and captured values.
131 *
132 * @return the map of operands and captured values, as an unmodifiable map.
133 */
134 public SequencedMap<Value, Object> operandsAndCapturedValues() {
135 return operandsAndCapturedValues;
136 }
137
138 /**
139 * Embeds the given operation into a quoting code model whose behaviour quotes the operation.
140 * <p>
141 * The result is a {@link jdk.incubator.code.dialect.core.CoreOp.FuncOp func operation}
142 * that has one body with one block (<i>fblock</i>).
143 * <br>
144 * <i>fblock</i> will have a block parameter, in order, for every value in the key set of the map of operands and
145 * captured values.
146 * If the value is a result of a {@link jdk.incubator.code.dialect.core.CoreOp.VarOp var operation} then the
147 * parameter's type element is the var operation's value type, and <i>fblock</i> will have a var operation whose
148 * operand is the block parameter.
149 * Otherwise, the parameter's type element is the value's type element.
150 * <br>
151 * Then <i>fblock</i> has a {@link jdk.incubator.code.dialect.core.CoreOp.QuotedOp quoted operation}
152 * that has one body with one block (<i>qblock</i>). Inside <i>qblock</i> there is a copy of {@code op}
153 * and a {@link jdk.incubator.code.dialect.core.CoreOp.YieldOp yield operation} whose operand is the operation
154 * result of {@code op}'s copy.
155 * <br>
156 * <i>fblock</i> terminates with a {@link jdk.incubator.code.dialect.core.CoreOp.ReturnOp return operation},
157 * whose operand is the operation result of the quoted op.
158 *
159 * @param op the operation
160 * @return the quoting code model.
161 * @throws IllegalArgumentException if {@code op} is unbuilt.
162 */
163 public static CoreOp.FuncOp embedOp(Op op) {
164 if (op.result() == null) {
165 throw new IllegalArgumentException("Op is unbuilt");
166 }
167
168 // if we don't remove duplicate operands we will have unused params in the new model
169 // if we don't remove captured values that are operands we will have unused params in the new model
170 SequencedSet<Value> s = new LinkedHashSet<>(op.operands());
171 s.addAll(op.capturedValues());
172 List<Value> inputOperandsAndCaptures = s.stream().toList();
173
174 // Build the function type
175 List<TypeElement> params = inputOperandsAndCaptures.stream()
176 .map(v -> v.type() instanceof VarType vt ? vt.valueType() : v.type())
177 .toList();
178 FunctionType ft = CoreType.functionType(CoreOp.QuotedOp.QUOTED_OP_TYPE, params);
179
180 // Build the function that quotes the lambda
181 return CoreOp.func("q", ft).body(b -> {
182 // Create variables as needed and obtain the operands and captured values for the copied lambda
183 List<Value> outputOperandsAndCaptures = new ArrayList<>();
184 for (int i = 0; i < inputOperandsAndCaptures.size(); i++) {
185 Value inputValue = inputOperandsAndCaptures.get(i);
186 Value outputValue = b.parameters().get(i);
187 if (inputValue.type() instanceof VarType) {
188 outputValue = b.op(CoreOp.var(String.valueOf(i), outputValue));
189 }
190 outputOperandsAndCaptures.add(outputValue);
191 }
192
193 // Quoted the lambda expression
194 Value q = b.op(CoreOp.quoted(b.parentBody(), qb -> {
195 // Map the entry block of the op's ancestor body to the quoted block
196 // We are copying op in the context of the quoted block, the block mapping
197 // ensures the use of operands and captured values are reachable when building
198 qb.context().mapBlock(op.ancestorBody().entryBlock(), qb);
199 // Map the op's operands and captured values
200 qb.context().mapValues(inputOperandsAndCaptures, outputOperandsAndCaptures);
201 // Return the op to be copied in the quoted operation
202 return op;
203 }));
204 b.op(CoreOp.return_(q));
205 });
206 }
207
208 private static IllegalArgumentException invalidQuotedModel(CoreOp.FuncOp model) {
209 return new IllegalArgumentException("Invalid code model for quoted operation : " + model);
210 }
211
212 /**
213 * Extracts the quoted operation from a quoting code model with the given runtime arguments.
214 * <p>
215 * This method behaves as if the quoting code model is executed with the runtime arguments by interpreting the
216 * behaviour of the code elements, as specified, in the code model, but it may be implemented more efficiently.
217 *
218 * @param funcOp the quoting code model
219 * @param args the runtime arguments
220 * @return quoted form of the extracted operation
221 * @throws IllegalArgumentException if the quoting code model is invalid and does not conform to that specified
222 * the method {@link #embedOp(Op)}.
223 * @throws IllegalArgumentException if the quoting code model's parameters differ in arity from the
224 * runtime arguments.
225 */
226 public static Quoted<Op> extractOp(CoreOp.FuncOp funcOp, List<Object> args) {
227 if (funcOp.body().blocks().size() != 1) {
228 throw invalidQuotedModel(funcOp);
229 }
230 Block fblock = funcOp.body().entryBlock();
231 if (fblock.ops().size() < 2) {
232 throw invalidQuotedModel(funcOp);
233 }
234 if (!(fblock.ops().get(fblock.ops().size() - 2) instanceof CoreOp.QuotedOp qop)) {
235 throw invalidQuotedModel(funcOp);
236 }
237 if (!(fblock.ops().getLast() instanceof CoreOp.ReturnOp returnOp)) {
238 throw invalidQuotedModel(funcOp);
239 }
240 if (returnOp.returnValue() == null) {
241 throw invalidQuotedModel(funcOp);
242 }
243 if (!returnOp.returnValue().equals(qop.result())) {
244 throw invalidQuotedModel(funcOp);
245 }
246
247 Op op = qop.quotedOp();
248
249 SequencedSet<Value> operandsAndCaptures = new LinkedHashSet<>();
250 operandsAndCaptures.addAll(op.operands());
251 operandsAndCaptures.addAll(op.capturedValues());
252
253 // validation rule of block params and ConstantOp result
254 // let v be a block param or ConstantOp result
255 // if v not used -> throw
256 // if v used once and user is VarOp and VarOp not used or VarOp used in funcOp entry block -> throw
257 // if v is used once and user is not a VarOp and usage isn't as operand or capture -> throw
258 // if v is used more than once and one of the uses is in funcOp entry block -> throw
259 Consumer<Value> validate = v -> {
260 if (v.uses().isEmpty()) {
261 throw invalidQuotedModel(funcOp);
262 } else if (v.uses().size() == 1 && v.uses().iterator().next().op() instanceof CoreOp.VarOp vop
263 && (vop.result().uses().isEmpty() ||
264 vop.result().uses().stream().anyMatch(u -> u.op().ancestorBlock() == fblock))) {
265 throw invalidQuotedModel(funcOp);
266 } else if (v.uses().size() == 1 && !(v.uses().iterator().next().op() instanceof CoreOp.VarOp)
267 && !operandsAndCaptures.contains(v)) {
268 throw invalidQuotedModel(funcOp);
269 } else if (v.uses().size() > 1 && v.uses().stream().anyMatch(u -> u.op().ancestorBlock() == fblock)) {
270 throw invalidQuotedModel(funcOp);
271 }
272 };
273
274 for (Block.Parameter p : fblock.parameters()) {
275 validate.accept(p);
276 }
277
278 List<Op> ops = fblock.ops().subList(0, fblock.ops().size() - 2);
279 for (Op o : ops) {
280 switch (o) {
281 case CoreOp.VarOp varOp -> {
282 if (varOp.isUninitialized()) {
283 throw invalidQuotedModel(funcOp);
284 }
285 if (varOp.initOperand() instanceof Op.Result opr && !(opr.op() instanceof CoreOp.ConstantOp)) {
286 throw invalidQuotedModel(funcOp);
287 }
288 }
289 case CoreOp.ConstantOp cop -> validate.accept(cop.result());
290 default -> throw invalidQuotedModel(funcOp);
291 }
292 }
293
294 // map operands and captures to their corresponding runtime values
295 // operand and capture can be:
296 // 1- block param
297 // 2- result of VarOp whose initial value is constant
298 // 3- result of VarOp whose initial value is block param
299 // 4- result of ConstantOp
300 List<Block.Parameter> params = funcOp.parameters();
301 if (params.size() != args.size()) {
302 throw invalidQuotedModel(funcOp);
303 }
304 SequencedMap<Value, Object> m = new LinkedHashMap<>();
305 for (Value v : operandsAndCaptures) {
306 switch (v) {
307 case Block.Parameter p -> {
308 Object rv = args.get(p.index());
309 m.put(v, rv);
310 }
311 case Op.Result opr when opr.op() instanceof CoreOp.VarOp varOp -> {
312 if (varOp.initOperand() instanceof Op.Result r && r.op() instanceof CoreOp.ConstantOp cop) {
313 m.put(v, CoreOp.Var.of(cop.value()));
314 } else if (varOp.initOperand() instanceof Block.Parameter p) {
315 Object rv = args.get(p.index());
316 m.put(v, CoreOp.Var.of(rv));
317 }
318 }
319 case Op.Result opr when opr.op() instanceof CoreOp.ConstantOp cop -> {
320 m.put(v, cop.value());
321 }
322 default -> throw invalidQuotedModel(funcOp);
323 }
324 }
325
326 return new Quoted<>(op, m);
327 }
328
329 /**
330 * Extracts the quoted operation from {@code funcOp}
331 * and map its operands and captured values to the runtime values in {@code args}.
332 * <p>
333 * {@code funcOp} must have the same structure as if it's produced by {@link #embedOp(Op)}.
334 *
335 * @param funcOp Model to extract the quoted op from
336 * @param args Runtime values for {@code funcOp} parameters
337 * @return Quoted instance that wraps the quoted operation,
338 * plus the mapping of its operands and captured values to the given runtime values
339 * @throws RuntimeException If {@code funcOp} isn't a valid code model
340 * @throws RuntimeException If {@code funcOp} parameters size is different from {@code args} length
341 * @see Quoted#extractOp(CoreOp.FuncOp, List)
342 */
343 public static Quoted<Op> extractOp(CoreOp.FuncOp funcOp, Object... args) {
344 return extractOp(funcOp, Arrays.asList(args));
345 }
346 }