feat(wast): add support for map custom keyword and encoding in component model

Signed-off-by: Yordis Prieto <yordis.prieto@gmail.com>

Author: yordis

crates/wasmparser/src/validator/component.rs

@@ -3922,16 +3922,13 @@ impl ComponentState {
             )),
             crate::ComponentDefinedType::Map(key, value) => {
                 if !self.features.cm_map() {
-                    bail!(
-                        offset,
-                        "Maps require the component model map feature"
-                    )
+                    bail!(offset, "Maps require the component model map feature")
                 }
                 Ok(ComponentDefinedType::Map(
                     self.create_component_val_type(key, offset)?,
                     self.create_component_val_type(value, offset)?,
                 ))
-            },
+            }
             crate::ComponentDefinedType::FixedSizeList(ty, elements) => {
                 if !self.features.cm_fixed_size_list() {
                     bail!(

crates/wast/src/component/binary.rs

@@ -109,6 +109,9 @@ fn encode_defined_type(encoder: ComponentDefinedTypeEncoder, ty: &ComponentDefin
         ComponentDefinedType::List(l) => {
             encoder.list(l.element.as_ref());
         }
+        ComponentDefinedType::Map(m) => {
+            encoder.map(m.key.as_ref(), m.value.as_ref());
+        }
         ComponentDefinedType::FixedSizeList(l) => {
             encoder.fixed_size_list(l.element.as_ref(), l.elements);
         }

crates/wast/src/component/expand.rs

@@ -557,6 +557,10 @@ impl<'a> Expander<'a> {
             }) => {
                 self.expand_component_val_ty(t);
             }
+            ComponentDefinedType::Map(Map { key: k, value: v }) => {
+                self.expand_component_val_ty(k);
+                self.expand_component_val_ty(v);
+            }
             ComponentDefinedType::Tuple(t) => {
                 for field in t.fields.iter_mut() {
                     self.expand_component_val_ty(field);

crates/wast/src/component/resolve.rs

@@ -566,6 +566,10 @@ impl<'a> Resolver<'a> {
             }) => {
                 self.component_val_type(t)?;
             }
+            ComponentDefinedType::Map(Map { key: k, value: v }) => {
+                self.component_val_type(k)?;
+                self.component_val_type(v)?;
+            }
             ComponentDefinedType::Tuple(t) => {
                 for field in t.fields.iter_mut() {
                     self.component_val_type(field)?;

crates/wast/src/component/types.rs

@@ -386,6 +386,7 @@ pub enum ComponentDefinedType<'a> {
     Record(Record<'a>),
     Variant(Variant<'a>),
     List(List<'a>),
+    Map(Map<'a>),
     FixedSizeList(FixedSizeList<'a>),
     Tuple(Tuple<'a>),
     Flags(Flags<'a>),
@@ -407,6 +408,8 @@ impl<'a> ComponentDefinedType<'a> {
             Ok(Self::Variant(parser.parse()?))
         } else if l.peek::<kw::list>()? {
             parse_list(parser)
+        } else if l.peek::<kw::map>()? {
+            Ok(Self::Map(parser.parse()?))
         } else if l.peek::<kw::tuple>()? {
             Ok(Self::Tuple(parser.parse()?))
         } else if l.peek::<kw::flags>()? {
@@ -451,6 +454,7 @@ impl Peek for ComponentDefinedType<'_> {
                 Some(("record", _))
                     | Some(("variant", _))
                     | Some(("list", _))
+                    | Some(("map", _))
                     | Some(("tuple", _))
                     | Some(("flags", _))
                     | Some(("enum", _))
@@ -587,6 +591,15 @@ pub struct List<'a> {
     pub element: Box<ComponentValType<'a>>,
 }
 
+/// A map type.
+#[derive(Debug)]
+pub struct Map<'a> {
+    /// The key type of the map.
+    pub key: Box<ComponentValType<'a>>,
+    /// The value type of the map.
+    pub value: Box<ComponentValType<'a>>,
+}
+
 /// A fixed size list type.
 #[derive(Debug)]
 pub struct FixedSizeList<'a> {
@@ -612,6 +625,18 @@ fn parse_list<'a>(parser: Parser<'a>) -> Result<ComponentDefinedType<'a>> {
     }
 }
 
+impl<'a> Parse<'a> for Map<'a> {
+    fn parse(parser: Parser<'a>) -> Result<Self> {
+        parser.parse::<kw::map>()?;
+        let key = parser.parse()?;
+        let value = parser.parse()?;
+        Ok(Self {
+            key: Box::new(key),
+            value: Box::new(value),
+        })
+    }
+}
+
 /// A tuple type.
 #[derive(Debug)]
 pub struct Tuple<'a> {

crates/wast/src/lib.rs

@@ -533,6 +533,7 @@ pub mod kw {
     custom_keyword!(option);
     custom_keyword!(tuple);
     custom_keyword!(list);
+    custom_keyword!(map);
     custom_keyword!(error);
     custom_keyword!(canon);
     custom_keyword!(lift);

tests/cli/component-model/map.wast

@@ -0,0 +1,69 @@
+;; RUN: wast % --assert default --snapshot tests/snapshots -f cm-map
+
+(component
+  (core module $m
+    (memory (export "memory") 1)
+    (func (export "ret-map") (result i32 i32) unreachable)
+  )
+  (core instance $i (instantiate $m))
+
+  (func (export "ret-map") (result (map string u32))
+    (canon lift (core func $i "ret-map") (memory $i "memory"))
+  )
+)
+
+(component
+  (core module $m
+    (func (export "param-map") (param i32 i32) unreachable)
+  )
+  (core instance $i (instantiate $m))
+
+  (func (export "param-map") (param "m" (map string u32))
+    (canon lift (core func $i "param-map"))
+  )
+)
+
+(component
+  (type $map-type (map u32 string))
+  (func (export "f") (param "x" $map-type))
+)
+
+(component
+  (type $nested-map (map string (map string u32)))
+  (func (export "f") (param "x" $nested-map))
+)
+
+(component
+  (type $map-with-list (map string (list u32)))
+  (func (export "f") (param "x" $map-with-list))
+)
+
+(component
+  (type $map-with-option (map u32 (option string)))
+  (func (export "f") (param "x" $map-with-option))
+)
+
+(assert_invalid
+  (component
+    (import "y" (component $c
+      (type $t (map string u32))
+      (import "x" (type (eq $t)))
+    ))
+
+    (type $x (map u32 string))
+    (instance (instantiate $c (with "x" (type $x))))
+  )
+  "type mismatch for import `x`")
+
+(assert_invalid
+  (component
+    (import "y" (component $c
+      (type $t (map string u32))
+      (import "x" (type (eq $t)))
+    ))
+
+    (type $x (list u32))
+    (instance (instantiate $c (with "x" (type $x))))
+  )
+  "type mismatch for import `x`")
+

tests/cli/missing-features/component-model/map.wast

@@ -0,0 +1,39 @@
+;; RUN: wast % --assert default --snapshot tests/snapshots
+
+(assert_invalid
+  (component
+    (core module $m
+      (memory (export "memory") 1)
+      (func (export "ret-map") (result i32 i32) unreachable)
+    )
+    (core instance $i (instantiate $m))
+
+    (func (export "ret-map") (result (map string u32))
+      (canon lift (core func $i "ret-map") (memory $i "memory"))
+    )
+  )
+  "Maps require the component model map feature (at offset 0x54)"
+)
+
+(assert_invalid
+  (component
+    (core module $m
+      (func (export "param-map") (param i32 i32) unreachable)
+    )
+    (core instance $i (instantiate $m))
+
+    (func (export "param-map") (param "m" (map string u32))
+      (canon lift (core func $i "param-map"))
+    )
+  )
+  "Maps require the component model map feature"
+)
+
+(assert_invalid
+  (component
+    (type $map-type (map u32 string))
+    (func (export "f") (param "x" $map-type))
+  )
+  "Maps require the component model map feature"
+)
+