feat: Wasm data section parsing

frontend-wasm/Cargo.toml

@@ -22,6 +22,7 @@ log.workspace = true
 anyhow.workspace = true
 wasmparser = "0.107"
 itertools = "0.11"
+derive_more = "0.99"
 
 [dev-dependencies]
 wat = "1.0.69"

frontend-wasm/src/environ/module_env.rs

@@ -4,7 +4,8 @@ use crate::error::{WasmError, WasmResult};
 use crate::func_translator::FuncTranslator;
 use crate::translation_utils::sig_from_funct_type;
 use crate::wasm_types::{
-    DefinedFuncIndex, FuncIndex, Global, GlobalIndex, Memory, MemoryIndex, TypeIndex,
+    DataSegment, DataSegmentIndex, DefinedFuncIndex, FuncIndex, Global, GlobalIndex, Memory,
+    MemoryIndex, TypeIndex,
 };
 use miden_diagnostics::{DiagnosticsHandler, SourceSpan};
 use miden_hir::cranelift_entity::{EntityRef, PrimaryMap, SecondaryMap};
@@ -43,6 +44,12 @@ pub struct ModuleInfo {
     /// Global names.
     global_names: SecondaryMap<GlobalIndex, String>,
 
+    /// Data segments declared in the module.
+    pub data_segments: PrimaryMap<DataSegmentIndex, DataSegment>,
+
+    /// Data segment names.
+    pub data_segment_names: SecondaryMap<DataSegmentIndex, String>,
+
     /// The start function.
     pub start_func: Option<FuncIndex>,
 }
@@ -59,6 +66,8 @@ impl ModuleInfo {
             globals: PrimaryMap::new(),
             function_names: SecondaryMap::new(),
             global_names: SecondaryMap::new(),
+            data_segments: PrimaryMap::new(),
+            data_segment_names: SecondaryMap::new(),
         }
     }
 
@@ -120,6 +129,7 @@ impl<'a> ModuleEnvironment<'a> {
     ) -> WasmResult<Module> {
         let mut module_builder = ModuleBuilder::new(self.info.name.as_str());
         self.build_globals(&mut module_builder, diagnostics)?;
+        self.build_data_segments(&mut module_builder, diagnostics)?;
         let get_num_func_imports = self.get_num_func_imports();
         for (def_func_index, body) in &self.function_bodies {
             let func_index = FuncIndex::new(get_num_func_imports + def_func_index.index());
@@ -152,7 +162,7 @@ impl<'a> ModuleEnvironment<'a> {
     }
 
     fn build_globals(
-        &mut self,
+        &self,
         module_builder: &mut ModuleBuilder,
         diagnostics: &DiagnosticsHandler,
     ) -> Result<(), WasmError> {
@@ -178,6 +188,30 @@ impl<'a> ModuleEnvironment<'a> {
         })
     }
 
+    fn build_data_segments(
+        &self,
+        _module_builder: &mut ModuleBuilder,
+        diagnostics: &DiagnosticsHandler,
+    ) -> Result<(), WasmError> {
+        for (data_segment_idx, data_segment) in &self.info.data_segments {
+            let data_segment_name = self.info.data_segment_names[data_segment_idx].clone();
+            let _readonly = data_segment_name == ".rodata";
+            let _init = ConstantData::from(data_segment.data.clone());
+            let _offset = data_segment
+                .offset
+                .as_i32(&self.info.globals, diagnostics)?;
+            // if let Err(e) = module_builder.declare_data_segment(offset, size, init, readonly) {
+            //     let message = format!("Failed to declare data segment {init} for data segment size {size} at {offset} with error: {:?}", e);
+            //     diagnostics
+            //         .diagnostic(miden_diagnostics::Severity::Error)
+            //         .with_message(message.clone())
+            //         .emit();
+            //     return Err(WasmError::Unexpected(message));
+            // }
+        }
+        Ok(())
+    }
+
     /// Declares a function signature to the environment.
     pub fn declare_type_func(&mut self, func_type: FunctionType) {
         self.info.func_types.push(func_type);
@@ -237,6 +271,14 @@ impl<'a> ModuleEnvironment<'a> {
         self.info.function_names[func_index] = String::from(name);
     }
 
+    pub fn declare_data_segment(&mut self, segment: DataSegment) {
+        self.info.data_segments.push(segment);
+    }
+
+    pub fn declare_data_segment_name(&mut self, segment_index: DataSegmentIndex, name: &'a str) {
+        self.info.data_segment_names[segment_index] = String::from(name);
+    }
+
     /// Indicates that a custom section has been found in the wasm file
     pub fn custom_section(&mut self, _name: &'a str, _data: &'a [u8]) {
         // Do we need to support custom sections?

frontend-wasm/src/sections_translator.rs

@@ -6,14 +6,15 @@ use crate::{
     error::{WasmError, WasmResult},
     unsupported_diag,
     wasm_types::{
-        convert_func_type, convert_global_type, FuncIndex, GlobalIndex, GlobalInit, TypeIndex,
+        convert_func_type, convert_global_type, DataSegment, DataSegmentIndex, DataSegmentOffset,
+        FuncIndex, GlobalIndex, GlobalInit, TypeIndex,
     },
 };
 use miden_diagnostics::DiagnosticsHandler;
 use wasmparser::{
-    DataSectionReader, ElementSectionReader, FunctionSectionReader, GlobalSectionReader,
-    ImportSectionReader, MemorySectionReader, NameSectionReader, Naming, Operator, Type, TypeRef,
-    TypeSectionReader,
+    Data, DataKind, DataSectionReader, ElementSectionReader, FunctionSectionReader,
+    GlobalSectionReader, ImportSectionReader, MemorySectionReader, NameSectionReader, Naming,
+    Operator, Type, TypeRef, TypeSectionReader,
 };
 
 /// Parses the Type section of the wasm module.
@@ -139,11 +140,49 @@ pub fn parse_element_section<'a>(
 
 /// Parses the Data section of the wasm module.
 pub fn parse_data_section<'a>(
-    _data: DataSectionReader<'a>,
-    _environ: &mut ModuleEnvironment<'a>,
-    _diagnostics: &DiagnosticsHandler,
+    data: DataSectionReader<'a>,
+    environ: &mut ModuleEnvironment<'a>,
+    diagnostics: &DiagnosticsHandler,
 ) -> WasmResult<()> {
-    todo!("Data section are not yet implemented");
+    for (_index, entry) in data.into_iter().enumerate() {
+        let Data {
+            kind,
+            data,
+            range: _,
+        } = entry?;
+        match kind {
+            DataKind::Active {
+                // ignored, since for Wasm spec v1 it's always 0
+                memory_index: _,
+                offset_expr,
+            } => {
+                let mut offset_expr_reader = offset_expr.get_binary_reader();
+                let offset = match offset_expr_reader.read_operator()? {
+                    Operator::I32Const { value } => DataSegmentOffset::I32Const(value),
+                    Operator::GlobalGet { global_index } => {
+                        DataSegmentOffset::GetGlobal(GlobalIndex::from_u32(global_index))
+                    }
+                    ref s => {
+                        unsupported_diag!(
+                            diagnostics,
+                            "unsupported init expr in data section offset: {:?}",
+                            s
+                        );
+                    }
+                };
+                let segment = DataSegment {
+                    offset,
+                    data: data.to_owned(),
+                };
+                environ.declare_data_segment(segment);
+            }
+            DataKind::Passive => {
+                // Passive data segments type is added in Wasm spec 2.0
+                unsupported_diag!(diagnostics, "Passive data segments are not supported");
+            }
+        }
+    }
+    Ok(())
 }
 
 /// Parses the Name section of the wasm module.
@@ -185,12 +224,19 @@ pub fn parse_name_section<'a>(
                     }
                 }
             }
+            wasmparser::Name::Data(names) => {
+                for name in names {
+                    let Naming { index, name } = name?;
+                    if index != u32::max_value() {
+                        environ.declare_data_segment_name(DataSegmentIndex::from_u32(index), name);
+                    }
+                }
+            }
             wasmparser::Name::Label(_)
             | wasmparser::Name::Type(_)
             | wasmparser::Name::Table(_)
             | wasmparser::Name::Memory(_)
             | wasmparser::Name::Element(_)
-            | wasmparser::Name::Data(_)
             | wasmparser::Name::Unknown { .. } => {}
         }
     }

frontend-wasm/src/wasm_types.rs

@@ -1,5 +1,6 @@
 //! Internal types for parsed WebAssembly.
 
+use miden_diagnostics::DiagnosticsHandler;
 use miden_hir::cranelift_entity::entity_impl;
 use miden_hir::cranelift_entity::PrimaryMap;
 use miden_hir_type::FunctionType;
@@ -8,6 +9,7 @@ use miden_hir_type::Type;
 use crate::environ::ModuleInfo;
 use crate::error::WasmError;
 use crate::error::WasmResult;
+use crate::unsupported_diag;
 
 /// Index type of a function (imported or defined) inside the WebAssembly module.
 #[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug)]
@@ -20,7 +22,7 @@ pub struct DefinedFuncIndex(u32);
 entity_impl!(DefinedFuncIndex);
 
 /// Index type of a global variable (imported or defined) inside the WebAssembly module.
-#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug)]
+#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug, derive_more::Display)]
 pub struct GlobalIndex(u32);
 entity_impl!(GlobalIndex);
 
@@ -34,6 +36,11 @@ entity_impl!(MemoryIndex);
 pub struct TypeIndex(u32);
 entity_impl!(TypeIndex);
 
+/// Index type of a data segment inside the WebAssembly module.
+#[derive(Copy, Clone, PartialEq, Eq, Hash, PartialOrd, Ord, Debug)]
+pub struct DataSegmentIndex(u32);
+entity_impl!(DataSegmentIndex);
+
 /// A WebAssembly global.
 #[derive(Debug, Clone, Hash, Eq, PartialEq)]
 pub struct Global {
@@ -74,6 +81,23 @@ impl GlobalInit {
             }
         }
     }
+
+    pub fn as_i32(
+        &self,
+        globals: &PrimaryMap<GlobalIndex, Global>,
+        diagnostics: &DiagnosticsHandler,
+    ) -> WasmResult<i32> {
+        Ok(match self {
+            GlobalInit::I32Const(x) => *x,
+            GlobalInit::GetGlobal(global_idx) => {
+                let global = &globals[*global_idx];
+                global.init.as_i32(globals, diagnostics)?
+            }
+            g => {
+                unsupported_diag!(diagnostics, "Expected global init to be i32, got: {:?}", g);
+            }
+        })
+    }
 }
 
 /// WebAssembly linear memory.
@@ -97,6 +121,53 @@ impl From<wasmparser::MemoryType> for Memory {
     }
 }
 
+/// Offset of a data segment inside a linear memory.
+#[derive(Debug, Clone, Copy, Hash, Eq, PartialEq)]
+pub enum DataSegmentOffset {
+    /// An `i32.const` offset.
+    I32Const(i32),
+    /// An offset as a `global.get` of another global.
+    GetGlobal(GlobalIndex),
+}
+
+impl DataSegmentOffset {
+    /// Returns the offset as a i32, resolving the global if necessary.
+    pub fn as_i32(
+        &self,
+        globals: &PrimaryMap<GlobalIndex, Global>,
+        diagnostics: &DiagnosticsHandler,
+    ) -> WasmResult<i32> {
+        Ok(match self {
+            DataSegmentOffset::I32Const(x) => *x,
+            DataSegmentOffset::GetGlobal(global_idx) => {
+                let global = &globals[*global_idx];
+                match global.init.as_i32(globals, diagnostics) {
+                    Err(e) => {
+                        diagnostics
+                            .diagnostic(miden_diagnostics::Severity::Error)
+                            .with_message(format!(
+                                "Failed to get data segment offset from global init {:?} with global index {global_idx}",
+                                global.init,
+                            ))
+                            .emit();
+                        return Err(e);
+                    }
+                    Ok(v) => v,
+                }
+            }
+        })
+    }
+}
+
+/// A WebAssembly data segment.
+/// https://www.w3.org/TR/wasm-core-1/#data-segments%E2%91%A0
+pub struct DataSegment {
+    /// The offset of the data segment inside the linear memory.
+    pub offset: DataSegmentOffset,
+    /// The initialization data.
+    pub data: Vec<u8>,
+}
+
 #[derive(Debug, Clone, PartialEq, Eq, Hash, Default)]
 pub struct BlockType {
     pub params: Vec<Type>,

frontend-wasm/tests/rust_source/array.rs

@@ -0,0 +1,18 @@
+#![no_std]
+#![no_main]
+
+#[panic_handler]
+fn my_panic(_info: &core::panic::PanicInfo) -> ! {
+    loop {}
+}
+
+#[inline(never)]
+#[no_mangle]
+pub fn sum_arr(arr: &[u32]) -> u32 {
+    arr.iter().sum()
+}
+
+#[no_mangle]
+pub extern "C" fn __main() -> u32 {
+    sum_arr(&[1, 2, 3, 4, 5]) + sum_arr(&[6, 7, 8, 9, 10])
+}