Read files with memory mapping

Cargo.toml

@@ -19,6 +19,7 @@ anyhow = "1.0"
 lexopt = "0.3.0"
 libc = "0.2"
 rustc-hash = "1.1.0"
+memmap2 = "0.9.5"
 
 [target.'cfg(windows)'.dependencies.windows-sys]
 version = "0.48"
@@ -50,6 +51,10 @@ harness = false
 name = "canon"
 harness = false
 
+[[bench]]
+name = "loader"
+harness = false
+
 [features]
 default = ["jemalloc"]
 jemalloc = ["jemallocator"]

benches/loader.rs

@@ -0,0 +1,24 @@
+use divan::Bencher;
+
+mod loader {
+    use super::*;
+    use n2::load;
+
+    #[divan::bench(sample_size = 3, sample_count = 3)]
+    fn file_via_loader(bencher: Bencher) {
+        bencher.bench_local(|| {
+            load::testing::read_internal("benches/build.ninja").unwrap();
+        });
+    }
+
+    #[divan::bench(sample_size = 3, sample_count = 3)]
+    fn file_via_loader_slow(bencher: Bencher) {
+        bencher.bench_local(|| {
+            load::testing::read_internal_slow("benches/build.ninja").unwrap();
+        });
+    }
+}
+
+fn main() {
+    divan::main();
+}

benches/parse.rs

@@ -32,7 +32,7 @@ mod parser {
     }
 
     // This can take a while to run (~100ms per sample), so reduce total count.
-    #[divan::bench(sample_size = 3, max_time = 1)]
+    #[divan::bench(sample_size = 3, sample_count = 3)]
     fn file(bencher: Bencher) {
         let input = match n2::scanner::read_file_with_nul("benches/build.ninja".as_ref()) {
             Ok(input) => input,

rust-toolchain.toml

@@ -0,0 +1,3 @@
+[toolchain]
+channel = "nightly-2025-04-26"
+profile = "complete"

src/densemap.rs

@@ -50,6 +50,10 @@ impl<K: Index, V> DenseMap<K, V> {
         self.vec.push(val);
         id
     }
+
+    pub fn all_ids(&self) -> impl Iterator<Item = K> {
+        (0..self.vec.len()).map(|id| K::from(id))
+    }
 }
 
 impl<K: Index, V: Clone> DenseMap<K, V> {

src/file_loader.rs

@@ -0,0 +1,97 @@
+use std::{
+    cell::RefCell, fs::File, path::PathBuf, rc::Rc, sync::Arc
+};
+use libc::{
+    c_void, mmap, munmap, sysconf, MAP_ANONYMOUS, MAP_FAILED, MAP_FIXED, MAP_PRIVATE,
+    PROT_READ, PROT_WRITE, _SC_PAGESIZE,
+};
+
+use anyhow::{anyhow, bail, Result};
+use ref_cell_hash_map::RefCellHashMapFileMap;
+// use libc::{sysconf, _SC_PAGESIZE};
+use memmap2::{Mmap, MmapOptions};
+
+use crate::graph::{FileId, Graph};
+
+pub mod ref_cell_hash_map;
+
+pub trait FileMap {
+    fn new() -> Self where Self : Sized;
+    fn read_file(&self, file_id: FileId, file_path: &PathBuf) -> Result<FileHandle>;
+}
+
+pub struct FileLoader<FileMapT : FileMap = RefCellHashMapFileMap> {
+    graph: Rc<RefCell<Graph>>,
+    files: FileMapT,
+}
+
+impl FileLoader {
+    pub fn new(graph: Rc<RefCell<Graph>>) -> Self {
+        FileLoader { graph, files: RefCellHashMapFileMap::new() }
+    }
+
+    pub fn read_file(&self, file_id: FileId) -> Result<FileHandle> {
+        let file_path = self.graph.borrow().file(file_id).path().to_path_buf();
+
+        self.files.read_file(file_id, &file_path)
+            .map_err(|err| anyhow!("read {}: {}", file_path.display(), err))
+    }
+}
+
+#[derive(Clone)]
+pub struct FileHandle {
+    pub size: usize,
+    pub path: PathBuf,
+
+    mmap: Option<Arc<Mmap>>,
+}
+
+impl FileHandle {
+    fn from_path(path: &PathBuf) -> Result<FileHandle> {
+        let file = Arc::new(File::options().read(true).write(true).open(path)?);
+        let metadata = file.metadata()?;
+        let size = metadata.len() as usize;
+
+        let mmap = if size > 0 {
+            let mmap = unsafe {
+                let page_size = sysconf(_SC_PAGESIZE) as usize;
+                let mapping_size = (size + page_size).next_multiple_of(page_size);
+                let mut mmap = MmapOptions::new().len(mapping_size).map_copy(&file)?;
+
+                // TODO: upstream it to the memmap2 crate? Also what about other platforms?
+                let addr2 = libc::mmap(
+                    mmap.as_ptr_range().end.sub(page_size) as *mut c_void,
+                    page_size,
+                    PROT_READ | PROT_WRITE,
+                    MAP_PRIVATE | MAP_ANONYMOUS | MAP_FIXED,
+                    -1,
+                    0,
+                );
+                if addr2 == MAP_FAILED {
+                    bail!("mmap failed");
+                }
+
+                // Ensure we have a 0 byte at the end
+                mmap[size] = 0;
+
+                mmap.make_read_only()?
+            };
+
+            Some(Arc::new(mmap))
+        } else {
+            None
+        };
+
+        Ok(FileHandle { mmap, path: path.to_path_buf(), size })
+    }
+}
+
+impl AsRef<[u8]> for FileHandle {
+    #[inline]
+    fn as_ref(&self) -> &[u8] {
+        match self.mmap {
+            Some(ref mmap) => mmap[0 ..= self.size].as_ref(),
+            None => &[0],
+        }
+    }
+}

src/file_loader/ref_cell_hash_map.rs

@@ -0,0 +1,31 @@
+use std::{cell::RefCell, path::PathBuf};
+use std::collections::HashMap;
+
+use anyhow::Result;
+
+use crate::graph::{FileId};
+
+use super::{FileHandle, FileMap};
+
+#[derive(Default)]
+pub struct RefCellHashMapFileMap(RefCell<HashMap<FileId, FileHandle>>);
+
+impl FileMap for RefCellHashMapFileMap {
+    fn new() -> Self where Self : Sized {
+        RefCellHashMapFileMap(RefCell::new(HashMap::with_capacity(16)))
+    }
+
+    fn read_file(&self, file_id: FileId, file_path: &PathBuf) -> Result<FileHandle> {
+        let file = { self.0.borrow().get(&file_id).map(Clone::clone) };
+        match file {
+            Some(handle) => Ok(handle),
+            None => {
+                let handle = FileHandle::from_path(file_path)?;
+                let _ = self.0.borrow_mut().insert(file_id, handle.clone());
+
+                Ok(handle)
+            }
+        }
+    }
+}
+