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