allow passive wasm segments for manganis (#5163)

* fix passive segment using tls calculation!

* clean things up

* clean diff

* revert some pr changes for review

* cleanups

packages/cli/src/build/assets.rs

@@ -367,6 +367,7 @@ fn find_native_symbol_offsets<'a, R: ReadRef<'a>>(
     Ok(offsets)
 }
 
+/// Evaluate a walrus global expression to get its value.
 fn eval_walrus_global_expr(module: &walrus::Module, expr: &walrus::ConstExpr) -> Option<u64> {
     match expr {
         walrus::ConstExpr::Value(walrus::ir::Value::I32(value)) => Some(*value as u64),
@@ -383,7 +384,26 @@ fn eval_walrus_global_expr(module: &walrus::Module, expr: &walrus::ConstExpr) ->
     }
 }
 
+/// Find the value of a global export by name.
+fn find_global_export_value(module: &walrus::Module, name: &str) -> Option<u64> {
+    for export in module.exports.iter() {
+        if export.name == name {
+            if let walrus::ExportItem::Global(g) = export.item {
+                if let walrus::GlobalKind::Local(expr) = &module.globals.get(g).kind {
+                    return eval_walrus_global_expr(module, expr);
+                }
+            }
+        }
+    }
+    None
+}
+
 /// Find the offsets of any manganis symbols in the wasm file.
+///
+/// This handles both standard WASM builds and builds with advanced features like:
+/// - Bulk memory operations (passive data segments)
+/// - Thread Local Storage (TLS)
+/// - Atomics and shared memory
 fn find_wasm_symbol_offsets<'a, R: ReadRef<'a>>(
     file_contents: &[u8],
     file: &File<'a, R>,
@@ -395,57 +415,85 @@ fn find_wasm_symbol_offsets<'a, R: ReadRef<'a>>(
         tracing::error!("Failed to find <data> section in WASM file");
         return Ok(Vec::new());
     };
+
     let Some((_, section_range_end)) = section.file_range() else {
         tracing::error!("Failed to find file range for <data> section in WASM file");
         return Ok(Vec::new());
     };
+
     let section_size = section.data()?.len() as u64;
     let section_start = section_range_end - section_size;
 
-    // Translate the section_relative_address to the file offset
+    // Parse data segments with wasmparser to get file offsets.
-    // WASM files have a section address of 0 in object, reparse the data section with wasmparser
+    // Walrus doesn't expose file offset information, so we need wasmparser for this.
-    // to get the correct address and section start
+    // With bulk memory operations, there may be multiple data segments.
-    // Note: We need to reparse just the data section with wasmparser to get the file offset because walrus does
-    // not expose the file offset information
     let reader = wasmparser::DataSectionReader::new(wasmparser::BinaryReader::new(
         &file_contents[section_start as usize..section_range_end as usize],
         0,
     ))
     .context("Failed to create WASM data section reader")?;
-    let main_memory = reader
-        .into_iter()
-        .next()
-        .context("Failed find main memory from WASM data section")?
-        .context("Failed to read main memory from WASM data section")?;
-    // main_memory.data is a slice somewhere in file_contents. Find out the offset in the file
-    let data_start_offset = (main_memory.data.as_ptr() as u64)
-        .checked_sub(file_contents.as_ptr() as u64)
-        .expect("Data section start offset should be within the file contents");
 
-    // Parse the wasm file to find the globals
+    // Collect all data segments with their file offsets and sizes
-    let module = walrus::Module::from_buffer(file_contents).unwrap();
+    let mut segment_file_info: Vec<(u64, u64)> = Vec::new();
-    let mut offsets = Vec::new();
+    for segment in reader.into_iter() {
+        let segment = segment.context("Failed to read data segment")?;
+        segment_file_info.push((
+            (segment.data.as_ptr() as u64)
+                .checked_sub(file_contents.as_ptr() as u64)
+                .expect("Data segment should be within file contents"),
+            segment.data.len() as u64,
+        ));
+    }
 
-    // Find the main memory offset
+    if segment_file_info.is_empty() {
-    let main_memory = module
+        return Ok(Vec::new());
+    }
+
+    // Parse the wasm file with walrus to find globals and exports
+    let module = walrus::Module::from_buffer(file_contents)
+        .context("Failed to parse WASM module with walrus")?;
+
+    // Determine the memory base address for symbol lookup
+    let main_memory_walrus = module
         .data
         .iter()
         .next()
         .context("Failed to find main memory in WASM module")?;
 
-    let walrus::DataKind::Active {
+    let main_memory_offset = match &main_memory_walrus.kind {
-        offset: main_memory_offset,
+        walrus::DataKind::Active { offset, .. } => {
-        ..
+            // Active segments have an explicit offset expression
-    } = main_memory.kind
+            eval_walrus_global_expr(&module, offset).unwrap_or_default()
-    else {
+        }
-        tracing::error!("Failed to find main memory offset in WASM module");
+        walrus::DataKind::Passive => {
-        return Ok(Vec::new());
+            // For passive segments (bulk memory operations), there's no static offset.
+            // The memory.init instruction determines placement at runtime.
+            //
+            // Rust/LLVM places static data at 0x100000 (1MB) in linear memory.
+            //
+            // With TLS support, the linker calculates symbol addresses as if TLS data
+            // is at 0x100000 followed by main data. But at runtime, TLS is stored
+            // separately per-thread via __wasm_init_tls. The main data segment starts
+            // at 0x100000 in the file, but symbol addresses include the TLS offset.
+            //
+            // We detect TLS by looking for __tls_size and remove the TLS segment from
+            // our file info, then add the TLS size to the base address.
+            let tls_size = find_global_export_value(&module, "__tls_size");
+
+            // If TLS is present and segment 0 matches TLS size, remove TLS segment
+            if let Some(tls) = tls_size {
+                if !segment_file_info.is_empty() && segment_file_info[0].1 == tls {
+                    segment_file_info.remove(0);
+                }
+            }
+
+            // Base = 0x100000 + TLS size (to account for TLS offset in symbol addresses)
+            0x100000u64 + tls_size.unwrap_or(0)
+        }
     };
 
-    // In the hot patch build, the main memory offset is a global from the main module and each global
+    // Find all manganis symbols and calculate their file offsets
-    // is it's own global. Use an offset of 0 instead if we can't evaluate the global
+    let mut offsets = Vec::new();
-    let main_memory_offset =
-        eval_walrus_global_expr(&module, &main_memory_offset).unwrap_or_default();
 
     for export in module.exports.iter() {
         let Some(version) = looks_like_manganis_symbol(&export.name) else {
@@ -456,7 +504,8 @@ fn find_wasm_symbol_offsets<'a, R: ReadRef<'a>>(
             continue;
         };
 
-        let walrus::GlobalKind::Local(pointer) = module.globals.get(global).kind else {
+        let global_data = module.globals.get(global);
+        let walrus::GlobalKind::Local(pointer) = global_data.kind else {
             continue;
         };
 
@@ -468,11 +517,41 @@ fn find_wasm_symbol_offsets<'a, R: ReadRef<'a>>(
             continue;
         };
 
-        let section_relative_address: u64 = ((virtual_address as i128)
+        // Calculate offset relative to the data base address
-            - main_memory_offset as i128)
+        let data_relative_offset =
-            .try_into()
+            match (virtual_address as i128).checked_sub(main_memory_offset as i128) {
-            .expect("Virtual address should be greater than or equal to section address");
+                Some(offset) if offset >= 0 => offset as u64,
-        let file_offset = data_start_offset + section_relative_address;
+                _ => {
+                    tracing::error!(
+                        "Virtual address 0x{:x} is below main memory offset 0x{:x}",
+                        virtual_address,
+                        main_memory_offset
+                    );
+                    continue;
+                }
+            };
+
+        // Find which segment this offset falls into.
+        // Segments are laid out contiguously in memory.
+        let mut cumulative_offset = 0u64;
+        let mut file_offset = None;
+
+        for (seg_file_offset, seg_size) in segment_file_info.iter() {
+            if data_relative_offset < cumulative_offset + seg_size {
+                let offset_in_segment = data_relative_offset - cumulative_offset;
+                file_offset = Some(seg_file_offset + offset_in_segment);
+                break;
+            }
+            cumulative_offset += seg_size;
+        }
+
+        let Some(file_offset) = file_offset else {
+            tracing::error!(
+                "Virtual address 0x{:x} is beyond all data segments",
+                virtual_address
+            );
+            continue;
+        };
 
         offsets.push(ManganisSymbolOffset::new(version, file_offset));
     }
@@ -530,7 +609,6 @@ pub(crate) async fn extract_assets_from_file(path: impl AsRef<Path>) -> Result<A
         let new_data = version.serialize(asset);
 
         file.seek(std::io::SeekFrom::Start(offset))?;
-        // Write the modified binary data back to the file
         file.write_all(new_data.as_ref())?;
     }
 
@@ -548,7 +626,6 @@ pub(crate) async fn extract_assets_from_file(path: impl AsRef<Path>) -> Result<A
             .arg(path)
             .output()
             .context("Failed to run codesign - is `codesign` in your path?")?;
-
         if !output.status.success() {
             bail!(
                 "Failed to re-sign the binary with codesign after finalizing the assets: {}",

packages/cli/src/wasm_opt.rs

@@ -58,6 +58,7 @@ impl WasmOpt {
             "--enable-bulk-memory",
             "--enable-mutable-globals",
             "--enable-nontrapping-float-to-int",
+            "--enable-threads",
         ];
 
         if self.cfg.memory_packing {