RD-10562 Implement OSR for iterators

snapi-client/src/test/scala/raw/compiler/rql2/tests/builtin/collection/CollectionRangeTest.scala

@@ -21,6 +21,10 @@ trait CollectionRangeTest extends CompilerTestContext {
   test("""Long.Range(0, 10)""")(_ should evaluateTo("[0L,1L,2L,3L,4L,5L,6L,7L,8L,9L]"))
   test("""Long.Range(0, 1)""")(_ should evaluateTo("[0L]"))
 
+  test("""Collection.Filter(Long.Range(0, 1000000), x -> x == 999999)""")(
+    _ should evaluateTo("Collection.Build(999999L)")
+  )
+
   // end = start
   test("""Long.Range(0, 0)""")(_ should evaluateTo("[]"))
   test("""Long.Range(12, 12)""")(_ should evaluateTo("[]"))

snapi-truffle/src/main/java/module-info.java

@@ -303,8 +303,6 @@
 
   exports raw.runtime.truffle;
   exports raw.runtime.truffle.boundary;
-  exports raw.runtime.truffle.runtime.aggregation;
-  exports raw.runtime.truffle.runtime.aggregation.aggregator;
   exports raw.runtime.truffle.runtime.record;
   exports raw.runtime.truffle.runtime.operators;
   exports raw.runtime.truffle.runtime.function;
@@ -330,6 +328,7 @@
       .abstract_generator
       .compute_next
       .operations;
+  exports raw.runtime.truffle.runtime.generator.collection.abstract_generator.compute_next.osr;
   exports raw.runtime.truffle.runtime.generator.collection.off_heap_generator.record_shaper;
   exports raw.runtime.truffle.runtime.generator.collection.off_heap_generator.input_buffer;
   exports raw.runtime.truffle.runtime.generator.collection.off_heap_generator.off_heap;
@@ -340,6 +339,7 @@
   exports raw.runtime.truffle.runtime.iterable.operations;
   exports raw.runtime.truffle.runtime.iterable.list;
   exports raw.runtime.truffle.runtime.iterable.sources;
+  exports raw.runtime.truffle.runtime.iterable.osr;
   exports raw.runtime.truffle.runtime.kryo;
   exports raw.runtime.truffle.utils;
   exports raw.runtime.truffle.ast;
@@ -348,6 +348,7 @@
   exports raw.runtime.truffle.ast.io.jdbc;
   exports raw.runtime.truffle.ast.io.json.reader;
   exports raw.runtime.truffle.ast.io.json.reader.parser;
+  exports raw.runtime.truffle.ast.io.json.reader.parser.osr;
   exports raw.runtime.truffle.ast.io.json.writer;
   exports raw.runtime.truffle.ast.io.json.writer.internal;
   exports raw.runtime.truffle.ast.io.csv.reader;
@@ -358,7 +359,9 @@
   exports raw.runtime.truffle.ast.local;
   exports raw.runtime.truffle.ast.expressions.unary;
   exports raw.runtime.truffle.ast.expressions.iterable.collection;
+  exports raw.runtime.truffle.ast.expressions.iterable.collection.osr;
   exports raw.runtime.truffle.ast.expressions.iterable.list;
+  exports raw.runtime.truffle.ast.expressions.iterable.list.osr;
   exports raw.runtime.truffle.ast.expressions.record;
   exports raw.runtime.truffle.ast.expressions.option;
   exports raw.runtime.truffle.ast.expressions.function;
@@ -387,6 +390,7 @@
   exports raw.runtime.truffle.ast.expressions.builtin.string_package;
   exports raw.runtime.truffle.ast.expressions.builtin.location_package;
   exports raw.runtime.truffle.ast.expressions.builtin.binary_package;
+  exports raw.runtime.truffle.ast.expressions.aggregation;
   exports raw.runtime.truffle.ast.controlflow;
   exports raw.runtime.truffle.runtime.exceptions.validation;
   exports raw.compiler.snapi.truffle.compiler;

snapi-truffle/src/main/java/raw/compiler/snapi/truffle/builtin/collection_extension/TruffleCountCollectionEntry.java

@@ -19,12 +19,13 @@
 import raw.compiler.snapi.truffle.TruffleEntryExtension;
 import raw.runtime.truffle.ExpressionNode;
 import raw.runtime.truffle.RawLanguage;
-import raw.runtime.truffle.ast.expressions.iterable.collection.CollectionCountNodeGen;
+import raw.runtime.truffle.ast.expressions.aggregation.AggregateSingleNode;
+import raw.runtime.truffle.ast.expressions.aggregation.Aggregations;
 
 public class TruffleCountCollectionEntry extends CountCollectionEntry
     implements TruffleEntryExtension {
   @Override
   public ExpressionNode toTruffle(Type type, List<TruffleArg> args, RawLanguage rawLanguage) {
-    return CollectionCountNodeGen.create(args.get(0).exprNode());
+    return new AggregateSingleNode(args.get(0).exprNode(), Aggregations.COUNT);
   }
 }

snapi-truffle/src/main/java/raw/compiler/snapi/truffle/builtin/collection_extension/TruffleExistsCollectionEntry.java

@@ -15,7 +15,6 @@
 import java.util.List;
 import raw.compiler.base.source.Type;
 import raw.compiler.rql2.builtin.ExistsCollectionEntry;
-import raw.compiler.rql2.source.FunType;
 import raw.compiler.snapi.truffle.TruffleArg;
 import raw.compiler.snapi.truffle.TruffleEntryExtension;
 import raw.runtime.truffle.ExpressionNode;
@@ -26,8 +25,6 @@ public class TruffleExistsCollectionEntry extends ExistsCollectionEntry
     implements TruffleEntryExtension {
   @Override
   public ExpressionNode toTruffle(Type type, List<TruffleArg> args, RawLanguage rawLanguage) {
-    FunType funType = (FunType) args.get(1).type();
-
     return CollectionExistsNodeGen.create(args.get(0).exprNode(), args.get(1).exprNode());
   }
 }

snapi-truffle/src/main/java/raw/compiler/snapi/truffle/builtin/collection_extension/TruffleLastCollectionEntry.java

@@ -19,12 +19,13 @@
 import raw.compiler.snapi.truffle.TruffleEntryExtension;
 import raw.runtime.truffle.ExpressionNode;
 import raw.runtime.truffle.RawLanguage;
-import raw.runtime.truffle.ast.expressions.iterable.collection.CollectionLastNodeGen;
+import raw.runtime.truffle.ast.expressions.aggregation.AggregateSingleNode;
+import raw.runtime.truffle.ast.expressions.aggregation.Aggregations;
 
 public class TruffleLastCollectionEntry extends LastCollectionEntry
     implements TruffleEntryExtension {
   @Override
   public ExpressionNode toTruffle(Type type, List<TruffleArg> args, RawLanguage rawLanguage) {
-    return CollectionLastNodeGen.create(args.get(0).exprNode());
+    return new AggregateSingleNode(args.get(0).exprNode(), Aggregations.LAST);
   }
 }

snapi-truffle/src/main/java/raw/compiler/snapi/truffle/builtin/collection_extension/TruffleMaxCollectionEntry.java

@@ -19,11 +19,12 @@
 import raw.compiler.snapi.truffle.TruffleEntryExtension;
 import raw.runtime.truffle.ExpressionNode;
 import raw.runtime.truffle.RawLanguage;
-import raw.runtime.truffle.ast.expressions.iterable.collection.CollectionMaxNodeGen;
+import raw.runtime.truffle.ast.expressions.aggregation.AggregateSingleNode;
+import raw.runtime.truffle.ast.expressions.aggregation.Aggregations;
 
 public class TruffleMaxCollectionEntry extends MaxCollectionEntry implements TruffleEntryExtension {
   @Override
   public ExpressionNode toTruffle(Type type, List<TruffleArg> args, RawLanguage rawLanguage) {
-    return CollectionMaxNodeGen.create(args.get(0).exprNode());
+    return new AggregateSingleNode(args.get(0).exprNode(), Aggregations.MAX);
   }
 }

snapi-truffle/src/main/java/raw/compiler/snapi/truffle/builtin/collection_extension/TruffleMinCollectionEntry.java

@@ -19,11 +19,12 @@
 import raw.compiler.snapi.truffle.TruffleEntryExtension;
 import raw.runtime.truffle.ExpressionNode;
 import raw.runtime.truffle.RawLanguage;
-import raw.runtime.truffle.ast.expressions.iterable.collection.CollectionMinNodeGen;
+import raw.runtime.truffle.ast.expressions.aggregation.AggregateSingleNode;
+import raw.runtime.truffle.ast.expressions.aggregation.Aggregations;
 
 public class TruffleMinCollectionEntry extends MinCollectionEntry implements TruffleEntryExtension {
   @Override
   public ExpressionNode toTruffle(Type type, List<TruffleArg> args, RawLanguage rawLanguage) {
-    return CollectionMinNodeGen.create(args.get(0).exprNode());
+    return new AggregateSingleNode(args.get(0).exprNode(), Aggregations.MIN);
   }
 }

snapi-truffle/src/main/java/raw/compiler/snapi/truffle/builtin/collection_extension/TruffleSumCollectionEntry.java

@@ -19,11 +19,12 @@
 import raw.compiler.snapi.truffle.TruffleEntryExtension;
 import raw.runtime.truffle.ExpressionNode;
 import raw.runtime.truffle.RawLanguage;
-import raw.runtime.truffle.ast.expressions.iterable.collection.CollectionSumNodeGen;
+import raw.runtime.truffle.ast.expressions.aggregation.AggregateSingleNode;
+import raw.runtime.truffle.ast.expressions.aggregation.Aggregations;
 
 public class TruffleSumCollectionEntry extends SumCollectionEntry implements TruffleEntryExtension {
   @Override
   public ExpressionNode toTruffle(Type type, List<TruffleArg> args, RawLanguage rawLanguage) {
-    return CollectionSumNodeGen.create(args.get(0).exprNode());
+    return new AggregateSingleNode(args.get(0).exprNode(), Aggregations.SUM);
   }
 }

snapi-truffle/src/main/java/raw/compiler/snapi/truffle/builtin/collection_extension/TruffleTupleAvgCollectionEntry.java

@@ -19,12 +19,12 @@
 import raw.compiler.snapi.truffle.TruffleEntryExtension;
 import raw.runtime.truffle.ExpressionNode;
 import raw.runtime.truffle.RawLanguage;
-import raw.runtime.truffle.ast.expressions.iterable.collection.CollectionTupleAvgNodeGen;
+import raw.runtime.truffle.ast.expressions.iterable.collection.CollectionTupleAvgNode;
 
 public class TruffleTupleAvgCollectionEntry extends TupleAvgCollectionEntry
     implements TruffleEntryExtension {
   @Override
   public ExpressionNode toTruffle(Type type, List<TruffleArg> args, RawLanguage rawLanguage) {
-    return CollectionTupleAvgNodeGen.create(args.get(0).exprNode());
+    return new CollectionTupleAvgNode(args.get(0).exprNode());
   }
 }

snapi-truffle/src/main/java/raw/compiler/snapi/truffle/builtin/list_extension/TruffleFilterListEntry.java

@@ -15,6 +15,8 @@
 import java.util.List;
 import raw.compiler.base.source.Type;
 import raw.compiler.rql2.builtin.FilterListEntry;
+import raw.compiler.rql2.source.Rql2ListType;
+import raw.compiler.rql2.source.Rql2Type;
 import raw.compiler.snapi.truffle.TruffleArg;
 import raw.compiler.snapi.truffle.TruffleEntryExtension;
 import raw.runtime.truffle.ExpressionNode;
@@ -24,6 +26,8 @@
 public class TruffleFilterListEntry extends FilterListEntry implements TruffleEntryExtension {
   @Override
   public ExpressionNode toTruffle(Type type, List<TruffleArg> args, RawLanguage rawLanguage) {
-    return ListFilterNodeGen.create(args.get(0).exprNode(), args.get(1).exprNode());
+    Rql2ListType listType = (Rql2ListType) type;
+    return ListFilterNodeGen.create(
+        args.get(0).exprNode(), args.get(1).exprNode(), (Rql2Type) listType.innerType());
   }
 }

snapi-truffle/src/main/java/raw/compiler/snapi/truffle/builtin/list_extension/TruffleMaxListEntry.java

@@ -19,11 +19,14 @@
 import raw.compiler.snapi.truffle.TruffleEntryExtension;
 import raw.runtime.truffle.ExpressionNode;
 import raw.runtime.truffle.RawLanguage;
-import raw.runtime.truffle.ast.expressions.iterable.list.ListMaxNodeGen;
+import raw.runtime.truffle.ast.expressions.aggregation.AggregateSingleNode;
+import raw.runtime.truffle.ast.expressions.aggregation.Aggregations;
+import raw.runtime.truffle.ast.expressions.iterable.collection.CollectionFromNodeGen;
 
 public class TruffleMaxListEntry extends MaxListEntry implements TruffleEntryExtension {
   @Override
   public ExpressionNode toTruffle(Type type, List<TruffleArg> args, RawLanguage rawLanguage) {
-    return ListMaxNodeGen.create(args.get(0).exprNode());
+    return new AggregateSingleNode(
+        CollectionFromNodeGen.create(args.get(0).exprNode()), Aggregations.MAX);
   }
 }

snapi-truffle/src/main/java/raw/compiler/snapi/truffle/builtin/list_extension/TruffleMinListEntry.java

@@ -19,11 +19,14 @@
 import raw.compiler.snapi.truffle.TruffleEntryExtension;
 import raw.runtime.truffle.ExpressionNode;
 import raw.runtime.truffle.RawLanguage;
-import raw.runtime.truffle.ast.expressions.iterable.list.ListMinNodeGen;
+import raw.runtime.truffle.ast.expressions.aggregation.AggregateSingleNode;
+import raw.runtime.truffle.ast.expressions.aggregation.Aggregations;
+import raw.runtime.truffle.ast.expressions.iterable.collection.CollectionFromNodeGen;
 
 public class TruffleMinListEntry extends MinListEntry implements TruffleEntryExtension {
   @Override
   public ExpressionNode toTruffle(Type type, List<TruffleArg> args, RawLanguage rawLanguage) {
-    return ListMinNodeGen.create(args.get(0).exprNode());
+    return new AggregateSingleNode(
+        CollectionFromNodeGen.create(args.get(0).exprNode()), Aggregations.MIN);
   }
 }

snapi-truffle/src/main/java/raw/compiler/snapi/truffle/builtin/list_extension/TruffleSumListEntry.java

@@ -19,11 +19,14 @@
 import raw.compiler.snapi.truffle.TruffleEntryExtension;
 import raw.runtime.truffle.ExpressionNode;
 import raw.runtime.truffle.RawLanguage;
-import raw.runtime.truffle.ast.expressions.iterable.list.ListSumNodeGen;
+import raw.runtime.truffle.ast.expressions.aggregation.AggregateSingleNode;
+import raw.runtime.truffle.ast.expressions.aggregation.Aggregations;
+import raw.runtime.truffle.ast.expressions.iterable.collection.CollectionFromNodeGen;
 
 public class TruffleSumListEntry extends SumListEntry implements TruffleEntryExtension {
   @Override
   public ExpressionNode toTruffle(Type type, List<TruffleArg> args, RawLanguage rawLanguage) {
-    return ListSumNodeGen.create(args.get(0).exprNode());
+    return new AggregateSingleNode(
+        CollectionFromNodeGen.create(args.get(0).exprNode()), Aggregations.SUM);
   }
 }

snapi-truffle/src/main/java/raw/runtime/truffle/ast/expressions/aggregation/AggregateMultipleNode.java

@@ -0,0 +1,79 @@
+/*
+ * Copyright 2023 RAW Labs S.A.
+ *
+ * Use of this software is governed by the Business Source License
+ * included in the file licenses/BSL.txt.
+ *
+ * As of the Change Date specified in that file, in accordance with
+ * the Business Source License, use of this software will be governed
+ * by the Apache License, Version 2.0, included in the file
+ * licenses/APL.txt.
+ */
+
+package raw.runtime.truffle.ast.expressions.aggregation;
+
+import com.oracle.truffle.api.Truffle;
+import com.oracle.truffle.api.frame.VirtualFrame;
+import com.oracle.truffle.api.nodes.LoopNode;
+import raw.runtime.truffle.ExpressionNode;
+import raw.runtime.truffle.runtime.exceptions.RawTruffleRuntimeException;
+import raw.runtime.truffle.runtime.generator.collection.GeneratorNodes;
+import raw.runtime.truffle.runtime.generator.collection.GeneratorNodesFactory;
+import raw.runtime.truffle.runtime.iterable.IterableNodes;
+import raw.runtime.truffle.runtime.iterable.IterableNodesFactory;
+import raw.runtime.truffle.runtime.primitives.ErrorObject;
+
+public class AggregateMultipleNode extends ExpressionNode {
+
+  @Child ExpressionNode iterableNode;
+
+  @Child private LoopNode loop;
+
+  @Child
+  GeneratorNodes.GeneratorInitNode initNode = GeneratorNodesFactory.GeneratorInitNodeGen.create();
+
+  @Child
+  GeneratorNodes.GeneratorCloseNode closeNode =
+      GeneratorNodesFactory.GeneratorCloseNodeGen.create();
+
+  @Child
+  private GeneratorNodes.GeneratorHasNextNode hasNextNode =
+      GeneratorNodesFactory.GeneratorHasNextNodeGen.create();
+
+  @Child
+  IterableNodes.GetGeneratorNode getGeneratorNode =
+      IterableNodesFactory.GetGeneratorNodeGen.create();
+
+  @Child AggregatorNodes.Zero zeroNode = AggregatorNodesFactory.ZeroNodeGen.create();
+
+  private final byte[] aggregationTypes;
+
+  public AggregateMultipleNode(ExpressionNode iterableNode, byte[] aggregationTypes) {
+    this.iterableNode = iterableNode;
+    loop = Truffle.getRuntime().createLoopNode(new OSRMultiAggregationNode(aggregationTypes));
+    this.aggregationTypes = aggregationTypes;
+  }
+
+  @Override
+  public Object executeGeneric(VirtualFrame virtualFrame) {
+    Object generator = getGeneratorNode.execute(this, iterableNode.executeGeneric(virtualFrame));
+    try {
+      initNode.execute(this, generator);
+      Object[] results = new Object[aggregationTypes.length];
+      for (int i = 0; i < aggregationTypes.length; i++) {
+        results[i] = zeroNode.execute(this, aggregationTypes[i]);
+      }
+      if (!hasNextNode.execute(this, generator)) {
+        return results;
+      }
+      OSRMultiAggregationNode OSRNode = (OSRMultiAggregationNode) loop.getRepeatingNode();
+      OSRNode.init(generator, results);
+      loop.execute(virtualFrame);
+      return OSRNode.getResults();
+    } catch (RawTruffleRuntimeException e) {
+      return new ErrorObject(e.getMessage());
+    } finally {
+      closeNode.execute(this, generator);
+    }
+  }
+}

snapi-truffle/src/main/java/raw/runtime/truffle/ast/expressions/aggregation/AggregateSingleNode.java

@@ -0,0 +1,76 @@
+/*
+ * Copyright 2023 RAW Labs S.A.
+ *
+ * Use of this software is governed by the Business Source License
+ * included in the file licenses/BSL.txt.
+ *
+ * As of the Change Date specified in that file, in accordance with
+ * the Business Source License, use of this software will be governed
+ * by the Apache License, Version 2.0, included in the file
+ * licenses/APL.txt.
+ */
+
+package raw.runtime.truffle.ast.expressions.aggregation;
+
+import com.oracle.truffle.api.Truffle;
+import com.oracle.truffle.api.frame.VirtualFrame;
+import com.oracle.truffle.api.nodes.LoopNode;
+import raw.runtime.truffle.ExpressionNode;
+import raw.runtime.truffle.runtime.exceptions.RawTruffleRuntimeException;
+import raw.runtime.truffle.runtime.generator.collection.GeneratorNodes;
+import raw.runtime.truffle.runtime.generator.collection.GeneratorNodesFactory;
+import raw.runtime.truffle.runtime.iterable.IterableNodes;
+import raw.runtime.truffle.runtime.iterable.IterableNodesFactory;
+import raw.runtime.truffle.runtime.primitives.ErrorObject;
+
+public class AggregateSingleNode extends ExpressionNode {
+
+  @Child ExpressionNode iterableNode;
+
+  @Child private LoopNode loop;
+
+  @Child
+  GeneratorNodes.GeneratorInitNode initNode = GeneratorNodesFactory.GeneratorInitNodeGen.create();
+
+  @Child
+  GeneratorNodes.GeneratorCloseNode closeNode =
+      GeneratorNodesFactory.GeneratorCloseNodeGen.create();
+
+  @Child
+  private GeneratorNodes.GeneratorHasNextNode hasNextNode =
+      GeneratorNodesFactory.GeneratorHasNextNodeGen.create();
+
+  @Child
+  IterableNodes.GetGeneratorNode getGeneratorNode =
+      IterableNodesFactory.GetGeneratorNodeGen.create();
+
+  @Child AggregatorNodes.Zero zeroNode = AggregatorNodesFactory.ZeroNodeGen.create();
+
+  private final byte aggregationType;
+
+  public AggregateSingleNode(ExpressionNode iterableNode, byte aggregationType) {
+    this.iterableNode = iterableNode;
+    loop = Truffle.getRuntime().createLoopNode(new OSRSingleAggregationNode(aggregationType));
+    this.aggregationType = aggregationType;
+  }
+
+  @Override
+  public Object executeGeneric(VirtualFrame virtualFrame) {
+    Object generator = getGeneratorNode.execute(this, iterableNode.executeGeneric(virtualFrame));
+    try {
+      initNode.execute(this, generator);
+      if (!hasNextNode.execute(this, generator)) {
+        return zeroNode.execute(this, aggregationType);
+      }
+      Object result = zeroNode.execute(this, aggregationType);
+      OSRSingleAggregationNode OSRNode = (OSRSingleAggregationNode) loop.getRepeatingNode();
+      OSRNode.init(generator, result);
+      loop.execute(virtualFrame);
+      return OSRNode.getResult();
+    } catch (RawTruffleRuntimeException e) {
+      return new ErrorObject(e.getMessage());
+    } finally {
+      closeNode.execute(this, generator);
+    }
+  }
+}