Merge branch 'mk/system_api64_4' into 'master'

Feat: RUN-979: Support for wasm64 in native stable64

This MR modifies stable64 instrumentation to support writing to 64bit heap.
It also enables 64bit addresses for active data segments. 

See merge request dfinity-lab/public/ic!19304

rs/embedders/src/wasm_utils/instrumentation.rs

@@ -142,6 +142,16 @@ pub(crate) enum WasmMemoryType {
     Wasm64,
 }
 
+pub(crate) fn main_memory_type(module: &Module<'_>) -> WasmMemoryType {
+    let mut mem_type = WasmMemoryType::Wasm32;
+    if let Some(mem) = module.memories.first() {
+        if mem.memory64 {
+            mem_type = WasmMemoryType::Wasm64;
+        }
+    }
+    mem_type
+}
+
 // The indices of injected function imports.
 pub(crate) enum InjectedImports {
     OutOfInstructions = 0,
@@ -908,12 +918,7 @@ pub(super) fn instrument(
     subnet_type: SubnetType,
     dirty_page_overhead: NumInstructions,
 ) -> Result<InstrumentationOutput, WasmInstrumentationError> {
-    let mut main_memory_type = WasmMemoryType::Wasm32;
-    if let Some(mem) = module.memories.first() {
-        if mem.memory64 {
-            main_memory_type = WasmMemoryType::Wasm64;
-        }
-    }
+    let main_memory_type = main_memory_type(&module);
     let stable_memory_index;
     let mut module = inject_helper_functions(module, wasm_native_stable_memory, main_memory_type);
     module = export_table(module);
@@ -1012,6 +1017,7 @@ pub(super) fn instrument(
             special_indices,
             subnet_type,
             dirty_page_overhead,
+            main_memory_type,
         )
     }
 
@@ -1095,6 +1101,7 @@ fn replace_system_api_functions(
     special_indices: SpecialIndices,
     subnet_type: SubnetType,
     dirty_page_overhead: NumInstructions,
+    main_memory_type: WasmMemoryType,
 ) {
     let api_indexes = calculate_api_indexes(module);
     let number_of_func_imports = module
@@ -1106,9 +1113,12 @@ fn replace_system_api_functions(
     // Collect a single map of all the function indexes that need to be
     // replaced.
     let mut func_index_replacements = BTreeMap::new();
-    for (api, (ty, body)) in
-        replacement_functions(special_indices, subnet_type, dirty_page_overhead)
-    {
+    for (api, (ty, body)) in replacement_functions(
+        special_indices,
+        subnet_type,
+        dirty_page_overhead,
+        main_memory_type,
+    ) {
         if let Some(old_index) = api_indexes.get(&api) {
             let type_idx = add_func_type(module, ty);
             let new_index = (number_of_func_imports + module.functions.len()) as u32;
@@ -1825,6 +1835,7 @@ fn get_data(
                     offset_expr,
                 } => match offset_expr {
                     Operator::I32Const { value } => *value as usize,
+                    Operator::I64Const { value } => *value as usize,
                     _ => return Some(Err(WasmInstrumentationError::WasmDeserializeError(WasmError::new(
                         "complex initialization expressions for data segments are not supported!".into()
                     )))),

rs/embedders/src/wasm_utils/system_api_replacements.rs

@@ -12,8 +12,9 @@
 //!
 
 use crate::{
-    wasm_utils::instrumentation::InjectedImports,
-    wasmtime_embedder::system_api_complexity::overhead_native, InternalErrorCode,
+    wasm_utils::instrumentation::{InjectedImports, WasmMemoryType},
+    wasmtime_embedder::system_api_complexity::overhead_native,
+    InternalErrorCode,
 };
 use ic_interfaces::execution_environment::StableMemoryApi;
 use ic_registry_subnet_type::SubnetType;
@@ -31,6 +32,7 @@ pub(super) fn replacement_functions(
     special_indices: SpecialIndices,
     subnet_type: SubnetType,
     dirty_page_overhead: NumInstructions,
+    main_memory_type: WasmMemoryType,
 ) -> Vec<(SystemApiFunc, (FuncType, Body<'static>))> {
     let count_clean_pages_fn_index = special_indices.count_clean_pages_fn.unwrap();
     let dirty_pages_counter_index = special_indices.dirty_pages_counter_ix.unwrap();
@@ -41,6 +43,12 @@ pub(super) fn replacement_functions(
     use Operator::*;
     let page_size_shift = PAGE_SIZE.trailing_zeros() as i32;
     let stable_memory_bytemap_index = stable_memory_index + 1;
+
+    let cast_to_heap_addr_type = match main_memory_type {
+        WasmMemoryType::Wasm32 => I32WrapI64,
+        WasmMemoryType::Wasm64 => Nop,
+    };
+
     vec![
         (
             SystemApiFunc::StableSize,
@@ -731,10 +739,10 @@ pub(super) fn replacement_functions(
                             },
                             Else,
                             LocalGet { local_index: DST },
-                            I32WrapI64,
+                            cast_to_heap_addr_type.clone(),
                             LocalGet { local_index: SRC },
                             LocalGet { local_index: LEN },
-                            I32WrapI64,
+                            cast_to_heap_addr_type.clone(),
                             MemoryCopy {
                                 dst_mem: 0,
                                 src_mem: stable_memory_index,
@@ -1205,9 +1213,9 @@ pub(super) fn replacement_functions(
                             // copy memory contents
                             LocalGet { local_index: DST },
                             LocalGet { local_index: SRC },
-                            I32WrapI64,
+                            cast_to_heap_addr_type.clone(),
                             LocalGet { local_index: LEN },
-                            I32WrapI64,
+                            cast_to_heap_addr_type,
                             MemoryCopy {
                                 dst_mem: stable_memory_index,
                                 src_mem: 0,

rs/embedders/src/wasm_utils/validation.rs

@@ -20,7 +20,8 @@ use crate::wasmtime_embedder::{
 };
 use crate::{
     wasm_utils::instrumentation::{
-        ACCESSED_PAGES_COUNTER_GLOBAL_NAME, DIRTY_PAGES_COUNTER_GLOBAL_NAME,
+        main_memory_type, WasmMemoryType, ACCESSED_PAGES_COUNTER_GLOBAL_NAME,
+        DIRTY_PAGES_COUNTER_GLOBAL_NAME,
     },
     MAX_WASM_STACK_SIZE, MIN_GUARD_REGION_SIZE,
 };
@@ -954,24 +955,45 @@ fn validate_export_section(
 // expression. Required because of OP. See also:
 // instrumentation.rs
 fn validate_data_section(module: &Module) -> Result<(), WasmValidationError> {
-    fn validate_segment(s: &DataSegment) -> Result<(), WasmValidationError> {
-        match &s.kind {
-            DataSegmentKind::Passive => Ok(()),
-            DataSegmentKind::Active {
-                memory_index: _,
-                offset_expr,
-            } => match offset_expr {
+    fn validate_segment(
+        s: &DataSegment,
+        mem_type: WasmMemoryType,
+    ) -> Result<(), WasmValidationError> {
+        match (&s.kind, mem_type) {
+            (DataSegmentKind::Passive, _) => Ok(()),
+            (
+                DataSegmentKind::Active {
+                    memory_index: _,
+                    offset_expr,
+                },
+                WasmMemoryType::Wasm32,
+            ) => match offset_expr {
                 Operator::I32Const { .. } => Ok(()),
                 _ => Err(WasmValidationError::InvalidDataSection(format!(
-                    "Invalid offset expression in data segment: {:?}",
+                    "Invalid offset expression in data segment for 32bit memory: {:?}",
+                    offset_expr
+                ))),
+            },
+            (
+                DataSegmentKind::Active {
+                    memory_index: _,
+                    offset_expr,
+                },
+                WasmMemoryType::Wasm64,
+            ) => match offset_expr {
+                Operator::I64Const { .. } => Ok(()),
+                _ => Err(WasmValidationError::InvalidDataSection(format!(
+                    "Invalid offset expression in data segment for 64bit memory: {:?}",
                     offset_expr
                 ))),
             },
         }
     }
 
+    let mem_type = main_memory_type(module);
+
     for d in &module.data {
-        validate_segment(d)?;
+        validate_segment(d, mem_type)?;
     }
     Ok(())
 }

rs/embedders/tests/wasmtime_random_memory_writes.rs

@@ -259,6 +259,57 @@ fn make_module_wat_for_api_calls(heap_size: usize) -> String {
     )
 }
 
+fn make_module64_wat_for_api_calls(heap_size: usize) -> String {
+    format!(
+        r#"
+    (module
+      (import "ic0" "msg_reply" (func $msg_reply))
+      (import "ic0" "msg_reply_data_append"
+        (func $msg_reply_data_append (param i32) (param i32)))
+      (import "ic0" "msg_arg_data_copy"
+        (func $ic0_msg_arg_data_copy (param i32) (param i32) (param i32)))
+      (import "ic0" "msg_arg_data_size"
+        (func $ic0_msg_arg_data_size (result i32)))
+      (import "ic0" "msg_caller_copy"
+        (func $ic0_msg_caller_copy (param i32) (param i32) (param i32)))
+      (import "ic0" "msg_caller_size"
+        (func $ic0_msg_caller_size (result i32)))
+      (import "ic0" "canister_self_copy"
+        (func $ic0_canister_self_copy (param i32) (param i32) (param i32)))
+      (import "ic0" "canister_self_size"
+        (func $ic0_canister_self_size (result i32)))
+
+      (import "ic0" "canister_cycle_balance128"
+        (func $ic0_canister_cycle_balance128 (param i32)))
+
+      (import "ic0" "stable_grow"
+        (func $ic0_stable_grow (param $pages i32) (result i32)))
+      (import "ic0" "stable64_read"
+        (func $ic0_stable64_read (param $dst i64) (param $offset i64) (param $size i64)))
+      (import "ic0" "stable64_write"
+        (func $ic0_stable64_write (param $offset i64) (param $src i64) (param $size i64)))
+
+      (func $touch_heap_with_api_calls
+        (call $ic0_msg_caller_copy (i32.const 4096) (i32.const 0) (call $ic0_msg_caller_size))
+        (call $ic0_msg_arg_data_copy (i32.const 12288) (i32.const 0) (call $ic0_msg_arg_data_size))
+        (call $ic0_canister_self_copy (i32.const 20480) (i32.const 0) (call $ic0_canister_self_size))
+        (call $ic0_canister_cycle_balance128 (i32.const 36864))
+
+        (; Write some data to page 10 using stable_read, by first copying 4
+        bytes from the second page to stable memory, then copying back ;)
+        (drop (call $ic0_stable_grow (i32.const 1)))
+        (call $ic0_stable64_write (i64.const 0) (i64.const 4096) (i64.const 4))
+        (call $ic0_stable64_read (i64.const 40960) (i64.const 0) (i64.const 4))
+      )
+
+      (memory $memory i64 {HEAP_SIZE})
+      (export "memory" (memory $memory))
+      (export "canister_update touch_heap_with_api_calls" (func $touch_heap_with_api_calls))
+    )"#,
+        HEAP_SIZE = heap_size
+    )
+}
+
 fn make_module_wat_with_write_fun(heap_size: usize, write_fun: &str) -> String {
     format!(
         r#"
@@ -1385,6 +1436,62 @@ mod tests {
         });
     }
 
+    #[test]
+    fn touch_heap64_with_api_calls() {
+        with_test_replica_logger(|log| {
+            let wat = make_module64_wat_for_api_calls(TEST_NUM_PAGES);
+            let wasm = wat2wasm(&wat).unwrap();
+            let mut config = EmbeddersConfig::default();
+            config.feature_flags.wasm64 = FlagStatus::Enabled;
+            let embedder = WasmtimeEmbedder::new(config, log);
+            let (embedder_cache, result) = compile(&embedder, &wasm);
+            result.unwrap();
+
+            let mut dirty_pages: BTreeSet<u64> = BTreeSet::new();
+
+            let payload = vec![0, 1, 2, 3, 4, 5, 6, 7];
+
+            let api = test_api_for_update(
+                no_op_logger(),
+                None,
+                payload,
+                SubnetType::Application,
+                MAX_NUM_INSTRUCTIONS,
+            );
+            let instruction_limit = api.slice_instruction_limit();
+            let mut instance = embedder
+                .new_instance(
+                    canister_test_id(1),
+                    &embedder_cache,
+                    None,
+                    &Memory::new(PageMap::new_for_testing(), NumWasmPages::from(0)),
+                    &Memory::new(PageMap::new_for_testing(), NumWasmPages::from(0)),
+                    ModificationTracking::Track,
+                    Some(api),
+                )
+                .map_err(|r| r.0)
+                .expect("Failed to create instance");
+            instance.set_instruction_counter(i64::try_from(instruction_limit.get()).unwrap());
+
+            let result = instance
+                .run(FuncRef::Method(WasmMethod::Update(
+                    "touch_heap_with_api_calls".to_string(),
+                )))
+                .expect("call to touch_heap_with_api_calls failed");
+            dirty_pages.extend(result.dirty_pages.iter().map(|x| x.get()));
+
+            let mut expected_dirty_pages: BTreeSet<u64> = BTreeSet::new();
+            expected_dirty_pages.insert(1); // caller_copy
+            expected_dirty_pages.insert(3); // data_copy
+            expected_dirty_pages.insert(5); // canister_self_copy
+            expected_dirty_pages.insert(9); // canister_cycle_balance128
+            expected_dirty_pages.insert(9); // msg_cycles_available128
+            expected_dirty_pages.insert(10); // stable_read
+
+            assert_eq!(expected_dirty_pages, dirty_pages);
+        });
+    }
+
     #[test]
     fn wasmtime_random_memory_writes_ignore_dirty_pages() {
         // The seed value will always be the same for a particular version of

rs/execution_environment/src/hypervisor/tests.rs

@@ -1438,6 +1438,42 @@ fn ic0_msg_reject_works() {
     assert_eq!(WasmResult::Reject("panic!".to_string()), result);
 }
 
+#[test]
+fn wasm64_active_data_segments() {
+    let mut test = ExecutionTestBuilder::new().with_wasm64().build();
+    let wat = r#"
+        (module
+            (import "ic0" "msg_reply" (func $msg_reply))
+            (func (export "canister_update test")
+                (if (i64.ne
+                         (i64.load8_u (i64.const 0))
+                         (i64.const 112) ;; p
+                    )
+                    (then (unreachable))
+                )
+                (if (i64.ne
+                         (i64.load8_u (i64.const 1))
+                         (i64.const 0)
+                    )
+                    (then (unreachable))
+                )
+                (if (i64.ne
+                         (i64.load8_u (i64.const 2))
+                         (i64.const 97) ;; a
+                    )
+                    (then (unreachable))
+                )
+                (call $msg_reply)
+            )
+            (memory i64 1 1)
+            (data (i64.const 0) "p")
+            (data (i64.const 2) "a")
+        )"#;
+    let canister_id = test.canister_from_wat(wat).unwrap();
+    let result = test.ingress(canister_id, "test", vec![]).unwrap();
+    assert_eq!(WasmResult::Reply(vec![]), result);
+}
+
 #[test]
 fn ic0_msg_caller_size_works_in_reply_callback() {
     let mut test = ExecutionTestBuilder::new().build();

rs/test_utilities/execution_environment/src/lib.rs

@@ -1929,6 +1929,11 @@ impl ExecutionTestBuilder {
         self
     }
 
+    pub fn with_wasm64(mut self) -> Self {
+        self.execution_config.embedders_config.feature_flags.wasm64 = FlagStatus::Enabled;
+        self
+    }
+
     pub fn with_metering_type(mut self, metering_type: MeteringType) -> Self {
         self.execution_config.embedders_config.metering_type = metering_type;
         self