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 java.lang.reflect.code.*;
 25 import java.lang.reflect.code.op.CoreOp;
 26 import java.lang.reflect.code.op.ExtendedOp.JavaEnhancedForOp;
 27 import java.lang.reflect.code.type.ClassType;
 28 import java.lang.reflect.code.type.FunctionType;
 29 import java.lang.reflect.code.type.JavaType;
 30 import java.util.ArrayList;
 31 import java.util.List;
 32 import java.util.function.BiConsumer;
 33 import java.util.function.BiFunction;
 34 import java.util.function.Consumer;
 35 import java.util.function.Function;
 36 
 37 import static java.lang.reflect.code.op.CoreOp.*;
 38 import static java.lang.reflect.code.op.ExtendedOp._continue;
 39 import static java.lang.reflect.code.op.ExtendedOp.enhancedFor;
 40 import static java.lang.reflect.code.type.JavaType.parameterized;
 41 import static java.lang.reflect.code.type.JavaType.type;
 42 
 43 public final class StreamFuser {
 44 
 45     StreamFuser() {}
 46 
 47     public static StreamExprBuilder fromList(JavaType elementType) {
 48         // java.util.List<E>
 49         JavaType listType = parameterized(type(List.class), elementType);
 50         return new StreamExprBuilder(listType, elementType,
 51                 (b, v) -> StreamExprBuilder.enhancedForLoop(b, elementType, v)::body);
 52     }
 53 
 54     public static class StreamExprBuilder {
 55         static class StreamOp {
 56             final Quoted quotedClosure;
 57 
 58             StreamOp(Quoted quotedClosure) {
 59                 if (!(quotedClosure.op() instanceof CoreOp.ClosureOp)) {
 60                     throw new IllegalArgumentException("Quoted operation is not closure operation");
 61                 }
 62                 this.quotedClosure = quotedClosure;
 63             }
 64 
 65             CoreOp.ClosureOp op() {
 66                 return (CoreOp.ClosureOp) quotedClosure.op();
 67             }
 68         }
 69 
 70         static class MapStreamOp extends StreamOp {
 71             public MapStreamOp(Quoted quotedClosure) {
 72                 super(quotedClosure);
 73                 // @@@ Check closure signature
 74             }
 75         }
 76 
 77         static class FlatMapStreamOp extends StreamOp {
 78             public FlatMapStreamOp(Quoted quotedClosure) {
 79                 super(quotedClosure);
 80                 // @@@ Check closure signature
 81             }
 82         }
 83 
 84         static class FilterStreamOp extends StreamOp {
 85             public FilterStreamOp(Quoted quotedClosure) {
 86                 super(quotedClosure);
 87                 // @@@ Check closure signature
 88             }
 89         }
 90 
 91         final JavaType sourceType;
 92         final JavaType sourceElementType;
 93         final BiFunction<Body.Builder, Value, Function<Consumer<Block.Builder>, Op>> loopSupplier;
 94         final List<StreamOp> streamOps;
 95 
 96         StreamExprBuilder(JavaType sourceType, JavaType sourceElementType,
 97                           BiFunction<Body.Builder, Value, Function<Consumer<Block.Builder>, Op>> loopSupplier) {
 98             this.sourceType = sourceType;
 99             this.sourceElementType = sourceElementType;
100             this.loopSupplier = loopSupplier;
101             this.streamOps = new ArrayList<>();
102         }
103 
104         static JavaEnhancedForOp.BodyBuilder enhancedForLoop(Body.Builder ancestorBody, JavaType elementType,
105                                                              Value iterable) {
106             return enhancedFor(ancestorBody, iterable.type(), elementType)
107                     .expression(b -> {
108                         b.op(_yield(iterable));
109                     })
110                     .definition(b -> {
111                         b.op(_yield(b.parameters().get(0)));
112                     });
113         }
114 
115         public StreamExprBuilder map(Quoted f) {
116             streamOps.add(new MapStreamOp(f));
117             return this;
118         }
119 
120         public StreamExprBuilder flatMap(Quoted f) {
121             streamOps.add(new FlatMapStreamOp(f));
122             return this;
123         }
124 
125         public StreamExprBuilder filter(Quoted f) {
126             streamOps.add(new FilterStreamOp(f));
127             return this;
128         }
129 
130         void fuseIntermediateOperations(Block.Builder body, BiConsumer<Block.Builder, Value> terminalConsumer) {
131             fuseIntermediateOperation(0, body, body.parameters().get(0), null, terminalConsumer);
132         }
133 
134         void fuseIntermediateOperation(int i, Block.Builder body, Value element, Block.Builder continueBlock,
135                                        BiConsumer<Block.Builder, Value> terminalConsumer) {
136             if (i == streamOps.size()) {
137                 terminalConsumer.accept(body, element);
138                 return;
139             }
140 
141             StreamOp sop = streamOps.get(i);
142             if (sop instanceof MapStreamOp) {
143                 body.inline(sop.op(), List.of(element), (block, value) -> {
144                     fuseIntermediateOperation(i + 1, block, value, continueBlock, terminalConsumer);
145                 });
146             } else if (sop instanceof FilterStreamOp) {
147                 body.inline(sop.op(), List.of(element), (block, p) -> {
148                     Block.Builder _if = block.block();
149                     Block.Builder _else = continueBlock;
150                     if (continueBlock == null) {
151                         _else = block.block();
152                         _else.op(_continue());
153                     }
154 
155                     block.op(conditionalBranch(p, _if.successor(), _else.successor()));
156 
157                     fuseIntermediateOperation(i + 1, _if, element, _else, terminalConsumer);
158                 });
159             } else if (sop instanceof FlatMapStreamOp) {
160                 body.inline(sop.op(), List.of(element), (block, iterable) -> {
161                     JavaEnhancedForOp forOp = enhancedFor(block.parentBody(),
162                             iterable.type(), ((ClassType) iterable.type()).typeArguments().get(0))
163                             .expression(b -> {
164                                 b.op(_yield(iterable));
165                             })
166                             .definition(b -> {
167                                 b.op(_yield(b.parameters().get(0)));
168                             })
169                             .body(b -> {
170                                 fuseIntermediateOperation(i + 1,
171                                         b,
172                                         b.parameters().get(0),
173                                         null, terminalConsumer);
174                             });
175 
176                     block.op(forOp);
177                     block.op(_continue());
178                 });
179             }
180         }
181 
182         public FuncOp forEach(Quoted quotedConsumer) {
183             if (!(quotedConsumer.op() instanceof CoreOp.ClosureOp consumer)) {
184                 throw new IllegalArgumentException("Quoted consumer is not closure operation");
185             }
186 
187             return func("fused.forEach", FunctionType.functionType(JavaType.VOID, sourceType))
188                     .body(b -> {
189                         Value source = b.parameters().get(0);
190 
191                         Op sourceLoop = loopSupplier.apply(b.parentBody(), source)
192                                 .apply(loopBlock -> {
193                                     fuseIntermediateOperations(loopBlock, (terminalBlock, resultValue) -> {
194                                         terminalBlock.inline(consumer, List.of(resultValue),
195                                                 (_, _) -> {
196                                                 });
197                                         terminalBlock.op(_continue());
198                                     });
199 
200                                 });
201                         b.op(sourceLoop);
202                         b.op(_return());
203                     });
204         }
205 
206         // Supplier<C> supplier, BiConsumer<C, T> accumulator
207         public FuncOp collect(Quoted quotedSupplier, Quoted quotedAccumulator) {
208             if (!(quotedSupplier.op() instanceof CoreOp.ClosureOp supplier)) {
209                 throw new IllegalArgumentException("Quoted supplier is not closure operation");
210             }
211             if (!(quotedAccumulator.op() instanceof CoreOp.ClosureOp accumulator)) {
212                 throw new IllegalArgumentException("Quoted accumulator is not closure operation");
213             }
214 
215             JavaType collectType = (JavaType) supplier.invokableType().returnType();
216             return func("fused.collect", FunctionType.functionType(collectType, sourceType))
217                     .body(b -> {
218                         Value source = b.parameters().get(0);
219 
220                         b.inline(supplier, List.of(), (block, collect) -> {
221                             Op sourceLoop = loopSupplier.apply(block.parentBody(), source)
222                                     .apply(loopBlock -> {
223                                         fuseIntermediateOperations(loopBlock, (terminalBlock, resultValue) -> {
224                                             terminalBlock.inline(accumulator, List.of(collect, resultValue),
225                                                     (_, _) -> {
226                                                     });
227                                             terminalBlock.op(_continue());
228                                         });
229                                     });
230                             block.op(sourceLoop);
231                             block.op(_return(collect));
232                         });
233                     });
234         }
235 
236     }
237 }