1//! Implements an elasticlunr.js document store. Most users do not need to use this module directly.
2
3use std::collections::BTreeMap;
4
5/// The document store saves the complete text of each item saved to the index, if enabled.
6/// Most users do not need to use this type directly.
7#[derive(Serialize, Deserialize, Debug, Clone)]
8#[serde(rename_all = "camelCase")]
9pub struct DocumentStore {
10 pub save: bool,
11 pub docs: BTreeMap<String, BTreeMap<String, String>>,
12 pub doc_info: BTreeMap<String, BTreeMap<String, usize>>,
13 // Redundant with docs.len(), but needed for serialization
14 pub length: usize,
15}
16
17impl DocumentStore {
18 pub fn new(save: bool) -> Self {
19 DocumentStore {
20 save,
21 docs: BTreeMap::new(),
22 doc_info: BTreeMap::new(),
23 length: 0,
24 }
25 }
26
27 pub fn len(&self) -> usize {
28 self.docs.len()
29 }
30
31 pub fn is_empty(&self) -> bool {
32 self.len() == 0
33 }
34
35 pub fn is_stored(&self) -> bool {
36 self.save
37 }
38
39 pub fn has_doc(&self, doc_ref: &str) -> bool {
40 self.docs.contains_key(doc_ref)
41 }
42
43 pub fn add_doc(&mut self, doc_ref: &str, doc: BTreeMap<String, String>) {
44 if !self.has_doc(doc_ref) {
45 self.length += 1;
46 }
47
48 self.docs.insert(
49 doc_ref.into(),
50 if self.save { doc } else { BTreeMap::new() },
51 );
52 }
53
54 pub fn get_doc(&self, doc_ref: &str) -> Option<BTreeMap<String, String>> {
55 self.docs.get(doc_ref).cloned()
56 }
57
58 pub fn remove_doc(&mut self, doc_ref: &str) {
59 if self.has_doc(doc_ref) {
60 self.length -= 1;
61 }
62
63 self.docs.remove(doc_ref);
64 }
65
66 pub fn add_field_length(&mut self, doc_ref: &str, field: &str, length: usize) {
67 self.doc_info
68 .entry(doc_ref.into())
69 .or_insert_with(BTreeMap::new)
70 .insert(field.into(), length);
71 }
72
73 pub fn get_field_length(&self, doc_ref: &str, field: &str) -> usize {
74 if self.has_doc(doc_ref) {
75 self.doc_info
76 .get(doc_ref)
77 .and_then(|e| e.get(field))
78 .cloned()
79 .unwrap_or(0)
80 } else {
81 0
82 }
83 }
84}
85
86#[cfg(test)]
87mod tests {
88 use super::*;
89
90 #[test]
91 fn add_doc_tokens() {
92 let mut store = DocumentStore::new(true);
93 let doc = btreemap! { "title".into() => "eggs bread".into() };
94
95 store.add_doc("1", doc.clone());
96 assert_eq!(store.get_doc("1").unwrap(), doc);
97 }
98
99 #[test]
100 fn create_doc_no_store() {
101 let mut store = DocumentStore::new(false);
102 let doc = btreemap! { "title".into() => "eggs bread".into() };
103
104 store.add_doc("1", doc);
105 assert_eq!(store.len(), 1);
106 assert_eq!(store.is_stored(), false);
107 assert_eq!(store.has_doc("1"), true);
108 }
109
110 #[test]
111 fn add_doc_no_store() {
112 let mut store = DocumentStore::new(false);
113 let doc1 = btreemap! { "title".into() => "eggs bread".into() };
114 let doc2 = btreemap! { "title".into() => "hello world".into() };
115
116 store.add_doc("1", doc1);
117 store.add_doc("2", doc2);
118 assert_eq!(store.len(), 2);
119 assert_eq!(store.is_stored(), false);
120 assert_eq!(store.has_doc("1"), true);
121 assert_eq!(store.has_doc("2"), true);
122 }
123
124 #[test]
125 fn is_stored_true() {
126 let store = DocumentStore::new(true);
127 assert_eq!(store.is_stored(), true);
128 }
129
130 #[test]
131 fn is_stored_false() {
132 let store = DocumentStore::new(false);
133 assert_eq!(store.is_stored(), false);
134 }
135
136 #[test]
137 fn get_doc_no_store() {
138 let mut store = DocumentStore::new(false);
139 let doc1 = btreemap! { "title".into() => "eggs bread".into() };
140 let doc2 = btreemap! { "title".into() => "hello world".into() };
141
142 store.add_doc("1", doc1);
143 store.add_doc("2", doc2);
144 assert_eq!(store.len(), 2);
145 assert_eq!(store.is_stored(), false);
146 assert_eq!(store.get_doc("1").unwrap(), BTreeMap::new());
147 assert_eq!(store.get_doc("2").unwrap(), BTreeMap::new());
148 }
149
150 #[test]
151 fn get_nonexistant_doc_no_store() {
152 let mut store = DocumentStore::new(false);
153 let doc1 = btreemap! { "title".into() => "eggs bread".into() };
154 let doc2 = btreemap! { "title".into() => "hello world".into() };
155
156 store.add_doc("1", doc1);
157 store.add_doc("2", doc2);
158 assert_eq!(store.len(), 2);
159 assert_eq!(store.is_stored(), false);
160 assert_eq!(store.get_doc("6"), None);
161 assert_eq!(store.get_doc("2").unwrap(), BTreeMap::new());
162 }
163
164 #[test]
165 fn remove_doc_no_store() {
166 let mut store = DocumentStore::new(false);
167 let doc1 = btreemap! { "title".into() => "eggs bread".into() };
168 let doc2 = btreemap! { "title".into() => "hello world".into() };
169
170 store.add_doc("1", doc1);
171 store.add_doc("2", doc2);
172 store.remove_doc("1");
173 assert_eq!(store.len(), 1);
174 assert_eq!(store.is_stored(), false);
175 assert_eq!(store.get_doc("2").unwrap(), BTreeMap::new());
176 assert_eq!(store.get_doc("1"), None);
177 }
178
179 #[test]
180 fn remove_nonexistant_doc() {
181 let mut store = DocumentStore::new(false);
182 let doc1 = btreemap! { "title".into() => "eggs bread".into() };
183 let doc2 = btreemap! { "title".into() => "hello world".into() };
184
185 store.add_doc("1", doc1);
186 store.add_doc("2", doc2);
187 store.remove_doc("8");
188 assert_eq!(store.len(), 2);
189 assert_eq!(store.is_stored(), false);
190 assert_eq!(store.get_doc("2").unwrap(), BTreeMap::new());
191 assert_eq!(store.get_doc("1").unwrap(), BTreeMap::new());
192 }
193
194 #[test]
195 fn get_num_docs() {
196 let mut store = DocumentStore::new(true);
197
198 assert_eq!(store.len(), 0);
199 store.add_doc("1", btreemap! { "title".into() => "eggs bread".into() });
200 assert_eq!(store.len(), 1);
201 }
202
203 #[test]
204 fn get_doc() {
205 let mut store = DocumentStore::new(true);
206
207 assert_eq!(store.len(), 0);
208 store.add_doc("1", btreemap! { "title".into() => "eggs bread".into() });
209 assert_eq!(
210 store.get_doc("1").unwrap(),
211 btreemap! { "title".into() => "eggs bread".into() }
212 );
213 }
214
215 #[test]
216 fn get_doc_many_fields() {
217 let mut store = DocumentStore::new(true);
218
219 assert_eq!(store.len(), 0);
220 store.add_doc(
221 "1",
222 btreemap! {
223 "title".into() => "eggs bread".into()
224 },
225 );
226 store.add_doc(
227 "2",
228 btreemap! {
229 "title".into() => "boo bar".into()
230 },
231 );
232 store.add_doc(
233 "3",
234 btreemap! {
235 "title".into() => "oracle".into(),
236 "body".into() => "Oracle is demonspawn".into()
237 },
238 );
239 assert_eq!(
240 store.get_doc("3").unwrap(),
241 btreemap! {
242 "title".into() => "oracle".into(),
243 "body".into() => "Oracle is demonspawn".into()
244 }
245 );
246 assert_eq!(store.len(), 3);
247 }
248
249 #[test]
250 fn get_nonexistant_doc() {
251 let mut store = DocumentStore::new(true);
252
253 assert_eq!(store.len(), 0);
254 store.add_doc(
255 "1",
256 btreemap! {
257 "title".into() => "eggs bread".into()
258 },
259 );
260 store.add_doc(
261 "2",
262 btreemap! {
263 "title".into() => "boo bar".into()
264 },
265 );
266 store.add_doc(
267 "3",
268 btreemap! {
269 "title".into() => "oracle".into(),
270 "body".into() => "Oracle is demonspawn".into()
271 },
272 );
273 assert_eq!(store.get_doc("4"), None);
274 assert_eq!(store.get_doc("0"), None);
275 assert_eq!(store.len(), 3);
276 }
277
278 #[test]
279 fn check_store_has_key() {
280 let mut store = DocumentStore::new(true);
281
282 assert!(!store.has_doc("foo"));
283 store.add_doc("foo", btreemap! { "title".into() => "eggs bread".into() });
284 assert!(store.has_doc("foo"));
285 }
286
287 #[test]
288 fn remove_doc() {
289 let mut store = DocumentStore::new(true);
290
291 store.add_doc("foo", btreemap! { "title".into() => "eggs bread".into() });
292 assert!(store.has_doc("foo"));
293 assert_eq!(store.len(), 1);
294 store.remove_doc("foo");
295 assert!(!store.has_doc("foo"));
296 assert_eq!(store.len(), 0);
297 }
298
299 #[test]
300 fn remove_nonexistant_store() {
301 let mut store = DocumentStore::new(true);
302
303 store.add_doc("foo", btreemap! { "title".into() => "eggs bread".into() });
304 assert!(store.has_doc("foo"));
305 assert_eq!(store.len(), 1);
306 store.remove_doc("bar");
307 assert!(store.has_doc("foo"));
308 assert_eq!(store.len(), 1);
309 }
310
311 #[test]
312 fn add_field_len() {
313 let mut store = DocumentStore::new(true);
314
315 store.add_doc("foo", btreemap! { "title".into() => "eggs bread".into() });
316 store.add_field_length("foo", "title", 2);
317 assert_eq!(store.get_field_length("foo", "title"), 2);
318 }
319
320 #[test]
321 fn add_field_length_multiple() {
322 let mut store = DocumentStore::new(true);
323
324 store.add_doc("foo", btreemap! { "title".into() => "eggs bread".into() });
325 store.add_field_length("foo", "title", 2);
326 store.add_field_length("foo", "body", 10);
327 assert_eq!(store.get_field_length("foo", "title"), 2);
328 assert_eq!(store.get_field_length("foo", "body"), 10);
329 }
330}
331