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 import java.util.stream.Stream;
 36 
 37 /**
 38  * The quoted form of an operation.
 39  * <p>
 40  * The quoted form is utilized when the code model of some code is to be obtained rather than obtaining the result of
 41  * executing that code. For example passing the of a lambda expression in quoted form rather than the expression being
 42  * targeted to a functional interface from which it can be invoked.
 43  */
 44 public final class Quoted {
 45     private final Op op;
 46     private final SequencedMap<Value, Object> capturedValues;
 47 
 48     static final SequencedMap<Value, Object> EMPTY_SEQUENCED_MAP = new LinkedHashMap<>();
 49     /**
 50      * Constructs the quoted form of a given operation.
 51      *
 52      * @param op the invokable operation.
 53      */
 54     public Quoted(Op op) {
 55         this(op, EMPTY_SEQUENCED_MAP);
 56     }
 57 
 58     /**
 59      * Constructs the quoted form of a given operation.
 60      * <p>
 61      * The captured values key set must have the same elements and same encounter order as
 62      * operation's captured values, specifically the following expression should evaluate to true:
 63      * {@snippet lang=java :
 64      * op.capturedValues().equals(new ArrayList<>(capturedValues.keySet()));
 65      * }
 66      *
 67      * @param op             the operation.
 68      * @param capturedValues the captured values referred to by the operation
 69      * @see Op#capturedValues()
 70      */
 71     public Quoted(Op op, SequencedMap<Value, Object> capturedValues) {
 72         // @@@ This check is potentially expensive, remove or keep as assert?
 73         // @@@ Or make Quoted an interface, with a module private implementation?
 74         assert Stream.concat(op.operands().stream(), op.capturedValues().stream()).toList()
 75                 .equals(new ArrayList<>(capturedValues.keySet()));
 76         this.op = op;
 77         this.capturedValues = Collections.unmodifiableSequencedMap(new LinkedHashMap<>(capturedValues));
 78     }
 79 
 80     /**
 81      * Returns the operation.
 82      *
 83      * @return the operation.
 84      */
 85     public Op op() {
 86         return op;
 87     }
 88 
 89     /**
 90      * Returns the captured values.
 91      * <p>
 92      * The captured values key set has the same elements and same encounter order as
 93      * operation's captured values, specifically the following expression evaluates to true:
 94      * {@snippet lang=java :
 95      * op().capturedValues().equals(new ArrayList<>(capturedValues().keySet()));
 96      * }
 97      *
 98      * @return the captured values, as an unmodifiable map.
 99      */
100     public SequencedMap<Value, Object> capturedValues() {
101         return capturedValues;
102     }
103 
104     // Take an op from its original context to a new one where operands and captured values are parameters
105     public static CoreOp.FuncOp quoteOp(Op op) {
106 
107         if (op.result() == null) {
108             throw new IllegalArgumentException("Op not bound");
109         }
110 
111         List<Value> inputOperandsAndCaptures = Stream.concat(op.operands().stream(), op.capturedValues().stream()).toList();
112 
113         // Build the function type
114         List<TypeElement> params = inputOperandsAndCaptures.stream()
115                 .map(v -> v.type() instanceof VarType vt ? vt.valueType() : v.type())
116                 .toList();
117         FunctionType ft = CoreType.functionType(CoreOp.QuotedOp.QUOTED_TYPE, params);
118 
119         // Build the function that quotes the lambda
120         return CoreOp.func("q", ft).body(b -> {
121             // Create variables as needed and obtain the operands and captured values for the copied lambda
122             List<Value> outputOperandsAndCaptures = new ArrayList<>();
123             for (int i = 0; i < inputOperandsAndCaptures.size(); i++) {
124                 Value inputValue = inputOperandsAndCaptures.get(i);
125                 Value outputValue = b.parameters().get(i);
126                 if (inputValue.type() instanceof VarType) {
127                     outputValue = b.op(CoreOp.var(String.valueOf(i), outputValue));
128                 }
129                 outputOperandsAndCaptures.add(outputValue);
130             }
131 
132             // Quoted the lambda expression
133             Value q = b.op(CoreOp.quoted(b.parentBody(), qb -> {
134                 // Map the entry block of the op's ancestor body to the quoted block
135                 // We are copying op in the context of the quoted block, the block mapping
136                 // ensures the use of operands and captured values are reachable when building
137                 qb.context().mapBlock(op.ancestorBody().entryBlock(), qb);
138                 // Map the op's operands and captured values
139                 qb.context().mapValues(inputOperandsAndCaptures, outputOperandsAndCaptures);
140                 // Return the op to be copied in the quoted operation
141                 return op;
142             }));
143             b.op(CoreOp.return_(q));
144         });
145     }
146 
147     private static RuntimeException invalidQuotedModel(CoreOp.FuncOp model) {
148         return new RuntimeException("Invalid code model for quoted operation : " + model);
149     }
150 
151     // Extract the quoted operation from funcOp and maps the operands and captured values to the runtime values
152     // @@@ Add List<Object> accepting method, varargs array defers to it
153     public static Quoted quotedOp(CoreOp.FuncOp funcOp, Object... args) {
154 
155         if (funcOp.body().blocks().size() != 1) {
156             throw invalidQuotedModel(funcOp);
157         }
158         Block fblock = funcOp.body().entryBlock();
159 
160         if (fblock.ops().size() < 2) {
161             throw invalidQuotedModel(funcOp);
162         }
163 
164         if (!(fblock.ops().get(fblock.ops().size() - 2) instanceof CoreOp.QuotedOp qop)) {
165             throw invalidQuotedModel(funcOp);
166         }
167 
168         if (!(fblock.ops().getLast() instanceof CoreOp.ReturnOp returnOp)) {
169             throw invalidQuotedModel(funcOp);
170         }
171         if (returnOp.returnValue() == null) {
172             throw invalidQuotedModel(funcOp);
173         }
174         if (!returnOp.returnValue().equals(qop.result())) {
175             throw invalidQuotedModel(funcOp);
176         }
177 
178         Op op = qop.quotedOp();
179 
180         SequencedSet<Value> operandsAndCaptures = new LinkedHashSet<>();
181         operandsAndCaptures.addAll(op.operands());
182         operandsAndCaptures.addAll(op.capturedValues());
183 
184         // validation rule of block params and constant op result
185         Consumer<Value> validate = v -> {
186             if (v.uses().isEmpty()) {
187                 throw invalidQuotedModel(funcOp);
188             } else if (v.uses().size() == 1
189                     && !(v.uses().iterator().next().op() instanceof CoreOp.VarOp vop && vop.result().uses().size() >= 1
190                     && vop.result().uses().stream().noneMatch(u -> u.op().parentBlock() == fblock))
191                     && !operandsAndCaptures.contains(v)) {
192                 throw invalidQuotedModel(funcOp);
193             } else if (v.uses().size() > 1 && v.uses().stream().anyMatch(u -> u.op().parentBlock() == fblock)) {
194                 throw invalidQuotedModel(funcOp);
195             }
196         };
197 
198         for (Block.Parameter p : fblock.parameters()) {
199             validate.accept(p);
200         }
201 
202         List<Op> ops = fblock.ops().subList(0, fblock.ops().size() - 2);
203         for (Op o : ops) {
204             switch (o) {
205                 case CoreOp.VarOp varOp -> {
206                     if (varOp.isUninitialized()) {
207                         throw invalidQuotedModel(funcOp);
208                     }
209                     if (varOp.initOperand() instanceof Op.Result opr && !(opr.op() instanceof CoreOp.ConstantOp)) {
210                         throw invalidQuotedModel(funcOp);
211                     }
212                 }
213                 case CoreOp.ConstantOp cop -> validate.accept(cop.result());
214                 default -> throw invalidQuotedModel(funcOp);
215             }
216         }
217 
218         // map captured values to their corresponding runtime values
219         // captured value can be:
220         // 1- block param
221         // 2- result of VarOp whose initial value is constant
222         // 3- result of VarOp whose initial value is block param
223         // 4- result of ConstantOp
224         List<Block.Parameter> params = funcOp.parameters();
225         if (params.size() != args.length) {
226             throw invalidQuotedModel(funcOp);
227         }
228         SequencedMap<Value, Object> m = new LinkedHashMap<>();
229         for (Value v : operandsAndCaptures) {
230             switch (v) {
231                 case Block.Parameter p -> {
232                     Object rv = args[p.index()];
233                     m.put(v, rv);
234                 }
235                 case Op.Result opr when opr.op() instanceof CoreOp.VarOp varOp -> {
236                     if (varOp.initOperand() instanceof Op.Result r && r.op() instanceof CoreOp.ConstantOp cop) {
237                         m.put(v, CoreOp.Var.of(cop.value()));
238                     } else if (varOp.initOperand() instanceof Block.Parameter p) {
239                         Object rv = args[p.index()];
240                         m.put(v, CoreOp.Var.of(rv));
241                     }
242                 }
243                 case Op.Result opr when opr.op() instanceof CoreOp.ConstantOp cop -> {
244                     m.put(v, cop.value());
245                 }
246                 default -> throw invalidQuotedModel(funcOp);
247             }
248         }
249 
250         return new Quoted(op, m);
251     }
252 }