1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159

use std::collections::BTreeMap;
use bson::{Document, DateTime};
use bson::oid::ObjectId;
use anyhow::{anyhow, Result};

use crate::storage::record::*;
use crate::storage::sstable::*;


/// The in-memory buffer for an LSM Tree.
/// 
/// This buffer is comprised of a red-black tree of records, sorted by key.
pub struct MemTable {
    pub records: BTreeMap<ObjectId, Value<Document>>,
}

impl MemTable {
    /// Creates a new MemTable.
    pub fn new() -> Self {
        Self::default()
    }

    pub fn insert(&mut self, key: ObjectId, value: Value<Document>) {
        self.records.insert(key, value);
    }

    pub fn set(&mut self, key: ObjectId, doc: Document) {
        self.insert(key, Value::Data(doc));
    }

    pub fn del(&mut self, key: ObjectId) {
        self.insert(key, Value::Tombstone);
    }

    pub fn get(&self, key: ObjectId) -> Option<Value<Document>> {
        self.records
            .get(&key)
            .map(|value| value.clone())
    }

    /// Flushes the contents of the MemTable to disk, returning an SSTable.
    pub fn flush(&mut self) -> Result<SSTable> {
        // Create a vector of records from the BTreeMap...
        let records: Vec<_> = self.records
            .iter()
            .map(|(key, value)| {
                Record {
                    key: *key,
                    value: value.clone(),
                }
            })
            .collect();

        // Get the min/max keys and count from the records...
        let min_key = records
            .first()
            .ok_or(anyhow!("records vec was empty"))?
            .key;
        let max_key = records
            .last()
            .ok_or(anyhow!("records vec was empty"))?
            .key;
        let num_records = records.len();
        let meta = SSTableMeta {
            id: ObjectId::new(),
            created_at: DateTime::now(),
            min_key,
            max_key,
            num_records,
        };

        // Create and return!
        Ok(SSTable {
            meta,
            records,
        })
    }
}

impl Default for MemTable {
    fn default() -> Self {
        MemTable {
            records: BTreeMap::new(),
        }
    }
}


#[cfg(test)]
mod test {
    use super::*;
    use bson::doc;

    #[test]
    fn create() {
        let _mt = MemTable::new();
    }

    #[test]
    fn set_and_get() {
        // Define a key/value pair to add to the memtable...
        let k = ObjectId::new();
        let v = doc! {
            "msg": "Hello, World!",
            "num": 42,
            "skyIsBlue": true,
        };
        let exp = Some(Value::Data(v.clone()));

        // Create an empty memtable...
        let mut mt = MemTable::new();

        // Add it to the memtable...
        mt.set(k, v);

        // Check that the BTree contains the key...
        assert!(mt.records.contains_key(&k), "Key doesn't exist in the btree");

        // Get the value back from the memtable...
        let res = mt.get(k);

        // Check that it matches...
        assert_eq!(res, exp);
    }

    #[test]
    fn set_del_get() {
        // Define a key/value pair to add to the memtable...
        let k = ObjectId::new();
        let v = doc! {
            "msg": "Hello, World!",
            "num": 42,
            "skyIsBlue": true,
        };

        // Create an empty memtable...
        let mut mt = MemTable::new();

        // Add it to the memtable...
        mt.set(k, v);

        // Check that the BTree contains the key...
        assert!(mt.records.contains_key(&k), "Key doesn't exist in the btree");

        // Delete the value from the memtable...
        mt.del(k);

        // Check that the key is still in the btree...
        assert!(mt.records.contains_key(&k), "Key should still exist in the btree after 'deletion'");
        
        // Check that the returned value is a tombstone...
        let res = mt.get(k);
        let exp = Some(Value::<Document>::Tombstone);
        assert_eq!(res, exp, "Expecting a present tombstone");
    }
}