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 * The quoted form of an operation.
38 * <p>
39 * The quoted form is utilized when the code model of some code is to be obtained rather than obtaining the result of
40 * executing that code. For example passing the of a lambda expression in quoted form rather than the expression being
41 * targeted to a functional interface from which it can be invoked.
42 */
43 public final class Quoted {
44 private final Op op;
45 private final SequencedMap<Value, Object> operandsAndCapturedValues;
46
47 static final SequencedMap<Value, Object> EMPTY_SEQUENCED_MAP = new LinkedHashMap<>();
48 /**
49 * Constructs the quoted form of a given operation.
50 *
51 * @param op the invokable operation.
52 */
53 public Quoted(Op op) {
54 this(op, EMPTY_SEQUENCED_MAP);
55 }
56
57 /**
58 * Constructs the quoted form of a given operation.
59 * <p>
60 * The {@code operandsAndCapturedValues} key set must be equal to
61 * the sequenced set of operation's operands + captured values, in order.
62 *
63 * @param op the operation.
64 * @param operandsAndCapturedValues sequenced map of {@link Value} to {@link Object}, with the requirement defined above
65 * @throws IllegalArgumentException If {@code operandsAndCapturedValues} doesn't satisfy the requirement described above
66 * @see Op#capturedValues()
67 * @see Op#operands()
68 */
69 public Quoted(Op op, SequencedMap<Value, Object> operandsAndCapturedValues) {
70 // @@@ This check is potentially expensive, remove or keep ?
71 // @@@ Or make Quoted an interface, with a module private implementation?
72 SequencedSet<Value> s = new LinkedHashSet<>(op.operands());
73 s.addAll(op.capturedValues());
74 if (!s.stream().toList().equals(operandsAndCapturedValues.keySet().stream().toList())) {
75 throw new IllegalArgumentException("The map key set isn't equal to the sequenced set of operands + captured values");
76 }
77 this.op = op;
78 this.operandsAndCapturedValues = Collections.unmodifiableSequencedMap(new LinkedHashMap<>(operandsAndCapturedValues));
79 }
80
81 /**
82 * Returns the operation.
83 *
84 * @return the operation.
85 */
86 public Op op() {
87 return op;
88 }
89
90 /**
91 * Returns the captured values.
92 * <p>
93 * The captured values key set has the same elements and same encounter order as
94 * operation's captured values, specifically the following expression evaluates to true:
95 * {@snippet lang=java :
96 * op().capturedValues().equals(new ArrayList<>(capturedValues().keySet()));
97 * }
98 *
99 * @return the captured values.
100 */
101 public SequencedMap<Value, Object> capturedValues() {
102 SequencedMap<Value, Object> m = new LinkedHashMap<>();
103 for (Value cv : op.capturedValues()) {
104 m.put(cv, operandsAndCapturedValues.get(cv));
105 }
106 return m;
107 }
108
109 /**
110 * Returns the operands.
111 * <p>
112 * The result key set has the same elements and same encounter order as the sequenced set of operation's operands,
113 * specifically the following expression evaluates to true:
114 * {@snippet lang = java:
115 * new LinkedHashSet<>(op.operands()).equals(operands().keySet());
116 *}
117 *
118 * @return the operands.
119 */
120 public SequencedMap<Value, Object> operands() {
121 SequencedMap<Value, Object> m = new LinkedHashMap<>();
122 for (Value operand : op.operands()) {
123 // putIfAbsent is used because a value may be used as operand more than once
124 m.putIfAbsent(operand, operandsAndCapturedValues.get(operand));
125 }
126 return m;
127 }
128
129 /**
130 * Returns the operands and captured values.
131 * The result key set is equal to the sequenced set of operands + captured values.
132 *
133 * @return the operands + captured values, as an unmodifiable map.
134 */
135 public SequencedMap<Value, Object> operandsAndCapturedValues() {
136 return operandsAndCapturedValues;
137 }
138
139 /**
140 * Embeds the given {@code op}, copying it from its original context to a new one,
141 * where its operands and captured values will be parameters.
142 * <p>
143 * The result is a {@link jdk.incubator.code.dialect.core.CoreOp.FuncOp FuncOp}
144 * that has one body with one block (<i>fblock</i>).
145 * <br>
146 * <i>fblock</i> will have a parameter for every element in the sequenced set of {@code op}'s operands + captured values.
147 * If the operand or capture is a result of a {@link jdk.incubator.code.dialect.core.CoreOp.VarOp VarOp},
148 * <i>fblock</i> will have a {@link jdk.incubator.code.dialect.core.CoreOp.VarOp VarOp}
149 * whose initial value is the parameter.
150 * <br>
151 * Then <i>fblock</i> has a {@link jdk.incubator.code.dialect.core.CoreOp.QuotedOp QuotedOp}
152 * that has one body with one block (<i>qblock</i>).
153 * Inside <i>qblock</i> we find a copy of {@code op}
154 * and a {@link jdk.incubator.code.dialect.core.CoreOp.YieldOp YieldOp}
155 * whose yield value is the result of the {@code op}'s copy.
156 * <br>
157 * <i>fblock</i> terminates with a {@link jdk.incubator.code.dialect.core.CoreOp.ReturnOp ReturnOp},
158 * the returned value is the result of the {@link jdk.incubator.code.dialect.core.CoreOp.QuotedOp QuotedOp}
159 * object described previously.
160 *
161 * @param op The operation to embed
162 * @return The model that represent the quoting of {@code op}
163 * @throws IllegalArgumentException if {@code op} is not bound
164 */
165 public static CoreOp.FuncOp embedOp(Op op) {
166 if (op.result() == null) {
167 throw new IllegalArgumentException("Op not bound");
168 }
169
170 // if we don't remove duplicate operands we will have unused params in the new model
171 // if we don't remove captured values that are operands we will have unused params in the new model
172 SequencedSet<Value> s = new LinkedHashSet<>(op.operands());
173 s.addAll(op.capturedValues());
174 List<Value> inputOperandsAndCaptures = s.stream().toList();
175
176 // Build the function type
177 List<TypeElement> params = inputOperandsAndCaptures.stream()
178 .map(v -> v.type() instanceof VarType vt ? vt.valueType() : v.type())
179 .toList();
180 FunctionType ft = CoreType.functionType(CoreOp.QuotedOp.QUOTED_TYPE, params);
181
182 // Build the function that quotes the lambda
183 return CoreOp.func("q", ft).body(b -> {
184 // Create variables as needed and obtain the operands and captured values for the copied lambda
185 List<Value> outputOperandsAndCaptures = new ArrayList<>();
186 for (int i = 0; i < inputOperandsAndCaptures.size(); i++) {
187 Value inputValue = inputOperandsAndCaptures.get(i);
188 Value outputValue = b.parameters().get(i);
189 if (inputValue.type() instanceof VarType) {
190 outputValue = b.op(CoreOp.var(String.valueOf(i), outputValue));
191 }
192 outputOperandsAndCaptures.add(outputValue);
193 }
194
195 // Quoted the lambda expression
196 Value q = b.op(CoreOp.quoted(b.parentBody(), qb -> {
197 // Map the entry block of the op's ancestor body to the quoted block
198 // We are copying op in the context of the quoted block, the block mapping
199 // ensures the use of operands and captured values are reachable when building
200 qb.context().mapBlock(op.ancestorBody().entryBlock(), qb);
201 // Map the op's operands and captured values
202 qb.context().mapValues(inputOperandsAndCaptures, outputOperandsAndCaptures);
203 // Return the op to be copied in the quoted operation
204 return op;
205 }));
206 b.op(CoreOp.return_(q));
207 });
208 }
209
210 private static RuntimeException invalidQuotedModel(CoreOp.FuncOp model) {
211 return new RuntimeException("Invalid code model for quoted operation : " + model);
212 }
213
214 /**
215 * Extracts the quoted operation from {@code funcOp}
216 * and map its operands and captured values to the runtime values in {@code args}.
217 * <p>
218 * {@code funcOp} must have the same structure as if it's produced by {@link #embedOp(Op)}.
219 *
220 * @param funcOp Model to extract the quoted op from
221 * @param args Runtime values for {@code funcOp} parameters
222 * @return Quoted instance that wraps the quoted operation,
223 * plus the mapping of its operands and captured values to the given runtime values
224 * @throws RuntimeException If {@code funcOp} isn't a valid code model
225 * @throws RuntimeException If {@code funcOp} parameters size is different from {@code args} length
226 */
227 public static Quoted 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 extractOp(CoreOp.FuncOp funcOp, Object... args) {
345 return extractOp(funcOp, List.of(args));
346 }
347 }