fix(untyped_cst): parse local functions in expression fragments

moon.mod

@@ -1,6 +1,6 @@
 name = "moonbitlang/parser"
 
-version = "0.3.17"
+version = "0.3.18"
 
 import {
   "moonbitlang/x@0.4.39",

untyped_cst/ast_equiv_test.mbt

@@ -333,6 +333,29 @@ test "expression conversion produces expression AST" {
   guard expr is Some(_) else { fail("expected expression") }
 }
 
+///|
+test "expression fragments lower complete statement sequences" {
+  let cases = [
+    "fn helper(x) { x }", "fn helper(x) { x }; helper(1)", "fn helper(x) { x };",
+    "1; 2", "let x = 1; x", "async fn helper(x) { x }; helper(1)",
+  ]
+  for source in cases {
+    assert_expr_ast_equiv(source)
+  }
+  let (local_expr, local_reports) = @untyped_cst.parse_expression(
+    "fn helper(x) { x }",
+  ).to_expr()
+  @test.assert_eq(local_reports.length(), 0)
+  guard local_expr is Some(@syntax.Expr::LetFn(..)) else {
+    fail("expected local function expression root")
+  }
+  let (sequence, sequence_reports) = @untyped_cst.parse_expression("1; 2").to_expr()
+  @test.assert_eq(sequence_reports.length(), 0)
+  guard sequence is Some(@syntax.Expr::Sequence(..)) else {
+    fail("expected sequence expression root")
+  }
+}
+
 ///|
 test "string and labelled argument AST equivalence cases" {
   let cases = [

untyped_cst/diagnostics_equiv_test.mbt

@@ -170,6 +170,36 @@ test "expression diagnostics come from CST parser recovery" {
   assert_expression_acceptance("line break separator", "1\n2", true)
 }
 
+///|
+test "local function expression fragment diagnostics match handrolled" {
+  let valid = [
+    ("local function", "fn helper(x) { x }"),
+    ("local function sequence", "fn helper(x) { x }; helper(1)"),
+    ("async local function", "async fn helper(x) { x }"),
+    ("newline after fn", "fn\nhelper(x) { x }"),
+    ("comment after fn", "fn // between fn and name\nhelper(x) { x }"),
+    ("newline after async", "async\nfn helper(x) { x }"),
+    ("comment after async", "async // between async and fn\nfn helper(x) { x }"),
+    ("newline after async fn", "async fn\nhelper(x) { x }"),
+    ("anonymous function", "fn(x) { x }"),
+  ]
+  for case in valid {
+    let (label, source) = case
+    assert_expression_acceptance(label, source, true)
+  }
+  let invalid = [
+    ("local function missing parameters", "fn helper { x }"),
+    ("local function missing parameter close", "fn helper(x"),
+    ("local function malformed parameter close", "fn helper(x { x }"),
+    ("local function missing body", "fn helper(x)"),
+    ("nested named function", "consume(fn helper(x) { x })"),
+  ]
+  for case in invalid {
+    let (label, source) = case
+    assert_expression_acceptance(label, source, false)
+  }
+}
+
 ///|
 test "acceptance equivalence: labelled arguments" {
   let invalid = [

untyped_cst/lower.mbt

@@ -1335,11 +1335,14 @@ fn lower_using_names(
 ///|
 fn lower_expr(node : CstNode) -> @syntax.Expr {
   match node.kind {
-    Expression =>
-      node
-      .first_child_with_label("expr")
-      .map(lower_expr)
-      .unwrap_or(lower_hole_expr(node.loc))
+    Expression => {
+      let items = node.children_with_label("expr")
+      if items.is_empty() {
+        lower_hole_expr(node.loc)
+      } else {
+        lower_statement_sequence(items, 0, node.loc)
+      }
+    }
     Expr_Ident =>
       @syntax.Expr::Ident(
         id=@syntax.Var::{

untyped_cst/parse.mbt

@@ -531,7 +531,9 @@ fn parse_parameter_list(state : ParserState) -> CstNode {
     parse_parameter_item,
     List_ParameterList,
     sep_kind=TokenKind::TK_COMMA,
+    follow_set=[TokenKind::TK_LBRACE],
     item_expected="parameter name or label",
+    post_item_expected=[TokenKind::TK_COMMA, TokenKind::TK_RPAREN],
   )
 }
 
@@ -3506,7 +3508,7 @@ fn parse_invalid_guard_letrec(
 fn parse_expr_fragment(state : ParserState) -> CstNode {
   let reports_before_expr = state.reports.length()
   let expr = state.required_expr_node(
-    parse_expr_until(state, [TokenKind::TK_EOF]),
+    parse_expr_fragment_item(state),
     reports_before_expr,
   )
   let children : Array[(String?, CstNode)] = []
@@ -3520,7 +3522,7 @@ fn parse_expr_fragment(state : ParserState) -> CstNode {
     }
     push_children(children, state.consume_trivia())
     if !state.current_is(TokenKind::TK_EOF) && !(state.peek() is None) {
-      let next = parse_expr_until(state, [TokenKind::TK_EOF])
+      let next = parse_expr_fragment_item(state)
       if !next.children.is_empty() {
         push_child(children, Some("expr"), next)
       }
@@ -3556,6 +3558,19 @@ fn parse_expr_fragment(state : ParserState) -> CstNode {
   root
 }
 
+///|
+fn parse_expr_fragment_item(state : ParserState) -> CstNode {
+  if local_function_statement_can_start(state) {
+    discard_expr_trivia(state)
+    parse_local_function_statement(state, [
+      TokenKind::TK_SEMI,
+      TokenKind::TK_EOF,
+    ])
+  } else {
+    parse_expr_until(state, [TokenKind::TK_EOF])
+  }
+}
+
 ///|
 fn parse_expr_until(
   state : ParserState,
@@ -8685,7 +8700,12 @@ fn parse_statement(state : ParserState) -> CstNode {
     Some((LETREC, _, _)) => parse_letrec_statement(state)
     Some((FN, _, _)) | Some((ASYNC, _, _)) if local_function_statement_can_start(
         state,
-      ) => parse_local_function_statement(state)
+      ) =>
+      parse_local_function_statement(state, [
+        TokenKind::TK_SEMI,
+        TokenKind::TK_RBRACE,
+        TokenKind::TK_EOF,
+      ])
     Some((GUARD, _, _)) | Some((GUARD_EXCLAMATION, _, _)) =>
       parse_guard_statement(state)
     Some((DEFER, _, _)) => parse_defer_statement(state)
@@ -9051,21 +9071,26 @@ fn parse_letrec_binding(state : ParserState) -> CstNode {
 }
 
 ///|
-fn parse_local_function_statement(state : ParserState) -> CstNode {
+fn parse_local_function_statement(
+  state : ParserState,
+  stop_kinds : Array[TokenKind],
+) -> CstNode {
   let children : Array[(String?, CstNode)] = []
   if state.current_is(TokenKind::TK_ASYNC) {
     if state.consume_any() is Some(kw) {
       push_child(children, Some("kw"), kw)
     } else {
       ()
     }
+    push_children(children, state.consume_trivia())
   }
   if state.current_is(TokenKind::TK_FN) {
     if state.consume_any() is Some(kw) {
       push_child(children, Some("kw"), kw)
     } else {
       ()
     }
+    push_children(children, state.consume_trivia())
   }
   match state.peek() {
     Some((LIDENT(_), _, _)) | Some((UIDENT(_), _, _)) | Some((EXTEND, _, _)) =>
@@ -9076,42 +9101,42 @@ fn parse_local_function_statement(state : ParserState) -> CstNode {
       }
     _ => ()
   }
-  if state.current_is(TokenKind::TK_EXCLAMATION) {
-    match state.peek() {
-      Some((_, start, end)) =>
-        state.report(
-          mk_loc(start, end),
-          "Failed to parse func: expected '(' or '{', got EXCLAMATION",
-        )
-      None => ()
-    }
-    while !state.current_is(TokenKind::TK_RBRACE) &&
-          !state.current_is(TokenKind::TK_EOF) &&
-          !(state.peek() is None) {
-      if state.consume_any() is Some(bad) {
-        push_child(
-          children,
-          Some("error"),
-          error_node(bad.loc, bad.source_span, [(Some("item"), bad)]),
-        )
-      } else {
-        ()
+  match state.peek() {
+    Some((tok, start, end)) if !(tok is LPAREN) => {
+      state.report(
+        mk_loc(start, end),
+        "Failed to parse func: expected '(' or '{', got \{@debug.render(tok.to_repr())}",
+      )
+      while !state.current_in(stop_kinds) &&
+            !state.current_is(TokenKind::TK_EOF) &&
+            !(state.peek() is None) {
+        if state.consume_any() is Some(bad) {
+          push_child(
+            children,
+            Some("error"),
+            error_node(bad.loc, bad.source_span, [(Some("item"), bad)]),
+          )
+        } else {
+          ()
+        }
       }
+      return state
+        .node_from_children(Expr_LetFn, children)
+        .with_bool_marker(
+          "async",
+          Flag_Async,
+          children
+          .iter()
+          .any(fn(entry) {
+            entry.0 is Some("kw") && entry.1.has_child_kind(Flag_Async)
+          }),
+        )
+        .with_async_location_anchor()
+        .with_decl_like_function_location_anchor()
     }
-    return state
-      .node_from_children(Expr_LetFn, children)
-      .with_bool_marker(
-        "async",
-        Flag_Async,
-        children
-        .iter()
-        .any(fn(entry) {
-          entry.0 is Some("kw") && entry.1.has_child_kind(Flag_Async)
-        }),
-      )
-      .with_async_location_anchor()
-      .with_decl_like_function_location_anchor()
+    _ => ()
   }
+  let reports_before_signature = state.reports.length()
   if state.current_is(TokenKind::TK_LPAREN) {
     push_child(children, Some("params"), parse_parameter_list(state))
   }
@@ -9122,6 +9147,15 @@ fn parse_local_function_statement(state : ParserState) -> CstNode {
   }
   if state.current_is(TokenKind::TK_LBRACE) {
     push_child(children, Some("body"), parse_block_expr(state))
+  } else {
+    let body = if state.reports.length() > reports_before_signature {
+      state.missing_token_node_silent(TokenKind::TK_LBRACE)
+    } else {
+      state.missing_token_node_with_expected(TokenKind::TK_LBRACE, [
+        TokenKind::TK_LBRACE,
+      ])
+    }
+    push_child(children, Some("body"), body)
   }
   state
   .node_from_children(Expr_LetFn, children)

untyped_cst/parse_test.mbt

@@ -76,6 +76,20 @@ fn root_contains_expr(root : @untyped_cst.CstNode) -> Bool {
   false
 }
 
+///|
+fn root_expr_children(
+  root : @untyped_cst.CstNode,
+) -> Array[@untyped_cst.CstNode] {
+  let exprs : Array[@untyped_cst.CstNode] = []
+  for entry in root.children {
+    match entry {
+      (Some("expr"), child) => exprs.push(child)
+      _ => ()
+    }
+  }
+  exprs
+}
+
 ///|
 fn find_node(
   root : @untyped_cst.CstNode,
@@ -397,6 +411,42 @@ test "expression entry returns expression" {
   @test.assert_eq(result.root().text(source), source[:])
 }
 
+///|
+test "expression fragments recognize top-level local functions" {
+  let single_cases = [
+    "fn helper(x) { x }", "async fn helper(x) { x }", "// leading trivia\nfn helper(x) { x }",
+    "fn\nhelper(x) { x }", "fn // between fn and name\nhelper(x) { x }", "async\nfn helper(x) { x }",
+    "async // between async and fn\nfn helper(x) { x }", "async fn\nhelper(x) { x }",
+  ]
+  for source in single_cases {
+    let result = @untyped_cst.parse_expression(source)
+    @test.assert_eq(result.diagnostics_view().length(), 0)
+    let exprs = root_expr_children(result.root())
+    @test.assert_eq(exprs.length(), 1)
+    @test.assert_eq(exprs[0].kind is @untyped_cst.Expr_LetFn, true)
+    @test.assert_eq(result.root().text(source), source[:])
+  }
+  let sequence_cases = [
+    "fn helper(x) { x }; helper(1)", "fn helper(x) { x }\nhelper(1)",
+  ]
+  for source in sequence_cases {
+    let result = @untyped_cst.parse_expression(source)
+    @test.assert_eq(result.diagnostics_view().length(), 0)
+    let exprs = root_expr_children(result.root())
+    @test.assert_eq(exprs.length(), 2)
+    @test.assert_eq(exprs[0].kind is @untyped_cst.Expr_LetFn, true)
+    @test.assert_eq(count_root_label(result.root(), "sep"), 1)
+    @test.assert_eq(result.root().text(source), source[:])
+  }
+  let anonymous = @untyped_cst.parse_expression("fn(x) { x }")
+  @test.assert_eq(anonymous.diagnostics_view().length(), 0)
+  let anonymous_exprs = root_expr_children(anonymous.root())
+  @test.assert_eq(anonymous_exprs.length(), 1)
+  @test.assert_eq(anonymous_exprs[0].kind is @untyped_cst.Expr_Function, true)
+  let nested = @untyped_cst.parse_expression("consume(fn helper(x) { x })")
+  @test.assert_eq(nested.diagnostics_view().is_empty(), false)
+}
+
 ///|
 test "successful expressions are grammar shaped without placeholders" {
   let cases = [