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 }