Various fixes for the IML formatter

imlformat/iml-prettier.ts

@@ -178,15 +178,28 @@ function trim_parentheses(s: string): string {
   return s;
 }
 
-function comments(cur, options: Options): Doc[] {
+interface Position {
+  pos_fname: string;
+  pos_cnum: number;
+  pos_lnum: number;
+  pos_bol: number;
+}
+
+interface Location {
+  loc_start: Position;
+  loc_ghost: boolean;
+  loc_end: Position;
+}
+
+function comments(cur: Location, options: Options): Doc[] {
   if (cur && cur.loc_start.pos_cnum >= 0) {
     const last = options.last_loc;
     if (last) {
       let src = get_source_between(last, cur, options);
       const cstart = src.indexOf("(*");
       if (cstart != -1) {
         // Could be a docstring right at the beginning of the file.
-        if (src.length > 4 && src[cstart + 2] != '*' && src[cstart + 3] != ' ') {
+        if (src.length > 4 && src[cstart + 2] != '*') {
           const had_newline = src.endsWith('\n') || src.endsWith('\r');
           src = trim(src.substring(cstart), ['\n', '\r', ';', ' ']);
           let cend = src.lastIndexOf("*)") + 2;
@@ -203,6 +216,43 @@ function comments(cur, options: Options): Doc[] {
   return [];
 }
 
+function gobble_line_comment(loc: Location, options: Options): Doc[] {
+  let i = loc.loc_end.pos_cnum;
+  const src = options.originalText as string;
+  let comment_from = undefined;
+  while (src[i] != '\n' && i < src.length - 1) {
+    if (src[i] == '(' && src[i + 1] == "*") {
+      comment_from = i;
+      break;
+    }
+    i++;
+  }
+  if (comment_from) {
+    i = comment_from;
+    const r = [];
+    let tmp = "";
+    while (i < src.length - 1 && src[i] != '\n') {
+      if (src[i] == "(" && src[i + 1] == "*") {
+        while (i < src.length - 1 && (src[i] != "*" || src[i + 1] != ")")) {
+          tmp += src[i];
+          i++;
+        }
+        if (i < src.length - 1)
+          tmp += src[i] + src[i + 1];
+        r.push(tmp);
+        tmp = "";
+        options.last_loc = loc;
+        (options.last_loc as Location).loc_end.pos_cnum = i;
+      }
+      else
+        i++;
+    }
+    return [line, ...join(line, r)];
+  }
+  else
+    return [];
+}
+
 enum Notation { None, Infix, Prefix }
 enum Associativity { None, Left, Right }
 
@@ -378,6 +428,10 @@ function par_if(c: boolean, x: Doc): Doc[] {
   return c ? ["(", l, x, l, ")"] : [x];
 }
 
+function bracketize(d: Doc[]): Doc[] {
+  return ["[", ...d, "]"];
+}
+
 function print_constant_desc(node: AST, options: Options): Doc {
   const constructor = node[0];
   const args = node.slice(1);
@@ -1201,20 +1255,25 @@ function print_list(e: AST, options: Options): Doc {
   return f([join([softline, "::", softline], es)]);
 }
 
-function print_flat_list_pattern_elems(p: AST, options: Options): Doc[] {
+function print_flat_list_pattern_elems(p: AST, options: Options): [Doc[], boolean] {
   const d0 = print_pattern(p.ppat_desc[1][0], options);
   const e1 = p.ppat_desc[1][1];
   if (e1.ppat_desc[0] == "Ppat_construct" && e1.ppat_desc[1].txt[1] == "[]")
-    return [d0];
-  else if (e1.ppat_desc[0] == "Ppat_construct" && e1.ppat_desc[1].txt[1] == "::")
-    return [d0, ...print_flat_list_pattern_elems(e1.ppat_desc[2][1], options)];
+    return [[d0], true];
+  else if (e1.ppat_desc[0] == "Ppat_construct" && e1.ppat_desc[1].txt[1] == "::") {
+    const [x, b] = print_flat_list_pattern_elems(e1.ppat_desc[2][1], options);
+    return [[d0, ...x], b];
+  }
   else
-    return [d0, print_pattern(e1, options)];
+    return [[d0, print_pattern(e1, options)], false];
 }
 
 function print_pattern_list(p: AST, options: Options): Doc {
-  const ps = print_flat_list_pattern_elems(p, options);
-  return f([join([softline, "::", softline], ps)]);
+  const [ps, b] = print_flat_list_pattern_elems(p, options);
+  if (b)
+    return f(bracketize(ps));
+  else
+    return f(join([softline, "::", softline], ps));
 }
 
 function print_class_structure(node: AST, options: Options): Doc {
@@ -1738,9 +1797,8 @@ function print_expression_desc(node: AST, options: Options): Doc {
       }
       else
         return f([
-          ...par_if(true, [
-            ...print_open_declaration(args[0], options, false), ".(", softline,
-            ...print_expression(args[1], options)]), softline, ")"]);
+          ...print_open_declaration(args[0], options, false), ".(", softline,
+          ...print_expression(args[1], options), softline, ")"]);
     }
     case "Pexp_letop":
       // | Pexp_letop of letop
@@ -1794,13 +1852,13 @@ function print_constructor_declaration(node: AST, options: Options): Doc {
   //  pcd_loc: Location.t;
   //  pcd_attributes: attributes;  (** [C of ... [\@id1] [\@id2]] *)
   // }
-  let r: Doc = comments(node.pcd_loc, options);;
+  let r: Doc = comments(node.pcd_loc, options);
   if (node.pcd_args[1].length == 0)
-    r = [node.pcd_name.txt];
+    r.push(node.pcd_name.txt);
   else
-    r = [
+    r = r.concat([
       print_string_loc(node.pcd_name, options), line, "of", line,
-      print_constructor_arguments(node.pcd_args, options)];
+      ...print_constructor_arguments(node.pcd_args, options)]);
   return [...r, ...ifnonempty(line, print_attributes(node.pcd_attributes, 1, options))]
 }
 
@@ -1830,8 +1888,12 @@ function print_type_kind(node: AST, options: Options): Doc {
       return [];
     case "Ptype_variant":
       // | Ptype_variant of constructor_declaration list
-      return g([ifBreak("| ", ""), join([line, "| "], args[0].map(x =>
-        f([print_constructor_declaration(x, options)])
+      return g([ifBreak("| ", ""), join([line, "| "], args[0].map((x, i) =>
+        g([
+          print_constructor_declaration(x, options),
+          ...((i + 1 < args[0].length) ?
+            ifnonempty([line], comments(args[0][i + 1].pcd_loc, options)) :
+            gobble_line_comment(x.pcd_loc, options))])
       ))]);
     case "Ptype_record":
       // | Ptype_record of label_declaration list  (** Invariant: non-empty list *)
@@ -2008,6 +2070,24 @@ function print_type_exception(node: AST, options: Options): Doc {
       ...ifnonempty(line, print_attributes(node.ptyexn_attributes, 2, options))])]);
 }
 
+function print_type_constraint(node: AST, options: Options): Doc {
+  const constructor = node[0];
+  const args = node.slice(1);
+  switch (constructor) {
+    case "Pconstraint": {
+      // | Pconstraint of core_type
+      return print_core_type(args[0], options);
+    }
+    case "Pcoerce": {
+      // | Pcoerce of core_type option * core_type
+      niy();
+      break;
+    }
+    default:
+      throw new Error(`Unexpected node type: ${constructor}`);
+  }
+}
+
 function print_structure_item_desc(node: AST, options: Options): Doc {
   const constructor = node[0];
   const args = node.slice(1);
@@ -2082,17 +2162,20 @@ function print_structure_item_desc(node: AST, options: Options): Doc {
       else if (args[1].length > 0 && pvb.pvb_expr.pexp_desc[0] == "Pexp_function") {
         // For function definitions we want to hoist the arguments
         const params = pvb.pvb_expr.pexp_desc[1];
+        const type_cnstrnt_opt = pvb.pvb_expr.pexp_desc[2];
         const fundef = pvb.pvb_expr.pexp_desc[3];
-        return [...r,
-        f([indent([
-          line,
-          ...par_if(is_infix_op_pattern(pvb.pvb_pat), print_pattern(pvb.pvb_pat, options)),
-          line,
-          ...join(line, params.map(x => print_function_param(x, options))),
-          ...(params && params.length > 0 ? [line] : []),
-          "="]),
-        f([indent([line, print_function_body(fundef, options)]),
-        ...ifnonempty(line, print_attributes(attrs, 2, options))])])];
+        return [
+          ...r,
+          f([indent([
+            line,
+            ...par_if(is_infix_op_pattern(pvb.pvb_pat), print_pattern(pvb.pvb_pat, options)),
+            line,
+            ...join(line, params.map(x => print_function_param(x, options))),
+            ...(params && params.length > 0 ? [line] : []),
+            ...(type_cnstrnt_opt ? [":", line, print_type_constraint(type_cnstrnt_opt, options), line] : []),
+            "="]),
+          f([indent([line, print_function_body(fundef, options)]),
+          ...ifnonempty(line, print_attributes(attrs, 2, options))])])];
       }
       // Generic version
       return [

imlformat/test/basic.test.ts

@@ -111,4 +111,15 @@ let f x = 1;;
 let g y = 0
 
 let h z = 2;;`))
+})
+
+test("open expr", () => {
+  return format(`
+let f e x =
+  match e with
+  | Some _ -> Real.(g x = 0.0)
+  | None -> Real.(g x >= 0.0)
+`).then(x => expect(x).toEqual(`\
+let f e x = match e with Some _ -> Real.(g x = 0.0) | None -> Real.(g x >= 0.0)`
+))
 })

imlformat/test/comments.test.ts

@@ -49,5 +49,28 @@ let f = 1
 
 (** This is a docstring *)
 
-let g = 1`))
+let g = 1`
+    ))
 })
+
+test("line comments in variants", () => {
+  return format(`
+(* Comment 1 *)
+
+type expr =
+  | Eq of poly  (*  = 0 *)
+  | Geq of poly (* >= 0 *)
+  | Gt of poly  (*  > 0 *) (* a second line commment *)
+
+(* Comment 2 *)
+
+type something_else = int option
+`).then(x =>
+    expect(x).toEqual(`\
+(* Comment 1 *)
+type expr = Eq of poly (*  = 0 *)  | Geq of poly (* >= 0 *)  | Gt of poly (*  > 0 *) (* a second line commment *)
+
+(* Comment 2 *)
+type something_else = int option`
+    ))
+})

imlformat/test/functions.test.ts

@@ -0,0 +1,39 @@
+import { expect, test } from '@jest/globals';
+
+import { format } from "../imlformat.format";
+
+test("output_type", () => {
+  return format(`let f x : bool = x = 0`).then(x => expect(x).toEqual(`let f x : bool = x = 0`));
+})
+
+test("no_output_type", () => {
+  return format(`let f x = x = 0`).then(x => expect(x).toEqual(`let f x = x = 0`));
+})
+
+test("eval_poly", () => {
+  return format(`\
+let rec eval_poly (p:poly) (x:Real.t list) : Real.t =
+  match p, x with
+  | a :: p, b :: x ->
+    Real.(a*b + eval_poly p x)
+  | [a], [] -> a
+  | _ -> 0.0`
+  ).then(x => expect(x).toEqual(`\
+let rec eval_poly (p : poly) (x : Real.t list) : Real.t =
+  match p, x with
+  | a::p, b::x -> Real.(a * b + eval_poly p x)
+  | [a], [] -> a
+  | _ -> 0.0`
+  ))
+});
+
+test("eval_system", () => {
+  return format(`\
+let rec eval_system (es:system) (x:Real.t list) : bool =
+  match es with
+  | [] -> true
+  | e::es -> eval_expr e x && eval_system es x`
+  ).then(x => expect(x).toEqual(`\
+let rec eval_system (es : system) (x : Real.t list) : bool =
+  match es with [] -> true | e::es -> eval_expr e x && eval_system es x`))
+});

package.json

@@ -10,7 +10,7 @@
     "theme": "dark"
   },
   "license": "MIT",
-  "version": "0.0.40",
+  "version": "0.0.41",
   "repository": {
     "type": "git",
     "url": "https://github.com/imandra-ai/imandrax-vscode"