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(NAME, 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(AnfLetRecOp.NAME, 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(NAME, 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(NAME,
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(AnfApply.NAME, 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(AnfApplyStub.NAME, 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 }