Make attributes like other places, include module attrs in well-forme…

…dness check

compiler/src/compile.re

@@ -127,9 +127,14 @@ let next_state = (~is_root_file=false, {cstate_desc, cstate_filename} as cs) =>
       cleanup();
       Parsed(parsed);
     | Parsed(p) =>
+      let has_attr = name =>
+        List.exists(
+          attr => attr.Asttypes.attr_name.txt == name,
+          p.attributes,
+        );
       Grain_utils.Config.apply_attribute_flags(
-        ~no_pervasives=Option.is_some(p.attributes.no_pervasives),
-        ~runtime_mode=Option.is_some(p.attributes.runtime_mode),
+        ~no_pervasives=has_attr("noPervasives"),
+        ~runtime_mode=has_attr("runtimeMode"),
       );
       if (is_root_file) {
         Grain_utils.Config.set_root_config();

compiler/src/formatting/fmt.re

@@ -3921,12 +3921,8 @@ let print_comment_range =
 };
 
 let print_program = (fmt, parsed_program) => {
-  let attrs =
-    Option.to_list(parsed_program.attributes.no_pervasives)
-    @ Option.to_list(parsed_program.attributes.runtime_mode);
-
   let first_loc =
-    switch (attrs) {
+    switch (parsed_program.attributes) {
     | [fst, ..._] => fst.attr_loc
     | _ => parsed_program.prog_core_loc
     };
@@ -3995,38 +3991,37 @@ let print_program = (fmt, parsed_program) => {
     };
 
   group @@
-  fmt.print_comment_range(
-    fmt,
-    enclosing_start_location(parsed_program.prog_loc),
-    first_loc,
+  concat_map(
+    ~lead=
+      _ =>
+        fmt.print_comment_range(
+          fmt,
+          enclosing_start_location(parsed_program.prog_loc),
+          first_loc,
+        ),
+    ~sep=
+      (prev, next) =>
+        fmt.print_comment_range(
+          fmt,
+          ~none=hardline,
+          ~lead=space,
+          ~trail=hardline,
+          prev.Asttypes.attr_loc,
+          next.attr_loc,
+        ),
+    ~trail=
+      prev =>
+        fmt.print_comment_range(
+          fmt,
+          ~none=hardline,
+          ~lead=space,
+          ~trail=hardline,
+          prev.Asttypes.attr_loc,
+          parsed_program.prog_core_loc,
+        ),
+    ~f=(~final, a) => fmt.print_attribute(fmt, a),
+    parsed_program.attributes,
   )
-  ++ group(
-       concat_map(
-         ~lead=_ => empty,
-         ~sep=
-           (prev, next) =>
-             fmt.print_comment_range(
-               fmt,
-               ~none=hardline,
-               ~lead=space,
-               ~trail=hardline,
-               prev.Asttypes.attr_loc,
-               next.attr_loc,
-             ),
-         ~trail=
-           prev =>
-             fmt.print_comment_range(
-               fmt,
-               ~none=hardline,
-               ~lead=space,
-               ~trail=hardline,
-               prev.Asttypes.attr_loc,
-               parsed_program.prog_core_loc,
-             ),
-         ~f=(~final, a) => fmt.print_attribute(fmt, a),
-         attrs,
-       ),
-     )
   ++ string("module ")
   ++ string(parsed_program.module_name.txt)
   ++ toplevel;

compiler/src/parsing/driver.re

@@ -150,6 +150,11 @@ let parse = (~name=?, lexbuf, source): Parsetree.parsed_program => {
 let read_imports = (program: Parsetree.parsed_program) => {
   open Parsetree_iter;
 
+  let module_has_attr = name =>
+    List.exists(
+      attr => attr.Asttypes.attr_name.txt == name,
+      program.attributes,
+    );
   let implicit_opens =
     List.map(
       o => {
@@ -159,9 +164,8 @@ let read_imports = (program: Parsetree.parsed_program) => {
         }
       },
       Grain_utils.Config.with_attribute_flags(
-        ~on_error=_ => (),
-        ~no_pervasives=Option.is_some(program.attributes.no_pervasives),
-        ~runtime_mode=Option.is_some(program.attributes.runtime_mode),
+        ~no_pervasives=module_has_attr("noPervasives"),
+        ~runtime_mode=module_has_attr("runtimeMode"),
         Grain_utils.Config.get_implicit_opens,
       ),
     );

compiler/src/parsing/parser_header.re

@@ -121,36 +121,9 @@ let make_program = (~loc, ~core_loc, ~attributes, module_name, statements) => {
   // there's whitespace or comments
   let loc_start = {...loc.loc_start, pos_lnum: 1, pos_cnum: 0, pos_bol: 0};
   let prog_loc = {...loc, loc_start};
-  let (no_pervasives, runtime_mode) =
-    List.fold_left(
-      ((no_pervasives, runtime_mode), attr) => {
-        switch (attr) {
-        | Asttypes.{attr_name: {txt: "noPervasives"}, attr_args: []} => (
-            Some(attr),
-            runtime_mode,
-          )
-        | {attr_name: {txt: "runtimeMode"}, attr_args: []} => (
-            no_pervasives,
-            Some(attr),
-          )
-        | {attr_name: {loc}} =>
-          raise(
-            SyntaxError(
-              loc,
-              "@noPervasives and @runtimeMode are the only valid top-level module attributes.",
-            ),
-          )
-        }
-      },
-      (None, None),
-      attributes,
-    );
 
   fix_blocks({
-    attributes: {
-      no_pervasives,
-      runtime_mode,
-    },
+    attributes,
     module_name,
     statements,
     comments: [],

compiler/src/parsing/parsetree.re

@@ -501,6 +501,11 @@ type attribute = Asttypes.attribute;
 [@deriving (sexp, yojson)]
 type attributes = Asttypes.attributes;
 
+type attribute_context =
+  | ModuleAttribute
+  | ToplevelAttribute
+  | ExpressionAttribute;
+
 /** Type for expressions (i.e. things which evaluate to something) */
 
 [@deriving (sexp, yojson)]
@@ -691,15 +696,9 @@ type comment =
 
 /** The type for parsed programs */
 
-[@deriving (sexp, yojson)]
-type module_attributes = {
-  no_pervasives: option(attribute),
-  runtime_mode: option(attribute),
-};
-
 [@deriving (sexp, yojson)]
 type parsed_program = {
-  attributes: module_attributes,
+  attributes,
   module_name: loc(string),
   statements: list(toplevel_stmt),
   comments: list(comment),

compiler/src/parsing/parsetree_iter.re

@@ -3,8 +3,8 @@ open Parsetree;
 type hooks = {
   enter_location: Location.t => unit,
   leave_location: Location.t => unit,
-  enter_attribute: attribute => unit,
-  leave_attribute: attribute => unit,
+  enter_attribute: (attribute, attribute_context) => unit,
+  leave_attribute: (attribute, attribute_context) => unit,
   enter_parsed_program: parsed_program => unit,
   leave_parsed_program: parsed_program => unit,
   enter_include: include_declaration => unit,
@@ -73,20 +73,21 @@ let iter_ident = (hooks, id) => {
 };
 
 let iter_attribute =
-    (hooks, {Asttypes.attr_name, attr_args, attr_loc} as attr) => {
-  hooks.enter_attribute(attr);
+    (hooks, attr_context, {Asttypes.attr_name, attr_args, attr_loc} as attr) => {
+  hooks.enter_attribute(attr, attr_context);
   iter_loc(hooks, attr_name);
   List.iter(iter_loc(hooks), attr_args);
   iter_location(hooks, attr_loc);
-  hooks.leave_attribute(attr);
+  hooks.leave_attribute(attr, attr_context);
 };
 
-let iter_attributes = (hooks, attrs) => {
-  List.iter(iter_attribute(hooks), attrs);
+let iter_attributes = (hooks, attr_context, attrs) => {
+  List.iter(iter_attribute(hooks, attr_context), attrs);
 };
 
 let rec iter_parsed_program = (hooks, {statements} as program) => {
   hooks.enter_parsed_program(program);
+  iter_attributes(hooks, ModuleAttribute, program.attributes);
   iter_toplevel_stmts(hooks, statements);
   hooks.leave_parsed_program(program);
 }
@@ -108,7 +109,7 @@ and iter_toplevel_stmt =
   hooks.enter_toplevel_stmt(top);
   iter_location(hooks, loc);
   iter_location(hooks, core_loc);
-  iter_attributes(hooks, attrs);
+  iter_attributes(hooks, ToplevelAttribute, attrs);
   switch (desc) {
   | PTopInclude(id) => iter_include(hooks, id)
   | PTopProvide(ex) => iter_provide(hooks, ex)
@@ -254,7 +255,7 @@ and iter_expression =
   hooks.enter_expression(expr);
   iter_location(hooks, loc);
   iter_location(hooks, core_loc);
-  iter_attributes(hooks, attrs);
+  iter_attributes(hooks, ExpressionAttribute, attrs);
   switch (desc) {
   | PExpId(i) => iter_ident(hooks, i)
   | PExpConstant(c) => iter_constant(hooks, c)
@@ -524,8 +525,8 @@ let default_hooks = {
   enter_location: _ => (),
   leave_location: _ => (),
 
-  enter_attribute: _ => (),
-  leave_attribute: _ => (),
+  enter_attribute: (_, _) => (),
+  leave_attribute: (_, _) => (),
 
   enter_parsed_program: _ => (),
   leave_parsed_program: _ => (),

compiler/src/parsing/parsetree_iter.rei

@@ -3,8 +3,8 @@ open Parsetree;
 type hooks = {
   enter_location: Location.t => unit,
   leave_location: Location.t => unit,
-  enter_attribute: attribute => unit,
-  leave_attribute: attribute => unit,
+  enter_attribute: (attribute, attribute_context) => unit,
+  leave_attribute: (attribute, attribute_context) => unit,
   enter_parsed_program: parsed_program => unit,
   leave_parsed_program: parsed_program => unit,
   enter_include: include_declaration => unit,

compiler/src/parsing/well_formedness.re

@@ -11,7 +11,7 @@ type wferr =
   | RHSLetRecMayOnlyBeFunction(Location.t)
   | NoLetRecMut(Location.t)
   | RationalZeroDenominator(Location.t)
-  | UnknownAttribute(string, Location.t)
+  | UnknownAttribute(string, string, Location.t)
   | InvalidAttributeArity(string, int, Location.t)
   | AttributeDisallowed(string, Location.t)
   | LoopControlOutsideLoop(string, Location.t)
@@ -49,8 +49,8 @@ let prepare_error =
         errorf(~loc, "let rec may not be used with the `mut` keyword.")
       | RationalZeroDenominator(loc) =>
         errorf(~loc, "Rational numbers may not have a denominator of zero.")
-      | UnknownAttribute(attr, loc) =>
-        errorf(~loc, "Unknown attribute `%s`.", attr)
+      | UnknownAttribute(attr_context, attr, loc) =>
+        errorf(~loc, "Unknown %s attribute `%s`.", attr_context, attr)
       | InvalidAttributeArity(attr, arity, loc) =>
         switch (arity) {
         | 0 => errorf(~loc, "Attribute `%s` expects no arguments.", attr)
@@ -300,22 +300,40 @@ type known_attribute = {
   arity: int,
 };
 
-let known_attributes = [
-  {name: "disableGC", arity: 0},
-  {name: "unsafe", arity: 0},
-  {name: "externalName", arity: 1},
-];
-
 let valid_attributes = (errs, super) => {
   let enter_attribute =
-      ({Asttypes.attr_name: {txt, loc}, attr_args: args} as attr) => {
+      (
+        {Asttypes.attr_name: {txt, loc}, attr_args: args} as attr,
+        attr_context,
+      ) => {
+    let known_attributes =
+      switch (attr_context) {
+      | ModuleAttribute => [
+          {name: "runtimeMode", arity: 0},
+          {name: "noPervasives", arity: 0},
+        ]
+      | ToplevelAttribute
+      | ExpressionAttribute => [
+          {name: "disableGC", arity: 0},
+          {name: "unsafe", arity: 0},
+          {name: "externalName", arity: 1},
+        ]
+      };
+
     switch (List.find_opt(({name}) => name == txt, known_attributes)) {
     | Some({arity}) when List.length(args) != arity =>
       errs := [InvalidAttributeArity(txt, arity, loc), ...errs^]
-    | None => errs := [UnknownAttribute(txt, loc), ...errs^]
+    | None =>
+      let context_string =
+        switch (attr_context) {
+        | ModuleAttribute => "module"
+        | ToplevelAttribute => "top-level"
+        | ExpressionAttribute => "expression"
+        };
+      errs := [UnknownAttribute(context_string, txt, loc), ...errs^];
     | _ => ()
     };
-    super.enter_attribute(attr);
+    super.enter_attribute(attr, attr_context);
   };
 
   {

compiler/src/utils/config.re

@@ -533,7 +533,7 @@ let apply_attribute_flags = (~no_pervasives as np, ~runtime_mode as rm) => {
   };
 };
 
-let with_attribute_flags = (~on_error, ~no_pervasives, ~runtime_mode, thunk) => {
+let with_attribute_flags = (~no_pervasives, ~runtime_mode, thunk) => {
   preserve_config(() => {
     apply_attribute_flags(~no_pervasives, ~runtime_mode);
     thunk();

compiler/src/utils/config.rei

@@ -183,13 +183,7 @@ let with_cli_options: 'a => Cmdliner.Term.t('a);
 let apply_attribute_flags: (~no_pervasives: bool, ~runtime_mode: bool) => unit;
 
 let with_attribute_flags:
-  (
-    ~on_error: [> | `Help | `Message(string)] => unit,
-    ~no_pervasives: bool,
-    ~runtime_mode: bool,
-    unit => 'a
-  ) =>
-  'a;
+  (~no_pervasives: bool, ~runtime_mode: bool, unit => 'a) => 'a;
 
 type implicit_opens =
   | Pervasives_mod