RD-10562 Implement OSR for iterators

snapi-truffle/src/main/java/raw/runtime/truffle/ast/expressions/iterable/list/ListFilterNode.java

@@ -18,17 +18,18 @@
 import com.oracle.truffle.api.nodes.LoopNode;
 import com.oracle.truffle.api.nodes.Node;
 import com.oracle.truffle.api.nodes.NodeInfo;
+import java.util.ArrayList;
 import raw.compiler.rql2.source.Rql2Type;
 import raw.runtime.truffle.ExpressionNode;
 import raw.runtime.truffle.ast.TypeGuards;
-import raw.runtime.truffle.ast.expressions.iterable.list.osr.OSRListFilterNode;
 import raw.runtime.truffle.ast.expressions.iterable.list.osr.OSRToArrayNode;
+import raw.runtime.truffle.ast.osr.OSRGeneratorNode;
+import raw.runtime.truffle.ast.osr.filter.OSRListFilterBodyNode;
+import raw.runtime.truffle.ast.osr.filter.OSRListFilterConditionNode;
 import raw.runtime.truffle.runtime.generator.collection.GeneratorNodes;
 import raw.runtime.truffle.runtime.iterable.IterableNodes;
 import raw.runtime.truffle.runtime.list.*;
 
-import java.util.ArrayList;
-
 @ImportStatic(value = TypeGuards.class)
 @NodeInfo(shortName = "List.Filter")
 @NodeChild("list")
@@ -52,7 +53,11 @@ public static int[] getSlots(VirtualFrame frame) {
   }
 
   public static LoopNode getFilterLoopNode(int[] slots) {
-    return Truffle.getRuntime().createLoopNode(new OSRListFilterNode(slots[0], slots[1], slots[2]));
+    return Truffle.getRuntime()
+        .createLoopNode(
+            new OSRGeneratorNode(
+                new OSRListFilterConditionNode(slots[0]),
+                new OSRListFilterBodyNode(slots[0], slots[1], slots[2])));
   }
 
   public static LoopNode getToArrayLoopNode(Rql2Type resultType) {

snapi-truffle/src/main/java/raw/runtime/truffle/ast/expressions/iterable/list/osr/OSRListFilterNode.java

@@ -12,13 +12,10 @@
 
 package raw.runtime.truffle.ast.expressions.iterable.list.osr;
 
-import com.oracle.truffle.api.CompilerDirectives;
 import com.oracle.truffle.api.frame.VirtualFrame;
 import com.oracle.truffle.api.nodes.Node;
 import com.oracle.truffle.api.nodes.RepeatingNode;
 import java.util.ArrayList;
-
-import raw.runtime.truffle.runtime.function.Function;
 import raw.runtime.truffle.runtime.function.FunctionExecuteNodes;
 import raw.runtime.truffle.runtime.function.FunctionExecuteNodesFactory;
 import raw.runtime.truffle.runtime.generator.collection.GeneratorNodes;

snapi-truffle/src/main/java/raw/runtime/truffle/ast/osr/OSRGeneratorNode.java

@@ -0,0 +1,44 @@
+/*
+ * 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.osr;
+
+import com.oracle.truffle.api.frame.VirtualFrame;
+import com.oracle.truffle.api.nodes.Node;
+import com.oracle.truffle.api.nodes.RepeatingNode;
+import com.oracle.truffle.api.nodes.UnexpectedResultException;
+import raw.runtime.truffle.ExpressionNode;
+import raw.runtime.truffle.runtime.exceptions.RawTruffleInternalErrorException;
+
+public class OSRGeneratorNode extends Node implements RepeatingNode {
+
+  @Child private ExpressionNode conditionNode;
+
+  @Child private ExpressionNode bodyNode;
+
+  public OSRGeneratorNode(ExpressionNode conditionNode, ExpressionNode bodyNode) {
+    this.conditionNode = conditionNode;
+    this.bodyNode = bodyNode;
+  }
+
+  public boolean executeRepeating(VirtualFrame frame) {
+    try {
+      if (conditionNode.executeBoolean(frame)) {
+        bodyNode.executeVoid(frame);
+        return true;
+      }
+      return false;
+    } catch (UnexpectedResultException e) {
+      throw new RawTruffleInternalErrorException(e);
+    }
+  }
+}

snapi-truffle/src/main/java/raw/runtime/truffle/ast/osr/filter/OSRListFilterBodyNode.java

@@ -0,0 +1,64 @@
+/*
+ * 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.osr.filter;
+
+import com.oracle.truffle.api.frame.VirtualFrame;
+import java.util.ArrayList;
+import raw.runtime.truffle.ExpressionNode;
+import raw.runtime.truffle.runtime.function.FunctionExecuteNodes;
+import raw.runtime.truffle.runtime.function.FunctionExecuteNodesFactory;
+import raw.runtime.truffle.runtime.generator.collection.GeneratorNodes;
+import raw.runtime.truffle.runtime.generator.collection.GeneratorNodesFactory;
+import raw.runtime.truffle.tryable_nullable.TryableNullable;
+
+public class OSRListFilterBodyNode extends ExpressionNode {
+
+  @Child
+  private GeneratorNodes.GeneratorNextNode nextNode =
+      GeneratorNodesFactory.GeneratorNextNodeGen.create();
+
+  @Child
+  FunctionExecuteNodes.FunctionExecuteOne functionExecuteOneNode =
+      FunctionExecuteNodesFactory.FunctionExecuteOneNodeGen.create();
+
+  private final int generatorSlot;
+  private final int functionSlot;
+  private final int llistSlot;
+
+  public OSRListFilterBodyNode(int generatorSlot, int functionSlot, int llistSlot) {
+    this.generatorSlot = generatorSlot;
+    this.functionSlot = functionSlot;
+    this.llistSlot = llistSlot;
+  }
+
+  @Override
+  public Object executeGeneric(VirtualFrame frame) {
+    Object generator = frame.getAuxiliarySlot(generatorSlot);
+    Object v = nextNode.execute(this, generator);
+    Boolean predicate = null;
+    Object function = frame.getAuxiliarySlot(functionSlot);
+    predicate =
+        TryableNullable.handlePredicate(functionExecuteOneNode.execute(this, function, v), false);
+    if (predicate) {
+      @SuppressWarnings("unchecked")
+      ArrayList<Object> llist = (ArrayList<Object>) frame.getAuxiliarySlot(llistSlot);
+      llist.add(v);
+    }
+    return null;
+  }
+
+  @Override
+  public void executeVoid(VirtualFrame virtualFrame) {
+    executeGeneric(virtualFrame);
+  }
+}

snapi-truffle/src/main/java/raw/runtime/truffle/ast/osr/filter/OSRListFilterConditionNode.java

@@ -0,0 +1,46 @@
+/*
+ * 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.osr.filter;
+
+import com.oracle.truffle.api.frame.VirtualFrame;
+import com.oracle.truffle.api.nodes.UnexpectedResultException;
+import raw.runtime.truffle.ExpressionNode;
+import raw.runtime.truffle.runtime.generator.collection.GeneratorNodes;
+import raw.runtime.truffle.runtime.generator.collection.GeneratorNodesFactory;
+
+public class OSRListFilterConditionNode extends ExpressionNode {
+
+  @Child
+  private GeneratorNodes.GeneratorHasNextNode hasNextNode =
+      GeneratorNodesFactory.GeneratorHasNextNodeGen.create();
+
+  private final int generatorSlot;
+
+  public OSRListFilterConditionNode(int generatorSlot) {
+    this.generatorSlot = generatorSlot;
+  }
+
+  @Override
+  public Object executeGeneric(VirtualFrame frame) {
+    Object generator = frame.getAuxiliarySlot(generatorSlot);
+    if (!hasNextNode.execute(this, generator)) {
+      return false;
+    }
+    return true;
+  }
+
+  @Override
+  public boolean executeBoolean(VirtualFrame virtualFrame) throws UnexpectedResultException {
+    return (boolean) executeGeneric(virtualFrame);
+  }
+}