feat: (CM64) Make realloc signature match the memory's address type (#2501)

* realloc

* fix task.return

* files

* address comments

Author: michael-weigelt

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

@@ -18,7 +18,8 @@ pub enum CanonicalOption {
     /// The realloc function to use if the lifting or lowering of a function requires memory
     /// allocation.
     ///
-    /// The value is an index to a core function of type `(func (param i32 i32 i32 i32) (result i32))`.
+    /// The value is an index to a core function of type `(func (param $T $T $T $T) (result $T))` where
+    /// `$T` is the index type of the memory, i.e., either `i32` or `i64`.
     Realloc(u32),
     /// The post-return function to use if the lifting of a function requires
     /// cleanup after the function returns.

crates/wasmparser/src/validator/component.rs

@@ -2635,17 +2635,7 @@ impl ComponentState {
                 CanonicalOption::Realloc(idx) => {
                     realloc = match realloc {
                         None => {
-                            let ty_id = self.core_function_at(*idx, offset)?;
-                            let func_ty = types[ty_id].unwrap_func();
-                            if func_ty.params()
-                                != [ValType::I32, ValType::I32, ValType::I32, ValType::I32]
-                                || func_ty.results() != [ValType::I32]
-                            {
-                                return Err(BinaryReaderError::new(
-                                    "canonical option `realloc` uses a core function with an incorrect signature",
-                                    offset,
-                                ));
-                            }
+                            // Validation deferred because it may depend on the memory option.
                             Some(*idx)
                         }
                         Some(_) => {
@@ -2771,6 +2761,29 @@ impl ComponentState {
             bail!(offset, "cannot specify `core-type` without `gc`")
         }
 
+        // Validate `realloc`
+        if let Some(realloc_idx) = realloc {
+            let addr_type = match memory {
+                // If a memory was specified, `realloc` must match its address type.
+                Some(memory_idx) => match self.memory_at(memory_idx, offset)?.memory64 {
+                    true => ValType::I64,
+                    false => ValType::I32,
+                },
+                // Backwards compatibility: Assume `i32` memory if none was specified.
+                None => ValType::I32,
+            };
+            let ty_id = self.core_function_at(realloc_idx, offset)?;
+            let func_ty = types[ty_id].unwrap_func();
+            if func_ty.params() != [addr_type, addr_type, addr_type, addr_type]
+                || func_ty.results() != [addr_type]
+            {
+                return Err(BinaryReaderError::new(
+                    "canonical option `realloc` uses a core function with an incorrect signature",
+                    offset,
+                ));
+            }
+        }
+
         Ok(CanonicalOptions {
             string_encoding,
             memory,

tests/cli/component-model/memory64/realloc.wast

@@ -0,0 +1,50 @@
+;; RUN: wast --assert default --snapshot tests/snapshots -f cm64 %
+
+(component
+  (core module $m
+    (memory (export "m") i64 1)
+    (func (export "f") (result i32) unreachable)
+    (func (export "realloc") (param i64 i64 i64 i64) (result i64) unreachable)
+  )
+  (core instance $i (instantiate $m))
+  (func (result string)
+    (canon lift (core func $i "f")
+      (memory $i "m")
+      (realloc (func $i "realloc"))
+    )
+  )
+)
+
+(assert_invalid
+  (component
+    (core module $m
+      (memory (export "m") i64 1)
+      (func (export "f") (param i32 i32))
+      (func (export "realloc") (param i32 i32 i32 i32) (result i32) unreachable)
+    )
+    (core instance $i (instantiate $m))
+    (func (param "p1" (list u8))
+      (canon lift (core func $i "f")
+        (memory $i "m")
+        (realloc (func $i "realloc"))
+      )
+    )
+  )
+  "canonical option `realloc` uses a core function with an incorrect signature")
+
+(assert_invalid
+  (component
+    (core module $m
+      (memory (export "m") i32 1)
+      (func (export "f") (param i32 i32))
+      (func (export "realloc") (param i64 i64 i64 i64) (result i64) unreachable)
+    )
+    (core instance $i (instantiate $m))
+    (func (param "p1" (list u8))
+      (canon lift (core func $i "f")
+        (memory $i "m")
+        (realloc (func $i "realloc"))
+      )
+    )
+  )
+  "canonical option `realloc` uses a core function with an incorrect signature")

tests/snapshots/cli/component-model/memory64/realloc.wast.json

@@ -0,0 +1,25 @@
+{
+  "source_filename": "tests/cli/component-model/memory64/realloc.wast",
+  "commands": [
+    {
+      "type": "module",
+      "line": 3,
+      "filename": "realloc.0.wasm",
+      "module_type": "binary"
+    },
+    {
+      "type": "assert_invalid",
+      "line": 19,
+      "filename": "realloc.1.wasm",
+      "module_type": "binary",
+      "text": "canonical option `realloc` uses a core function with an incorrect signature"
+    },
+    {
+      "type": "assert_invalid",
+      "line": 36,
+      "filename": "realloc.2.wasm",
+      "module_type": "binary",
+      "text": "canonical option `realloc` uses a core function with an incorrect signature"
+    }
+  ]
+}

tests/snapshots/cli/component-model/memory64/realloc.wast/0.print

@@ -0,0 +1,22 @@
+(component
+  (core module $m (;0;)
+    (type (;0;) (func (result i32)))
+    (type (;1;) (func (param i64 i64 i64 i64) (result i64)))
+    (memory (;0;) i64 1)
+    (export "m" (memory 0))
+    (export "f" (func 0))
+    (export "realloc" (func 1))
+    (func (;0;) (type 0) (result i32)
+      unreachable
+    )
+    (func (;1;) (type 1) (param i64 i64 i64 i64) (result i64)
+      unreachable
+    )
+  )
+  (core instance $i (;0;) (instantiate $m))
+  (type (;0;) (func (result string)))
+  (alias core export $i "f" (core func (;0;)))
+  (alias core export $i "m" (core memory (;0;)))
+  (alias core export $i "realloc" (core func (;1;)))
+  (func (;0;) (type 0) (canon lift (core func 0) (memory 0) (realloc 1)))
+)