1use fs4::fs_std::FileExt;
2use serde::de::DeserializeOwned;
3use serde::Serialize;
4use std::collections::btree_map::Keys;
5use std::collections::BTreeMap;
6use std::fs::{File, OpenOptions};
7use std::io::Read;
8use std::io::Seek;
9use std::io::{BufRead, Write};
10use std::marker::PhantomData;
11use std::path::{Path, PathBuf};
12use std::vec;
13
14pub mod index;
15pub mod rtree;
16
17#[derive(Debug)]
18pub struct Ref<V> {
19 start: usize,
20 length: usize,
21 p: PhantomData<V>,
22}
23
24#[derive(Debug)]
25pub struct JSStore<K, V>(BTreeMap<K, Ref<V>>, File, PathBuf);
26
27impl<K: DeserializeOwned + Eq + Ord, V: DeserializeOwned> JSStore<K, V> {
28 pub fn new<P: AsRef<Path>>(filename: P) -> Result<JSStore<K, V>, std::io::Error> {
29 let mut map = BTreeMap::new();
30 if !Path::exists(filename.as_ref()) {
31 let file = OpenOptions::new()
32 .read(true)
33 .create(true)
34 .write(true)
35 .open(&filename)?;
36 file.try_lock_exclusive()?;
37 return Ok(JSStore(map, file, filename.as_ref().to_path_buf()));
38 }
39
40 let file = OpenOptions::new()
41 .read(true)
42 .create(true)
43 .write(true)
44 .open(filename.as_ref())?;
45 file.try_lock_exclusive()?;
46 let mut buf = std::io::BufReader::new(file);
47 let mut offset = 0;
48 loop {
49 let mut line = String::new();
50 let size = buf.read_line(&mut line)?;
51 if size == 0 {
52 break;
53 }
54 let data: Vec<(K, Option<V>)> = match serde_json::from_str(&line) {
55 Ok(data) => data,
56 Err(_e) => {
57 offset = offset + size;
58 continue;
59 }
60 };
61 for (key, value) in data {
62 match value {
63 Some(_value) => {
64 map.insert(
65 key,
66 Ref {
67 start: offset,
68 length: size,
69 p: PhantomData,
70 },
71 );
72 }
73 None => {
74 map.remove(&key);
75 }
76 }
77 }
78 offset = offset + size;
79 }
80 let file = buf.into_inner();
81
82 let result = JSStore(map, file, filename.as_ref().to_path_buf());
83 Ok(result)
84 }
85}
86
87impl<K: Eq + Ord + Serialize, V: Serialize> JSStore<K, V> {
88 pub fn insert<Key: Into<K>>(&mut self, k: Key, v: V) -> Result<(), std::io::Error> {
89 let k = k.into();
90 let mut file = &self.1;
91 file.seek(std::io::SeekFrom::End(0))?;
92 let offset = file.stream_position()?;
93 file.write_all(b"\n")?;
94 let buf = serde_json::to_vec(&[(&k, &v)])?;
95 file.write_all(&buf)?;
96 self.0.insert(
97 k,
98 Ref {
99 start: offset as usize,
100 length: buf.len() + 1,
101 p: PhantomData,
102 },
103 );
104 Ok(())
105 }
106 pub fn remove<Key: Into<K>>(&mut self, k: Key) -> Result<(), std::io::Error> {
107 let k = k.into();
108 let mut file = &self.1;
109 file.seek(std::io::SeekFrom::End(0))?;
110 file.write_all(b"\n")?;
111 let buf = serde_json::to_vec::<[(&K, Option<V>)]>(&[(&k, None)])?;
112 file.write_all(&buf)?;
113 self.0.remove(&k);
114 Ok(())
115 }
116 pub fn batch(&mut self) -> Batch<JSStore<K, V>, K, V> {
117 Batch::new(self)
118 }
119 pub fn dump<P: AsRef<Path>>(&self, filename: P) -> Result<Self, std::io::Error>
120 where
121 V: DeserializeOwned,
122 K: DeserializeOwned + Clone,
123 {
124 let mut store = Self::new(filename)?;
125 for (k, _v) in &self.0 {
126 let v = self.get(k.clone())?.unwrap();
127 store.insert(k.clone(), v)?;
128 }
129 Ok(store)
130 }
131 pub fn compress<P: AsRef<Path>>(&mut self, tempfilename: P) -> Result<(), std::io::Error>
132 where
133 V: DeserializeOwned,
134 K: DeserializeOwned + Clone,
135 {
136 let mut newstore = self.dump(tempfilename)?;
137 std::mem::swap(&mut newstore.0, &mut self.0);
138 std::mem::swap(&mut newstore.1, &mut self.1);
139 let tempfilename = newstore.2.clone();
140 drop(newstore);
141 let targetname = self.2.as_path();
142 std::fs::remove_file(&targetname)?;
143 std::fs::rename(&tempfilename, &targetname)?;
144 Ok(())
145 }
146 pub fn get<Key: Into<K>>(&self, key: Key) -> Result<Option<V>, std::io::Error>
147 where
148 V: DeserializeOwned,
149 K: DeserializeOwned,
150 {
151 let key = key.into();
152 let r = match self.0.get(&key) {
153 Some(r) => r,
154 None => {
155 return Ok(None);
156 }
157 };
158 let mut file = &self.1;
159 file.seek(std::io::SeekFrom::Start(r.start.try_into().unwrap()))?;
160 let mut buf = vec![0u8; r.length];
161 file.read_exact(&mut buf)?;
162 let s = unsafe { String::from_utf8_unchecked(buf) };
163 let data: Vec<(K, Option<V>)> = serde_json::from_str(&s).unwrap();
164 let data = data.into_iter().rev();
165 for (k, v) in data {
166 if k == key {
167 return Ok(v);
168 }
169 }
170 unreachable!()
171 }
172 pub fn contains_key<Key: Into<K>>(&self, key: Key) -> bool {
173 self.0.contains_key(&key.into())
174 }
175 pub fn iter<'a>(&'a self) -> Iter<'a, K, V> {
176 Iter(self.0.iter(), self)
177 }
178 pub fn keys(&self)->Keys<K, Ref<V>>{
179 self.0.keys()
180 }
181}
182
183pub struct Iter<'a, K, V>(
184 std::collections::btree_map::Iter<'a, K, Ref<V>>,
185 &'a JSStore<K, V>,
186);
187
188impl<'a, K, V> Iterator for Iter<'a, K, V>
189where
190 K: Eq + Ord + Serialize + Clone + DeserializeOwned,
191 V: Serialize + DeserializeOwned,
192{
193 type Item = (K, V);
194
195 fn next(&mut self) -> Option<Self::Item> {
196 self.0.next().and_then(|(k, _v)| {
197 let v = match self.1.get(k.to_owned()) {
198 Ok(Some(v)) => v,
199 _ => return None,
200 };
201 Some((k.to_owned(), v))
202 })
203 }
204}
205
206pub struct Batch<'a, B, K, V>(&'a mut B, Vec<(K, Option<V>)>);
207impl<'a, B, K, V> Batch<'a, B, K, V>
208where
209 B: BatchCommit<K, V>,
210{
211 fn new(store: &'a mut B) -> Batch<'a, B, K, V> {
212 Batch(store, vec![])
213 }
214 pub fn insert<Key: Into<K>>(mut self, k: Key, v: V) -> Self {
215 let k = k.into();
216 self.1.push((k, Some(v)));
217 self
218 }
219 pub fn remove<Key: Into<K>>(mut self, k: Key) -> Self {
220 let k = k.into();
221 self.1.push((k, None));
222 self
223 }
224 pub fn commit(self) -> Result<(), std::io::Error>
225 where
226 K: Eq + Ord + Serialize,
227 V: Serialize,
228 {
229 self.0.commit(self.1)
230 }
231}
232
233pub trait BatchCommit<K, V> {
234 fn commit(&mut self, oplist: Vec<(K, Option<V>)>) -> Result<(), std::io::Error>;
235}
236
237impl<K, V> BatchCommit<K, V> for JSStore<K, V>
238where
239 K: Eq + Ord + Serialize,
240 V: Serialize,
241{
242 fn commit(&mut self, oplist: Vec<(K, Option<V>)>) -> Result<(), std::io::Error> {
243 let mut file = &self.1;
244 let v: Vec<(K, Option<V>)> = oplist;
245 file.seek(std::io::SeekFrom::End(0))?;
246 let offset = file.stream_position()?;
247 file.write_all(b"\n")?;
248 let buf = serde_json::to_vec(&v).unwrap();
249 file.write_all(&buf)?;
250 for (key, value) in v {
251 match value {
252 Some(_value) => {
253 self.0.insert(
254 key,
255 Ref {
256 start: offset as usize,
257 length: buf.len() + 1,
258 p: PhantomData,
259 },
260 );
261 }
262 None => {
263 self.0.remove(&key);
264 }
265 }
266 }
267 Ok(())
268 }
269}
270
271#[cfg(test)]
272mod tests {
273 use crate::JSStore;
274 #[test]
275 fn test() {
276 let mut x: JSStore<String, i32> = JSStore::new("test.json").unwrap();
277 x.insert("x", 100).unwrap();
278 x.insert("z", 200).unwrap();
279 drop(x);
280 let mut x: JSStore<String, i32> = JSStore::new("test.json").unwrap();
281 assert!(x.get("x").unwrap() == Some(100));
282 x.batch()
283 .insert("x", 10)
284 .insert("y", 20)
285 .remove("z")
286 .commit()
287 .unwrap();
288 drop(x);
289 let mut x: JSStore<String, i32> = JSStore::new("test.json").unwrap();
290 assert!(x.get("z").unwrap() == None);
291 assert!(x.get("x").unwrap() == Some(10));
292 assert!(x.get("y").unwrap() == Some(20));
293 x.compress("tem.json").unwrap();
294 assert!(x.get("z").unwrap() == None);
295 assert!(x.get("x").unwrap() == Some(10));
296 assert!(x.get("y").unwrap() == Some(20));
297 assert!(x.contains_key("y"));
298 let n: Vec<_> = x.iter().collect();
299 assert!(n.len() == 2);
300 let x: Result<JSStore<String, i32>, _> = JSStore::new("test.json");
301 assert!(x.is_err());
302 }
303}