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