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.
  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.Op;
 25 import jdk.incubator.code.Quoted;
 26 import jdk.incubator.code.TypeElement;
 27 import jdk.incubator.code.Value;
 28 import jdk.incubator.code.analysis.Inliner;
 29 import jdk.incubator.code.dialect.java.JavaOp;
 30 import jdk.incubator.code.dialect.java.JavaType;
 31 import jdk.incubator.code.dialect.java.MethodRef;
 32 import jdk.incubator.code.interpreter.Interpreter;
 33 import org.junit.jupiter.api.Assertions;
 34 import org.junit.jupiter.api.Test;
 35 
 36 import java.lang.invoke.MethodHandles;
 37 import java.util.stream.Stream;
 38 
 39 import static jdk.incubator.code.dialect.core.CoreOp.*;
 40 import static jdk.incubator.code.dialect.core.CoreType.functionType;
 41 import static jdk.incubator.code.dialect.java.MethodRef.method;
 42 
 43 /*
 44  * @test
 45  * @modules jdk.incubator.code
 46  * @run junit TestLinqUsingQuoted
 47  */
 48 
 49 public class TestLinqUsingQuoted {
 50 
 51     // Query interfaces
 52 
 53     public interface Queryable {
 54         TypeElement elementType();
 55 
 56         // Queryable<T> -> Queryable<U>
 57         FuncOp expression();
 58 
 59         QueryProvider provider();
 60 
 61         // T -> boolean
 62         // Predicate<T>
 63         default Queryable where(Quoted f) {
 64             // @@@@ validate
 65             ClosureOp c = (ClosureOp) f.op();
 66             return insertQuery(elementType(), "where", c);
 67         }
 68 
 69         // T -> R
 70         // Function<T, R>
 71         default Queryable select(Quoted f) {
 72             // @@@@ validate
 73             ClosureOp c = (ClosureOp) f.op();
 74             return insertQuery(c.invokableType().returnType(), "select", c);
 75         }
 76 
 77         private Queryable insertQuery(TypeElement et, String name, ClosureOp c) {
 78             QueryProvider qp = provider();
 79 
 80             FuncOp currentQueryExpression = expression();
 81             FuncOp nextQueryExpression = currentQueryExpression.transform((block, op) -> {
 82                 if (op instanceof ReturnOp rop && rop.ancestorBody() == currentQueryExpression.body()) {
 83                     Value query = block.context().getValue(rop.returnValue());
 84 
 85                     Op.Result quotedLambda = block.op(quoted(block.parentBody(), qblock -> c));
 86 
 87                     MethodRef md = method(qp.queryableType(), name,
 88                             functionType(qp.queryableType(), QuotedOp.QUOTED_TYPE));
 89                     Op.Result queryable = block.op(JavaOp.invoke(md, query, quotedLambda));
 90 
 91                     block.op(return_(queryable));
 92                 } else {
 93                     block.op(op);
 94                 }
 95                 return block;
 96             });
 97 
 98             return qp.createQuery(et, nextQueryExpression);
 99         }
100 
101         // Iterate
102         // Queryable -> Stream
103         default QueryResult elements() {
104             TypeElement resultType = JavaType.parameterized(JavaType.type(Stream.class), (JavaType) elementType());
105             return insertQueryResult("elements", resultType);
106         }
107 
108         // Count
109         // Queryable -> Long
110         default QueryResult count() {
111             return insertQueryResult("count", JavaType.LONG);
112         }
113 
114         private QueryResult insertQueryResult(String name, TypeElement resultType) {
115             QueryProvider qp = provider();
116 
117             // Copy function expression, replacing return type
118             FuncOp currentQueryExpression = expression();
119             FuncOp nextQueryExpression = func("queryresult",
120                     functionType(qp.queryResultType(), currentQueryExpression.invokableType().parameterTypes()))
121                     .body(b -> Inliner.inline(b, currentQueryExpression, b.parameters(), (block, query) -> {
122                         MethodRef md = method(qp.queryableType(), name, functionType(qp.queryResultType()));
123                         Op.Result queryResult = block.op(JavaOp.invoke(md, query));
124 
125                         block.op(return_(queryResult));
126                     }));
127             return qp.createQueryResult(resultType, nextQueryExpression);
128         }
129     }
130 
131     public interface QueryResult {
132         TypeElement resultType();
133 
134         // Queryable -> QueryResult
135         FuncOp expression();
136 
137         Object execute();
138     }
139 
140     public interface QueryProvider {
141         TypeElement queryableType();
142 
143         TypeElement queryResultType();
144 
145         Queryable createQuery(TypeElement elementType, FuncOp expression);
146 
147         QueryResult createQueryResult(TypeElement resultType, FuncOp expression);
148 
149         Queryable newQuery(TypeElement elementType);
150     }
151 
152 
153     // Query implementation
154 
155     public static final class TestQueryable implements Queryable {
156         final TypeElement elementType;
157         final TestQueryProvider provider;
158         final FuncOp expression;
159 
160         TestQueryable(TypeElement elementType, TestQueryProvider provider) {
161             this.elementType = elementType;
162             this.provider = provider;
163 
164             // Initial expression is an identity function
165             var funType = functionType(provider().queryableType(), provider().queryableType());
166             this.expression = func("query", funType)
167                     .body(b -> b.op(return_(b.parameters().get(0))));
168         }
169 
170         TestQueryable(TypeElement elementType, TestQueryProvider provider, FuncOp expression) {
171             this.elementType = elementType;
172             this.provider = provider;
173             this.expression = expression;
174         }
175 
176         @Override
177         public TypeElement elementType() {
178             return elementType;
179         }
180 
181         @Override
182         public FuncOp expression() {
183             return expression;
184         }
185 
186         @Override
187         public QueryProvider provider() {
188             return provider;
189         }
190     }
191 
192     public record TestQueryResult(TypeElement resultType, FuncOp expression) implements QueryResult {
193         @Override
194         public Object execute() {
195             // @@@ Compile/translate the expression and execute it
196             throw new UnsupportedOperationException();
197         }
198     }
199 
200     public static final class TestQueryProvider implements QueryProvider {
201         final TypeElement queryableType;
202         final TypeElement queryResultType;
203 
204         TestQueryProvider() {
205             this.queryableType = JavaType.type(Queryable.class);
206             this.queryResultType = JavaType.type(QueryResult.class);
207         }
208 
209         @Override
210         public TypeElement queryableType() {
211             return queryableType;
212         }
213 
214         @Override
215         public TypeElement queryResultType() {
216             return queryResultType;
217         }
218 
219         @Override
220         public TestQueryable createQuery(TypeElement elementType, FuncOp expression) {
221             return new TestQueryable(elementType, this, expression);
222         }
223 
224         @Override
225         public QueryResult createQueryResult(TypeElement resultType, FuncOp expression) {
226             return new TestQueryResult(resultType, expression);
227         }
228 
229         @Override
230         public Queryable newQuery(TypeElement elementType) {
231             return new TestQueryable(elementType, this);
232         }
233     }
234 
235 
236     static class Customer {
237         // 1st column
238         String contactName;
239         // 2nd column
240         String phone;
241         // 3rd column
242         String city;
243     }
244 
245     @Test
246     public void testSimpleQuery() {
247         QueryProvider qp = new TestQueryProvider();
248 
249         QueryResult qr = qp.newQuery(JavaType.type(Customer.class))
250                 // c -> c.city.equals("London")
251                 .where((Customer c) -> c.city.equals("London"))
252                 // c -> c.contactName
253                 .select((Customer c) -> c.contactName).elements();
254 
255         Op op1 = qr.expression();
256         System.out.println(op1.toText());
257 
258         QueryResult qr2 = (QueryResult) Interpreter.invoke(MethodHandles.lookup(),
259                 qr.expression(), qp.newQuery(JavaType.type(Customer.class)));
260 
261         Op op = qr2.expression();
262         System.out.println(op.toText());
263 
264         Assertions.assertEquals(qr2.expression().toText(), qr.expression().toText());
265     }
266 }