From 6b78816c1dd0f1c6670362c2d74b725895217138 Mon Sep 17 00:00:00 2001 From: Jonathan Kelley Date: Wed, 7 Jan 2026 02:35:08 -0800 Subject: [PATCH] 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 | 153 +++++++++++++++++++++++-------- packages/cli/src/wasm_opt.rs | 1 + 2 files changed, 116 insertions(+), 38 deletions(-) diff --git a/packages/cli/src/build/assets.rs b/packages/cli/src/build/assets.rs index 7ccc50b65..dc9df2078 100644 --- a/packages/cli/src/build/assets.rs +++ b/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 { 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 { + 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 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 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 - // WASM files have a section address of 0 in object, reparse the data section with wasmparser - // to get the correct address and section start - // 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 + // Parse data segments with wasmparser to get file offsets. + // Walrus doesn't expose file offset information, so we need wasmparser for this. + // With bulk memory operations, there may be multiple data segments. 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 - let module = walrus::Module::from_buffer(file_contents).unwrap(); - let mut offsets = Vec::new(); + // Collect all data segments with their file offsets and sizes + let mut segment_file_info: Vec<(u64, u64)> = 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 - let main_memory = module + if segment_file_info.is_empty() { + 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 { - offset: main_memory_offset, - .. - } = main_memory.kind - else { - tracing::error!("Failed to find main memory offset in WASM module"); - return Ok(Vec::new()); + let main_memory_offset = match &main_memory_walrus.kind { + walrus::DataKind::Active { offset, .. } => { + // Active segments have an explicit offset expression + eval_walrus_global_expr(&module, offset).unwrap_or_default() + } + walrus::DataKind::Passive => { + // 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 - // is it's own global. Use an offset of 0 instead if we can't evaluate the global - let main_memory_offset = - eval_walrus_global_expr(&module, &main_memory_offset).unwrap_or_default(); + // Find all manganis symbols and calculate their file offsets + let mut offsets = Vec::new(); 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) - - main_memory_offset as i128) - .try_into() - .expect("Virtual address should be greater than or equal to section address"); - let file_offset = data_start_offset + section_relative_address; + // Calculate offset relative to the data base address + let data_relative_offset = + match (virtual_address as i128).checked_sub(main_memory_offset as i128) { + Some(offset) if offset >= 0 => offset as u64, + _ => { + 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) -> Result) -> Result