Flat container storage abstraction

container/Cargo.toml

@@ -7,5 +7,6 @@ license = "MIT"
 
 [dependencies]
 columnation = { git = "https://github.com/frankmcsherry/columnation" }
-flatcontainer = "0.5"
+# flatcontainer = "0.5"
+flatcontainer = { git = "https://github.com/antiguru/flatcontainer" }
 serde = { version = "1.0"}

container/src/flatcontainer.rs

@@ -1,9 +1,14 @@
 //! Present a [`FlatStack`] as a timely container.
 
 pub use flatcontainer::*;
-use crate::{buffer, Container, SizableContainer, PushInto};
+use flatcontainer::impls::index::IndexContainer;
+use crate::{buffer, Container, SizableContainer, PushInto, CapacityContainer};
 
-impl<R: Region + Clone + 'static> Container for FlatStack<R> {
+impl<R, S> Container for FlatStack<R, S>
+where
+    R: Region + Clone + 'static,
+    S: IndexContainer<<R as Region>::Index> + Clone + 'static,
+{
     type ItemRef<'a> = R::ReadItem<'a>  where Self: 'a;
     type Item<'a> = R::ReadItem<'a> where Self: 'a;
 
@@ -28,6 +33,7 @@ impl<R: Region + Clone + 'static> Container for FlatStack<R> {
     }
 }
 
+// Only implemented for `FlatStack` with `Vec` offsets.
 impl<R: Region + Clone + 'static> SizableContainer for FlatStack<R> {
     fn capacity(&self) -> usize {
         self.capacity()
@@ -42,7 +48,31 @@ impl<R: Region + Clone + 'static> SizableContainer for FlatStack<R> {
     }
 }
 
-impl<R: Region + Push<T>, T> PushInto<T> for FlatStack<R> {
+impl<R: Region> CapacityContainer for FlatStack<R>
+{
+    fn preferred_capacity() -> usize {
+        // We don't have a good way to present any pre-defined capacity here, since it's a
+        // concept foreign to flat containers. Each region might have a capacity, but overall
+        // the concept of capacity does not exist. For this reason, we just hardcode a number,
+        // which seems to work reasonably well.
+        //
+        // We should revisit this if/once we have an abstraction that can express a capacity
+        // for `FlatStack`, but we aren't there yet.
+        1024
+    }
+
+    fn ensure_preferred_capacity(&mut self) {
+        if self.capacity() < Self::preferred_capacity() {
+            self.reserve(Self::preferred_capacity() - self.capacity());
+        }
+    }
+}
+
+impl<R, S, T> PushInto<T> for FlatStack<R, S>
+where
+    R: Region + Push<T>,
+    S: IndexContainer<R::Index> + Clone + 'static,
+{
     #[inline]
     fn push_into(&mut self, item: T) {
         self.copy(item);

container/src/lib.rs

@@ -73,6 +73,18 @@ pub trait SizableContainer: Container {
     fn reserve(&mut self, additional: usize);
 }
 
+/// A container that has a preferred capacity and can ensure it has said capacity.
+///
+// TODO: This trait shouldn't exist, but @antiguru cannot come up with a better way to encode that
+// we might want to preallocate a buffer :(
+pub trait CapacityContainer {
+    /// The preferred capacity
+    fn preferred_capacity() -> usize;
+
+    /// Ensure that the container has sufficient capacity to absorb `preferred_capacity` elements.
+    fn ensure_preferred_capacity(&mut self);
+}
+
 /// A container that can absorb items of a specific type.
 pub trait PushInto<T> {
     /// Push item into self.
@@ -245,6 +257,20 @@ impl<T: Clone + 'static> SizableContainer for Vec<T> {
     }
 }
 
+impl<T> CapacityContainer for Vec<T> {
+    fn preferred_capacity() -> usize {
+        buffer::default_capacity::<T>()
+    }
+
+    #[inline]
+    fn ensure_preferred_capacity(&mut self) {
+        if self.capacity() < <Self as CapacityContainer>::preferred_capacity() {
+            self.reserve(<Self as CapacityContainer>::preferred_capacity() - self.capacity());
+        }
+    }
+}
+
+
 impl<T> PushInto<T> for Vec<T> {
     #[inline]
     fn push_into(&mut self, item: T) {

timely/src/order.rs

@@ -280,6 +280,18 @@ mod product {
             }
         }
 
+        impl<'a, 'b, RO, RI> ReserveItems<&'b Product<RO::ReadItem<'a>, RI::ReadItem<'a>>> for ProductRegion<RO, RI>
+        where
+            RO: Region + ReserveItems<&'b <RO as Region>::ReadItem<'a>> + 'a,
+            RI: Region + ReserveItems<&'b <RI as Region>::ReadItem<'a>> + 'a,
+        {
+            #[inline]
+            fn reserve_items<I>(&mut self, items: I) where I: Iterator<Item=&'b Product<RO::ReadItem<'a>, RI::ReadItem<'a>>> + Clone {
+                self.outer_region.reserve_items(items.clone().map(|i| &i.outer));
+                self.inner_region.reserve_items(items.clone().map(|i| &i.inner));
+            }
+        }
+
         impl<TO, TI, RO, RI> Push<Product<TO, TI>> for ProductRegion<RO, RI>
         where
             RO: Region + Push<TO>,