Update *.wast parser with recently added directives (#1762)

* Update `*.wast` parser with recently added directives

Parsing them is a bit wonky because it looks like this crate isn't
architected the same way as the reference interpreter, but this so far
otherwise seems to work. I've copied the upstream test added in
WebAssembly/spec#1796 here for reference.

* Fix compile of benches

* Minor fixes

crates/wasmparser/benches/benchmark.rs

@@ -58,7 +58,8 @@ fn collect_test_files(path: &Path, list: &mut Vec<BenchmarkInput>) -> Result<()>
                 };
                 for directive in wast.directives {
                     match directive {
-                        wast::WastDirective::Wat(mut module) => {
+                        wast::WastDirective::Module(mut module)
+                        | wast::WastDirective::ModuleDefinition(mut module) => {
                             let wasm = module.encode()?;
                             list.push(BenchmarkInput::new(path.clone(), wasm));
                         }

crates/wast/src/component/component.rs

@@ -103,16 +103,11 @@ impl<'a> Component<'a> {
         }
         Ok(())
     }
-}
-
-impl<'a> Parse<'a> for Component<'a> {
-    fn parse(parser: Parser<'a>) -> Result<Self> {
-        let _r = parser.register_annotation("custom");
-        let _r = parser.register_annotation("producers");
-        let _r = parser.register_annotation("name");
-        let _r = parser.register_annotation("metadata.code.branch_hint");
 
-        let span = parser.parse::<kw::component>()?.0;
+    pub(crate) fn parse_without_component_keyword(
+        component_keyword_span: Span,
+        parser: Parser<'a>,
+    ) -> Result<Self> {
         let id = parser.parse()?;
         let name = parser.parse()?;
 
@@ -127,14 +122,23 @@ impl<'a> Parse<'a> for Component<'a> {
             ComponentKind::Text(ComponentField::parse_remaining(parser)?)
         };
         Ok(Component {
-            span,
+            span: component_keyword_span,
             id,
             name,
             kind,
         })
     }
 }
 
+impl<'a> Parse<'a> for Component<'a> {
+    fn parse(parser: Parser<'a>) -> Result<Self> {
+        parser.with_standard_annotations_registered(|parser| {
+            let span = parser.parse::<kw::component>()?.0;
+            Component::parse_without_component_keyword(span, parser)
+        })
+    }
+}
+
 /// A listing of all possible fields that can make up a WebAssembly component.
 #[allow(missing_docs)]
 #[derive(Debug)]

crates/wast/src/core/module.rs

@@ -103,17 +103,11 @@ impl<'a> Module<'a> {
         }
         Ok(())
     }
-}
 
-impl<'a> Parse<'a> for Module<'a> {
-    fn parse(parser: Parser<'a>) -> Result<Self> {
-        let _r = parser.register_annotation("custom");
-        let _r = parser.register_annotation("producers");
-        let _r = parser.register_annotation("name");
-        let _r = parser.register_annotation("dylink.0");
-        let _r = parser.register_annotation("metadata.code.branch_hint");
-
-        let span = parser.parse::<kw::module>()?.0;
+    pub(crate) fn parse_without_module_keyword(
+        module_keyword_span: Span,
+        parser: Parser<'a>,
+    ) -> Result<Self> {
         let id = parser.parse()?;
         let name = parser.parse()?;
 
@@ -128,14 +122,23 @@ impl<'a> Parse<'a> for Module<'a> {
             ModuleKind::Text(ModuleField::parse_remaining(parser)?)
         };
         Ok(Module {
-            span,
+            span: module_keyword_span,
             id,
             name,
             kind,
         })
     }
 }
 
+impl<'a> Parse<'a> for Module<'a> {
+    fn parse(parser: Parser<'a>) -> Result<Self> {
+        parser.with_standard_annotations_registered(|parser| {
+            let span = parser.parse::<kw::module>()?.0;
+            Self::parse_without_module_keyword(span, parser)
+        })
+    }
+}
+
 /// A listing of all possible fields that can make up a WebAssembly module.
 #[allow(missing_docs)]
 #[derive(Debug)]

crates/wast/src/lib.rs

@@ -544,6 +544,7 @@ pub mod kw {
     custom_keyword!(import_info = "import-info");
     custom_keyword!(thread);
     custom_keyword!(wait);
+    custom_keyword!(definition);
 }
 
 /// Common annotations used to parse WebAssembly text files.

crates/wast/src/parser.rs

@@ -993,6 +993,19 @@ impl<'a> Parser<'a> {
     pub(crate) fn track_instr_spans(&self) -> bool {
         self.buf.track_instr_spans
     }
+
+    #[cfg(feature = "wasm-module")]
+    pub(crate) fn with_standard_annotations_registered<R>(
+        self,
+        f: impl FnOnce(Self) -> Result<R>,
+    ) -> Result<R> {
+        let _r = self.register_annotation("custom");
+        let _r = self.register_annotation("producers");
+        let _r = self.register_annotation("name");
+        let _r = self.register_annotation("dylink.0");
+        let _r = self.register_annotation("metadata.code.branch_hint");
+        f(self)
+    }
 }
 
 impl<'a> Cursor<'a> {

crates/wast/src/wast.rs

@@ -20,17 +20,19 @@ impl<'a> Parse<'a> for Wast<'a> {
     fn parse(parser: Parser<'a>) -> Result<Self> {
         let mut directives = Vec::new();
 
-        // If it looks like a directive token is in the stream then we parse a
-        // bunch of directives, otherwise assume this is an inline module.
-        if parser.peek2::<WastDirectiveToken>()? {
-            while !parser.is_empty() {
-                directives.push(parser.parens(|p| p.parse())?);
+        parser.with_standard_annotations_registered(|parser| {
+            // If it looks like a directive token is in the stream then we parse a
+            // bunch of directives, otherwise assume this is an inline module.
+            if parser.peek2::<WastDirectiveToken>()? {
+                while !parser.is_empty() {
+                    directives.push(parser.parens(|p| p.parse())?);
+                }
+            } else {
+                let module = parser.parse::<Wat>()?;
+                directives.push(WastDirective::Module(QuoteWat::Wat(module)));
             }
-        } else {
-            let module = parser.parse::<Wat>()?;
-            directives.push(WastDirective::Wat(QuoteWat::Wat(module)));
-        }
-        Ok(Wast { directives })
+            Ok(Wast { directives })
+        })
     }
 }
 
@@ -56,67 +58,103 @@ impl Peek for WastDirectiveToken {
 
 /// The different kinds of directives found in a `*.wast` file.
 ///
-/// It's not entirely clear to me what all of these are per se, but they're only
-/// really interesting to test harnesses mostly.
+///
+/// Some more information about these various branches can be found at
+/// <https://github.com/WebAssembly/spec/blob/main/interpreter/README.md#scripts>.
 #[allow(missing_docs)]
 #[derive(Debug)]
 pub enum WastDirective<'a> {
-    Wat(QuoteWat<'a>),
+    /// The provided module is defined, validated, and then instantiated.
+    Module(QuoteWat<'a>),
+
+    /// The provided module is defined and validated.
+    ///
+    /// This module is not instantiated automatically.
+    ModuleDefinition(QuoteWat<'a>),
+
+    /// The named module is instantiated under the instance name provided.
+    ModuleInstance {
+        span: Span,
+        instance: Option<Id<'a>>,
+        module: Option<Id<'a>>,
+    },
+
+    /// Asserts the module cannot be decoded with the given error.
     AssertMalformed {
         span: Span,
         module: QuoteWat<'a>,
         message: &'a str,
     },
+
+    /// Asserts the module cannot be validated with the given error.
     AssertInvalid {
         span: Span,
         module: QuoteWat<'a>,
         message: &'a str,
     },
+
+    /// Registers the `module` instance with the given `name` to be available
+    /// for importing in future module instances.
     Register {
         span: Span,
         name: &'a str,
         module: Option<Id<'a>>,
     },
+
+    /// Invokes the specified export.
     Invoke(WastInvoke<'a>),
+
+    /// The invocation provided should trap with the specified error.
     AssertTrap {
         span: Span,
         exec: WastExecute<'a>,
         message: &'a str,
     },
+
+    /// The invocation provided should succeed with the specified results.
     AssertReturn {
         span: Span,
         exec: WastExecute<'a>,
         results: Vec<WastRet<'a>>,
     },
+
+    /// The invocation provided should exhaust system resources (e.g. stack
+    /// overflow).
     AssertExhaustion {
         span: Span,
         call: WastInvoke<'a>,
         message: &'a str,
     },
+
+    /// The provided module should fail to link when instantiation is attempted.
     AssertUnlinkable {
         span: Span,
         module: Wat<'a>,
         message: &'a str,
     },
-    AssertException {
-        span: Span,
-        exec: WastExecute<'a>,
-    },
+
+    /// The invocation provided should throw an exception.
+    AssertException { span: Span, exec: WastExecute<'a> },
+
+    /// Creates a new system thread which executes the given commands.
     Thread(WastThread<'a>),
-    Wait {
-        span: Span,
-        thread: Id<'a>,
-    },
+
+    /// Waits for the specified thread to exit.
+    Wait { span: Span, thread: Id<'a> },
 }
 
 impl WastDirective<'_> {
     /// Returns the location in the source that this directive was defined at
     pub fn span(&self) -> Span {
         match self {
-            WastDirective::Wat(QuoteWat::Wat(w)) => w.span(),
-            WastDirective::Wat(QuoteWat::QuoteModule(span, _)) => *span,
-            WastDirective::Wat(QuoteWat::QuoteComponent(span, _)) => *span,
-            WastDirective::AssertMalformed { span, .. }
+            WastDirective::Module(QuoteWat::Wat(w))
+            | WastDirective::ModuleDefinition(QuoteWat::Wat(w)) => w.span(),
+            WastDirective::Module(QuoteWat::QuoteModule(span, _))
+            | WastDirective::ModuleDefinition(QuoteWat::QuoteModule(span, _)) => *span,
+            WastDirective::Module(QuoteWat::QuoteComponent(span, _))
+            | WastDirective::ModuleDefinition(QuoteWat::QuoteComponent(span, _)) => *span,
+            WastDirective::ModuleInstance { span, .. }
+            | WastDirective::AssertMalformed { span, .. }
             | WastDirective::Register { span, .. }
             | WastDirective::AssertTrap { span, .. }
             | WastDirective::AssertReturn { span, .. }
@@ -135,7 +173,7 @@ impl<'a> Parse<'a> for WastDirective<'a> {
     fn parse(parser: Parser<'a>) -> Result<Self> {
         let mut l = parser.lookahead1();
         if l.peek::<kw::module>()? || l.peek::<kw::component>()? {
-            Ok(WastDirective::Wat(parser.parse()?))
+            parse_wast_module(parser)
         } else if l.peek::<kw::assert_malformed>()? {
             let span = parser.parse::<kw::assert_malformed>()?.0;
             Ok(WastDirective::AssertMalformed {
@@ -296,6 +334,52 @@ impl<'a> Parse<'a> for WastInvoke<'a> {
     }
 }
 
+fn parse_wast_module<'a>(parser: Parser<'a>) -> Result<WastDirective<'a>> {
+    if parser.peek2::<kw::quote>()? {
+        QuoteWat::parse(parser).map(WastDirective::Module)
+    } else if parser.peek2::<kw::definition>()? {
+        fn parse_module(span: Span, parser: Parser<'_>) -> Result<Wat<'_>> {
+            Ok(Wat::Module(
+                crate::core::Module::parse_without_module_keyword(span, parser)?,
+            ))
+        }
+        fn parse_component(span: Span, parser: Parser<'_>) -> Result<Wat<'_>> {
+            Ok(Wat::Component(
+                crate::component::Component::parse_without_component_keyword(span, parser)?,
+            ))
+        }
+        let (span, ctor) = if parser.peek::<kw::component>()? {
+            (
+                parser.parse::<kw::component>()?.0,
+                parse_component as fn(_, _) -> _,
+            )
+        } else {
+            (
+                parser.parse::<kw::module>()?.0,
+                parse_module as fn(_, _) -> _,
+            )
+        };
+        parser.parse::<kw::definition>()?;
+        Ok(WastDirective::ModuleDefinition(QuoteWat::Wat(ctor(
+            span, parser,
+        )?)))
+    } else if parser.peek2::<kw::instance>()? {
+        let span = if parser.peek::<kw::component>()? {
+            parser.parse::<kw::component>()?.0
+        } else {
+            parser.parse::<kw::module>()?.0
+        };
+        parser.parse::<kw::instance>()?;
+        Ok(WastDirective::ModuleInstance {
+            span,
+            instance: parser.parse()?,
+            module: parser.parse()?,
+        })
+    } else {
+        QuoteWat::parse(parser).map(WastDirective::Module)
+    }
+}
+
 #[allow(missing_docs)]
 #[derive(Debug)]
 pub enum QuoteWat<'a> {
@@ -304,7 +388,7 @@ pub enum QuoteWat<'a> {
     QuoteComponent(Span, Vec<(Span, &'a [u8])>),
 }
 
-impl QuoteWat<'_> {
+impl<'a> QuoteWat<'a> {
     /// Encodes this module to bytes, either by encoding the module directly or
     /// parsing the contents and then encoding it.
     pub fn encode(&mut self) -> Result<Vec<u8>, Error> {
@@ -342,6 +426,15 @@ impl QuoteWat<'_> {
         Ok(QuoteWatTest::Text(ret))
     }
 
+    /// Returns the identifier, if registered, for this module.
+    pub fn name(&self) -> Option<Id<'a>> {
+        match self {
+            QuoteWat::Wat(Wat::Module(m)) => m.id,
+            QuoteWat::Wat(Wat::Component(m)) => m.id,
+            QuoteWat::QuoteModule(..) | QuoteWat::QuoteComponent(..) => None,
+        }
+    }
+
     /// Returns the defining span of this module.
     pub fn span(&self) -> Span {
         match self {