feat: (CM64) context.{set, get} accept i64 immediate

crates/wasm-encoder/src/component/builder.rs

@@ -507,14 +507,14 @@ impl ComponentBuilder {
     }
 
     /// Declares a new `context.get` intrinsic.
-    pub fn context_get(&mut self, i: u32) -> u32 {
-        self.canonical_functions().context_get(i);
+    pub fn context_get(&mut self, ty: ValType, i: u32) -> u32 {
+        self.canonical_functions().context_get(ty, i);
         self.core_funcs.add(Some(&format!("context.get {i}")))
     }
 
     /// Declares a new `context.set` intrinsic.
-    pub fn context_set(&mut self, i: u32) -> u32 {
-        self.canonical_functions().context_set(i);
+    pub fn context_set(&mut self, ty: ValType, i: u32) -> u32 {
+        self.canonical_functions().context_set(ty, i);
         self.core_funcs.add(Some(&format!("context.set {i}")))
     }
 

crates/wasm-encoder/src/component/canonicals.rs

@@ -1,4 +1,6 @@
-use crate::{ComponentSection, ComponentSectionId, ComponentValType, Encode, encode_section};
+use crate::{
+    ComponentSection, ComponentSectionId, ComponentValType, Encode, ValType, encode_section,
+};
 use alloc::vec::Vec;
 
 /// Represents options for canonical function definitions.
@@ -234,19 +236,21 @@ impl CanonicalFunctionSection {
         self
     }
 
-    /// Defines a new `context.get` intrinsic of the ith slot.
-    pub fn context_get(&mut self, i: u32) -> &mut Self {
+    /// Defines a new `context.get` intrinsic of the ith slot with the given
+    /// value type.
+    pub fn context_get(&mut self, ty: ValType, i: u32) -> &mut Self {
         self.bytes.push(0x0a);
-        self.bytes.push(0x7f);
+        ty.encode(&mut self.bytes);
         i.encode(&mut self.bytes);
         self.num_added += 1;
         self
     }
 
-    /// Defines a new `context.set` intrinsic of the ith slot.
-    pub fn context_set(&mut self, i: u32) -> &mut Self {
+    /// Defines a new `context.set` intrinsic of the ith slot with the given
+    /// value type.
+    pub fn context_set(&mut self, ty: ValType, i: u32) -> &mut Self {
         self.bytes.push(0x0b);
-        self.bytes.push(0x7f);
+        ty.encode(&mut self.bytes);
         i.encode(&mut self.bytes);
         self.num_added += 1;
         self

crates/wasm-encoder/src/reencode/component.rs

@@ -996,11 +996,11 @@ pub mod component_utils {
             wasmparser::CanonicalFunction::TaskCancel => {
                 section.task_cancel();
             }
-            wasmparser::CanonicalFunction::ContextGet(i) => {
-                section.context_get(i);
+            wasmparser::CanonicalFunction::ContextGet { ty, slot } => {
+                section.context_get(reencoder.val_type(ty)?, slot);
             }
-            wasmparser::CanonicalFunction::ContextSet(i) => {
-                section.context_set(i);
+            wasmparser::CanonicalFunction::ContextSet { ty, slot } => {
+                section.context_set(reencoder.val_type(ty)?, slot);
             }
             wasmparser::CanonicalFunction::ThreadYield { cancellable } => {
                 section.thread_yield(cancellable);

crates/wasmparser/src/readers/component/canonicals.rs

@@ -1,6 +1,6 @@
 use crate::limits::MAX_WASM_CANONICAL_OPTIONS;
 use crate::prelude::*;
-use crate::{BinaryReader, ComponentValType, FromReader, Result, SectionLimited};
+use crate::{BinaryReader, ComponentValType, FromReader, Result, SectionLimited, ValType};
 
 /// Represents options for component functions.
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
@@ -108,9 +108,23 @@ pub enum CanonicalFunction {
     /// A function to acknowledge cancellation of the current task.
     TaskCancel,
     /// A `context.get` intrinsic for the `i`th slot of task-local storage.
-    ContextGet(u32),
+    ContextGet {
+        /// The type of the slot. Currently only `ValType::I32` and
+        /// `ValType::I64` are accepted by the validator (with `I64` gated on
+        /// the component-model 64-bit feature).
+        ty: ValType,
+        /// The index of the task-local storage slot.
+        slot: u32,
+    },
     /// A `context.set` intrinsic for the `i`th slot of task-local storage.
-    ContextSet(u32),
+    ContextSet {
+        /// The type of the slot. Currently only `ValType::I32` and
+        /// `ValType::I64` are accepted by the validator (with `I64` gated on
+        /// the component-model 64-bit feature).
+        ty: ValType,
+        /// The index of the task-local storage slot.
+        slot: u32,
+    },
     /// A function which yields control to the host so that other tasks are able
     /// to make progress, if any.
     ThreadYield {
@@ -337,13 +351,13 @@ impl<'a> FromReader<'a> for CanonicalFunction {
                 result: crate::read_resultlist(reader)?,
                 options: read_opts(reader)?,
             },
-            0x0a => match reader.read_u8()? {
-                0x7f => CanonicalFunction::ContextGet(reader.read_var_u32()?),
-                x => return reader.invalid_leading_byte(x, "context.get intrinsic type"),
+            0x0a => CanonicalFunction::ContextGet {
+                ty: reader.read()?,
+                slot: reader.read_var_u32()?,
             },
-            0x0b => match reader.read_u8()? {
-                0x7f => CanonicalFunction::ContextSet(reader.read_var_u32()?),
-                x => return reader.invalid_leading_byte(x, "context.set intrinsic type"),
+            0x0b => CanonicalFunction::ContextSet {
+                ty: reader.read()?,
+                slot: reader.read_var_u32()?,
             },
             0x0c => CanonicalFunction::ThreadYield {
                 cancellable: reader.read()?,

crates/wasmparser/src/validator/component.rs

@@ -1203,8 +1203,8 @@ impl ComponentState {
                 self.task_return(&result, &options, types, offset)
             }
             CanonicalFunction::TaskCancel => self.task_cancel(types, offset),
-            CanonicalFunction::ContextGet(i) => self.context_get(i, types, offset),
-            CanonicalFunction::ContextSet(i) => self.context_set(i, types, offset),
+            CanonicalFunction::ContextGet { ty, slot } => self.context_get(ty, slot, types, offset),
+            CanonicalFunction::ContextSet { ty, slot } => self.context_set(ty, slot, types, offset),
             CanonicalFunction::ThreadYield { cancellable: _ } => self.thread_yield(types, offset),
             CanonicalFunction::SubtaskDrop => self.subtask_drop(types, offset),
             CanonicalFunction::SubtaskCancel { async_ } => {
@@ -1517,34 +1517,64 @@ impl ComponentState {
         Ok(())
     }
 
-    fn context_get(&mut self, i: u32, types: &mut TypeAlloc, offset: usize) -> Result<()> {
+    fn context_get(
+        &mut self,
+        ty: ValType,
+        i: u32,
+        types: &mut TypeAlloc,
+        offset: usize,
+    ) -> Result<()> {
         if !self.features.cm_async() {
             bail!(
                 offset,
                 "`context.get` requires the component model async feature"
             )
         }
+        self.validate_context_type(ty, "context.get", offset)?;
         self.validate_context_immediate(i, "context.get", offset)?;
 
         self.core_funcs
-            .push(types.intern_func_type(FuncType::new([], [ValType::I32]), offset));
+            .push(types.intern_func_type(FuncType::new([], [ty]), offset));
         Ok(())
     }
 
-    fn context_set(&mut self, i: u32, types: &mut TypeAlloc, offset: usize) -> Result<()> {
+    fn context_set(
+        &mut self,
+        ty: ValType,
+        i: u32,
+        types: &mut TypeAlloc,
+        offset: usize,
+    ) -> Result<()> {
         if !self.features.cm_async() {
             bail!(
                 offset,
                 "`context.set` requires the component model async feature"
             )
         }
+        self.validate_context_type(ty, "context.set", offset)?;
         self.validate_context_immediate(i, "context.set", offset)?;
 
         self.core_funcs
-            .push(types.intern_func_type(FuncType::new([ValType::I32], []), offset));
+            .push(types.intern_func_type(FuncType::new([ty], []), offset));
         Ok(())
     }
 
+    fn validate_context_type(&self, ty: ValType, intrinsic: &str, offset: usize) -> Result<()> {
+        match ty {
+            ValType::I32 => Ok(()),
+            ValType::I64 => {
+                if !self.features.cm64() {
+                    bail!(
+                        offset,
+                        "64-bit `{intrinsic}` requires the component model 64-bit feature"
+                    )
+                }
+                Ok(())
+            }
+            _ => bail!(offset, "`{intrinsic}` only supports `i32` or `i64`"),
+        }
+    }
+
     fn thread_yield(&mut self, types: &mut TypeAlloc, offset: usize) -> Result<()> {
         if !self.features.cm_async() {
             bail!(

crates/wasmprinter/src/component.rs

@@ -977,15 +977,19 @@ impl Printer<'_, '_> {
                 CanonicalFunction::TaskCancel => {
                     self.print_intrinsic(state, "canon task.cancel", &|_, _| Ok(()))?;
                 }
-                CanonicalFunction::ContextGet(i) => {
-                    self.print_intrinsic(state, "canon context.get", &|me, _state| {
-                        write!(me.result, " i32 {i}")?;
+                CanonicalFunction::ContextGet { ty, slot } => {
+                    self.print_intrinsic(state, "canon context.get", &|me, state| {
+                        me.result.write_str(" ")?;
+                        me.print_valtype(state, ty)?;
+                        write!(me.result, " {slot}")?;
                         Ok(())
                     })?;
                 }
-                CanonicalFunction::ContextSet(i) => {
-                    self.print_intrinsic(state, "canon context.set", &|me, _state| {
-                        write!(me.result, " i32 {i}")?;
+                CanonicalFunction::ContextSet { ty, slot } => {
+                    self.print_intrinsic(state, "canon context.set", &|me, state| {
+                        me.result.write_str(" ")?;
+                        me.print_valtype(state, ty)?;
+                        write!(me.result, " {slot}")?;
                         Ok(())
                     })?;
                 }

crates/wast/src/component/binary.rs

@@ -392,13 +392,13 @@ impl<'a> Encoder<'a> {
                     self.core_func_names.push(name);
                     self.funcs.task_cancel();
                 }
-                CoreFuncKind::ContextGet(i) => {
+                CoreFuncKind::ContextGet(ty, i) => {
                     self.core_func_names.push(name);
-                    self.funcs.context_get(*i);
+                    self.funcs.context_get((*ty).into(), *i);
                 }
-                CoreFuncKind::ContextSet(i) => {
+                CoreFuncKind::ContextSet(ty, i) => {
                     self.core_func_names.push(name);
-                    self.funcs.context_set(*i);
+                    self.funcs.context_set((*ty).into(), *i);
                 }
                 CoreFuncKind::ThreadYield(info) => {
                     self.core_func_names.push(name);

crates/wast/src/component/func.rs

@@ -58,8 +58,8 @@ pub enum CoreFuncKind<'a> {
     BackpressureDec,
     TaskReturn(CanonTaskReturn<'a>),
     TaskCancel,
-    ContextGet(u32),
-    ContextSet(u32),
+    ContextGet(crate::core::ValType<'a>, u32),
+    ContextSet(crate::core::ValType<'a>, u32),
     ThreadYield(CanonThreadYield),
     SubtaskDrop,
     SubtaskCancel(CanonSubtaskCancel),
@@ -139,12 +139,12 @@ impl<'a> CoreFuncKind<'a> {
             Ok(CoreFuncKind::TaskCancel)
         } else if l.peek::<kw::context_get>()? {
             parser.parse::<kw::context_get>()?;
-            parser.parse::<kw::i32>()?;
-            Ok(CoreFuncKind::ContextGet(parser.parse()?))
+            let ty = parser.parse()?;
+            Ok(CoreFuncKind::ContextGet(ty, parser.parse()?))
         } else if l.peek::<kw::context_set>()? {
             parser.parse::<kw::context_set>()?;
-            parser.parse::<kw::i32>()?;
-            Ok(CoreFuncKind::ContextSet(parser.parse()?))
+            let ty = parser.parse()?;
+            Ok(CoreFuncKind::ContextSet(ty, parser.parse()?))
         } else if l.peek::<kw::thread_yield>()? {
             Ok(CoreFuncKind::ThreadYield(parser.parse()?))
         } else if l.peek::<kw::subtask_drop>()? {

crates/wast/src/component/resolve.rs

@@ -408,7 +408,8 @@ impl<'a> Resolver<'a> {
                     }
                     self.canon_opts(&mut info.opts)?;
                 }
-                CoreFuncKind::ContextGet(_) | CoreFuncKind::ContextSet(_) => {}
+                CoreFuncKind::ContextGet(ty, _) => self.ref_type(ty)?,
+                CoreFuncKind::ContextSet(ty, _) => self.ref_type(ty)?,
                 CoreFuncKind::StreamNew(info) => {
                     self.resolve_ns(&mut info.ty, Ns::Type)?;
                 }
@@ -486,6 +487,18 @@ impl<'a> Resolver<'a> {
         Ok(())
     }
 
+    fn ref_type(&mut self, ty: &mut ValType<'a>) -> Result<(), Error> {
+        Ok(match ty {
+            ValType::I32 | ValType::I64 | ValType::F32 | ValType::F64 | ValType::V128 => {}
+            ValType::Ref(r) => match &mut r.heap {
+                core::HeapType::Abstract { .. } => {}
+                core::HeapType::Concrete(id) | core::HeapType::Exact(id) => {
+                    self.resolve_ns(id, Ns::Type)?;
+                }
+            },
+        })
+    }
+
     fn canon_opts(&mut self, opts: &mut [CanonOpt<'a>]) -> Result<(), Error> {
         for opt in opts {
             match opt {
@@ -658,15 +671,7 @@ impl<'a> Resolver<'a> {
                 self.stack.pop();
             }
             TypeDef::Resource(r) => {
-                match &mut r.rep {
-                    ValType::I32 | ValType::I64 | ValType::F32 | ValType::F64 | ValType::V128 => {}
-                    ValType::Ref(r) => match &mut r.heap {
-                        core::HeapType::Abstract { .. } => {}
-                        core::HeapType::Concrete(id) | core::HeapType::Exact(id) => {
-                            self.resolve_ns(id, Ns::Type)?;
-                        }
-                    },
-                }
+                self.ref_type(&mut r.rep)?;
                 if let Some(dtor) = &mut r.dtor {
                     self.core_item_ref(dtor)?;
                 }

crates/wit-component/src/encoding.rs

@@ -1977,12 +1977,12 @@ impl<'a> EncodingState<'a> {
                 let index = self.component.waitable_join();
                 Ok((ExportKind::Func, index))
             }
-            Import::ContextGet(n) => {
-                let index = self.component.context_get(*n);
+            Import::ContextGet { ty, slot } => {
+                let index = self.component.context_get((*ty).try_into()?, *slot);
                 Ok((ExportKind::Func, index))
             }
-            Import::ContextSet(n) => {
-                let index = self.component.context_set(*n);
+            Import::ContextSet { ty, slot } => {
+                let index = self.component.context_set((*ty).try_into()?, *slot);
                 Ok((ExportKind::Func, index))
             }
             Import::ExportedTaskCancel => {
@@ -2630,8 +2630,8 @@ impl<'a> Shims<'a> {
                 | Import::WaitableSetNew
                 | Import::WaitableSetDrop
                 | Import::WaitableJoin
-                | Import::ContextGet(_)
-                | Import::ContextSet(_)
+                | Import::ContextGet { .. }
+                | Import::ContextSet { .. }
                 | Import::ThreadIndex
                 | Import::ThreadSuspendToSuspended { .. }
                 | Import::ThreadSuspend { .. }

crates/wit-component/src/encoding/world.rs

@@ -415,8 +415,8 @@ impl<'a> ComponentWorld<'a> {
                 | Import::MainModuleMemory
                 | Import::MainModuleExport { .. }
                 | Import::Item(_)
-                | Import::ContextGet(_)
-                | Import::ContextSet(_)
+                | Import::ContextGet { .. }
+                | Import::ContextSet { .. }
                 | Import::BackpressureInc
                 | Import::BackpressureDec
                 | Import::WaitableSetNew