Fix map adapter trampoline compilation and alignment bugs

The translate_map function had two categories of bugs preventing map
adapter trampolines from working:

1. Wasm stack discipline: local_set_new_tmp emits LocalSet which pops
   from the stack, but was called when the stack was empty (to
   "pre-allocate" locals). Fixed by computing values first, then
   calling local_set_new_tmp to consume them—matching translate_list's
   pattern. Also removed an erroneous LocalTee that left an orphan
   value on the stack. Affected: src_byte_len, dst_byte_len,
   cur_src_ptr, cur_dst_ptr.

2. Pointer advancement: after value translation, the pointer still
   points at the value start. The code only advanced by trailing
   padding instead of value_size + trailing_padding, causing every
   loop iteration to re-read the same memory.

Also fixes entry layout to use proper record alignment rules (entry
align = max(key_align, value_align), value at aligned offset).

Author: yordis

crates/environ/src/fact/trampoline.rs

@@ -2699,8 +2699,6 @@ impl<'a, 'b> Compiler<'a, 'b> {
         dst_ty: &InterfaceType,
         dst: &Destination,
     ) {
-        // Extract memory configuration for source and destination
-        // Get linear memory options (32-bit vs 64-bit pointers, which memory, etc.)
         let src_mem_opts = match &src.opts().data_model {
             DataModel::Gc {} => todo!("CM+GC"),
             DataModel::LinearMemory(opts) => opts,
@@ -2710,84 +2708,66 @@ impl<'a, 'b> Compiler<'a, 'b> {
             DataModel::LinearMemory(opts) => opts,
         };
 
-        // Get type information for source and destination maps
-        // Look up the TypeMap structs which contain the key and value InterfaceTypes
         let src_map_ty = &self.types[src_ty];
         let dst_map_ty = match dst_ty {
             InterfaceType::Map(r) => &self.types[*r],
             _ => panic!("expected a map"),
         };
 
-        // Load the map's pointer and length into temporary locals
-        // A map is represented as (ptr, len) - we need both values in locals
-        // for later use in the translation loop.
         match src {
             Source::Stack(s) => {
-                // If map descriptor is passed on the stack (as 2 locals: ptr, len)
                 assert_eq!(s.locals.len(), 2);
-                self.stack_get(&s.slice(0..1), src_mem_opts.ptr()); // Push ptr to wasm stack
-                self.stack_get(&s.slice(1..2), src_mem_opts.ptr()); // Push len to wasm stack
+                self.stack_get(&s.slice(0..1), src_mem_opts.ptr());
+                self.stack_get(&s.slice(1..2), src_mem_opts.ptr());
             }
             Source::Memory(mem) => {
-                // If map descriptor is stored in memory, load ptr and len from there
-                self.ptr_load(mem); // Load ptr
-                self.ptr_load(&mem.bump(src_mem_opts.ptr_size().into())); // Load len (next field)
+                self.ptr_load(mem);
+                self.ptr_load(&mem.bump(src_mem_opts.ptr_size().into()));
             }
             Source::Struct(_) | Source::Array(_) => todo!("CM+GC"),
         }
-        // Pop values from wasm stack into named locals (note: len is on top, then ptr)
         let src_len = self.local_set_new_tmp(src_mem_opts.ptr());
         let src_ptr = self.local_set_new_tmp(src_mem_opts.ptr());
 
-        // Calculate tuple sizes with proper alignment
-        // Each map entry is a (key, value) tuple. We need to know:
-        // - Size of key and value in bytes
-        // - Alignment requirements
-        // - Total tuple size including any padding
         let src_opts = src.opts();
         let dst_opts = dst.opts();
 
-        // Source tuple layout
+        // Each map entry is a tuple<K, V> following record layout rules:
+        // key at offset 0, padding to value_align, value, trailing padding
+        // to entry_align.
         let (src_key_size, src_key_align) = self.types.size_align(src_mem_opts, &src_map_ty.key);
-        let (src_value_size, _) = self.types.size_align(src_mem_opts, &src_map_ty.value);
-        // Total tuple size = key + value + padding to alignment
-        // e.g., if key is 4 bytes, value is 8 bytes, align is 4:
-        // tuple_size = (4 + 8 + 3) & ~3 = 12 bytes
+        let (src_value_size, src_value_align) =
+            self.types.size_align(src_mem_opts, &src_map_ty.value);
+        let src_entry_align = src_key_align.max(src_value_align);
+        let src_value_offset =
+            (src_key_size + src_value_align - 1) & !(src_value_align - 1);
         let src_tuple_size =
-            (src_key_size + src_value_size + src_key_align - 1) & !(src_key_align - 1);
+            (src_value_offset + src_value_size + src_entry_align - 1) & !(src_entry_align - 1);
 
-        // Destination tuple layout (may differ if types are converted)
         let (dst_key_size, dst_key_align) = self.types.size_align(dst_mem_opts, &dst_map_ty.key);
-        let (dst_value_size, _) = self.types.size_align(dst_mem_opts, &dst_map_ty.value);
+        let (dst_value_size, dst_value_align) =
+            self.types.size_align(dst_mem_opts, &dst_map_ty.value);
+        let dst_entry_align = dst_key_align.max(dst_value_align);
+        let dst_value_offset =
+            (dst_key_size + dst_value_align - 1) & !(dst_value_align - 1);
         let dst_tuple_size =
-            (dst_key_size + dst_value_size + dst_key_align - 1) & !(dst_key_align - 1);
+            (dst_value_offset + dst_value_size + dst_entry_align - 1) & !(dst_entry_align - 1);
 
-        // Create source memory operand and verify alignment
-        // This creates a Memory operand and verifies the source pointer is properly aligned
-        let src_mem = self.memory_operand(src_opts, src_ptr, src_key_align);
+        let src_mem = self.memory_operand(src_opts, src_ptr, src_entry_align);
 
-        // Calculate total byte lengths for source and destination
-        // total_bytes = count * tuple_size
+        // Calculate source/destination byte lengths and allocate destination.
+        self.instruction(LocalGet(src_len.idx));
+        self.ptr_uconst(src_mem_opts, src_tuple_size);
+        self.ptr_mul(src_mem_opts);
         let src_byte_len = self.local_set_new_tmp(src_mem_opts.ptr());
-        self.instruction(LocalGet(src_len.idx)); // Push len
-        self.ptr_uconst(src_mem_opts, src_tuple_size); // Push tuple_size
-        self.ptr_mul(src_mem_opts); // len * tuple_size
-        self.instruction(LocalSet(src_byte_len.idx)); // Save to local
 
+        self.instruction(LocalGet(src_len.idx));
+        self.ptr_uconst(dst_mem_opts, dst_tuple_size);
+        self.ptr_mul(dst_mem_opts);
         let dst_byte_len = self.local_set_new_tmp(dst_mem_opts.ptr());
-        self.instruction(LocalGet(src_len.idx)); // Push len (same count)
-        self.ptr_uconst(dst_mem_opts, dst_tuple_size); // Push dst tuple_size
-        self.ptr_mul(dst_mem_opts); // len * tuple_size
-        self.instruction(LocalTee(dst_byte_len.idx)); // Save AND keep on stack for malloc
-
-        // Allocate destination buffer
-        // Call realloc in the destination component to allocate space.
-        // dst_byte_len is still on the stack from LocalTee above.
-        let dst_mem = self.malloc(dst_opts, MallocSize::Local(dst_byte_len.idx), dst_key_align);
-
-        // Validate memory bounds
-        // Verify that ptr + byte_len doesn't overflow or exceed memory bounds.
-        // Trap if invalid.
+
+        let dst_mem = self.malloc(dst_opts, MallocSize::Local(dst_byte_len.idx), dst_entry_align);
+
         self.validate_memory_inbounds(
             src_mem_opts,
             src_mem.addr.idx,
@@ -2801,37 +2781,25 @@ impl<'a, 'b> Compiler<'a, 'b> {
             Trap::ListOutOfBounds,
         );
 
-        // Done with byte length locals
         self.free_temp_local(src_byte_len);
         self.free_temp_local(dst_byte_len);
 
-        // Main translation loop - copy each (key, value) tuple
-        // Skip loop entirely if tuples are zero-sized (nothing to copy)
         if src_tuple_size > 0 || dst_tuple_size > 0 {
-            // Loop setup
-            // Create counter for remaining elements
-            let remaining = self.local_set_new_tmp(src_mem_opts.ptr());
+            self.instruction(Block(BlockType::Empty));
+
             self.instruction(LocalGet(src_len.idx));
-            self.instruction(LocalSet(remaining.idx));
+            let remaining = self.local_tee_new_tmp(src_mem_opts.ptr());
+            self.ptr_eqz(src_mem_opts);
+            self.instruction(BrIf(0));
 
-            // Create pointer to current position in source
-            let cur_src_ptr = self.local_set_new_tmp(src_mem_opts.ptr());
             self.instruction(LocalGet(src_mem.addr.idx));
-            self.instruction(LocalSet(cur_src_ptr.idx));
+            let cur_src_ptr = self.local_set_new_tmp(src_mem_opts.ptr());
 
-            // Create pointer to current position in destination
-            let cur_dst_ptr = self.local_set_new_tmp(dst_mem_opts.ptr());
             self.instruction(LocalGet(dst_mem.addr.idx));
-            self.instruction(LocalSet(cur_dst_ptr.idx));
+            let cur_dst_ptr = self.local_set_new_tmp(dst_mem_opts.ptr());
 
-            // WebAssembly loop structure
-            // Block is the outer container (to break out of loop)
-            // Loop is what we branch back to for iteration
-            self.instruction(Block(BlockType::Empty));
             self.instruction(Loop(BlockType::Empty));
 
-            // Translate the key
-            // Create Source pointing to current key location
             let key_src = Source::Memory(self.memory_operand(
                 src_opts,
                 TempLocal {
@@ -2841,7 +2809,6 @@ impl<'a, 'b> Compiler<'a, 'b> {
                 },
                 src_key_align,
             ));
-            // Create Destination pointing to where key should go
             let key_dst = Destination::Memory(self.memory_operand(
                 dst_opts,
                 TempLocal {
@@ -2851,32 +2818,30 @@ impl<'a, 'b> Compiler<'a, 'b> {
                 },
                 dst_key_align,
             ));
-            // Recursively translate the key (handles any type: primitives, strings, etc.)
             self.translate(&src_map_ty.key, &key_src, &dst_map_ty.key, &key_dst);
 
-            // Advance pointers past the key to point at value
-            if src_key_size > 0 {
+            // Advance pointers from key start to value start
+            if src_value_offset > 0 {
                 self.instruction(LocalGet(cur_src_ptr.idx));
-                self.ptr_uconst(src_mem_opts, src_key_size);
+                self.ptr_uconst(src_mem_opts, src_value_offset);
                 self.ptr_add(src_mem_opts);
                 self.instruction(LocalSet(cur_src_ptr.idx));
             }
-            if dst_key_size > 0 {
+            if dst_value_offset > 0 {
                 self.instruction(LocalGet(cur_dst_ptr.idx));
-                self.ptr_uconst(dst_mem_opts, dst_key_size);
+                self.ptr_uconst(dst_mem_opts, dst_value_offset);
                 self.ptr_add(dst_mem_opts);
                 self.instruction(LocalSet(cur_dst_ptr.idx));
             }
 
-            // Translate the value
             let value_src = Source::Memory(self.memory_operand(
                 src_opts,
                 TempLocal {
                     idx: cur_src_ptr.idx,
                     ty: src_mem_opts.ptr(),
                     needs_free: false,
                 },
-                src_key_align,
+                src_value_align,
             ));
             let value_dst = Destination::Memory(self.memory_operand(
                 dst_opts,
@@ -2885,52 +2850,49 @@ impl<'a, 'b> Compiler<'a, 'b> {
                     ty: dst_mem_opts.ptr(),
                     needs_free: false,
                 },
-                dst_key_align,
+                dst_value_align,
             ));
-            // Recursively translate the value
             self.translate(&src_map_ty.value, &value_src, &dst_map_ty.value, &value_dst);
 
-            // Advance pointers past the value (including any alignment padding)
-            // If tuple_size > key_size + value_size, there's padding we need to skip
-            if src_tuple_size > src_key_size + src_value_size {
+            // Advance pointers past the value and any trailing padding to
+            // reach the next entry. After value translation, cur_ptr still
+            // points at the value start, so advance by value_size + trailing
+            // padding (i.e. tuple_size - value_offset).
+            let src_advance_to_next = src_tuple_size - src_value_offset;
+            if src_advance_to_next > 0 {
                 self.instruction(LocalGet(cur_src_ptr.idx));
-                self.ptr_uconst(src_mem_opts, src_tuple_size - src_key_size - src_value_size);
+                self.ptr_uconst(src_mem_opts, src_advance_to_next);
                 self.ptr_add(src_mem_opts);
                 self.instruction(LocalSet(cur_src_ptr.idx));
             }
-            if dst_tuple_size > dst_key_size + dst_value_size {
+            let dst_advance_to_next = dst_tuple_size - dst_value_offset;
+            if dst_advance_to_next > 0 {
                 self.instruction(LocalGet(cur_dst_ptr.idx));
-                self.ptr_uconst(dst_mem_opts, dst_tuple_size - dst_key_size - dst_value_size);
+                self.ptr_uconst(dst_mem_opts, dst_advance_to_next);
                 self.ptr_add(dst_mem_opts);
                 self.instruction(LocalSet(cur_dst_ptr.idx));
             }
 
-            // Loop continuation: decrement counter and branch if not done
             self.instruction(LocalGet(remaining.idx));
-            self.ptr_iconst(src_mem_opts, -1); // Push -1
-            self.ptr_add(src_mem_opts); // remaining - 1
-            self.instruction(LocalTee(remaining.idx)); // Save back AND keep on stack
-            self.ptr_br_if(src_mem_opts, 0); // If remaining != 0, branch back to Loop
-            self.instruction(End); // End Loop
-            self.instruction(End); // End Block
-
-            // Release loop locals
+            self.ptr_iconst(src_mem_opts, -1);
+            self.ptr_add(src_mem_opts);
+            self.instruction(LocalTee(remaining.idx));
+            self.ptr_br_if(src_mem_opts, 0);
+            self.instruction(End); // end of loop
+            self.instruction(End); // end of block
             self.free_temp_local(cur_dst_ptr);
             self.free_temp_local(cur_src_ptr);
             self.free_temp_local(remaining);
         }
 
-        // Store the result (ptr, len) to the destination
         match dst {
             Destination::Stack(s, _) => {
-                // Put ptr and len on the wasm stack for return
                 self.instruction(LocalGet(dst_mem.addr.idx));
                 self.stack_set(&s[..1], dst_mem_opts.ptr());
                 self.convert_src_len_to_dst(src_len.idx, src_mem_opts.ptr(), dst_mem_opts.ptr());
                 self.stack_set(&s[1..], dst_mem_opts.ptr());
             }
             Destination::Memory(mem) => {
-                // Store ptr and len to destination memory location
                 self.instruction(LocalGet(mem.addr.idx));
                 self.instruction(LocalGet(dst_mem.addr.idx));
                 self.ptr_store(mem);
@@ -2941,7 +2903,6 @@ impl<'a, 'b> Compiler<'a, 'b> {
             Destination::Struct(_) | Destination::Array(_) => todo!("CM+GC"),
         }
 
-        // Cleanup - release all temporary locals
         self.free_temp_local(src_len);
         self.free_temp_local(src_mem.addr);
         self.free_temp_local(dst_mem.addr);