fix(compiler): Prevent stack overflows when compiling long blocks (#1534

)

cli/bin/grain.js

@@ -11,8 +11,6 @@ const v8 = require("v8");
  * This seems to work for our needs with Node 16, but we should be cautious when updating.
  */
 v8.setFlagsFromString("--experimental-wasm-return-call");
-// TODO(#1580): Remove this flag
-v8.setFlagsFromString("--stack-size=2048");
 
 const commander = require("commander");
 const exec = require("./exec.js");

compiler/graindoc/graindoc.re

@@ -49,7 +49,11 @@ type params = {
 
 let compile_typed = (input: Fp.t(Fp.absolute)) => {
   switch (
-    Compile.compile_file(~hook=stop_after_typed, Filepath.to_string(input))
+    Compile.compile_file(
+      ~is_root_file=true,
+      ~hook=stop_after_typed,
+      Filepath.to_string(input),
+    )
   ) {
   | exception exn =>
     let bt =

compiler/src/compile.re

@@ -14,6 +14,7 @@ type compilation_state_desc =
   | Initial(input_source)
   | Parsed(Parsetree.parsed_program)
   | WellFormed(Parsetree.parsed_program)
+  | DependenciesCompiled(Parsetree.parsed_program)
   | TypeChecked(Typedtree.typed_program)
   | TypedWellFormed(Typedtree.typed_program)
   | Linearized(Anftree.anf_program)
@@ -68,6 +69,7 @@ let log_state = state =>
       prerr_string("\nParsed program:\n");
       prerr_sexp(Grain_parsing.Parsetree.sexp_of_parsed_program, p);
     | WellFormed(_) => prerr_string("\nWell-Formedness passed")
+    | DependenciesCompiled(_) => prerr_string("\nDependencies compiled")
     | TypeChecked(typed_mod) =>
       prerr_string("\nTyped program:\n");
       prerr_sexp(Grain_typed.Typedtree.sexp_of_typed_program, typed_mod);
@@ -153,7 +155,16 @@ let next_state = (~is_root_file=false, {cstate_desc, cstate_filename} as cs) =>
       };
       Well_formedness.check_well_formedness(p);
       WellFormed(p);
-    | WellFormed(p) => TypeChecked(Typemod.type_implementation(p))
+    | WellFormed(p) =>
+      if (is_root_file) {
+        let base_file = Option.value(~default="", cstate_filename);
+        Module_resolution.compile_dependency_graph(
+          ~base_file,
+          Driver.read_imports(p),
+        );
+      };
+      DependenciesCompiled(p);
+    | DependenciesCompiled(p) => TypeChecked(Typemod.type_implementation(p))
     | TypeChecked(typed_mod) =>
       Typed_well_formedness.check_well_formedness(typed_mod);
       TypedWellFormed(typed_mod);

compiler/src/compile.rei

@@ -12,6 +12,7 @@ type compilation_state_desc =
   | Initial(input_source)
   | Parsed(Parsetree.parsed_program)
   | WellFormed(Parsetree.parsed_program)
+  | DependenciesCompiled(Parsetree.parsed_program)
   | TypeChecked(Typedtree.typed_program)
   | TypedWellFormed(Typedtree.typed_program)
   | Linearized(Anftree.anf_program)

compiler/src/language_server/code_file.re

@@ -54,6 +54,7 @@ let compile_source = (uri, source) => {
 
   switch (
     Compile.compile_string(
+      ~is_root_file=true,
       ~hook=stop_after_typed_well_formed,
       ~name=filename,
       source,

compiler/src/middle_end/linearize.re

@@ -955,12 +955,19 @@ and transl_comp_expression =
       cond_setup,
     );
   | TExpBlock([]) => (Comp.imm(~allocation_type, Imm.const(Const_void)), [])
-  | TExpBlock([stmt]) => transl_comp_expression(stmt)
-  | TExpBlock([fst, ...rest]) =>
-    let (fst_ans, fst_setup) = transl_comp_expression(fst);
-    let (rest_ans, rest_setup) =
-      transl_comp_expression({...e, exp_desc: TExpBlock(rest)});
-    (rest_ans, fst_setup @ [BSeq(fst_ans)] @ rest_setup);
+  | TExpBlock(stmts) =>
+    let stmts = List.rev_map(transl_comp_expression, stmts);
+    // stmts is non-empty, so this cannot fail
+    let (last_ans, last_setup) = List.hd(stmts);
+    let stmts = List.tl(stmts);
+    let setup =
+      List.concat @@
+      List.fold_left(
+        (acc, (ans, setup)) => {[setup, [BSeq(ans)], ...acc]},
+        [last_setup],
+        stmts,
+      );
+    (last_ans, setup);
   | TExpLet(Nonrecursive, _, []) => (
       Comp.imm(~allocation_type, Imm.const(Const_void)),
       [],

compiler/src/parsing/driver.re

@@ -126,6 +126,37 @@ let parse = (~name=?, lexbuf, source): Parsetree.parsed_program => {
   };
 };
 
+let read_imports = ({Parsetree.comments, Parsetree.statements}) => {
+  let implicit_opens =
+    List.map(
+      o => {
+        switch (o) {
+        | Grain_utils.Config.Pervasives_mod => Location.mknoloc("pervasives")
+        | Grain_utils.Config.Gc_mod => Location.mknoloc("runtime/gc")
+        }
+      },
+      switch (comments) {
+      | [Block({cmt_content}), ..._] =>
+        Grain_utils.Config.with_inline_flags(
+          ~on_error=_ => (),
+          cmt_content,
+          Grain_utils.Config.get_implicit_opens,
+        )
+      | _ => Grain_utils.Config.get_implicit_opens()
+      },
+    );
+  let found_imports = ref([]);
+  let iter_mod = (self, import) =>
+    found_imports := [import.Parsetree.pimp_path, ...found_imports^];
+  let iterator = {...Ast_iterator.default_iterator, import: iter_mod};
+  List.iter(iterator.toplevel(iterator), statements);
+
+  List.sort_uniq(
+    (a, b) => String.compare(a.txt, b.txt),
+    List.append(implicit_opens, found_imports^),
+  );
+};
+
 let scan_for_imports =
     (~defer_errors=true, filename: string): list(loc(string)) => {
   let ic = open_in(filename);
@@ -137,37 +168,9 @@ let scan_for_imports =
       close_in(ic);
       source;
     };
-    let {Parsetree.comments, Parsetree.statements} =
-      parse(~name=filename, lexbuf, source);
-    let implicit_opens =
-      List.map(
-        o => {
-          switch (o) {
-          | Grain_utils.Config.Pervasives_mod =>
-            Location.mknoloc("pervasives")
-          | Grain_utils.Config.Gc_mod => Location.mknoloc("runtime/gc")
-          }
-        },
-        switch (comments) {
-        | [Block({cmt_content}), ..._] =>
-          Grain_utils.Config.with_inline_flags(
-            ~on_error=_ => (),
-            cmt_content,
-            Grain_utils.Config.get_implicit_opens,
-          )
-        | _ => Grain_utils.Config.get_implicit_opens()
-        },
-      );
-    let found_imports = ref([]);
-    let iter_mod = (self, import) =>
-      found_imports := [import.Parsetree.pimp_path, ...found_imports^];
-    let iterator = {...Ast_iterator.default_iterator, import: iter_mod};
-    List.iter(iterator.toplevel(iterator), statements);
+    let prog = parse(~name=filename, lexbuf, source);
     close_in(ic);
-    List.sort_uniq(
-      (a, b) => String.compare(a.txt, b.txt),
-      List.append(implicit_opens, found_imports^),
-    );
+    read_imports(prog);
   }) {
   | e =>
     close_in(ic);

compiler/src/parsing/driver.rei

@@ -4,5 +4,7 @@ let parse:
   (~name: string=?, Sedlexing.lexbuf, unit => string) =>
   Parsetree.parsed_program;
 
+let read_imports: Parsetree.parsed_program => list(Location.loc(string));
+
 let scan_for_imports:
   (~defer_errors: bool=?, string) => list(Location.loc(string));

compiler/src/typed/dependency_graph.re

@@ -133,15 +133,11 @@ module Make = (DV: Dependency_value) => {
     ret;
   };
 
-  let compile_dependencies = (~loc=?, filename) => {
-    switch (lookup_filename(filename)) {
-    | None => raise(Not_found)
-    | Some(vtx) =>
-      let to_compile = ref(solve_next_out_of_date(~stop=vtx, ()));
-      while (Option.is_some(to_compile^)) {
-        DV.compile_module(~loc?, Option.get(to_compile^));
-        to_compile := solve_next_out_of_date(~stop=vtx, ());
-      };
+  let compile_graph = () => {
+    let to_compile = ref(solve_next_out_of_date());
+    while (Option.is_some(to_compile^)) {
+      DV.compile_module(Option.get(to_compile^));
+      to_compile := solve_next_out_of_date();
     };
   };
 

compiler/src/typed/dependency_graph.rei

@@ -26,10 +26,9 @@ module Make:
     let lookup_filename: string => option(DV.t);
 
     /**
-   Compiles the dependencies of the given filename.
+   Compiles the full dependency graph.
    */
-    let compile_dependencies:
-      (~loc: Grain_parsing.Location.t=?, string) => unit;
+    let compile_graph: unit => unit;
 
     /**
    Dumps the edges in this graph to stderr.

compiler/src/typed/module_resolution.re

@@ -454,17 +454,24 @@ module Dependency_graph =
   });
 
 let locate_module_file = (~loc, ~disable_relpath=false, unit_name) => {
-  /* NOTE: We need to take care here to *not* wrap get_up_to_date with this try/with, since
-     it will falsely raise No_module_file if a Not_found is raised during the compilation */
   let base_dir = Filepath.String.dirname(current_filename^());
   let path = Config.module_search_path();
   let located =
     try(locate_module(~disable_relpath, base_dir, path, unit_name)) {
     | Not_found => error(No_module_file(loc, unit_name, None))
     };
+  located_to_out_file_name(located);
+};
+
+let process_dependency = (~loc, ~base_file, unit_name) => {
+  let base_dir = Filepath.String.dirname(base_file);
+  let path = Config.module_search_path();
+  let located =
+    try(locate_module(~disable_relpath=false, base_dir, path, unit_name)) {
+    | Not_found => error(No_module_file(loc, unit_name, None))
+    };
   let out_file = located_to_out_file_name(located);
-  let current_dep_node =
-    Dependency_graph.lookup_filename(current_filename^());
+  let current_dep_node = Dependency_graph.lookup_filename(base_file);
   let existing_dependency = Dependency_graph.lookup_filename(out_file);
   let dn =
     switch (existing_dependency) {
@@ -482,9 +489,16 @@ let locate_module_file = (~loc, ~disable_relpath=false, unit_name) => {
       };
     };
   Dependency_graph.register(dn);
-  Dependency_graph.compile_dependencies(~loc, out_file);
-  let ret = located_to_out_file_name(located);
-  ret;
+};
+
+let compile_dependency_graph = (~base_file, dependencies) => {
+  open Location;
+  List.iter(
+    ({txt: dependency, loc}) =>
+      process_dependency(~loc, ~base_file, dependency),
+    dependencies,
+  );
+  Dependency_graph.compile_graph();
 };
 
 let clear_dependency_graph = () => {

compiler/src/typed/module_resolution.rei

@@ -15,6 +15,9 @@ let resolve_unit:
 
 let compile_module_dependency: ref((string, string) => unit);
 
+let compile_dependency_graph:
+  (~base_file: string, list(Grain_parsing.Location.loc(string))) => unit;
+
 let read_file_cmi: string => Cmi_format.cmi_infos;
 
 let clear_dependency_graph: unit => unit;