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

Code before:
//!
//! Comprehensive unit tests are included to verify the correct routing of memory operations
//! to different regions, handling of access permissions, and proper fallback to variable memory.
//!
//! # Performance Considerations
//!
//! The use of `RangeMap` for memory layout allows for efficient lookup of the correct memory
//! region for a given address. However, the performance may vary depending on the number and
//! size of fixed memory regions.
use nexus_common::error::MemoryError;
use num_derive::FromPrimitive;
use num_traits::FromPrimitive;
use rangemap::RangeMap;
use std::{
    collections::BTreeMap,
    fmt::{Debug, Display, Formatter, Result as FmtResult},
};

use super::{
    FixedMemory, LoadOp, MemAccessSize, MemoryProcessor, StoreOp, VariableMemory, NA, RO, RW, WO,
};
use crate::memory::fixed::FixedMemoryAddrValBytesIter;

// Helper to cast &FixedMemory<RW> to &FixedMemory<RO>
fn as_ro(mem: &FixedMemory<RW>) -> &FixedMemory<RO> {
    unsafe { &*(mem as *const FixedMemory<RW> as *const FixedMemory<RO>) }
}
// Helper to cast &FixedMemory<WO> to &FixedMemory<NA>
fn as_na_from_wo(mem: &FixedMemory<WO>) -> &FixedMemory<NA> {
    unsafe { &*(mem as *const FixedMemory<WO> as *const FixedMemory<NA>) }
}

#[derive(Debug, Clone, Eq, PartialEq, FromPrimitive)]
pub enum Modes {
    NA = 0,
    RO = 1,
    WO = 2,
    RW = 3,
}

// nb: we store outside the map becaues `rangemap::RangeMap` does not support a `get_mut` interface (https://github.com/jeffparsons/rangemap/issues/85)
#[derive(Default, Clone)]
pub struct UnifiedMemory {
    // lookup for correct fixed memory, if any
    meta: RangeMap<u32, Modes>,
    // lookup and storage for fixed read-write memories
    frw: RangeMap<u32, usize>,
    frw_store: Vec<FixedMemory<RW>>,
    // lookup and storage for fixed read-only memories
    fro: RangeMap<u32, usize>,
    fro_store: Vec<FixedMemory<RO>>,

Review comment (performance):
As cryptographic software we have a very strong aversion to using `unsafe` wherever possible.

Moreover, is it not possible to just do

```
self.frw_store[idx].clone().into<FixedMemory<NA>>()
```

which will implicitly use the `From` implementation in `vm/src/memory/fixed.rs`? The compiler should do a good job of making the cloning/conversion quite fast, and though slower than the recast approach is I think worth doing to avoid any `unsafe` issues.

Code after:
//!
//! Comprehensive unit tests are included to verify the correct routing of memory operations
//! to different regions, handling of access permissions, and proper fallback to variable memory.
//!
//! # Performance Considerations
//!
//! The use of `RangeMap` for memory layout allows for efficient lookup of the correct memory
//! region for a given address. However, the performance may vary depending on the number and
//! size of fixed memory regions.
use nexus_common::error::MemoryError;
use num_derive::FromPrimitive;
use num_traits::FromPrimitive;
use rangemap::RangeMap;
use std::{
    collections::BTreeMap,
    fmt::{Debug, Display, Formatter, Result as FmtResult},
};

use super::{
    FixedMemory, LoadOp, MemAccessSize, MemoryProcessor, StoreOp, VariableMemory, NA, RO, RW, WO,
};

#[derive(Debug, Clone, Eq, PartialEq, FromPrimitive)]
pub enum Modes {
    NA = 0,
    RO = 1,
    WO = 2,
    RW = 3,
}

// nb: we store outside the map becaues `rangemap::RangeMap` does not support a `get_mut` interface (https://github.com/jeffparsons/rangemap/issues/85)
#[derive(Default, Clone)]
pub struct UnifiedMemory {
    // lookup for correct fixed memory, if any
    meta: RangeMap<u32, Modes>,
    // lookup and storage for fixed read-write memories
    frw: RangeMap<u32, usize>,
    pub frw_store: Vec<FixedMemory<RW>>,
    // lookup and storage for fixed read-only memories
    fro: RangeMap<u32, usize>,
    pub fro_store: Vec<FixedMemory<RO>>,
    // lookup and storage for fixed write-only memories
    fwo: RangeMap<u32, usize>,
    pub fwo_store: Vec<FixedMemory<WO>>,
    // lookup and storage for fixed no-access memories
    fna: RangeMap<u32, usize>,
    pub fna_store: Vec<FixedMemory<NA>>,
    // fallback variable read-write memory for all other addresses
    vrw: Option<VariableMemory<RW>>,
}