1 /*
2 * Copyright (c) 2024, 2025, 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.CoreType;
27 import jdk.incubator.code.dialect.core.FunctionType;
28 import jdk.incubator.code.extern.ExternalizedOp;
29 import jdk.incubator.code.extern.OpFactory;
30
31 import java.util.*;
32 import java.util.function.Consumer;
33
34 public final class AnfDialect {
35
36 private AnfDialect() {
37 }
38
39 public static final class AnfLetOp extends Op implements Op.Terminating, Op.Nested {
40 public static final String NAME = "anf.let";
41
42 public static class Builder {
43 final Body.Builder ancestorBody;
44 final TypeElement yieldType;
45
46 Builder(Body.Builder ancestorBody, TypeElement yieldType) {
47 this.ancestorBody = ancestorBody;
48 this.yieldType = yieldType;
49 }
50
51 public AnfLetOp body(Consumer<Block.Builder> c) {
52 Body.Builder body = Body.Builder.of(ancestorBody, CoreType.functionType(yieldType));
53 c.accept(body.entryBlock());
54 return new AnfLetOp(body);
55 }
56 }
57
58 // Terminating operation is the in expression that yields the result for this operation
59 // Any operation result in the entry block used by the terminating operation
60 // or a descendant operation is a let binding
61 final Body bindings;
62
63 public AnfLetOp(ExternalizedOp def) {
64 this(def.bodyDefinitions().get(0));
65 }
66
67 public AnfLetOp(AnfLetOp that, CopyContext cc, OpTransformer ot) {
68 super(that, cc);
69
70 this.bindings = that.bindings.transform(cc, ot).build(this);
71 }
72
73 @Override
74 public Op transform(CopyContext cc, OpTransformer ot) {
75 return new AnfLetOp(this, cc, ot);
76 }
77
78 public AnfLetOp(Body.Builder bodyBuilder) {
79 super(List.of());
80
81 this.bindings = bodyBuilder.build(this);
82 }
83
84 @Override
85 public List<Body> bodies() {
86 return List.of(bindings);
87 }
88
89 @Override
90 public TypeElement resultType() {
91 return this.bindings.yieldType();
92 }
93 }
94
95
96 public static final class AnfLetRecOp extends Op implements Op.Terminating, Op.Nested {
97 public static final String NAME = "anf.letrec";
98
99 public static class Builder {
100 final Body.Builder ancestorBody;
101 final TypeElement yieldType;
102
103 Builder(Body.Builder ancestorBody, TypeElement yieldType) {
104 this.ancestorBody = ancestorBody;
105 this.yieldType = yieldType;
106 }
107
108 public AnfLetRecOp body(Consumer<Block.Builder> c) {
109 Body.Builder body = Body.Builder.of(ancestorBody, CoreType.functionType(yieldType));
110 c.accept(body.entryBlock());
111 return new AnfLetRecOp(body);
112 }
113 }
114
115 // Terminating operation is the in expression that yields the result for this operation
116 // Any operation result in the entry block used by the terminating operation
117 // or a descendant operation is a letrec binding
118 final Body bindings;
119
120 public AnfLetRecOp(ExternalizedOp def) {
121 this(def.bodyDefinitions().get(0));
122 }
123
124 public AnfLetRecOp(AnfLetRecOp that, CopyContext cc, OpTransformer ot) {
125 super(that, cc);
126
127 this.bindings = that.bindings.transform(cc, ot).build(this);
128 }
129
130 @Override
131 public Op transform(CopyContext cc, OpTransformer ot) {
132 return new AnfLetRecOp(this, cc, ot);
133 }
134
135 public AnfLetRecOp(Body.Builder bodyBuilder) {
136 super(List.of());
137
138 this.bindings = bodyBuilder.build(this);
139 }
140
141 @Override
142 public List<Body> bodies() {
143 return List.of(bindings);
144 }
145
146 @Override
147 public TypeElement resultType() {
148 return this.bindings.yieldType();
149 }
150
151 public List<CoreOp.FuncOp> funcOps() {
152 List<Op> ops = bindings.entryBlock().ops();
153 return ops.subList(0, ops.size() - 1).stream()
154 .<CoreOp.FuncOp>mapMulti((op, objectConsumer) -> {
155 if (op instanceof CoreOp.FuncOp fop) {
156 objectConsumer.accept(fop);
157 }
158 }).toList();
159 }
160 }
161
162 public static final class AnfIfOp extends Op implements Op.Terminating, Op.Nested {
163 public static final String NAME = "anf.if";
164
165 public static class ThenBuilder {
166 final Body.Builder ancestorBody;
167 final TypeElement yieldType;
168 final Value test;
169
170 ThenBuilder(Body.Builder ancestorBody, TypeElement yieldType, Value test) {
171 this.ancestorBody = ancestorBody;
172 this.yieldType = yieldType;
173 this.test = test;
174 }
175
176 public ElseBuilder if_(Consumer<Block.Builder> c) {
177 Body.Builder then_ = Body.Builder.of(ancestorBody,
178 CoreType.functionType(yieldType));
179 c.accept(then_.entryBlock());
180
181 return new ElseBuilder(this, then_);
182 }
183 }
184
185 public static class ElseBuilder {
186 final ThenBuilder thenBuilder;
187 final Body.Builder then_;
188
189 public ElseBuilder(ThenBuilder thenBuilder, Body.Builder then_) {
190 this.thenBuilder = thenBuilder;
191 this.then_ = then_;
192 }
193
194 public AnfIfOp else_(Consumer<Block.Builder> c) {
195 Body.Builder else_ = Body.Builder.of(thenBuilder.ancestorBody,
196 CoreType.functionType(thenBuilder.yieldType));
197 c.accept(else_.entryBlock());
198
199 return new AnfIfOp(thenBuilder.test, then_, else_);
200 }
201 }
202
203 final Body then_, else_;
204
205 public AnfIfOp(ExternalizedOp def) {
206 this(def.operands().getFirst(),
207 def.bodyDefinitions().get(0),
208 def.bodyDefinitions().get(1));
209 }
210
211 public AnfIfOp(AnfIfOp that, CopyContext cc, OpTransformer ot) {
212 super(that, cc);
213
214 this.then_ = that.then_.transform(cc, ot).build(this);
215 this.else_ = that.else_.transform(cc, ot).build(this);
216 }
217
218 @Override
219 public Op transform(CopyContext cc, OpTransformer ot) {
220 return new AnfIfOp(this, cc, ot);
221 }
222
223 AnfIfOp(Value test, Body.Builder thenBodyBuilder, Body.Builder elseBodyBuilder) {
224 super(List.of(test));
225
226 this.then_ = thenBodyBuilder.build(this);
227 this.else_ = elseBodyBuilder.build(this);
228 }
229
230 public Value getTest() {
231 return this.operands().get(0);
232 }
233
234 public Body _then() {
235 return then_;
236 }
237
238 public Body _else() {
239 return else_;
240 }
241
242 @Override
243 public List<Body> bodies() {
244 return List.of(then_, else_);
245 }
246
247 @Override
248 public TypeElement resultType() {
249 return this.then_.yieldType();
250 }
251 }
252
253 public static final class AnfFuncOp extends Op implements Op.Nested {
254
255 public static class Builder {
256 final Body.Builder ancestorBody;
257 final String funcName;
258 final FunctionType funcType;
259
260 Builder(Body.Builder ancestorBody, String funcName, FunctionType funcType) {
261 this.ancestorBody = ancestorBody;
262 this.funcName = funcName;
263 this.funcType = funcType;
264 }
265
266 public AnfFuncOp body(Consumer<Block.Builder> c) {
267 Body.Builder body = Body.Builder.of(ancestorBody, funcType);
268 c.accept(body.entryBlock());
269 return new AnfFuncOp(funcName, body);
270 }
271 }
272
273 public static final String NAME = "anf.func";
274 public static final String ATTRIBUTE_FUNC_NAME = NAME + ".name";
275
276 final String funcName;
277 final Body body;
278
279 public static AnfFuncOp create(ExternalizedOp def) {
280 if (!def.operands().isEmpty()) {
281 throw new IllegalStateException("Bad op " + def.name());
282 }
283
284 String funcName = def.extractAttributeValue(ATTRIBUTE_FUNC_NAME, true,
285 v -> switch (v) {
286 case String s -> s;
287 case null, default -> throw new UnsupportedOperationException("Unsupported func name value:" + v);
288 });
289 return new AnfFuncOp(funcName, def.bodyDefinitions().get(0));
290 }
291
292 AnfFuncOp(AnfFuncOp that, CopyContext cc, OpTransformer oa) {
293 this(that, that.funcName, cc, oa);
294 }
295
296 AnfFuncOp(AnfFuncOp that, String funcName, CopyContext cc, OpTransformer ot) {
297 super(that, cc);
298
299 this.funcName = funcName;
300 this.body = that.body.transform(cc, ot).build(this);
301 }
302
303 @Override
304 public AnfFuncOp transform(CopyContext cc, OpTransformer ot) {
305 return new AnfFuncOp(this, cc, ot);
306 }
307
308 AnfFuncOp(String funcName, Body.Builder bodyBuilder) {
309 super(
310 List.of());
311
312 this.funcName = funcName;
313 this.body = bodyBuilder.build(this);
314 }
315
316 @Override
317 public List<Body> bodies() {
318 return List.of(body);
319 }
320
321 @Override
322 public Map<String, Object> externalize() {
323 return Map.of("", funcName);
324 }
325
326 public FunctionType invokableType() {
327 return body.bodyType();
328 }
329
330 public String funcName() {
331 return funcName;
332 }
333
334 public Body body() {
335 return body;
336 }
337
338 @Override
339 public TypeElement resultType() {
340 return invokableType();
341 }
342 }
343
344 public static final class AnfApply extends Op implements Op.Terminating {
345 public static final String NAME = "anf.apply";
346
347 public AnfApply(ExternalizedOp def) {
348 this(def.operands());
349 }
350
351 public AnfApply(AnfApply that, CopyContext cc) {
352 super(that, cc);
353 }
354
355 @Override
356 public Op transform(CopyContext cc, OpTransformer ot) {
357 return new AnfApply(this, cc);
358 }
359
360 @Override
361 public TypeElement resultType() {
362 FunctionType ft = (FunctionType) operands().get(0).type();
363 return ft.returnType();
364 }
365
366 public AnfApply(List<Value> arguments) {
367 super(arguments);
368
369 // First argument is func value
370 // Subsequent arguments are func arguments
371 }
372
373
374 public List<Value> args() {
375 return operands().subList(1, this.operands().size());
376 }
377 }
378
379 public static final class AnfApplyStub extends Op implements Op.Terminating {
380 public static final String NAME = "anf.apply.stub";
381 public static final String ATTRIBUTE_RESULT_TYPE = ".resultType";
382 public static final String ATTRIBUTE_CALLSITE_NAME = ".callsiteName";
383
384 public final String callSiteName;
385 public final TypeElement resultType;
386
387 public static AnfApplyStub create(ExternalizedOp def) {
388 if (!def.operands().isEmpty()) {
389 throw new IllegalStateException("Bad op " + def.name());
390 }
391
392 String callsiteName = def.extractAttributeValue(ATTRIBUTE_CALLSITE_NAME, true,
393 v -> switch (v) {
394 case String s -> s;
395 case null, default -> throw new UnsupportedOperationException("Unsupported func name value:" + v);
396 });
397 return new AnfApplyStub(callsiteName, def.operands(), def.resultType());
398 }
399
400 public AnfApplyStub(AnfApplyStub that, CopyContext cc) {
401 super(that, cc);
402 this.callSiteName = that.callSiteName;
403 this.resultType = that.resultType;
404 }
405
406 @Override
407 public Map<String, Object> externalize() {
408 return Map.of("", callSiteName);
409 }
410
411 @Override
412 public Op transform(CopyContext cc, OpTransformer ot) {
413 return new AnfApplyStub(this, cc);
414 }
415
416 public AnfApplyStub(String callSiteName, List<Value> arguments, TypeElement resultType) {
417 super(arguments);
418 this.resultType = resultType;
419 this.callSiteName = callSiteName;
420
421 // First argument is func value
422 // Subsequent arguments are func arguments
423 }
424
425 @Override
426 public TypeElement resultType() {
427 return this.resultType;
428 }
429
430 public List<Value> args() {
431 return operands().subList(1, this.operands().size());
432 }
433 }
434
435 static Op createOp(ExternalizedOp def) {
436 Op op = switch (def.name()) {
437 case "anf.apply" -> new AnfApply(def);
438 case "anf.apply.stub" -> AnfApplyStub.create(def);
439 case "anf.func" -> AnfFuncOp.create(def);
440 case "anf.if" -> new AnfIfOp(def);
441 case "anf.let" -> new AnfLetOp(def);
442 case "anf.letrec" -> new AnfLetRecOp(def);
443 default -> null;
444 };
445 if (op != null) {
446 op.setLocation(def.location());
447 }
448 return op;
449 }
450
451 static final OpFactory FACTORY = AnfDialect::createOp;
452
453 public static AnfLetRecOp.Builder letrec(Body.Builder ancestorBody, TypeElement yieldType) {
454 return new AnfLetRecOp.Builder(ancestorBody, yieldType);
455 }
456
457 public static AnfLetRecOp letrec(Body.Builder body) {
458 return new AnfLetRecOp(body);
459 }
460
461 public static AnfLetOp.Builder let(Body.Builder ancestorBody, TypeElement yieldType) {
462 return new AnfLetOp.Builder(ancestorBody, yieldType);
463 }
464
465 public static AnfLetOp let(Body.Builder body) {
466 return new AnfLetOp(body);
467 }
468
469 public static AnfIfOp.ThenBuilder if_(Body.Builder ancestorBody, TypeElement yieldType, Value test) {
470 return new AnfIfOp.ThenBuilder(ancestorBody, yieldType, test);
471 }
472
473 public static AnfIfOp if_(Body.Builder then_, Body.Builder else_, Value test) {
474 return new AnfIfOp(test, then_, else_);
475 }
476
477 public static AnfFuncOp.Builder func(Body.Builder ancestorBody, String funcName, FunctionType funcType) {
478 List<TypeElement> params = new ArrayList<>();
479 params.add(funcType.returnType());
480 params.addAll(funcType.parameterTypes());
481 return new AnfFuncOp.Builder(ancestorBody, funcName, CoreType.functionType(funcType.returnType(), params));
482 }
483
484 public static AnfFuncOp func(String funcName, Body.Builder body) {
485 return new AnfFuncOp(funcName, body);
486 }
487
488 public static AnfApply apply(List<Value> arguments) {
489 return new AnfApply(arguments);
490 }
491 //public static AnfApplyStub applyStub(String name, List<Value> arguments, TypeElement type) { return new AnfApplyStub(name, arguments, type);}
492 }