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 behavior 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 code type 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 code type is the value's code type.
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 unattached or a root operation.
162 * @throws IllegalStateException if an encountered block is being built and is not observable.
163 */
164 public static CoreOp.FuncOp embedOp(Op op) {
165 if (op.result() == null) {
166 throw new IllegalArgumentException("Operation is unattached or a root operation");
167 }
168
169 // if we don't remove duplicate operands we will have unused params in the new model
170 // if we don't remove captured values that are operands we will have unused params in the new model
171 SequencedSet<Value> s = new LinkedHashSet<>(op.operands());
172 s.addAll(op.capturedValues());
173 List<Value> inputOperandsAndCaptures = s.stream().toList();
174
175 // Build the function type
176 List<CodeType> params = inputOperandsAndCaptures.stream()
177 .map(v -> v.type() instanceof VarType vt ? vt.valueType() : v.type())
178 .toList();
179 FunctionType ft = CoreType.functionType(CoreOp.QuotedOp.QUOTED_OP_TYPE, params);
180
181 // Build the function that quotes the lambda
182 return CoreOp.func("q", ft).body(b -> {
183 // Create variables as needed and obtain the operands and captured values for the copied lambda
184 List<Value> outputOperandsAndCaptures = new ArrayList<>();
185 for (int i = 0; i < inputOperandsAndCaptures.size(); i++) {
186 Value inputValue = inputOperandsAndCaptures.get(i);
187 Value outputValue = b.parameters().get(i);
188 if (inputValue.type() instanceof VarType) {
189 outputValue = b.op(CoreOp.var(String.valueOf(i), outputValue));
190 }
191 outputOperandsAndCaptures.add(outputValue);
192 }
193
194 // Quoted the lambda expression
195 Value q = b.op(CoreOp.quoted(b.parentBody(), qb -> {
196 // Map the entry block of the op's ancestor body to the quoted block
197 // We are copying op in the context of the quoted block, the block mapping
198 // ensures the use of operands and captured values are reachable when building
199 qb.context().mapBlock(op.ancestorBody().entryBlock(), qb);
200 // Map the op's operands and captured values
201 qb.context().mapValues(inputOperandsAndCaptures, outputOperandsAndCaptures);
202 // Return the op to be copied in the quoted operation
203 return op;
204 }));
205 b.op(CoreOp.return_(q));
206 });
207 }
208
209 private static IllegalArgumentException invalidQuotedModel(CoreOp.FuncOp model) {
210 return new IllegalArgumentException("Invalid code model for quoted operation : " + model);
211 }
212
213 /**
214 * Extracts the quoted operation from a quoting code model with the given runtime arguments.
215 * <p>
216 * This method behaves as if the quoting code model is executed with the runtime arguments by interpreting the
217 * behavior of the code elements, as specified, in the code model, but it may be implemented more efficiently.
218 *
219 * @param funcOp the quoting code model
220 * @param args the runtime arguments
221 * @return quoted form of the extracted operation
222 * @throws IllegalArgumentException if the quoting code model is invalid and does not conform to that specified
223 * the method {@link #embedOp(Op)}.
224 * @throws IllegalArgumentException if the quoting code model's parameters differ in arity from the
225 * runtime arguments.
226 */
227 public static Quoted<Op> extractOp(CoreOp.FuncOp funcOp, List<Object> args) {
228 if (funcOp.body().blocks().size() != 1) {
229 throw invalidQuotedModel(funcOp);
230 }
231 Block fblock = funcOp.body().entryBlock();
232 if (fblock.ops().size() < 2) {
233 throw invalidQuotedModel(funcOp);
234 }
235 if (!(fblock.ops().get(fblock.ops().size() - 2) instanceof CoreOp.QuotedOp qop)) {
236 throw invalidQuotedModel(funcOp);
237 }
238 if (!(fblock.ops().getLast() instanceof CoreOp.ReturnOp returnOp)) {
239 throw invalidQuotedModel(funcOp);
240 }
241 if (returnOp.returnValue() == null) {
242 throw invalidQuotedModel(funcOp);
243 }
244 if (!returnOp.returnValue().equals(qop.result())) {
245 throw invalidQuotedModel(funcOp);
246 }
247
248 Op op = qop.quotedOp();
249
250 SequencedSet<Value> operandsAndCaptures = new LinkedHashSet<>();
251 operandsAndCaptures.addAll(op.operands());
252 operandsAndCaptures.addAll(op.capturedValues());
253
254 // validation rule of block params and ConstantOp result
255 // let v be a block param or ConstantOp result
256 // if v not used -> throw
257 // if v used once and user is VarOp and VarOp not used or VarOp used in funcOp entry block -> throw
258 // if v is used once and user is not a VarOp and usage isn't as operand or capture -> throw
259 // if v is used more than once and one of the uses is in funcOp entry block -> throw
260 Consumer<Value> validate = v -> {
261 if (v.uses().isEmpty()) {
262 throw invalidQuotedModel(funcOp);
263 } else if (v.uses().size() == 1 && v.uses().iterator().next().op() instanceof CoreOp.VarOp vop
264 && (vop.result().uses().isEmpty() ||
265 vop.result().uses().stream().anyMatch(u -> u.op().ancestorBlock() == fblock))) {
266 throw invalidQuotedModel(funcOp);
267 } else if (v.uses().size() == 1 && !(v.uses().iterator().next().op() instanceof CoreOp.VarOp)
268 && !operandsAndCaptures.contains(v)) {
269 throw invalidQuotedModel(funcOp);
270 } else if (v.uses().size() > 1 && v.uses().stream().anyMatch(u -> u.op().ancestorBlock() == fblock)) {
271 throw invalidQuotedModel(funcOp);
272 }
273 };
274
275 for (Block.Parameter p : fblock.parameters()) {
276 validate.accept(p);
277 }
278
279 List<Op> ops = fblock.ops().subList(0, fblock.ops().size() - 2);
280 for (Op o : ops) {
281 switch (o) {
282 case CoreOp.VarOp varOp -> {
283 if (varOp.isUninitialized()) {
284 throw invalidQuotedModel(funcOp);
285 }
286 if (varOp.initOperand() instanceof Op.Result opr && !(opr.op() instanceof CoreOp.ConstantOp)) {
287 throw invalidQuotedModel(funcOp);
288 }
289 }
290 case CoreOp.ConstantOp cop -> validate.accept(cop.result());
291 default -> throw invalidQuotedModel(funcOp);
292 }
293 }
294
295 // map operands and captures to their corresponding runtime values
296 // operand and capture can be:
297 // 1- block param
298 // 2- result of VarOp whose initial value is constant
299 // 3- result of VarOp whose initial value is block param
300 // 4- result of ConstantOp
301 List<Block.Parameter> params = funcOp.parameters();
302 if (params.size() != args.size()) {
303 throw invalidQuotedModel(funcOp);
304 }
305 SequencedMap<Value, Object> m = new LinkedHashMap<>();
306 for (Value v : operandsAndCaptures) {
307 switch (v) {
308 case Block.Parameter p -> {
309 Object rv = args.get(p.index());
310 m.put(v, rv);
311 }
312 case Op.Result opr when opr.op() instanceof CoreOp.VarOp varOp -> {
313 if (varOp.initOperand() instanceof Op.Result r && r.op() instanceof CoreOp.ConstantOp cop) {
314 m.put(v, CoreOp.Var.of(cop.value()));
315 } else if (varOp.initOperand() instanceof Block.Parameter p) {
316 Object rv = args.get(p.index());
317 m.put(v, CoreOp.Var.of(rv));
318 }
319 }
320 case Op.Result opr when opr.op() instanceof CoreOp.ConstantOp cop -> {
321 m.put(v, cop.value());
322 }
323 default -> throw invalidQuotedModel(funcOp);
324 }
325 }
326
327 return new Quoted<>(op, m);
328 }
329
330 /**
331 * Extracts the quoted operation from {@code funcOp}
332 * and map its operands and captured values to the runtime values in {@code args}.
333 * <p>
334 * {@code funcOp} must have the same structure as if it's produced by {@link #embedOp(Op)}.
335 *
336 * @param funcOp Model to extract the quoted op from
337 * @param args Runtime values for {@code funcOp} parameters
338 * @return Quoted instance that wraps the quoted operation,
339 * plus the mapping of its operands and captured values to the given runtime values
340 * @throws RuntimeException If {@code funcOp} isn't a valid code model
341 * @throws RuntimeException If {@code funcOp} parameters size is different from {@code args} length
342 * @see Quoted#extractOp(CoreOp.FuncOp, List)
343 */
344 public static Quoted<Op> extractOp(CoreOp.FuncOp funcOp, Object... args) {
345 return extractOp(funcOp, Arrays.asList(args));
346 }
347 }