More buf builders

src/layer/internal/base.rs

@@ -175,7 +175,7 @@
     pub fn add_node(&mut self, node: &str) -> u64 {
         let id = self.builder.add_node(node);
 
         id
     }
 
     pub fn add_node_bytes(&mut self, node: Bytes) -> u64 {
@@ -188,7 +188,7 @@
     pub fn add_predicate(&mut self, predicate: &str) -> u64 {
         let id = self.builder.add_predicate(predicate);
 
         id
     }
 
     pub fn add_predicate_bytes(&mut self, predicate: Bytes) -> u64 {
@@ -201,7 +201,7 @@
     pub fn add_value(&mut self, value: TypedDictEntry) -> u64 {
         let id = self.builder.add_value(value);
 
         id
     }
 
     /// Add nodes from an iterable.
@@ -217,7 +217,7 @@
     {
         let ids = self.builder.add_nodes(nodes);
 
         ids
     }
 
     pub fn add_nodes_bytes<I: 'static + IntoIterator<Item = Bytes> + Send>(
@@ -230,7 +230,7 @@
     {
         let ids = self.builder.add_nodes_bytes(nodes);
 
         ids
     }
 
     /// Add predicates from an iterable.
@@ -246,7 +246,7 @@
     {
         let ids = self.builder.add_predicates(predicates);
 
         ids
     }
 
     pub fn add_predicates_bytes<I: 'static + IntoIterator<Item = Bytes> + Send>(
@@ -259,7 +259,7 @@
     {
         let ids = self.builder.add_predicates_bytes(predicates);
 
         ids
     }
 
     /// Add values from an iterable.
@@ -275,7 +275,7 @@
     {
         let ids = self.builder.add_values(values);
 
         ids
     }
 
     /// Turn this builder into a phase 2 builder that will take triple data.
@@ -490,7 +490,7 @@
 }
 
 #[cfg(test)]
-pub mod tests {
+pub mod base_tests {
     use super::*;
     use crate::storage::memory::*;
     use futures::stream::TryStreamExt;

src/layer/internal/child.rs

@@ -702,9 +702,9 @@ pub async fn open_child_triple_stream<F: 'static + FileLoad + FileStore>(
 }
 
 #[cfg(test)]
-pub mod tests {
+pub mod child_tests {
     use super::*;
-    use crate::layer::base::tests::*;
+    use crate::layer::base::base_tests::*;
     use crate::storage::memory::*;
     use futures::stream::TryStreamExt;
 

src/layer/internal/mod.rs

@@ -525,10 +525,10 @@
     }
 
     pub fn is_rollup(&self) -> bool {
         match self {
             Rollup(_) => true,
             _ => false,
         }
     }
 }
 
@@ -636,7 +636,7 @@
                 parent_count = parent_count
                     - current_layer.node_dict_len() as u64
                     - current_layer.value_dict_len() as u64;
                 if corrected_id > parent_count as u64 {
                     // subject, if it exists, is in this layer
                     corrected_id -= parent_count;
                 } else {
@@ -667,7 +667,7 @@
             let mut corrected_id = id;
             if let Some(parent) = current_layer.immediate_parent() {
                 parent_count -= current_layer.predicate_dict_len() as u64;
                 if corrected_id > parent_count as u64 {
                     // subject, if it exists, is in this layer
                     corrected_id -= parent_count;
                 } else {
@@ -1098,7 +1098,7 @@
         assert_eq!(1, layer.triple_layer_removal_count().unwrap());
     }
 
-    use crate::layer::base::tests::*;
+    use crate::layer::base::base_tests::*;
     #[tokio::test]
     async fn base_layer_with_gaps_addition_count() {
         let files = base_layer_files();

src/layer/internal/object_iterator.rs

@@ -22,10 +22,10 @@
         s_p_adjacency_list: AdjacencyList,
     ) -> Self {
         Self {
             subjects: subjects,
             objects: objects,
             o_ps_adjacency_list: o_ps_adjacency_list,
             s_p_adjacency_list: s_p_adjacency_list,
             o_position: 0,
             o_ps_position: 0,
             peeked: None,
@@ -227,7 +227,7 @@
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::layer::base::tests::*;
+    use crate::layer::base::base_tests::*;
     use crate::storage::memory::*;
     use crate::storage::*;
 

src/layer/internal/predicate_iterator.rs

@@ -83,13 +83,13 @@
         // set the sp_boundary flag, ensuring that upon the subsequent
         // next() call, we move on to the next predicate.
         let next = self.subject_iterator.peek();
         if next.is_none()
             || next.map(|t| (t.subject, t.predicate)) != result.map(|t| (t.subject, t.predicate))
         {
             self.sp_boundary = true;
         } else {
             self.sp_boundary = false;
         }
 
         result
     }
@@ -182,8 +182,8 @@
 
 #[cfg(test)]
 mod tests {
-    use crate::layer::base::tests::*;
-    use crate::layer::child::tests::*;
+    use crate::layer::base::base_tests::*;
+    use crate::layer::child::child_tests::*;
     use crate::layer::*;
 
     use std::sync::Arc;

src/layer/internal/subject_iterator.rs

@@ -22,9 +22,9 @@
         sp_o_adjacency_list: AdjacencyList,
     ) -> Self {
         Self {
             subjects: subjects,
             s_p_adjacency_list: s_p_adjacency_list,
             sp_o_adjacency_list: sp_o_adjacency_list,
             s_position: 0,
             s_p_position: 0,
             sp_o_position: 0,
@@ -422,7 +422,7 @@
                 (Some(lowest_pos_index), Some(lowest_neg_index)) => {
                     let lowest_pos = self.positives[lowest_pos_index].peek().unwrap();
                     let lowest_neg = self.negatives[lowest_neg_index].peek().unwrap();
                     match lowest_pos.cmp(&lowest_neg) {
                         Ordering::Less => {
                             // next change is an addition, and there's no matching removal
                             return Some((
@@ -452,8 +452,8 @@
 
 #[cfg(test)]
 mod tests {
-    use crate::layer::base::tests::*;
-    use crate::layer::child::tests::*;
+    use crate::layer::base::base_tests::*;
+    use crate::layer::child::child_tests::*;
     use crate::layer::*;
     use crate::structure::TdbDataType;
 

src/layer/layer.rs

@@ -98,8 +98,8 @@
         self.subject_id(&triple.subject).and_then(|subject| {
             self.predicate_id(&triple.predicate).and_then(|predicate| {
                 match &triple.object {
                     ObjectType::Node(node) => self.object_node_id(&node),
                     ObjectType::Value(value) => self.object_value_id(&value),
                 }
                 .map(|object| IdTriple {
                     subject,
@@ -127,11 +127,11 @@
                 .unwrap_or(PossiblyResolved::Unresolved(triple.predicate)),
             object: match &triple.object {
                 ObjectType::Node(node) => self
                     .object_node_id(&node)
                     .map(PossiblyResolved::Resolved)
                     .unwrap_or(PossiblyResolved::Unresolved(triple.object)),
                 ObjectType::Value(value) => self
                     .object_value_id(&value)
                     .map(PossiblyResolved::Resolved)
                     .unwrap_or(PossiblyResolved::Unresolved(triple.object)),
             },
@@ -270,7 +270,7 @@
     /// Return a PossiblyResolved with the inner value as a reference.
     pub fn as_ref(&self) -> PossiblyResolved<&T> {
         match self {
             Self::Unresolved(u) => PossiblyResolved::Unresolved(&u),
             Self::Resolved(id) => PossiblyResolved::Resolved(*id),
         }
     }
@@ -306,7 +306,7 @@
         &self,
         node_map: &HashMap<String, u64>,
         predicate_map: &HashMap<String, u64>,
         value_map: &HashMap<TypedDictEntry, u64>,
     ) -> Option<IdTriple> {
         let subject = match self.subject.as_ref() {
             PossiblyResolved::Unresolved(s) => *node_map.get(s)?,
@@ -405,9 +405,9 @@
 #[cfg(test)]
 mod tests {
     use super::*;
-    use crate::layer::internal::base::tests::base_layer_files;
+    use crate::layer::internal::base::base_tests::base_layer_files;
     use crate::layer::internal::base::BaseLayer;
-    use crate::layer::internal::child::tests::child_layer_files;
+    use crate::layer::internal::child::child_tests::child_layer_files;
     use crate::layer::internal::child::ChildLayer;
     use crate::layer::internal::InternalLayer;
     use crate::layer::simple_builder::{LayerBuilder, SimpleLayerBuilder};

src/structure/adjacencylist.rs

@@ -15,6 +15,7 @@ use std::io;
 use std::pin::Pin;
 
 use bytes::Bytes;
+use bytes::BytesMut;
 
 use super::bitarray::*;
 use super::bitindex::*;
@@ -36,6 +37,15 @@ impl AdjacencyList {
         AdjacencyList { nums, bits }
     }
 
+    pub fn from_buffers(buffers: AdjacencyListBuffers) -> AdjacencyList {
+        Self::parse(
+            buffers.nums,
+            buffers.bits,
+            buffers.bitindex_blocks,
+            buffers.bitindex_sblocks,
+        )
+    }
+
     pub fn parse(
         nums_slice: Bytes,
         bits_slice: Bytes,
@@ -224,6 +234,82 @@ pub async fn adjacency_list_stream_pairs<F: 'static + FileLoad>(
     )
 }
 
+pub struct UnindexedAdjacencyListBufBuilder {
+    bitarray: BitArrayBufBuilder<BytesMut>,
+    nums: LogArrayBufBuilder<BytesMut>,
+    last_left: u64,
+    last_right: u64,
+}
+
+impl UnindexedAdjacencyListBufBuilder {
+    pub fn new(width: u8) -> Self {
+        Self {
+            bitarray: BitArrayBufBuilder::new(BytesMut::new()),
+            nums: LogArrayBufBuilder::new(BytesMut::new(), width),
+            last_left: 0,
+            last_right: 0,
+        }
+    }
+
+    pub fn push(&mut self, left: u64, right: u64) {
+        // the tricky thing with this code is that the bitarray lags one entry behind the logarray.
+        // The reason for this is that at push time, we do not yet know if this entry is going to be
+        // the last entry for `left`, we only know this when we push a greater `left` later on.
+        if left < self.last_left || (left == self.last_left && right <= self.last_right) {
+            panic!("tried to push an unordered adjacent pair");
+        }
+
+        // the left hand side of the adjacencylist is expected to be a continuous range from 1 up to the max
+        // but when adding entries, there may be holes. We handle holes by writing a '0' to the logarray
+        // (which is otherwise an invalid right-hand side) and pushing a 1 onto the bitarray to immediately close the segment.
+        let skip = left - self.last_left;
+
+        if self.last_left == 0 && skip == 1 {
+            // this is the first entry. we can't push a bit yet
+        } else if skip == 0 {
+            // same `left` as before. so the previous entry was not the last one, and the bitarray gets a 0 appended.
+            self.bitarray.push(false);
+        } else {
+            // if this is the first element, but we do need to skip, make sure we write one less bit than we'd usually do
+            let bitskip = if self.last_left == 0 { skip - 1 } else { skip };
+            // there's a different `left`. we push a bunch of 1s to the bitarray, and 0s to the num array.
+            for _ in 0..bitskip {
+                self.bitarray.push(true);
+            }
+            for _ in 0..(skip - 1) {
+                self.nums.push(0);
+            }
+        }
+
+        // finally push right to the logarray
+        self.nums.push(right);
+        self.last_left = left;
+        self.last_right = right;
+    }
+
+    pub fn push_all<I: Iterator<Item = (u64, u64)>>(&mut self, mut iter: I) {
+        while let Some((left, right)) = iter.next() {
+            self.push(left, right);
+        }
+    }
+
+    pub fn finalize(mut self) -> (Bytes, Bytes) {
+        if self.nums.count() != 0 {
+            // push last bit to bitarray
+            self.bitarray.push(true);
+        }
+
+        let ba = self.bitarray.finalize();
+        let nums = self.nums.finalize();
+
+        (ba.freeze(), nums.freeze())
+    }
+
+    pub fn count(&self) -> u64 {
+        self.bitarray.count()
+    }
+}
+
 pub struct UnindexedAdjacencyListBuilder<W1: SyncableFile, W2: SyncableFile> {
     bitarray: BitArrayFileBuilder<W1>,
     nums: LogArrayFileBuilder<W2>,
@@ -307,6 +393,59 @@ impl<W1: SyncableFile, W2: SyncableFile> UnindexedAdjacencyListBuilder<W1, W2> {
     }
 }
 
+pub struct AdjacencyListBufBuilder {
+    builder: UnindexedAdjacencyListBufBuilder,
+    bitindex_blocks: BytesMut,
+    bitindex_sblocks: BytesMut,
+}
+
+impl AdjacencyListBufBuilder {
+    pub fn new(width: u8) -> AdjacencyListBufBuilder {
+        AdjacencyListBufBuilder {
+            builder: UnindexedAdjacencyListBufBuilder::new(width),
+            bitindex_blocks: BytesMut::new(),
+            bitindex_sblocks: BytesMut::new(),
+        }
+    }
+
+    pub fn push(&mut self, left: u64, right: u64) {
+        self.builder.push(left, right)
+    }
+
+    pub fn push_all<I: Iterator<Item = (u64, u64)>>(&mut self, iter: I) {
+        self.builder.push_all(iter)
+    }
+
+    pub fn finalize(self) -> AdjacencyListBuffers {
+        let AdjacencyListBufBuilder {
+            builder,
+            mut bitindex_blocks,
+            mut bitindex_sblocks,
+        } = self;
+        let (bitfile, nums) = builder.finalize();
+
+        build_bitindex_from_buf(&bitfile[..], &mut bitindex_blocks, &mut bitindex_sblocks);
+
+        AdjacencyListBuffers {
+            nums,
+            bits: bitfile,
+            bitindex_blocks: bitindex_blocks.freeze(),
+            bitindex_sblocks: bitindex_sblocks.freeze(),
+        }
+    }
+
+    pub fn count(&self) -> u64 {
+        self.builder.count()
+    }
+}
+
+pub struct AdjacencyListBuffers {
+    nums: Bytes,
+    bits: Bytes,
+    bitindex_blocks: Bytes,
+    bitindex_sblocks: Bytes,
+}
+
 pub struct AdjacencyListBuilder<F, W1, W2, W3>
 where
     F: 'static + FileLoad + FileStore,
@@ -778,4 +917,16 @@ mod tests {
             result
         );
     }
+
+    #[test]
+    fn adjacencylist_buf_builder_works() {
+        let adjacencies = [(1, 1), (1, 5), (2, 3), (2, 7), (4, 8)];
+        let mut builder = AdjacencyListBufBuilder::new(8);
+        builder.push_all(adjacencies.iter().copied());
+        let buffers = builder.finalize();
+        let aj = AdjacencyList::from_buffers(buffers);
+
+        let result: Vec<_> = aj.iter().collect();
+        assert_eq!(&adjacencies[..], &result[..]);
+    }
 }

src/structure/bitarray.rs

@@ -34,7 +34,7 @@ use super::util;
 use crate::storage::*;
 use crate::structure::bititer::BitIter;
 use byteorder::{BigEndian, ByteOrder};
-use bytes::{Buf, Bytes, BytesMut};
+use bytes::{Buf, BufMut, Bytes, BytesMut};
 use futures::io;
 use futures::stream::{Stream, StreamExt, TryStreamExt};
 use std::{convert::TryFrom, error, fmt};
@@ -200,6 +200,70 @@ impl BitArray {
     }
 }
 
+pub struct BitArrayBufBuilder<B> {
+    /// Destination of the bit array data.
+    dest: B,
+    /// Storage for the next word to be written.
+    current: u64,
+    /// Number of bits written to the buffer
+    count: u64,
+}
+
+impl<B: BufMut> BitArrayBufBuilder<B> {
+    pub fn new(dest: B) -> BitArrayBufBuilder<B> {
+        BitArrayBufBuilder {
+            dest,
+            current: 0,
+            count: 0,
+        }
+    }
+
+    pub fn push(&mut self, bit: bool) {
+        // Set the bit in the current word.
+        if bit {
+            // Determine the position of the bit to be set from `count`.
+            let pos = self.count & 0b11_1111;
+            self.current |= 0x8000_0000_0000_0000 >> pos;
+        }
+
+        // Advance the bit count.
+        self.count += 1;
+
+        // Check if the new `count` has reached a word boundary.
+        if self.count & 0b11_1111 == 0 {
+            // We have filled `current`, so write it to the destination.
+            self.dest.put_u64(self.current);
+            self.current = 0;
+        }
+    }
+
+    pub fn push_all<I: Iterator<Item = bool>>(&mut self, mut iter: I) {
+        while let Some(bit) = iter.next() {
+            self.push(bit);
+        }
+    }
+
+    fn finalize_data(&mut self) {
+        if self.count & 0b11_1111 != 0 {
+            self.dest.put_u64(self.current);
+        }
+    }
+
+    pub fn finalize(mut self) -> B {
+        let count = self.count;
+        // Write the final data word.
+        self.finalize_data();
+        // Write the control word.
+        self.dest.put_u64(count);
+
+        self.dest
+    }
+
+    pub fn count(&self) -> u64 {
+        self.count
+    }
+}
+
 pub struct BitArrayFileBuilder<W> {
     /// Destination of the bit array data.
     dest: W,
@@ -317,22 +381,22 @@ pub fn bitarray_stream_blocks<R: AsyncRead + Unpin>(r: R) -> FramedRead<R, BitAr
     FramedRead::new(r, BitArrayBlockDecoder { readahead: None })
 }
 
-pub fn bitarray_iter_blocks<B: Buf>(b: &mut B) -> BitArrayBlockIterator<B> {
+pub fn bitarray_iter_blocks<B: Buf>(b: B) -> BitArrayBlockIterator<B> {
     BitArrayBlockIterator {
         buf: b,
         readahead: None,
     }
 }
 
-pub struct BitArrayBlockIterator<'a, B: Buf> {
-    buf: &'a mut B,
+pub struct BitArrayBlockIterator<B: Buf> {
+    buf: B,
     readahead: Option<u64>,
 }
 
-impl<'a, B: Buf> Iterator for BitArrayBlockIterator<'a, B> {
+impl<B: Buf> Iterator for BitArrayBlockIterator<B> {
     type Item = u64;
     fn next(&mut self) -> Option<u64> {
-        decode_next_bitarray_block(self.buf, &mut self.readahead)
+        decode_next_bitarray_block(&mut self.buf, &mut self.readahead)
     }
 }
 

src/structure/bitindex.rs

@@ -437,10 +437,10 @@ pub async fn build_bitindex<
     Ok(())
 }
 
-pub fn build_bitindex_from_block_iter<'a, I: Iterator<Item = u64>, B1: BufMut, B2: BufMut>(
-    blocks_iter: &'a mut I,
-    blocks: &mut B1,
-    sblocks: &mut B2,
+pub fn build_bitindex_from_block_iter<I: Iterator<Item = u64>, B1: BufMut, B2: BufMut>(
+    blocks_iter: I,
+    blocks: B1,
+    sblocks: B2,
 ) {
     // the following widths are unoptimized, but should always be large enough
     let mut blocks_builder =
@@ -474,9 +474,9 @@ pub fn build_bitindex_from_block_iter<'a, I: Iterator<Item = u64>, B1: BufMut, B
 }
 
 pub fn build_bitindex_from_buf<B1: Buf, B2: BufMut, B3: BufMut>(
-    bitarray: &mut B1,
-    blocks: &mut B2,
-    sblocks: &mut B3,
+    bitarray: B1,
+    blocks: B2,
+    sblocks: B3,
 ) {
     let mut iter = bitarray_iter_blocks(bitarray);
     build_bitindex_from_block_iter(&mut iter, blocks, sblocks)