Remove print_assigment for some utilities and less dry code

compiler/grainformat/grainformat.re

@@ -28,7 +28,7 @@ let get_program_string = filename => {
 let compile_parsed = filename => {
   let filename = Filepath.to_string(filename);
   let program_str = get_program_string(filename);
-  switch (Format.parse_source(program_str)) {
+  switch (Fmt.parse_source(program_str)) {
   | Error(ParseError(exn)) =>
     let bt =
       if (Printexc.backtrace_status()) {
@@ -56,28 +56,28 @@ let format_code =
     (
       ~eol,
       ~output=?,
-      ~original_source: array(string),
+      ~source: array(string),
       program: Parsetree.parsed_program,
     ) => {
-  let formatted_code =
-    Grain_formatting.Format.format_ast(~original_source, ~eol, program);
-
-  let buf = Buffer.create(0);
-  Buffer.add_string(buf, formatted_code);
-
-  let contents = Buffer.to_bytes(buf);
   switch (output) {
   | Some(outfile) =>
     let outfile = Filepath.to_string(outfile);
     // TODO: This crashes if you do something weird like `-o stdout/map.gr/foo`
     // because `foo` doesn't exist so it tries to mkdir it and raises
     Fs_access.ensure_parent_directory_exists(outfile);
     let oc = Fs_access.open_file_for_writing(outfile);
-    output_bytes(oc, contents);
+    set_binary_mode_out(oc, true);
+    Grain_formatting.Fmt.format(
+      ~write=output_string(oc),
+      ~source,
+      ~eol,
+      program,
+    );
     close_out(oc);
   | None =>
     set_binary_mode_out(stdout, true);
-    print_bytes(contents);
+    Grain_formatting.Fmt.format(~write=print_string, ~source, ~eol, program);
+    flush(stdout);
   };
 };
 
@@ -145,8 +145,8 @@ let enumerate_runs = opts =>
 let grainformat = runs => {
   List.iter(
     ({input_path, output_path}) => {
-      let (program, original_source, eol) = compile_parsed(input_path);
-      try(format_code(~eol, ~output=?output_path, ~original_source, program)) {
+      let (program, source, eol) = compile_parsed(input_path);
+      try(format_code(~eol, ~output=?output_path, ~source, program)) {
       | exn =>
         Stdlib.Format.eprintf("@[%s@]@.", Printexc.to_string(exn));
         exit(2);

compiler/src/diagnostics/commenttree.re

@@ -0,0 +1,74 @@
+open Grain_parsing;
+
+// This structure isn't a tree at all, but we use a binary search algorithm to
+// efficiently find comments, which is tree-like in spirit.
+
+type t = {
+  comments: array(Parsetree.comment),
+  line_map: Hashtbl.t(int, Parsetree.comment),
+};
+
+let empty: t = {comments: [||], line_map: Hashtbl.create(0)};
+
+let loc = cmt => {
+  switch (cmt) {
+  | Parsetree.Doc({cmt_loc})
+  | Block({cmt_loc})
+  | Line({cmt_loc})
+  | Shebang({cmt_loc}) => cmt_loc
+  };
+};
+
+let from_comments = x => {
+  // The array allows us to do a binary search for comments within a range.
+  let comments = Array.of_list(x);
+  // This map stores the last comment on a line, allowing us to quickly check
+  // for formatter-ignore comments.
+  let line_map = Hashtbl.create(Array.length(comments));
+  List.iter(
+    comment =>
+      Hashtbl.add(line_map, loc(comment).loc_start.pos_lnum, comment),
+    x,
+  );
+  {comments, line_map};
+};
+
+let rec find_start_index = (array, point, ans, left, right) =>
+  if (left <= right) {
+    let middle = (left + right) / 2;
+    if (loc(array[middle]).loc_start.pos_cnum >= point) {
+      find_start_index(array, point, Some(middle), left, middle - 1);
+    } else {
+      find_start_index(array, point, ans, middle + 1, right);
+    };
+  } else {
+    ans;
+  };
+
+let rec collect_range = (array, start, stop) =>
+  if (start == Array.length(array)) {
+    [];
+  } else {
+    let elem = array[start];
+    if (loc(elem).loc_end.pos_cnum <= stop) {
+      [elem, ...collect_range(array, start + 1, stop)];
+    } else {
+      [];
+    };
+  };
+
+let query =
+    (
+      tree,
+      {Location.loc_start: {pos_cnum: start}, loc_end: {pos_cnum: finish}},
+    ) => {
+  let array = tree.comments;
+  let start_index =
+    find_start_index(array, start, None, 0, Array.length(array) - 1);
+  switch (start_index) {
+  | None => []
+  | Some(start_index) => collect_range(array, start_index, finish)
+  };
+};
+
+let query_line = (tree, line) => Hashtbl.find_opt(tree.line_map, line);

compiler/src/diagnostics/commenttree.rei

@@ -0,0 +1,10 @@
+open Grain_parsing;
+
+type t;
+
+let empty: t;
+
+let from_comments: list(Parsetree.comment) => t;
+
+let query: (t, Location.t) => list(Parsetree.comment);
+let query_line: (t, int) => option(Parsetree.comment);