add ExprQuery::captures_from_ast

Author: myfreess

query/moon.pkg

@@ -9,3 +9,9 @@ import {
   "moonbitlang/parser",
   "moonbitlang/parser/syntax",
 }
+
+import {
+  "moonbitlang/core/list",
+  "moonbitlang/lexer",
+  "moonbitlang/parser/handrolled_parser",
+} for "test"

query/pkg.generated.mbti

@@ -13,6 +13,7 @@ import {
 type ExprQuery
 pub fn ExprQuery::ExprQuery(String) -> Self raise
 pub fn ExprQuery::captures(Self, source_name~ : String, String) -> ArrayView[Map[String, @untyped_ast.Node]] raise
+pub fn ExprQuery::captures_from_ast(Self, ast~ : @untyped_ast.Node) -> ArrayView[Map[String, @untyped_ast.Node]]
 
 // Type aliases
 

query/query.mbt

@@ -54,12 +54,22 @@ pub fn ExprQuery::captures(
   for node in impls {
     let impl_source = source_text_for_impl(source, node)
     if @rule_prefilter.rule_is_relevant_to_source(self.rule, impl_source) {
-      collect_impl_matches(self.pattern, node, matches)
+      collect_node_matches(self.pattern, @untyped_ast.from_impl(node), matches)
     }
   }
   matches
 }
 
+///|
+pub fn ExprQuery::captures_from_ast(
+  self : ExprQuery,
+  ast~ : @untyped_ast.Node,
+) -> ArrayView[Map[String, @untyped_ast.Node]] {
+  let matches : Array[Map[String, @untyped_ast.Node]] = []
+  collect_node_matches(self.pattern, ast, matches)
+  matches
+}
+
 ///|
 fn anonymous_pattern_rule(pattern : String) -> RawRuleSpec {
   {
@@ -76,14 +86,12 @@ fn anonymous_pattern_rule(pattern : String) -> RawRuleSpec {
 }
 
 ///|
-fn collect_impl_matches(
+fn collect_node_matches(
   pattern : CompiledRulePattern,
-  node : @syntax.Impl,
+  ast : @untyped_ast.Node,
   matches : Array[Map[String, @untyped_ast.Node]],
 ) -> Unit {
-  @untyped_ast.visit_node_scoped_expr_roots(@untyped_ast.from_impl(node), fn(
-    scoped_expr,
-  ) {
+  @untyped_ast.visit_node_scoped_expr_roots(ast, fn(scoped_expr) {
     collect_scoped_node_matches(pattern, scoped_expr, matches)
     @untyped_ast.NodeVisitAction::NodeVisitPrune
   })

query/query_test.mbt

@@ -34,6 +34,35 @@ fn query_constant_value(node : @untyped_ast.Node) -> String? {
   query_leaf_string_child(node, "value")
 }
 
+///|
+fn parse_query_expr(source : String) -> @syntax.Expr raise {
+  let lex_result = @lexer.tokens_from_string(
+    source,
+    comment=true,
+    name="query expr test",
+  )
+  guard lex_result.errors.length() == 0 else {
+    fail("expected expression to lex: \{source}")
+  }
+  let (expr, reports) = @handrolled_parser.parse_expr(lex_result.tokens)
+  guard reports is [] else { fail("expected expression to parse: \{source}") }
+  expr
+}
+
+///|
+fn parse_query_impls(source : String) -> @list.List[@syntax.Impl] raise {
+  let (impls, reports) = @parser.parse_string(source, name="query impl test")
+  guard reports is [] else { fail("expected impls to parse: \{source}") }
+  impls
+}
+
+///|
+fn parse_query_impl(source : String) -> @syntax.Impl raise {
+  let impls = parse_query_impls(source).to_array()
+  guard impls.length() == 1 else { fail("expected one impl: \{source}") }
+  impls[0]
+}
+
 ///|
 test "query matches simple expression and returns empty capture map" {
   let query = ExprQuery::ExprQuery("target()")
@@ -145,3 +174,90 @@ test "query raises when pattern is invalid" {
     _ => fail("expected invalid pattern")
   }
 }
+
+///|
+test "query captures_from_ast finds matches under an impl list root" {
+  let query = ExprQuery::ExprQuery("target()")
+  let source =
+    #|fn sample {
+    #|  target();
+    #|  other()
+    #|}
+    #|
+  let ast = @untyped_ast.from_impls(parse_query_impls(source))
+  let matches = query.captures_from_ast(ast~)
+  assert_eq(matches.length(), 1)
+  assert_eq(matches[0].length(), 0)
+}
+
+///|
+test "query captures_from_ast matches a direct expression root" {
+  let query = ExprQuery::ExprQuery("target($(literal:const))")
+  let ast = @untyped_ast.from_expr(parse_query_expr("target(42)"))
+  let matches = query.captures_from_ast(ast~)
+  assert_eq(matches.length(), 1)
+  guard matches[0].get("literal") is Some(literal) else {
+    fail("expected literal capture")
+  }
+  guard query_constant_value(literal) is Some(value) else {
+    fail("expected literal value")
+  }
+  inspect(value, content="42")
+}
+
+///|
+test "query captures_from_ast finds nested expressions from a non-expression impl root" {
+  let query = ExprQuery::ExprQuery("target()")
+  let ast = @untyped_ast.from_impl(parse_query_impl("fn sample { target() }"))
+  assert_true(ast.kind == Impl_TopFuncDef)
+  let matches = query.captures_from_ast(ast~)
+  assert_eq(matches.length(), 1)
+}
+
+///|
+test "query captures_from_ast preserves expression id and constant captures" {
+  let query = ExprQuery::ExprQuery(
+    "sink($(callee:id), $(value:exp), $(literal:const))",
+  )
+  let source =
+    #|fn sample {
+    #|  sink(source, make(), "danger")
+    #|}
+    #|
+  let source_matches = query.captures(source_name="sample.mbt", source)
+  let ast_matches = query.captures_from_ast(
+    ast=@untyped_ast.from_impls(parse_query_impls(source)),
+  )
+  assert_eq(ast_matches.length(), source_matches.length())
+  assert_eq(ast_matches.length(), 1)
+  guard ast_matches[0].get("callee") is Some(callee) else {
+    fail("expected callee capture")
+  }
+  guard query_leaf_string(callee) is Some(callee_name) else {
+    fail("expected callee name")
+  }
+  inspect(callee_name, content="source")
+  guard ast_matches[0].get("value") is Some(value) else {
+    fail("expected expression capture")
+  }
+  assert_true(value.kind == Expr_Apply)
+  guard ast_matches[0].get("literal") is Some(literal) else {
+    fail("expected literal capture")
+  }
+  guard query_constant_value(literal) is Some(literal_value) else {
+    fail("expected literal value")
+  }
+  inspect(literal_value, content="danger")
+}
+
+///|
+test "query captures_from_ast prunes nested expressions after an outer match" {
+  let query = ExprQuery::ExprQuery("wrap($(value:exp))")
+  let ast = @untyped_ast.from_expr(parse_query_expr("wrap(wrap(1))"))
+  let matches = query.captures_from_ast(ast~)
+  assert_eq(matches.length(), 1)
+  guard matches[0].get("value") is Some(value) else {
+    fail("expected value capture")
+  }
+  assert_true(value.kind == Expr_Apply)
+}