Repository: nexus-xyz/nexus-zkvm
File: vm/src/memory/unified.rs (Rust)

Code before:
                    Err(MemoryError::UndefinedMemoryRegion)
                }
            }
            Some(Modes::WO) => {
                if idx < self.fwo_store.len() {
                    Ok(self.fwo_store[idx].addr_val_bytes())
                } else {
                    Err(MemoryError::UndefinedMemoryRegion)
                }
            }
            Some(Modes::NA) => {
                if idx < self.fna_store.len() {
                    Ok(self.fna_store[idx].addr_val_bytes())
                } else {
                    Err(MemoryError::UndefinedMemoryRegion)
                }
            }
            _ => Err(MemoryError::UndefinedMemoryRegion),
        }
    }

    /// Iterator over addresses and values in the fixed memory, given bytewise, to avoid BTreeMap allocation
    pub fn addr_val_bytes_iter(
        &self,
        uidx: (usize, usize),
    ) -> Result<
        itertools::Either<
            std::iter::Empty<(u32, u8)>,
            itertools::Either<
                impl Iterator<Item = (u32, u8)> + '_,
                itertools::Either<
                    impl Iterator<Item = (u32, u8)> + '_,
                    itertools::Either<
                        impl Iterator<Item = (u32, u8)> + '_,
                        impl Iterator<Item = (u32, u8)> + '_,
                    >,
                >,
            >,
        >,
        MemoryError,
    > {
        let (store, idx) = uidx;
        match num_traits::FromPrimitive::from_usize(store) {
            Some(Modes::RW) => {
                if idx < self.frw_store.len() {
                    Ok(itertools::Either::Right(itertools::Either::Left(
                        self.frw_store[idx].addr_val_bytes_iter(),
                    )))
                } else {
                    Err(MemoryError::UndefinedMemoryRegion)
                }

Review comment (security):
This works, but is very unclear and unwieldy, and so I'd prefer not to merge it.

I'd propose the following alternate approach. The four modes on `FixedMemory` form a security lattice:

```
          RW
        /    \
      RO      WO
        \    /
          NA
```

As such, it is perfectly valid to coerce the fixed memories "down the lattice", since a RW memory can, e.g., be safely handled as a RO memory (but not vice-versa).

So I suggest implementing `From` on `FixedMemory` that allows converting `RW` -> `WO/RO` and `WO` -> `NA` and `RO` -> `NA` (for efficiency you could also add `RW` -> `NA` directly). You should then just be able to do something like `self.frw_store[idx].into<FixedMemory<NA>>.addr_val_bytes_iter()` for each of them, without needing to use `Either` in such a convoluted way.

Code after:
                    Err(MemoryError::UndefinedMemoryRegion)
                }
            }
            Some(Modes::WO) => {
                if idx < self.fwo_store.len() {
                    Ok(self.fwo_store[idx].addr_val_bytes())
                } else {
                    Err(MemoryError::UndefinedMemoryRegion)
                }
            }
            Some(Modes::NA) => {
                if idx < self.fna_store.len() {
                    Ok(self.fna_store[idx].addr_val_bytes())
                } else {
                    Err(MemoryError::UndefinedMemoryRegion)
                }
            }
            _ => Err(MemoryError::UndefinedMemoryRegion),
        }
    }

    pub fn segment_words(
        &self,
        uidx: (usize, usize),
        start: u32,
        end: Option<u32>,
    ) -> Result<&[u32], MemoryError> {
        let (store, idx) = uidx;

        match FromPrimitive::from_usize(store) {
            Some(Modes::RW) => {
                if idx < self.frw_store.len() {
                    Ok(self.frw_store[idx].segment_words(start, end))
                } else {
                    Err(MemoryError::UndefinedMemoryRegion)
                }
            }
            Some(Modes::RO) => {
                if idx < self.fro_store.len() {
                    Ok(self.fro_store[idx].segment_words(start, end))
                } else {
                    Err(MemoryError::UndefinedMemoryRegion)
                }
            }
            Some(Modes::WO) => {
                if idx < self.fwo_store.len() {
                    Ok(self.fwo_store[idx].segment_words(start, end))
                } else {
                    Err(MemoryError::UndefinedMemoryRegion)
                }
            }