From 21bf361e92f64dc59d14cb303afa1b99532b0f6b Mon Sep 17 00:00:00 2001
From: Moritz Hoffmann <antiguru@gmail.com>
Date: Mon, 26 Feb 2024 09:55:53 -0500
Subject: [PATCH] Introduce storage trait to abstract storage of slices

Signed-off-by: Moritz Hoffmann <antiguru@gmail.com>
---
 src/impls/mod.rs        |   1 +
 src/impls/slice_copy.rs |  97 ++++++++++----------
 src/impls/storage.rs    | 192 ++++++++++++++++++++++++++++++++++++++++
 3 files changed, 245 insertions(+), 45 deletions(-)
 create mode 100644 src/impls/storage.rs

diff --git a/src/impls/mod.rs b/src/impls/mod.rs
index be5a6a1..6f71cca 100644
--- a/src/impls/mod.rs
+++ b/src/impls/mod.rs
@@ -9,5 +9,6 @@ pub mod option;
 pub mod result;
 pub mod slice;
 pub mod slice_copy;
+pub mod storage;
 pub mod string;
 pub mod tuple;
diff --git a/src/impls/slice_copy.rs b/src/impls/slice_copy.rs
index d5697b5..0aca1e2 100644
--- a/src/impls/slice_copy.rs
+++ b/src/impls/slice_copy.rs
@@ -1,8 +1,11 @@
 //! A region that stores slices of copy types.
 
+use std::marker::PhantomData;
+
 #[cfg(feature = "serde")]
 use serde::{Deserialize, Serialize};
 
+use crate::impls::storage::Storage;
 use crate::{CopyIter, CopyOnto, Region, ReserveItems};
 
 /// A container for [`Copy`] types.
@@ -24,11 +27,13 @@ use crate::{CopyIter, CopyOnto, Region, ReserveItems};
 /// ```
 #[derive(Debug, Clone)]
 #[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
-pub struct CopyRegion<T: Copy> {
-    slices: Vec<T>,
+pub struct CopyRegion<T: Copy, S: Storage<T> = Vec<T>> {
+    slices: S,
+    offset: usize,
+    _marker: PhantomData<T>,
 }
 
-impl<T: Copy> Region for CopyRegion<T> {
+impl<T: Copy, S: Storage<T>> Region for CopyRegion<T, S> {
     type ReadItem<'a> = &'a [T] where Self: 'a;
     type Index = (usize, usize);
 
@@ -37,13 +42,15 @@ impl<T: Copy> Region for CopyRegion<T> {
         Self: 'a,
     {
         Self {
-            slices: Vec::with_capacity(regions.map(|r| r.slices.len()).sum()),
+            slices: S::merge_regions(regions.map(|r| &r.slices)),
+            offset: 0,
+            _marker: PhantomData,
         }
     }
 
     #[inline]
     fn index(&self, (start, end): Self::Index) -> Self::ReadItem<'_> {
-        &self.slices[start..end]
+        self.slices.index(start, end)
     }
 
     fn reserve_regions<'a, I>(&mut self, regions: I)
@@ -51,7 +58,7 @@ impl<T: Copy> Region for CopyRegion<T> {
         Self: 'a,
         I: Iterator<Item = &'a Self> + Clone,
     {
-        self.slices.reserve(regions.map(|r| r.slices.len()).sum());
+        self.slices.reserve_regions(regions.map(|r| &r.slices));
     }
 
     #[inline]
@@ -59,57 +66,55 @@ impl<T: Copy> Region for CopyRegion<T> {
         self.slices.clear();
     }
 
-    fn heap_size<F: FnMut(usize, usize)>(&self, mut callback: F) {
-        let size_of_t = std::mem::size_of::<T>();
-        callback(
-            self.slices.len() * size_of_t,
-            self.slices.capacity() * size_of_t,
-        );
+    fn heap_size<F: FnMut(usize, usize)>(&self, callback: F) {
+        self.slices.heap_size(callback);
     }
 }
 
-impl<T: Copy> Default for CopyRegion<T> {
+impl<T: Copy, S: Storage<T>> Default for CopyRegion<T, S> {
     fn default() -> Self {
         Self {
-            slices: Vec::default(),
+            slices: S::default(),
+            offset: 0,
+            _marker: PhantomData,
         }
     }
 }
 
-impl<T, const N: usize> CopyOnto<CopyRegion<T>> for [T; N]
+impl<T, S: Storage<T>, const N: usize> CopyOnto<CopyRegion<T, S>> for [T; N]
 where
     T: Copy,
 {
     #[inline]
-    fn copy_onto(self, target: &mut CopyRegion<T>) -> <CopyRegion<T> as Region>::Index {
+    fn copy_onto(self, target: &mut CopyRegion<T, S>) -> <CopyRegion<T, S> as Region>::Index {
         (&self).copy_onto(target)
     }
 }
 
-impl<T, const N: usize> CopyOnto<CopyRegion<T>> for &[T; N]
+impl<T, S: Storage<T>, const N: usize> CopyOnto<CopyRegion<T, S>> for &[T; N]
 where
     T: Copy,
 {
     #[inline]
-    fn copy_onto(self, target: &mut CopyRegion<T>) -> <CopyRegion<T> as Region>::Index {
-        let start = target.slices.len();
-        target.slices.extend_from_slice(self);
-        (start, target.slices.len())
+    fn copy_onto(self, target: &mut CopyRegion<T, S>) -> <CopyRegion<T, S> as Region>::Index {
+        let start = target.offset;
+        target.offset = target.slices.extend_from_slice(self);
+        (start, target.offset)
     }
 }
 
-impl<T, const N: usize> CopyOnto<CopyRegion<T>> for &&[T; N]
+impl<T, S: Storage<T>, const N: usize> CopyOnto<CopyRegion<T, S>> for &&[T; N]
 where
     T: Copy,
 {
     #[inline]
-    fn copy_onto(self, target: &mut CopyRegion<T>) -> <CopyRegion<T> as Region>::Index {
+    fn copy_onto(self, target: &mut CopyRegion<T, S>) -> <CopyRegion<T, S> as Region>::Index {
         (*self).copy_onto(target)
     }
 }
 
-impl<T: Copy, const N: usize> ReserveItems<CopyRegion<T>> for &[T; N] {
-    fn reserve_items<I>(target: &mut CopyRegion<T>, items: I)
+impl<T: Copy, S: Storage<T>, const N: usize> ReserveItems<CopyRegion<T, S>> for &[T; N] {
+    fn reserve_items<I>(target: &mut CopyRegion<T, S>, items: I)
     where
         I: Iterator<Item = Self> + Clone,
     {
@@ -117,30 +122,30 @@ impl<T: Copy, const N: usize> ReserveItems<CopyRegion<T>> for &[T; N] {
     }
 }
 
-impl<T> CopyOnto<CopyRegion<T>> for &[T]
+impl<T, S: Storage<T>> CopyOnto<CopyRegion<T, S>> for &[T]
 where
     T: Copy,
 {
     #[inline]
-    fn copy_onto(self, target: &mut CopyRegion<T>) -> <CopyRegion<T> as Region>::Index {
-        let start = target.slices.len();
-        target.slices.extend_from_slice(self);
-        (start, target.slices.len())
+    fn copy_onto(self, target: &mut CopyRegion<T, S>) -> <CopyRegion<T, S> as Region>::Index {
+        let start = target.offset;
+        target.offset = target.slices.extend_from_slice(self);
+        (start, target.offset)
     }
 }
 
-impl<T> CopyOnto<CopyRegion<T>> for &&[T]
+impl<T, S: Storage<T>> CopyOnto<CopyRegion<T, S>> for &&[T]
 where
     T: Copy,
 {
     #[inline]
-    fn copy_onto(self, target: &mut CopyRegion<T>) -> <CopyRegion<T> as Region>::Index {
+    fn copy_onto(self, target: &mut CopyRegion<T, S>) -> <CopyRegion<T, S> as Region>::Index {
         (*self).copy_onto(target)
     }
 }
 
-impl<T: Copy> ReserveItems<CopyRegion<T>> for &[T] {
-    fn reserve_items<I>(target: &mut CopyRegion<T>, items: I)
+impl<T: Copy, S: Storage<T>> ReserveItems<CopyRegion<T, S>> for &[T] {
+    fn reserve_items<I>(target: &mut CopyRegion<T, S>, items: I)
     where
         I: Iterator<Item = Self> + Clone,
     {
@@ -148,18 +153,18 @@ impl<T: Copy> ReserveItems<CopyRegion<T>> for &[T] {
     }
 }
 
-impl<T> CopyOnto<CopyRegion<T>> for &Vec<T>
+impl<T, S: Storage<T>> CopyOnto<CopyRegion<T, S>> for &Vec<T>
 where
     T: Copy,
 {
     #[inline]
-    fn copy_onto(self, target: &mut CopyRegion<T>) -> <CopyRegion<T> as Region>::Index {
+    fn copy_onto(self, target: &mut CopyRegion<T, S>) -> <CopyRegion<T, S> as Region>::Index {
         self.as_slice().copy_onto(target)
     }
 }
 
-impl<T: Copy> ReserveItems<CopyRegion<T>> for &Vec<T> {
-    fn reserve_items<I>(target: &mut CopyRegion<T>, items: I)
+impl<T: Copy, S: Storage<T>> ReserveItems<CopyRegion<T, S>> for &Vec<T> {
+    fn reserve_items<I>(target: &mut CopyRegion<T, S>, items: I)
     where
         I: Iterator<Item = Self> + Clone,
     {
@@ -167,20 +172,22 @@ impl<T: Copy> ReserveItems<CopyRegion<T>> for &Vec<T> {
     }
 }
 
-impl<T, I: IntoIterator<Item = T>> CopyOnto<CopyRegion<T>> for CopyIter<I>
+impl<T, S: Storage<T>, I: IntoIterator<Item = T>> CopyOnto<CopyRegion<T, S>> for CopyIter<I>
 where
     T: Copy,
 {
     #[inline]
-    fn copy_onto(self, target: &mut CopyRegion<T>) -> <CopyRegion<T> as Region>::Index {
-        let start = target.slices.len();
-        target.slices.extend(self.0);
-        (start, target.slices.len())
+    fn copy_onto(self, target: &mut CopyRegion<T, S>) -> <CopyRegion<T, S> as Region>::Index {
+        let start = target.offset;
+        target.offset = target.slices.extend(self.0);
+        (start, target.offset)
     }
 }
 
-impl<T: Copy, J: IntoIterator<Item = T>> ReserveItems<CopyRegion<T>> for CopyIter<J> {
-    fn reserve_items<I>(target: &mut CopyRegion<T>, items: I)
+impl<T: Copy, S: Storage<T>, J: IntoIterator<Item = T>> ReserveItems<CopyRegion<T, S>>
+    for CopyIter<J>
+{
+    fn reserve_items<I>(target: &mut CopyRegion<T, S>, items: I)
     where
         I: Iterator<Item = Self> + Clone,
     {
diff --git a/src/impls/storage.rs b/src/impls/storage.rs
new file mode 100644
index 0000000..b46baaf
--- /dev/null
+++ b/src/impls/storage.rs
@@ -0,0 +1,192 @@
+//! Storage abstractions to represent slices of data.
+
+#[cfg(feature = "serde")]
+use serde::{Deserialize, Serialize};
+
+/// TODO
+pub trait Storage<T>: Default {
+    /// TODO
+    fn with_capacity(capacity: usize) -> Self;
+    /// TODO
+    fn merge_regions<'a>(regions: impl Iterator<Item = &'a Self> + Clone) -> Self
+    where
+        Self: 'a;
+    /// TODO
+    fn reserve(&mut self, additional: usize);
+    /// TODO
+    fn reserve_regions<'a, I>(&mut self, regions: I)
+    where
+        Self: 'a,
+        I: Iterator<Item = &'a Self> + Clone;
+    /// TODO
+    fn clear(&mut self);
+    /// TODO
+    fn heap_size<F: FnMut(usize, usize)>(&self, callback: F);
+    /// TODO
+    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) -> usize;
+    /// TODO
+    fn extend_from_slice(&mut self, slice: &[T]) -> usize
+    where
+        T: Clone;
+    /// TODO
+    fn index(&self, start: usize, end: usize) -> &[T];
+}
+
+impl<T> Storage<T> for Vec<T> {
+    fn with_capacity(capacity: usize) -> Self {
+        Vec::with_capacity(capacity)
+    }
+
+    fn merge_regions<'a>(regions: impl Iterator<Item = &'a Self> + Clone) -> Self
+    where
+        Self: 'a,
+    {
+        Self::with_capacity(regions.map(Vec::len).sum())
+    }
+
+    fn reserve(&mut self, additional: usize) {
+        Vec::reserve(self, additional);
+    }
+
+    fn reserve_regions<'a, I>(&mut self, regions: I)
+    where
+        Self: 'a,
+        I: Iterator<Item = &'a Self> + Clone,
+    {
+        self.reserve(regions.map(Vec::len).sum());
+    }
+
+    fn clear(&mut self) {
+        self.clear();
+    }
+
+    fn heap_size<F: FnMut(usize, usize)>(&self, mut callback: F) {
+        let size_of_t = std::mem::size_of::<T>();
+        callback(self.len() * size_of_t, self.capacity() * size_of_t);
+    }
+
+    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) -> usize {
+        Extend::extend(self, iter);
+        self.len()
+    }
+
+    fn extend_from_slice(&mut self, slice: &[T]) -> usize
+    where
+        T: Clone,
+    {
+        self.extend_from_slice(slice);
+        self.len()
+    }
+
+    fn index(&self, start: usize, end: usize) -> &[T] {
+        &self[start..end]
+    }
+}
+
+#[derive(Debug)]
+#[cfg_attr(feature = "serde", derive(Serialize, Deserialize))]
+struct Doubling<T> {
+    inner: Vec<Vec<T>>,
+    offsets: Vec<usize>,
+}
+
+impl<T> Default for Doubling<T> {
+    fn default() -> Self {
+        Self {
+            inner: Vec::default(),
+            offsets: Vec::default(),
+        }
+    }
+}
+
+impl<T: Clone> Storage<T> for Doubling<T> {
+    fn with_capacity(capacity: usize) -> Self {
+        Self {
+            inner: vec![Vec::with_capacity(capacity)],
+            offsets: Vec::default(),
+        }
+    }
+
+    fn merge_regions<'a>(regions: impl Iterator<Item = &'a Self> + Clone) -> Self
+    where
+        Self: 'a,
+    {
+        Self::with_capacity(regions.flat_map(|r| &r.inner).map(Vec::len).sum())
+    }
+
+    fn reserve(&mut self, additional: usize) {
+        let (remaining, last_len) = self
+            .inner
+            .last()
+            .map_or((0, 0), |last| (last.capacity() - last.len(), last.len()));
+        if remaining < additional {
+            let len = 2 * last_len;
+            let len = std::cmp::max(additional, len);
+            let len = len.next_power_of_two();
+            self.offsets
+                .push(last_len + *self.offsets.last().unwrap_or(&0));
+            self.inner.push(Vec::with_capacity(len));
+        }
+    }
+
+    fn reserve_regions<'a, I>(&mut self, regions: I)
+    where
+        Self: 'a,
+        I: Iterator<Item = &'a Self> + Clone,
+    {
+        self.reserve(regions.flat_map(|r| &r.inner).map(Vec::len).sum());
+    }
+
+    fn clear(&mut self) {
+        self.inner.clear();
+    }
+
+    fn heap_size<F: FnMut(usize, usize)>(&self, mut callback: F) {
+        let size_of_t = std::mem::size_of::<T>();
+        for inner in &self.inner {
+            callback(inner.len() * size_of_t, inner.capacity() * size_of_t);
+        }
+    }
+
+    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) -> usize {
+        let vec: Vec<T> = iter.into_iter().collect();
+        self.extend_from_slice(&vec)
+    }
+
+    fn extend_from_slice(&mut self, slice: &[T]) -> usize
+    where
+        T: Clone,
+    {
+        self.reserve(slice.len());
+        self.inner.last_mut().unwrap().extend_from_slice(slice);
+        *self.offsets.last().unwrap_or(&0) + self.inner.last().unwrap().len()
+    }
+
+    fn index(&self, start: usize, end: usize) -> &[T] {
+        let index = self
+            .offsets
+            .iter()
+            .position(|&o| o > start)
+            .unwrap_or_else(|| self.offsets.len().saturating_sub(1));
+        let start = start - self.offsets[index];
+        let end = end - self.offsets[index];
+        &self.inner[index][start..end]
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_doubling() {
+        let mut d: Doubling<usize> = Doubling::default();
+        let mut start = 0;
+
+        for i in 0..1000 {
+            let end = d.extend_from_slice(&[i, i + 1, i + 3]);
+            assert_eq!(&[i, i + 1, i + 3], d.index(start, end));
+            start = end;
+        }
+    }
+}