Parse comments

Kestrel.g4

@@ -4,7 +4,7 @@ SLASH_4: '////';
 SLASH_3: '///';
 SLASH_2: '//';
 
-LineComment: SLASH_2 ~[\r\n]* -> channel(HIDDEN);
+LINE_COMMENT: SLASH_2 ~[\r\n]* -> channel(HIDDEN);
 
 EXPOSING_NESTED: '(' '..' ')';
 INFIX_ID: '(' INFIX_CHAR+ ')';

src/formatter.test.ts

@@ -311,6 +311,77 @@ test("toplevel nested let expr", () => {
 `).toBeFormatted();
 });
 
+test("allows spaces in toplevel nested let expr", () => {
+  expect(`let f = {
+  let x = value;
+
+  let y = value2;
+
+  body
+}
+`).toBeFormatted(`let f = {
+  let x = value;
+
+  let y = value2;
+
+  body
+}
+`);
+});
+
+test("force at 1 space in toplevel nested let expr", () => {
+  expect(`let f = { let x = value; body }
+`).toBeFormatted(`let f = {
+  let x = value;
+  body
+}
+`);
+});
+
+test("allows at most 1 space in toplevel nested let expr", () => {
+  expect(`let f = {
+  let x = value;
+
+
+
+  let y = value2;
+
+
+
+  body
+}
+`).toBeFormatted(`let f = {
+  let x = value;
+
+  let y = value2;
+
+  body
+}
+`);
+});
+
+test("allow zero lines after struct", () => {
+  expect(`let f = {
+  let p = Person {
+    name: "hello",
+    age: 42,
+  };
+  body
+}
+`).toBeFormatted();
+});
+
+test("nested let", () => {
+  expect(`let a = {
+  let l1 = {
+    let l2 = value;
+    e
+  };
+  body
+}
+`).toBeFormatted();
+});
+
 test("toplevel nested let# expr", () => {
   expect(`let f = {
   let#and_then x = value;
@@ -396,7 +467,65 @@ test("order between type declrs and declrs", () => {
 });
 
 describe("comments", () => {
-  test.todo("regular comments");
+  test("doc comments on declrs", () => {
+    expect(`let f =
+      // c1
+      0 +
+      // c2
+      1
+`).toBeFormatted(`let f =
+  // c1
+  // c2
+  0 + 1
+`);
+  });
+
+  test("doc comments on declrs", () => {
+    expect(`let f = fn {
+        // c
+        42
+    }
+`).toBeFormatted(`let f = fn {
+  // c
+  42
+}
+`);
+  });
+
+  test("doc comments in if expr", () => {
+    expect(`let f = if b {
+    // c
+    42
+} else {
+  // d
+  100
+}
+`).toBeFormatted(`let f =
+  if b {
+    // c
+    42
+  } else {
+    // d
+    100
+  }
+
+`);
+  });
+
+  test.todo("doc comments in lists", () => {
+    expect(`let f = [
+  0,
+  // comment
+  1,
+]
+`).toBeFormatted(`let f = [
+  0,
+  // comment
+  1,
+]
+
+`);
+  });
 
   test("doc comments on declrs", () => {
     expect(`/// First line

src/formatter.ts

@@ -1,7 +1,9 @@
 import {
   ConstLiteral,
+  LineComment,
   MatchPattern,
   PolyTypeAst,
+  RangeMeta,
   TypeAst,
   UntypedDeclaration,
   UntypedExpr,
@@ -28,6 +30,23 @@ import {
 } from "./pretty";
 import { gtEqPos } from "./typecheck/typedAst/common";
 
+let currentLineComments: LineComment[] = [];
+function popComments(ast: RangeMeta): Doc[] {
+  const poppedComments: string[] = [];
+  // eslint-disable-next-line no-constant-condition
+  while (true) {
+    const comment = currentLineComments.at(-1);
+
+    if (comment === undefined || comment.range.end.line >= ast.range.end.line) {
+      break;
+    }
+
+    poppedComments.push(comment.comment);
+    currentLineComments.pop();
+  }
+  return poppedComments.map((comment) => concat(text(comment), lines()));
+}
+
 const ORDERED_PREFIX_SYMBOLS = [["!"]];
 
 const ORDERED_INFIX_SYMBOLS = [
@@ -114,7 +133,7 @@ function constToDoc(lit: ConstLiteral): Doc {
 }
 
 function indentWithSpaceBreak(docs: Doc[], unbroken?: string): Doc {
-  return concat(
+  return group(
     nest(
       //
       break_(""),
@@ -148,6 +167,10 @@ function asBlock(isBlock: boolean, docs: Doc[]): Doc {
   return block_(...docs);
 }
 
+function exprToDocWithComments(ast: UntypedExpr, block: boolean): Doc {
+  return concat(...popComments(ast), exprToDoc(ast, block));
+}
+
 function exprToDoc(ast: UntypedExpr, block: boolean): Doc {
   switch (ast.type) {
     /* v8 ignore next 2 */
@@ -163,7 +186,10 @@ function exprToDoc(ast: UntypedExpr, block: boolean): Doc {
             sepBy(
               concat(text(","), break_()),
               ast.values.map((expr) =>
-                exprToDoc(expr, expr.type !== "let" && expr.type !== "let#"),
+                exprToDocWithComments(
+                  expr,
+                  expr.type !== "let" && expr.type !== "let#",
+                ),
               ),
             ),
           ],
@@ -308,7 +334,7 @@ function exprToDoc(ast: UntypedExpr, block: boolean): Doc {
         text("fn"),
         sepByString(",", params),
         text(" "),
-        block_(exprToDoc(ast.body, true)),
+        block_(exprToDocWithComments(ast.body, true)),
       );
     }
 
@@ -317,10 +343,10 @@ function exprToDoc(ast: UntypedExpr, block: boolean): Doc {
         text("if "),
         exprToDoc(ast.condition, false),
         text(" "),
-        block_(exprToDoc(ast.then, true)),
+        block_(exprToDocWithComments(ast.then, true)),
 
         text(" else "),
-        block_(exprToDoc(ast.else, true)),
+        block_(exprToDocWithComments(ast.else, true)),
       );
 
     case "let#": {
@@ -330,8 +356,8 @@ function exprToDoc(ast: UntypedExpr, block: boolean): Doc {
       const inner = concat(
         text(`let#${ns}${ast.mapper.name} `),
         patternToDoc(ast.pattern),
-        text(` = `),
-        exprToDoc(ast.value, false),
+        text(` =`),
+        declarationValueToDoc(ast.value),
         text(";"),
         break_(),
         exprToDoc(ast.body, true),
@@ -345,13 +371,18 @@ function exprToDoc(ast: UntypedExpr, block: boolean): Doc {
     }
 
     case "let": {
+      const linesDiff = Math.min(
+        Math.max(ast.body.range.start.line - ast.value.range.end.line - 1, 0),
+        1,
+      );
+
       const inner = concat(
         text("let "),
         patternToDoc(ast.pattern),
-        text(" = "),
-        exprToDoc(ast.value, false),
+        text(" ="),
+        declarationValueToDoc(ast.value),
         text(";"),
-        break_(),
+        lines(linesDiff),
         exprToDoc(ast.body, true),
       );
 
@@ -488,6 +519,32 @@ function handleDocComment(content: string, init = "///") {
   );
 }
 
+function declarationValueToDoc(expr: UntypedExpr): Doc {
+  const exprDoc = exprToDoc(expr, false);
+
+  switch (expr.type) {
+    case "if":
+      return indentWithSpaceBreak([exprDoc]);
+
+    default: {
+      const poppedComments = popComments(expr);
+
+      if (poppedComments.length === 0) {
+        return concat(text(" "), exprDoc);
+      }
+
+      return broken(
+        nest(
+          //
+          break_(),
+          ...poppedComments,
+          exprDoc,
+        ),
+      );
+    }
+  }
+}
+
 function declToDoc(ast: UntypedDeclaration): Doc {
   const name =
     isInfix(ast.binding.name) || isPrefix(ast.binding.name)
@@ -503,14 +560,7 @@ function declToDoc(ast: UntypedDeclaration): Doc {
     ast.typeHint === undefined
       ? nil
       : concat(text(": "), typeHintToDoc(ast.typeHint)),
-    ast.extern
-      ? nil
-      : concat(
-          text(" ="),
-          ["if"].includes(ast.value.type)
-            ? indentWithSpaceBreak([exprToDoc(ast.value, false)])
-            : concat(text(" "), exprToDoc(ast.value, false)),
-        ),
+    ast.extern ? nil : concat(text(" ="), declarationValueToDoc(ast.value)),
   );
 }
 
@@ -631,6 +681,9 @@ export function formatExpr(expr: UntypedExpr): string {
 }
 
 export function format(ast: UntypedModule): string {
+  currentLineComments = [...(ast.lineComments ?? [])];
+  currentLineComments.reverse();
+
   const importsDocs = ast.imports
     .sort((i1, i2) => (i1.ns > i2.ns ? 1 : -1))
     .map(importToDoc)