Add iteration and standard impls in once crate

ext/crates/once/Cargo.toml

@@ -15,6 +15,7 @@ maybe-rayon = { path = "../maybe-rayon" }
 
 [dev-dependencies]
 criterion = "0.5"
+expect-test = "1.1.0"
 pprof = { version = "0.15", features = ["criterion", "flamegraph"] }
 proptest = "1.7.0"
 rand = "0.9"

ext/crates/once/benches/criterion/benchable_impl.rs

@@ -18,6 +18,13 @@ impl<T> Benchable<1, T> for OnceBiVec<T> {
     fn get(&self, coords: [i32; 1]) -> Option<&T> {
         self.get(coords[0])
     }
+
+    fn iter<'a>(&'a self) -> impl Iterator<Item = ([i32; 1], &'a T)>
+    where
+        T: 'a,
+    {
+        self.iter_enum().map(|(i, v)| ([i], v))
+    }
 }
 
 impl<T> Benchable<2, T> for OnceBiVec<OnceBiVec<T>> {
@@ -44,6 +51,14 @@ impl<T> Benchable<2, T> for OnceBiVec<OnceBiVec<T>> {
         }
         layer1.get(coords[1])
     }
+
+    fn iter<'a>(&'a self) -> impl Iterator<Item = ([i32; 2], &'a T)>
+    where
+        T: 'a,
+    {
+        Benchable::<1, _>::iter(self)
+            .flat_map(|(start, v)| v.iter_enum().map(move |(end, val)| ([start[0], end], val)))
+    }
 }
 
 impl<T> Benchable<3, T> for OnceBiVec<OnceBiVec<OnceBiVec<T>>> {
@@ -75,6 +90,16 @@ impl<T> Benchable<3, T> for OnceBiVec<OnceBiVec<OnceBiVec<T>>> {
         }
         layer2.get(coords[2])
     }
+
+    fn iter<'a>(&'a self) -> impl Iterator<Item = ([i32; 3], &'a T)>
+    where
+        T: 'a,
+    {
+        Benchable::<2, _>::iter(self).flat_map(|(start, v)| {
+            v.iter_enum()
+                .map(move |(end, val)| ([start[0], start[1], end], val))
+        })
+    }
 }
 
 impl<T> Benchable<4, T> for OnceBiVec<OnceBiVec<OnceBiVec<OnceBiVec<T>>>> {
@@ -111,6 +136,16 @@ impl<T> Benchable<4, T> for OnceBiVec<OnceBiVec<OnceBiVec<OnceBiVec<T>>>> {
         }
         layer3.get(coords[3])
     }
+
+    fn iter<'a>(&'a self) -> impl Iterator<Item = ([i32; 4], &'a T)>
+    where
+        T: 'a,
+    {
+        Benchable::<3, _>::iter(self).flat_map(|(start, v)| {
+            v.iter_enum()
+                .map(move |(end, val)| ([start[0], start[1], start[2], end], val))
+        })
+    }
 }
 
 impl<T> Benchable<5, T> for OnceBiVec<OnceBiVec<OnceBiVec<OnceBiVec<OnceBiVec<T>>>>> {
@@ -152,4 +187,14 @@ impl<T> Benchable<5, T> for OnceBiVec<OnceBiVec<OnceBiVec<OnceBiVec<OnceBiVec<T>
         }
         layer4.get(coords[4])
     }
+
+    fn iter<'a>(&'a self) -> impl Iterator<Item = ([i32; 5], &'a T)>
+    where
+        T: 'a,
+    {
+        Benchable::<4, _>::iter(self).flat_map(|(start, v)| {
+            v.iter_enum()
+                .map(move |(end, val)| ([start[0], start[1], start[2], start[3], end], val))
+        })
+    }
 }

ext/crates/once/benches/criterion/main.rs

@@ -3,7 +3,7 @@ use std::hint::black_box;
 use criterion::{
     BenchmarkGroup, Criterion, criterion_group, criterion_main, measurement::WallTime,
 };
-use once::{MultiIndexed, OnceBiVec, TwoEndedGrove};
+use once::{MultiIndexed, OnceBiVec, TwoEndedGrove, multiindexed::kdtrie::KdTrie};
 use pprof::criterion::{Output, PProfProfiler};
 use rand::{rng, seq::SliceRandom};
 
@@ -19,6 +19,8 @@ fn run_benchmarks(c: &mut Criterion) {
     run_insert_benchmarks(c, &|i| [i; 1000]);
     run_lookup_benchmarks(c, &|i| i as i32);
     run_lookup_benchmarks(c, &|i| [i; 1000]);
+    run_iter_benchmarks(c, &|i| i as i32);
+    run_iter_benchmarks(c, &|i| [i; 1000]);
 }
 
 /// A trait that matches OnceBiVec's semantics for benchmarking
@@ -34,6 +36,11 @@ trait Benchable<const K: usize, T> {
 
     /// Get a value at the given coordinates if it exists and is within bounds
     fn get(&self, coords: [i32; K]) -> Option<&T>;
+
+    /// Iterate over all items in the container
+    fn iter<'a>(&'a self) -> impl Iterator<Item = ([i32; K], &'a T)>
+    where
+        T: 'a;
 }
 
 impl<T, const K: usize> Benchable<K, T> for MultiIndexed<K, T> {
@@ -52,6 +59,13 @@ impl<T, const K: usize> Benchable<K, T> for MultiIndexed<K, T> {
     fn get(&self, coords: [i32; K]) -> Option<&T> {
         self.get(coords)
     }
+
+    fn iter<'a>(&'a self) -> impl Iterator<Item = ([i32; K], &'a T)>
+    where
+        T: 'a,
+    {
+        self.iter()
+    }
 }
 
 impl<T> Benchable<1, T> for TwoEndedGrove<T> {
@@ -70,6 +84,45 @@ impl<T> Benchable<1, T> for TwoEndedGrove<T> {
     fn get(&self, coords: [i32; 1]) -> Option<&T> {
         self.get(coords[0])
     }
+
+    fn iter<'a>(&'a self) -> impl Iterator<Item = ([i32; 1], &'a T)>
+    where
+        T: 'a,
+    {
+        self.enumerate().map(|(k, v)| ([k], v))
+    }
+}
+
+impl<const K: usize, T> Benchable<K, T> for KdTrie<T> {
+    fn name() -> &'static str {
+        "kd_trie"
+    }
+
+    fn new(_min: [i32; K]) -> Self {
+        Self::new(K)
+    }
+
+    fn push_checked(&self, coords: [i32; K], value: T) {
+        self.insert(&coords, value);
+    }
+
+    fn get(&self, coords: [i32; K]) -> Option<&T> {
+        self.get(&coords)
+    }
+
+    fn iter<'a>(&'a self) -> impl Iterator<Item = ([i32; K], &'a T)>
+    where
+        T: 'a,
+    {
+        self.iter().map(|(coords, value)| {
+            (
+                coords
+                    .try_into()
+                    .expect("coordinates are the right length by construction"),
+                value,
+            )
+        })
+    }
 }
 
 mod benchable_impl;
@@ -127,43 +180,55 @@ fn run_insert_benchmarks<T>(c: &mut Criterion, make_value: &dyn Fn(usize) -> T)
     bench_insert_k::<1, _, OnceBiVec<_>>(&mut g, [0], make_value);
     bench_insert_k::<1, _, TwoEndedGrove<_>>(&mut g, [0], make_value);
     bench_insert_k::<1, _, MultiIndexed<1, _>>(&mut g, [0], make_value);
+    bench_insert_k::<1, _, KdTrie<_>>(&mut g, [0], make_value);
     g.finish();
 
-    run_insert_benchmark::<2, _, OnceBiVec<OnceBiVec<_>>, MultiIndexed<2, _>>(
+    run_insert_benchmark::<2, _, OnceBiVec<OnceBiVec<_>>, MultiIndexed<2, _>, KdTrie<_>>(
         c,
         [0, 0],
         make_value,
     );
-    run_insert_benchmark::<3, _, OnceBiVec<OnceBiVec<OnceBiVec<_>>>, MultiIndexed<3, _>>(
+    run_insert_benchmark::<3, _, OnceBiVec<OnceBiVec<OnceBiVec<_>>>, MultiIndexed<3, _>, KdTrie<_>>(
         c,
         [0, 0, 0],
         make_value,
     );
-    run_insert_benchmark::<4, _, OnceBiVec<OnceBiVec<OnceBiVec<OnceBiVec<_>>>>, MultiIndexed<4, _>>(
-        c,
-        [0, 0, 0, 0],
-        make_value,
-    );
+    run_insert_benchmark::<
+        4,
+        _,
+        OnceBiVec<OnceBiVec<OnceBiVec<OnceBiVec<_>>>>,
+        MultiIndexed<4, _>,
+        KdTrie<_>,
+    >(c, [0, 0, 0, 0], make_value);
     run_insert_benchmark::<
         5,
         _,
         OnceBiVec<OnceBiVec<OnceBiVec<OnceBiVec<OnceBiVec<_>>>>>,
         MultiIndexed<5, _>,
+        KdTrie<_>,
     >(c, [0, 0, 0, 0, 0], make_value);
 
     let mut g = c.benchmark_group(format!("insert_dim6_{}", std::any::type_name::<T>()));
     bench_insert_k::<6, _, MultiIndexed<6, _>>(&mut g, [0, 0, 0, 0, 0, 0], make_value);
+    bench_insert_k::<6, _, KdTrie<_>>(&mut g, [0, 0, 0, 0, 0, 0], make_value);
     g.finish();
 }
 
-fn run_insert_benchmark<const K: usize, T, B1: Benchable<K, T>, B2: Benchable<K, T>>(
+fn run_insert_benchmark<
+    const K: usize,
+    T,
+    B1: Benchable<K, T>,
+    B2: Benchable<K, T>,
+    B3: Benchable<K, T>,
+>(
     c: &mut Criterion,
     min: [i32; K],
     make_value: &dyn Fn(usize) -> T,
 ) {
     let mut g = c.benchmark_group(format!("insert_dim{K}_{}", std::any::type_name::<T>()));
     bench_insert_k::<K, _, B1>(&mut g, min, make_value);
     bench_insert_k::<K, _, B2>(&mut g, min, make_value);
+    bench_insert_k::<K, _, B3>(&mut g, min, make_value);
     g.finish();
 }
 
@@ -199,43 +264,55 @@ fn run_lookup_benchmarks<T>(c: &mut Criterion, make_value: &dyn Fn(usize) -> T)
     bench_lookup_k::<1, _, OnceBiVec<_>>(&mut g, [0], make_value);
     bench_lookup_k::<1, _, TwoEndedGrove<_>>(&mut g, [0], make_value);
     bench_lookup_k::<1, _, MultiIndexed<1, _>>(&mut g, [0], make_value);
+    bench_lookup_k::<1, _, KdTrie<_>>(&mut g, [0], make_value);
     g.finish();
 
-    run_lookup_benchmark::<2, _, OnceBiVec<OnceBiVec<_>>, MultiIndexed<2, _>>(
+    run_lookup_benchmark::<2, _, OnceBiVec<OnceBiVec<_>>, MultiIndexed<2, _>, KdTrie<_>>(
         c,
         [0, 0],
         make_value,
     );
-    run_lookup_benchmark::<3, _, OnceBiVec<OnceBiVec<OnceBiVec<_>>>, MultiIndexed<3, _>>(
+    run_lookup_benchmark::<3, _, OnceBiVec<OnceBiVec<OnceBiVec<_>>>, MultiIndexed<3, _>, KdTrie<_>>(
         c,
         [0, 0, 0],
         make_value,
     );
-    run_lookup_benchmark::<4, _, OnceBiVec<OnceBiVec<OnceBiVec<OnceBiVec<_>>>>, MultiIndexed<4, _>>(
-        c,
-        [0, 0, 0, 0],
-        make_value,
-    );
+    run_lookup_benchmark::<
+        4,
+        _,
+        OnceBiVec<OnceBiVec<OnceBiVec<OnceBiVec<_>>>>,
+        MultiIndexed<4, _>,
+        KdTrie<_>,
+    >(c, [0, 0, 0, 0], make_value);
     run_lookup_benchmark::<
         5,
         _,
         OnceBiVec<OnceBiVec<OnceBiVec<OnceBiVec<OnceBiVec<_>>>>>,
         MultiIndexed<5, _>,
+        KdTrie<_>,
     >(c, [0, 0, 0, 0, 0], make_value);
 
     let mut g = c.benchmark_group(format!("lookup_dim6_{}", std::any::type_name::<T>()));
     bench_lookup_k::<6, _, MultiIndexed<6, _>>(&mut g, [0, 0, 0, 0, 0, 0], make_value);
+    bench_lookup_k::<6, _, KdTrie<_>>(&mut g, [0, 0, 0, 0, 0, 0], make_value);
     g.finish();
 }
 
-fn run_lookup_benchmark<const K: usize, T, B1: Benchable<K, T>, B2: Benchable<K, T>>(
+fn run_lookup_benchmark<
+    const K: usize,
+    T,
+    B1: Benchable<K, T>,
+    B2: Benchable<K, T>,
+    B3: Benchable<K, T>,
+>(
     c: &mut Criterion,
     min: [i32; K],
     make_value: &dyn Fn(usize) -> T,
 ) {
     let mut g = c.benchmark_group(format!("lookup_dim{K}_{}", std::any::type_name::<T>()));
     bench_lookup_k::<K, _, B1>(&mut g, min, make_value);
     bench_lookup_k::<K, _, B2>(&mut g, min, make_value);
+    bench_lookup_k::<K, _, B3>(&mut g, min, make_value);
     g.finish();
 }
 
@@ -266,3 +343,89 @@ fn bench_lookup_k<const K: usize, T, B: Benchable<K, T>>(
         },
     );
 }
+
+// Iteration benchmarks
+
+fn run_iter_benchmarks<T>(c: &mut Criterion, make_value: &dyn Fn(usize) -> T) {
+    // Dim 1
+    let mut g = c.benchmark_group(format!("iter_dim1_{}", std::any::type_name::<T>()));
+    bench_iter_k::<1, _, OnceBiVec<_>>(&mut g, [0], make_value);
+    bench_iter_k::<1, _, TwoEndedGrove<_>>(&mut g, [0], make_value);
+    bench_iter_k::<1, _, MultiIndexed<1, _>>(&mut g, [0], make_value);
+    bench_iter_k::<1, _, KdTrie<_>>(&mut g, [0], make_value);
+    g.finish();
+
+    run_iter_benchmark::<2, _, OnceBiVec<OnceBiVec<_>>, MultiIndexed<2, _>, KdTrie<_>>(
+        c,
+        [0, 0],
+        make_value,
+    );
+    run_iter_benchmark::<3, _, OnceBiVec<OnceBiVec<OnceBiVec<_>>>, MultiIndexed<3, _>, KdTrie<_>>(
+        c,
+        [0, 0, 0],
+        make_value,
+    );
+    run_iter_benchmark::<
+        4,
+        _,
+        OnceBiVec<OnceBiVec<OnceBiVec<OnceBiVec<_>>>>,
+        MultiIndexed<4, _>,
+        KdTrie<_>,
+    >(c, [0, 0, 0, 0], make_value);
+    run_iter_benchmark::<
+        5,
+        _,
+        OnceBiVec<OnceBiVec<OnceBiVec<OnceBiVec<OnceBiVec<_>>>>>,
+        MultiIndexed<5, _>,
+        KdTrie<_>,
+    >(c, [0, 0, 0, 0, 0], make_value);
+
+    let mut g = c.benchmark_group(format!("iter_dim6_{}", std::any::type_name::<T>()));
+    bench_iter_k::<6, _, MultiIndexed<6, _>>(&mut g, [0, 0, 0, 0, 0, 0], make_value);
+    bench_iter_k::<6, _, KdTrie<_>>(&mut g, [0, 0, 0, 0, 0, 0], make_value);
+    g.finish();
+}
+
+fn run_iter_benchmark<
+    const K: usize,
+    T,
+    B1: Benchable<K, T>,
+    B2: Benchable<K, T>,
+    B3: Benchable<K, T>,
+>(
+    c: &mut Criterion,
+    min: [i32; K],
+    make_value: &dyn Fn(usize) -> T,
+) {
+    let mut g = c.benchmark_group(format!("iter_dim{K}_{}", std::any::type_name::<T>()));
+    bench_iter_k::<K, _, B1>(&mut g, min, make_value);
+    bench_iter_k::<K, _, B2>(&mut g, min, make_value);
+    bench_iter_k::<K, _, B3>(&mut g, min, make_value);
+    g.finish();
+}
+
+// Benchmark iters for different dimensions
+fn bench_iter_k<const K: usize, T, B: Benchable<K, T>>(
+    c: &mut BenchmarkGroup<'_, WallTime>,
+    min: [i32; K],
+    make_value: &dyn Fn(usize) -> T,
+) {
+    let vec = B::new(min);
+    let coords = get_n_coords(NUM_ELEMENTS, min);
+
+    // Insert data
+    for (i, coord) in coords.iter().enumerate() {
+        vec.push_checked(*coord, make_value(i));
+    }
+
+    c.bench_function(
+        format!("{}_iter_k{}_{}", B::name(), K, std::any::type_name::<T>()),
+        |b| {
+            b.iter(|| {
+                for x in vec.iter() {
+                    black_box(x);
+                }
+            })
+        },
+    );
+}