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.analysis.Inliner;
26 import jdk.incubator.code.dialect.java.ClassType;
27 import jdk.incubator.code.dialect.java.JavaOp;
28 import jdk.incubator.code.dialect.java.MethodRef;
29 import jdk.incubator.code.dialect.java.JavaType;
30
31 import java.util.function.Function;
32 import java.util.function.Predicate;
33 import java.util.stream.Stream;
34
35 import static jdk.incubator.code.dialect.java.JavaType.parameterized;
36 import static jdk.incubator.code.dialect.java.MethodRef.method;
37 import static jdk.incubator.code.dialect.core.CoreOp.*;
38 import static jdk.incubator.code.dialect.core.CoreType.functionType;
39 import static jdk.incubator.code.dialect.java.JavaType.type;
40
41 public interface Queryable<T> {
42 JavaType TYPE = type(Queryable.class);
43
44 QueryProvider provider();
45
46 // T
47 JavaType elementType();
48
49 // Queryable<T> -> Queryable<U>
50 FuncOp expression();
51
52 @SuppressWarnings("unchecked")
53 default Queryable<T> where(Predicate<T> f) {
54 JavaOp.LambdaOp l = (JavaOp.LambdaOp) Op.ofQuotable(f).get().op();
55 return (Queryable<T>) insertQuery(elementType(), "where", l);
56 }
57
58 @SuppressWarnings("unchecked")
59 default <R> Queryable<R> select(Function<T, R> f) {
60 JavaOp.LambdaOp l = (JavaOp.LambdaOp) Op.ofQuotable(f).get().op();
61 return (Queryable<R>) insertQuery((JavaType) l.invokableType().returnType(), "select", l);
62 }
63
64 private Queryable<?> insertQuery(JavaType elementType, String methodName, JavaOp.LambdaOp lambdaOp) {
65 // Copy function expression, replacing return operation
66 FuncOp queryExpression = expression();
67 JavaType queryableType = parameterized(Queryable.TYPE, elementType);
68 FuncOp nextQueryExpression = func("query",
69 functionType(queryableType, queryExpression.invokableType().parameterTypes()))
70 .body(b -> Inliner.inline(b, queryExpression, b.parameters(), (block, query) -> {
71 Op.Result fi = block.op(lambdaOp);
72
73 MethodRef md = method(Queryable.TYPE, methodName,
74 functionType(Queryable.TYPE, ((ClassType) lambdaOp.functionalInterface()).rawType()));
75 Op.Result queryable = block.op(JavaOp.invoke(queryableType, md, query, fi));
76
77 block.op(return_(queryable));
78 }));
79
80 return provider().createQuery(elementType, nextQueryExpression);
81 }
82
83 @SuppressWarnings("unchecked")
84 default QueryResult<Stream<T>> elements() {
85 JavaType resultType = parameterized(type(Stream.class), elementType());
86 return (QueryResult<Stream<T>>) insertQueryResult(resultType, "elements");
87 }
88
89 @SuppressWarnings("unchecked")
90 default QueryResult<Long> count() {
91 JavaType resultType = JavaType.LONG;
92 return (QueryResult<Long>) insertQueryResult(resultType, "count");
93 }
94
95 private QueryResult<?> insertQueryResult(JavaType resultType, String methodName) {
96 // Copy function expression, replacing return operation
97 FuncOp queryExpression = expression();
98 JavaType queryResultType = parameterized(QueryResult.TYPE, resultType);
99 FuncOp queryResultExpression = func("queryResult",
100 functionType(queryResultType, queryExpression.invokableType().parameterTypes()))
101 .body(b -> Inliner.inline(b, queryExpression, b.parameters(), (block, query) -> {
102 MethodRef md = method(Queryable.TYPE, methodName,
103 functionType(QueryResult.TYPE));
104 Op.Result queryResult = block.op(JavaOp.invoke(queryResultType, md, query));
105
106 block.op(return_(queryResult));
107 }));
108
109 return provider().createQueryResult(resultType, queryResultExpression);
110 }
111 }