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