Implement selector translation

Author: Killian Perlin

lkql_jit/options/src/main/java/com/adacore/lkql_jit/options/Refactorings/LKQLToLkt.java

@@ -11,6 +11,9 @@
 
 public class LKQLToLkt implements Refactoring {
 
+    /** Pointer to last-entered selector during rewriting. */
+    private Liblkqllang.SelectorDecl currentSelector = Liblkqllang.SelectorDecl.NONE;
+
     @Override
     public void applyRefactor(Refactoring.State state) {
         final var root = state.unit.getRoot();
@@ -30,11 +33,15 @@ private String refactorNode(Liblkqllang.LkqlNode node) {
             case Liblkqllang.ParameterDecl paramDecl -> refactorParamDecl(paramDecl);
             case Liblkqllang.Match match -> refactorMatch(match);
             case Liblkqllang.MatchArm arm -> refactorArm(arm, arm.fPattern(), arm.fExpr());
+            case Liblkqllang.SelectorArm arm -> refactorArm(arm, arm.fPattern(), arm.fExpr());
             case Liblkqllang.SelectorDecl selectorDecl -> refactorSelectorDecl(selectorDecl);
+            case Liblkqllang.RecExpr recExpr -> refactorRecExpr(recExpr);
             case Liblkqllang.BlockBodyExpr bbe -> "var _ = " + refactorGeneric(bbe);
+            case Liblkqllang.UnitLiteral _ -> "Unit()";
             case Liblkqllang.Expr expr when (
                 expr.parent() instanceof Liblkqllang.TopLevelList
             ) -> "val _ = " + refactorGeneric(expr);
+            case Liblkqllang.UniversalPattern _ -> "_";
             default -> refactorGeneric(node);
         };
     }
@@ -185,4 +192,86 @@ private String refactorArm(
             textRange(expr.tokenEnd().next(), arm.tokenEnd())
         );
     }
+
+    /*
+     *
+     * <annotations> selector <name> <docstring> <arms>
+     *
+     * <docstring>\n
+     * <annotations> fun <name> (this : Any) : Any = match this { <arms> }
+     *
+     */
+    private String refactorSelectorDecl(Liblkqllang.SelectorDecl selectorDecl) {
+        // save selector state on entering this function
+        var previousSelector = currentSelector;
+        // set new state for nested refactors
+        currentSelector = selectorDecl;
+        var s = "";
+
+        // pull docstring before declaration
+        if (!selectorDecl.fDocNode().isNone()) s = refactorNode(selectorDecl.fDocNode()) + "\n";
+
+        if (!selectorDecl.fAnnotation().isNone()) s +=
+            refactorNode(selectorDecl.fAnnotation()) +
+            selectorDecl.fAnnotation().tokenEnd().next().getText();
+
+        final var whitespace = textRange(
+            (selectorDecl.fDocNode().isNone()
+                    ? selectorDecl.fName().tokenEnd().next()
+                    : selectorDecl.fDocNode().tokenEnd().next()),
+            selectorDecl.fArms().tokenStart().previous()
+        );
+
+        s +=
+            "fun " +
+            refactorNode(selectorDecl.fName()) +
+            "(this : Any) : Any = match this {" +
+            whitespace +
+            refactorNode(selectorDecl.fArms()) +
+            "\n}\n";
+
+        // restore previous state on exiting this function
+        currentSelector = previousSelector;
+        return s;
+    }
+
+    /*
+     *
+     * 1) Expansion of implicit argument
+     *
+     * rec(<expr>) --> rec(<expr>, <expr>)
+     *
+     * 2) Case disjonction
+     *
+     * rec( <left>,  <right>) --> <right>          ::  <selector>(<left>)
+     * rec(*<left>,  <right>) --> <right>          ::  <left>.iterator.flatMap(<selector>)
+     * rec( <left>, *<right>) --> <right>.iterator ::: <selector>(<left>)
+     * rec(*<left>, *<right>) --> <right>.iterator ::: <left>.iterator.flatMap(<selector>)
+     *
+     */
+    private String refactorRecExpr(Liblkqllang.RecExpr recExpr) {
+        final var hasRight = !recExpr.fResultExpr().isNone();
+
+        final var unpackLeft = recExpr.fRecurseUnpack().pAsBool();
+        final var unpackRight = hasRight ? recExpr.fResultUnpack().pAsBool() : unpackLeft;
+
+        final var left = recExpr.fRecurseExpr();
+        final var right = hasRight ? recExpr.fResultExpr() : left;
+
+        var s = unpackRight ? refactorNode(right) + ".iterator :::" : refactorNode(right) + " ::";
+
+        // try to preserve spacing after "," (any newline for example)
+        if (hasRight && left.tokenEnd().next().getText().equals(",")) {
+            for (var tok = left.tokenEnd().next().next(); tok.isTrivia(); tok = tok.next()) s +=
+                tok.getText();
+        } else {
+            s += " ";
+        }
+
+        s += unpackLeft
+            ? refactorNode(left) + ".iterator.flatMap(" + currentSelector.fName().getText() + ")"
+            : currentSelector.fName().getText() + "(" + refactorNode(left) + ")";
+
+        return s;
+    }
 }

testsuite/drivers/refactor_driver.py

@@ -27,3 +27,6 @@ def run(self) -> None:
                 "input.lkql",
             ]
         )
+
+        if self.test_env["refactoring"] == "TO_LKQL_V2":
+            self.check_run(["lkt_parse", "-s", "-f", "test.out"])

testsuite/tests/refactor/selectors/input.lkql

@@ -0,0 +1,22 @@
+selector parents
+    | AdaNode => rec(this.parent, this)
+    | * => ()
+
+selector super_types
+    | BaseTypeDecl => rec(*this.p_base_types())
+    | *            => ()
+
+selector children
+    | AdaNode => rec(*this.children ,this)
+    | * => ()
+
+selector foo
+    | Foo => rec(this, *[s, o, m, e, t, h, i, n, g])
+
+selector infinite_sequence
+    |" Infinite sequence generator
+    | nb => rec(
+        nb + 1, # Recurse with value nb + 1
+        nb # Add nb to the result list
+    )
+

testsuite/tests/refactor/selectors/test.out

@@ -0,0 +1,32 @@
+# lkql version: 2
+
+fun parents(this : Any) : Any = match this {
+    case AdaNode => this :: parents(this.parent)
+    case _ => Unit()
+}
+
+
+fun super_types(this : Any) : Any = match this {
+    case BaseTypeDecl => this.p_base_types().iterator ::: this.p_base_types().iterator.flatMap(super_types)
+    case _            => Unit()
+}
+
+
+fun children(this : Any) : Any = match this {
+    case AdaNode => this :: this.children.iterator.flatMap(children)
+    case _ => Unit()
+}
+
+
+fun foo(this : Any) : Any = match this {
+    case Foo => [s, o, m, e, t, h, i, n, g].iterator ::: foo(this)
+}
+
+
+|" Infinite sequence generator
+fun infinite_sequence(this : Any) : Any = match this {
+    case nb => nb :: # Recurse with value nb + 1
+        infinite_sequence(nb + 1)
+}
+
+

testsuite/tests/refactor/selectors/test.yaml

@@ -0,0 +1,2 @@
+driver: refactor
+refactoring: TO_LKQL_V2