1
use std::convert::Infallible;
2

            
3
use byteorder::BigEndian;
4
use nanorand::{Pcg64, Rng};
5
use nebari::{
6
    tree::{
7
        btree::{Indexer, Reducer},
8
        EmbeddedIndex, Root, ScanEvaluation, Serializable, VersionedTreeRoot,
9
    },
10
    ArcBytes, Error,
11
};
12

            
13
fn main() -> Result<(), Error> {
14
    let roots = nebari::Config::default_for("embedded-indexes.nebari").open()?;
15

            
16
    // Nebari provides a way to embed data within the B-Tree directly through a
17
    // `EmbeddedIndex` and `EmbeddedStats` generics on either
18
    // `VersionedTreeRoot` or `UnversionedTreeRoot`.
19
    let tree = roots.tree(VersionedTreeRoot::<Zeroes>::tree("one"))?;
20
    let mut rng = Pcg64::new();
21
    // This example counts the number of '0' characters and stores it in the index.
22
    for i in 0_u32..100 {
23
        tree.set(i.to_be_bytes(), rng.generate::<u64>().to_string())?;
24
    }
25

            
26
    // When scanning the tree, we can retrieve our embedded statistics during
27
    // the scan at every level of the tree, and we have full control over
28
    // iteration of the tree.
29
    tree.scan::<Infallible, _, _, _, _>(
30
        &(..),
31
        true,
32
        // The node evaluator is called for every interior node -- nodes that
33
        // don't contain values directly, but point to other nodes instead.
34
        // Returning true allows the scan to dive into the node.
35
        |max_key, index, depth| {
36
            println!(
37
                "Interior found with a maximum key stored of {:?} at depth {} with {} zeros",
38
                max_key, depth, index.embedded.0,
39
            );
40
            ScanEvaluation::ReadData
41
        },
42
        // The key evaluator is called for each key discovered in the tree, but
43
        // before the stored value is read.
44
        |key, index| {
45
            println!("Key {:?} has {} zeros", key, index.embedded.0);
46
            ScanEvaluation::Skip
47
        },
48
        // The data callback is invoked once data is read. In this example, we
49
        // always skip reading, so this is unreachable.
50
        |_key, _index, _value| unreachable!(),
51
    )?;
52

            
53
    Ok(())
54
}
55

            
56
#[derive(Clone, Debug)]
57
pub struct Zeroes(pub u32);
58

            
59
impl EmbeddedIndex<ArcBytes<'static>> for Zeroes {
60
    type Reduced = Self;
61
    type Indexer = ZeroesIndexer;
62
}
63

            
64
#[derive(Default, Clone, Debug)]
65
pub struct ZeroesIndexer;
66

            
67
impl Indexer<ArcBytes<'static>, Zeroes> for ZeroesIndexer {
68
    fn index(
69
        &self,
70
        _key: &nebari::ArcBytes<'_>,
71
        value: Option<&nebari::ArcBytes<'static>>,
72
    ) -> Zeroes {
73
        Zeroes(
74
            value
75
                .map(|bytes| bytes.iter().filter(|&b| b as char == '0').count())
76
                .unwrap_or_default() as u32,
77
        )
78
    }
79
}
80

            
81
impl Reducer<Zeroes> for ZeroesIndexer {
82
    fn reduce<'a, Indexes, IndexesIter>(&self, indexes: Indexes) -> Zeroes
83
    where
84
        Indexes: IntoIterator<Item = &'a Zeroes, IntoIter = IndexesIter> + ExactSizeIterator,
85
        IndexesIter: Iterator<Item = &'a Zeroes> + ExactSizeIterator + Clone,
86
    {
87
        Zeroes(indexes.into_iter().map(|i| i.0).sum())
88
    }
89

            
90
    fn rereduce<'a, ReducedIndexes, ReducedIndexesIter>(&self, values: ReducedIndexes) -> Zeroes
91
    where
92
        Self: 'a,
93
        ReducedIndexes:
94
            IntoIterator<Item = &'a Zeroes, IntoIter = ReducedIndexesIter> + ExactSizeIterator,
95
        ReducedIndexesIter: Iterator<Item = &'a Zeroes> + ExactSizeIterator + Clone,
96
    {
97
        // TODO change reduce to an iterator too
98
        self.reduce(values)
99
    }
100
}
101

            
102
impl Serializable for Zeroes {
103
    fn serialize_to<W: byteorder::WriteBytesExt>(&self, writer: &mut W) -> Result<usize, Error> {
104
        writer.write_u32::<BigEndian>(self.0)?;
105
        Ok(std::mem::size_of::<u32>())
106
    }
107

            
108
    fn deserialize_from<R: byteorder::ReadBytesExt>(reader: &mut R) -> Result<Self, Error> {
109
        Ok(Self(reader.read_u32::<BigEndian>()?))
110
    }
111
}