feat: implement support for wasm data segments

Cargo.toml

@@ -1,5 +1,5 @@
 [workspace]
-
+resolver = "2"
 members = [
   "hir",
   "hir-analysis",
@@ -12,7 +12,7 @@ members = [
 
 [workspace.package]
 version = "0.1.0"
-rust-version = "1.67"
+rust-version = "1.71"
 authors = ["Miden Team"]
 description = "An intermediate representation and compiler for Miden Assembly"
 repository = "https://github.com/0xPolygonMiden/miden-ir"

hir/src/constants.rs

@@ -125,6 +125,11 @@ impl ConstantPool {
         &self.constants[&id]
     }
 
+    /// Returns true if this pool contains the given constant data
+    pub fn contains(&self, data: &ConstantData) -> bool {
+        self.cache.contains_key(data)
+    }
+
     /// Insert constant data into the pool, returning a handle for later referencing; when constant
     /// data is inserted that is a duplicate of previous constant data, the existing handle will be
     /// returned.

hir/src/globals.rs

@@ -102,7 +102,6 @@ pub enum ConflictResolutionStrategy {
 
 /// This table is used to lay out and link together global variables for a [Program].
 ///
-///
 /// See the docs for [Linkage], [GlobalVariableData], and [GlobalVariableTable::declare] for more details.
 pub struct GlobalVariableTable {
     layout: LinkedList<GlobalVariableAdapter>,
@@ -129,6 +128,11 @@ impl GlobalVariableTable {
         }
     }
 
+    /// Returns true if the global variable table is empty
+    pub fn is_empty(&self) -> bool {
+        self.layout.is_empty()
+    }
+
     /// Get a double-ended iterator over the current table layout
     pub fn iter<'a, 'b: 'a>(
         &'b self,
@@ -175,18 +179,19 @@ impl GlobalVariableTable {
     }
 
     /// Computes the offset, in bytes, of the given [GlobalVariable] from the
-    /// start of the linear memory heap, assuming that the layout of the global
-    /// variable table up to and including `id` remains unchanged.
+    /// start of the segment in which globals are allocated, assuming that the
+    /// layout of the global variable table up to and including `id` remains
+    /// unchanged.
     ///
     /// # SAFETY
     ///
-    /// This should only be used once all global variables have been declared, and
-    /// the layout of the table has been decided. It is technically safe to use
-    /// offsets obtained before all global variables are declared, _IF_ the layout
-    /// up to and including those global variables remains unchanged after that
-    /// point.
+    /// This should only be used once all data segments and global variables have
+    /// been declared, and the layout of the table has been decided. It is technically
+    /// safe to use offsets obtained before all global variables are declared, _IF_ the
+    /// data segments and global variable layout up to and including those global variables
+    /// remains unchanged after that point.
     ///
-    /// If the offset for a given global variable is obtained, and the layout is
+    /// If the offset for a given global variable is obtained, and the heap layout is
     /// subsequently changed in such a way that the original offset is no longer
     /// accurate, bad things will happen.
     pub unsafe fn offset_of(&self, id: GlobalVariable) -> usize {
@@ -232,6 +237,21 @@ impl GlobalVariableTable {
         self.try_insert(name, ty, linkage)
     }
 
+    /// Get the constant data associated with `id`
+    pub fn get_constant(&self, id: Constant) -> &ConstantData {
+        self.data.get(id)
+    }
+
+    /// Inserts the given constant data into this table without allocating a global
+    pub fn insert_constant(&mut self, data: ConstantData) -> Constant {
+        self.data.insert(data)
+    }
+
+    /// Returns true if the given constant data is in the constant pool
+    pub fn contains_constant(&self, data: &ConstantData) -> bool {
+        self.data.contains(data)
+    }
+
     /// This sets the initializer for the given [GlobalVariable] to `init`.
     ///
     /// This function will return `Err` if any of the following occur:

hir/src/lib.rs

@@ -5,8 +5,12 @@ pub use miden_diagnostics::SourceSpan;
 pub use miden_hir_symbol::{symbols, Symbol};
 pub use miden_hir_type::{FunctionType, Type, TypeRepr};
 
+/// Represents a field element in Miden
 pub type Felt = winter_math::fields::f64::BaseElement;
 
+/// Represents an offset from the base of linear memory in Miden
+pub type Offset = u32;
+
 macro_rules! assert_matches {
     ($left:expr, $(|)? $( $pattern:pat_param )|+ $( if $guard: expr )? $(,)?) => {
         match $left {
@@ -63,6 +67,7 @@ mod layout;
 mod locals;
 mod module;
 mod program;
+mod segments;
 #[cfg(test)]
 mod tests;
 mod value;
@@ -86,6 +91,7 @@ pub use self::layout::{ArenaMap, LayoutAdapter, LayoutNode, OrderedArenaMap};
 pub use self::locals::{Local, LocalId};
 pub use self::module::*;
 pub use self::program::{Linker, LinkerError, Program};
+pub use self::segments::{DataSegment, DataSegmentAdapter, DataSegmentError, DataSegmentTable};
 pub use self::value::{Value, ValueData, ValueList, ValueListPool};
 pub use self::write::{write_external_function, write_function};
 

hir/src/module.rs

@@ -42,6 +42,8 @@ pub struct Module {
     /// Documentation attached to this module, to be passed through to
     /// Miden Assembly during code generation.
     pub docs: Option<String>,
+    /// The set of data segments allocated in this module
+    segments: DataSegmentTable,
     /// The set of global variables declared in this module
     globals: GlobalVariableTable,
     /// The set of functions which belong to this module, in the order
@@ -61,11 +63,48 @@ pub struct Module {
 }
 impl fmt::Display for Module {
     fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
+        use std::fmt::Write;
+
         if self.is_kernel {
             write!(f, "kernel {}\n", &self.name)?;
         } else {
             write!(f, "module {}\n", &self.name)?;
         }
+
+        if !self.segments.is_empty() {
+            f.write_char('\n')?;
+            f.write_str("memory {\n")?;
+            for segment in self.segments.iter() {
+                let data = self.globals.get_constant(segment.init);
+                writeln!(
+                    f,
+                    "    segment @{:#x} x {} = {};",
+                    segment.offset, segment.size, data
+                )?;
+            }
+            f.write_str("}\n\n")?;
+        }
+
+        if !self.globals.is_empty() {
+            for global in self.globals.iter() {
+                write!(
+                    f,
+                    "global {} {} : {}",
+                    global.linkage, global.name, global.ty
+                )?;
+                match global.init {
+                    Some(init) => {
+                        let data = self.globals.get_constant(init);
+                        writeln!(f, " = {} {{ id = {} }};", data, global.id())?;
+                    }
+                    None => {
+                        writeln!(f, " {{ id = {} }};", global.id())?;
+                    }
+                }
+            }
+            writeln!(f)?;
+        }
+
         let mut external_functions = BTreeMap::<FunctionIdent, Signature>::default();
         for function in self.functions.iter() {
             for import in function.dfg.imports() {
@@ -124,6 +163,7 @@ impl Module {
             link: Default::default(),
             name,
             docs: None,
+            segments: Default::default(),
             globals: GlobalVariableTable::new(ConflictResolutionStrategy::None),
             functions: Default::default(),
             is_kernel: false,
@@ -153,6 +193,41 @@ impl Module {
         !self.link.is_linked()
     }
 
+    /// Return an iterator over the data segments allocated in this module
+    ///
+    /// The iterator is double-ended, so can be used to traverse the segments in either direction.
+    ///
+    /// Data segments are ordered by the address at which are are allocated, in ascending order.
+    pub fn segments<'a, 'b: 'a>(
+        &'b self,
+    ) -> intrusive_collections::linked_list::Iter<'a, DataSegmentAdapter> {
+        self.segments.iter()
+    }
+
+    /// Declare a new [DataSegment] in this module, with the given offset, size, and data.
+    ///
+    /// Returns `Err` if the proposed segment overlaps with an existing segment.
+    ///
+    /// Data segments are ordered by the address at which they are allocated, at link-time, all
+    /// segments from all modules are linked together, and they must either be disjoint, or exactly
+    /// identical in order to overlap - it is not permitted to have partially overlapping segments
+    /// with different views of the memory represented by that segment.
+    pub fn declare_data_segment(
+        &mut self,
+        offset: Offset,
+        size: u32,
+        init: ConstantData,
+        readonly: bool,
+    ) -> Result<(), DataSegmentError> {
+        let init_size = init
+            .len()
+            .try_into()
+            .expect("invalid constant data: must be smaller than 2^32 bytes");
+        let size = core::cmp::max(size, init_size);
+        let init = self.globals.insert_constant(init);
+        self.segments.insert(offset, size, init, readonly)
+    }
+
     /// Return an iterator over the global variables declared in this module
     ///
     /// The iterator is double-ended, so can be used to traverse the globals table in either direction

hir/src/program.rs

@@ -28,6 +28,8 @@ pub struct Program {
     /// a program or just a collection of modules; and in the case of the former, what code
     /// to emit in the root code block.
     entrypoint: Option<FunctionIdent>,
+    /// The data segments gathered from all modules in the program, and laid out in address order.
+    segments: DataSegmentTable,
     /// The global variable table produced by linking the global variable tables of all
     /// modules in this program. The layout of this table corresponds to the layout of
     /// global variables in the linear memory heap at runtime.
@@ -60,6 +62,7 @@ impl Program {
         Self {
             modules: Default::default(),
             entrypoint: None,
+            segments: Default::default(),
             globals: Default::default(),
         }
     }
@@ -97,6 +100,17 @@ impl Program {
         entry
     }
 
+    /// Return an iterator over the data segments allocated in this program
+    ///
+    /// The iterator is double-ended, so can be used to traverse the segments in either direction.
+    ///
+    /// Data segments are ordered by the address at which are are allocated, in ascending order.
+    pub fn segments<'a, 'b: 'a>(
+        &'b self,
+    ) -> intrusive_collections::linked_list::Iter<'a, DataSegmentAdapter> {
+        self.segments.iter()
+    }
+
     /// Get a reference to the global variable table for this program
     pub fn globals(&self) -> &GlobalVariableTable {
         &self.globals