diff --git a/compiler/src/formatting/fmt.re b/compiler/src/formatting/fmt.re
index 163a6f010..5d7938eb0 100644
--- a/compiler/src/formatting/fmt.re
+++ b/compiler/src/formatting/fmt.re
@@ -3928,12 +3928,6 @@ let print_comment_range =
 };
 
 let print_program = (fmt, parsed_program) => {
-  let first_loc =
-    switch (parsed_program.attributes) {
-    | [fst, ..._] => fst.attr_loc
-    | _ => parsed_program.prog_core_loc
-    };
-
   let toplevel =
     switch (parsed_program.statements) {
     | [] =>
@@ -4000,11 +3994,11 @@ let print_program = (fmt, parsed_program) => {
   group @@
   concat_map(
     ~lead=
-      _ =>
+      first =>
         fmt.print_comment_range(
           fmt,
           enclosing_start_location(parsed_program.prog_loc),
-          first_loc,
+          first.attr_loc,
         ),
     ~sep=
       (prev, next) =>
diff --git a/compiler/src/parsing/parsetree.re b/compiler/src/parsing/parsetree.re
index a500e69b8..b6f4bee5d 100644
--- a/compiler/src/parsing/parsetree.re
+++ b/compiler/src/parsing/parsetree.re
@@ -501,11 +501,6 @@ 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)]
diff --git a/compiler/src/parsing/parsetree_iter.re b/compiler/src/parsing/parsetree_iter.re
index 57794ccf9..d46f9a0c1 100644
--- a/compiler/src/parsing/parsetree_iter.re
+++ b/compiler/src/parsing/parsetree_iter.re
@@ -3,8 +3,6 @@ open Parsetree;
 type hooks = {
   enter_location: Location.t => unit,
   leave_location: Location.t => 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,21 +71,19 @@ let iter_ident = (hooks, id) => {
 };
 
 let iter_attribute =
-    (hooks, attr_context, {Asttypes.attr_name, attr_args, attr_loc} as attr) => {
-  hooks.enter_attribute(attr, attr_context);
+    (hooks, {Asttypes.attr_name, attr_args, attr_loc} as attr) => {
   iter_loc(hooks, attr_name);
   List.iter(iter_loc(hooks), attr_args);
   iter_location(hooks, attr_loc);
-  hooks.leave_attribute(attr, attr_context);
 };
 
-let iter_attributes = (hooks, attr_context, attrs) => {
-  List.iter(iter_attribute(hooks, attr_context), attrs);
+let iter_attributes = (hooks, attrs) => {
+  List.iter(iter_attribute(hooks), attrs);
 };
 
 let rec iter_parsed_program = (hooks, {statements} as program) => {
   hooks.enter_parsed_program(program);
-  iter_attributes(hooks, ModuleAttribute, program.attributes);
+  iter_attributes(hooks, program.attributes);
   iter_toplevel_stmts(hooks, statements);
   hooks.leave_parsed_program(program);
 }
@@ -109,7 +105,7 @@ and iter_toplevel_stmt =
   hooks.enter_toplevel_stmt(top);
   iter_location(hooks, loc);
   iter_location(hooks, core_loc);
-  iter_attributes(hooks, ToplevelAttribute, attrs);
+  iter_attributes(hooks, attrs);
   switch (desc) {
   | PTopInclude(id) => iter_include(hooks, id)
   | PTopProvide(ex) => iter_provide(hooks, ex)
@@ -255,7 +251,7 @@ and iter_expression =
   hooks.enter_expression(expr);
   iter_location(hooks, loc);
   iter_location(hooks, core_loc);
-  iter_attributes(hooks, ExpressionAttribute, attrs);
+  iter_attributes(hooks, attrs);
   switch (desc) {
   | PExpId(i) => iter_ident(hooks, i)
   | PExpConstant(c) => iter_constant(hooks, c)
@@ -525,9 +521,6 @@ let default_hooks = {
   enter_location: _ => (),
   leave_location: _ => (),
 
-  enter_attribute: (_, _) => (),
-  leave_attribute: (_, _) => (),
-
   enter_parsed_program: _ => (),
   leave_parsed_program: _ => (),
 
diff --git a/compiler/src/parsing/parsetree_iter.rei b/compiler/src/parsing/parsetree_iter.rei
index 55e391533..dcad23605 100644
--- a/compiler/src/parsing/parsetree_iter.rei
+++ b/compiler/src/parsing/parsetree_iter.rei
@@ -3,8 +3,6 @@ open Parsetree;
 type hooks = {
   enter_location: Location.t => unit,
   leave_location: Location.t => 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,
diff --git a/compiler/src/parsing/well_formedness.re b/compiler/src/parsing/well_formedness.re
index 39cbe1f79..59c3b599f 100644
--- a/compiler/src/parsing/well_formedness.re
+++ b/compiler/src/parsing/well_formedness.re
@@ -300,53 +300,28 @@ type known_attribute = {
   arity: int,
 };
 
-let valid_attributes = (errs, super) => {
-  let enter_attribute =
-      (
-        {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 =>
-      let context_string =
-        switch (attr_context) {
-        | ModuleAttribute => "module"
-        | ToplevelAttribute => "top-level"
-        | ExpressionAttribute => "expression"
-        };
-      errs := [UnknownAttribute(context_string, txt, loc), ...errs^];
-    | _ => ()
-    };
-    super.enter_attribute(attr, attr_context);
+let disallowed_attributes = (errs, super) => {
+  let validate_against_known = (attrs, known_attributes, context) => {
+    List.iter(
+      ({Asttypes.attr_name: {txt, loc}, attr_args: args}) => {
+        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(context, txt, loc), ...errs^]
+        | _ => ()
+        }
+      },
+      attrs,
+    );
   };
 
-  {
-    errs,
-    iter_hooks: {
-      ...super,
-      enter_attribute,
-    },
-  };
-};
+  let known_expr_attributes = [
+    {name: "disableGC", arity: 0},
+    {name: "unsafe", arity: 0},
+    {name: "externalName", arity: 1},
+  ];
 
-let disallowed_attributes = (errs, super) => {
-  let enter_expression = ({pexp_desc: desc, pexp_attributes: attrs} as e) => {
+  let enter_expression = ({pexp_attributes: attrs} as e) => {
     switch (
       List.find_opt(
         ({Asttypes.attr_name: {txt}}) => txt == "externalName",
@@ -363,8 +338,10 @@ let disallowed_attributes = (errs, super) => {
         ]
     | None => ()
     };
+    validate_against_known(attrs, known_expr_attributes, "expression");
     super.enter_expression(e);
   };
+
   let enter_toplevel_stmt =
       ({ptop_desc: desc, ptop_attributes: attrs} as top) => {
     switch (
@@ -417,15 +394,26 @@ let disallowed_attributes = (errs, super) => {
       }
     | None => ()
     };
+    validate_against_known(attrs, known_expr_attributes, "top-level");
     super.enter_toplevel_stmt(top);
   };
 
+  let enter_parsed_program = ({attributes} as prog) => {
+    let known_module_attributes = [
+      {name: "runtimeMode", arity: 0},
+      {name: "noPervasives", arity: 0},
+    ];
+    validate_against_known(attributes, known_module_attributes, "module");
+    super.enter_parsed_program(prog);
+  };
+
   {
     errs,
     iter_hooks: {
       ...super,
       enter_expression,
       enter_toplevel_stmt,
+      enter_parsed_program,
     },
   };
 };
@@ -860,7 +848,6 @@ let well_formedness_checks = [
   only_functions_oh_rhs_letrec,
   no_letrec_mut,
   no_zero_denominator_rational,
-  valid_attributes,
   disallowed_attributes,
   no_loop_control_statement_outside_of_loop,
   malformed_return_statements,
diff --git a/compiler/test/runner.re b/compiler/test/runner.re
index af61b95a4..2000d1b02 100644
--- a/compiler/test/runner.re
+++ b/compiler/test/runner.re
@@ -216,7 +216,7 @@ let doc = (file, arguments) => {
 let module_header = "module Test; ";
 
 let makeSnapshotRunner =
-    (~config_fn=?, test, ~default_module_header=true, name, prog) => {
+    (~config_fn=?, test, ~module_header=module_header, name, prog) => {
   test(name, ({expect}) => {
     Config.preserve_all_configs(() => {
       ignore @@
@@ -224,7 +224,7 @@ let makeSnapshotRunner =
         ~hook=stop_after_object_file_emitted,
         ~config_fn?,
         name,
-        (if (default_module_header) {module_header} else {""}) ++ prog,
+        module_header ++ prog,
       );
       expect.file(watfile(name)).toMatchSnapshot();
     })
@@ -232,15 +232,10 @@ let makeSnapshotRunner =
 };
 
 let makeFilesizeRunner =
-    (test, ~config_fn=?, ~default_module_header=true, name, prog, size) => {
+    (test, ~config_fn=?, ~module_header=module_header, name, prog, size) => {
   test(name, ({expect}) => {
     Config.preserve_all_configs(() => {
-      ignore @@
-      compile(
-        ~config_fn?,
-        name,
-        (if (default_module_header) {module_header} else {""}) ++ prog,
-      );
+      ignore @@ compile(~config_fn?, name, module_header ++ prog);
       let ic = open_in_bin(wasmfile(name));
       let filesize = in_channel_length(ic);
       close_in(ic);
@@ -268,18 +263,14 @@ let makeSnapshotFileRunner = (test, ~config_fn=?, name, filename) => {
 };
 
 let makeCompileErrorRunner =
-    (test, ~default_module_header=true, name, prog, msg) => {
+    (test, ~module_header=module_header, name, prog, msg) => {
   test(
     name,
     ({expect}) => {
       let error =
         try(
           {
-            ignore @@
-            compile(
-              name,
-              (if (default_module_header) {module_header} else {""}) ++ prog,
-            );
+            ignore @@ compile(name, module_header ++ prog);
             "";
           }
         ) {
@@ -291,29 +282,21 @@ let makeCompileErrorRunner =
 };
 
 let makeWarningRunner =
-    (test, ~default_module_header=true, name, prog, warning) => {
+    (test, ~module_header=module_header, name, prog, warning) => {
   test(name, ({expect}) => {
     Config.preserve_all_configs(() => {
       Config.print_warnings := false;
-      ignore @@
-      compile(
-        name,
-        (if (default_module_header) {module_header} else {""}) ++ prog,
-      );
+      ignore @@ compile(name, module_header ++ prog);
       expect.ext.warning.toHaveTriggered(warning);
     })
   });
 };
 
-let makeNoWarningRunner = (test, ~default_module_header=true, name, prog) => {
+let makeNoWarningRunner = (test, ~module_header=module_header, name, prog) => {
   test(name, ({expect}) => {
     Config.preserve_all_configs(() => {
       Config.print_warnings := false;
-      ignore @@
-      compile(
-        name,
-        (if (default_module_header) {module_header} else {""}) ++ prog,
-      );
+      ignore @@ compile(name, module_header ++ prog);
       expect.ext.warning.toHaveTriggeredNoWarnings();
     })
   });
@@ -325,20 +308,14 @@ let makeRunner =
       ~num_pages=?,
       ~config_fn=?,
       ~extra_args=?,
-      ~default_module_header=true,
+      ~module_header=module_header,
       name,
       prog,
       expected,
     ) => {
   test(name, ({expect}) => {
     Config.preserve_all_configs(() => {
-      ignore @@
-      compile(
-        ~num_pages?,
-        ~config_fn?,
-        name,
-        (if (default_module_header) {module_header} else {""}) ++ prog,
-      );
+      ignore @@ compile(~num_pages?, ~config_fn?, name, module_header ++ prog);
       let (result, _) = run(~num_pages?, ~extra_args?, wasmfile(name));
       expect.string(result).toEqual(expected);
     })
@@ -351,20 +328,14 @@ let makeErrorRunner =
       ~check_exists=true,
       ~num_pages=?,
       ~config_fn=?,
-      ~default_module_header=true,
+      ~module_header=module_header,
       name,
       prog,
       expected,
     ) => {
   test(name, ({expect}) => {
     Config.preserve_all_configs(() => {
-      ignore @@
-      compile(
-        ~num_pages?,
-        ~config_fn?,
-        name,
-        (if (default_module_header) {module_header} else {""}) ++ prog,
-      );
+      ignore @@ compile(~num_pages?, ~config_fn?, name, module_header ++ prog);
       let (result, _) = run(~num_pages?, wasmfile(name));
       if (check_exists) {
         expect.string(result).toMatch(expected);
diff --git a/compiler/test/suites/functions.re b/compiler/test/suites/functions.re
index 6ade12f80..0d9f0136e 100644
--- a/compiler/test/suites/functions.re
+++ b/compiler/test/suites/functions.re
@@ -155,9 +155,9 @@ describe("functions", ({test, testSkip}) => {
     "Unknown top-level attribute",
   );
   assertCompileError(
-    ~default_module_header=false,
+    ~module_header="@unsafe module Test",
     "unknown_module_attribute",
-    "@unsafe module Test",
+    "",
     "Unknown module attribute",
   );
   assertSnapshot(