use fs4::fs_std::FileExt;
use serde::de::DeserializeOwned;
use serde::Serialize;
use std::collections::btree_map::Keys;
use std::collections::BTreeMap;
use std::fs::{File, OpenOptions};
use std::io::Read;
use std::io::Seek;
use std::io::{BufRead, Write};
use std::marker::PhantomData;
use std::path::{Path, PathBuf};
use std::vec;
pub mod index;
pub mod rtree;
#[derive(Debug)]
pub struct Ref<V> {
start: usize,
length: usize,
p: PhantomData<V>,
}
#[derive(Debug)]
pub struct JSStore<K, V>(BTreeMap<K, Ref<V>>, File, PathBuf);
impl<K: DeserializeOwned + Eq + Ord, V: DeserializeOwned> JSStore<K, V> {
pub fn new<P: AsRef<Path>>(filename: P) -> Result<JSStore<K, V>, std::io::Error> {
let mut map = BTreeMap::new();
if !Path::exists(filename.as_ref()) {
let file = OpenOptions::new()
.read(true)
.create(true)
.write(true)
.open(&filename)?;
file.try_lock_exclusive()?;
return Ok(JSStore(map, file, filename.as_ref().to_path_buf()));
}
let file = OpenOptions::new()
.read(true)
.create(true)
.write(true)
.open(filename.as_ref())?;
file.try_lock_exclusive()?;
let mut buf = std::io::BufReader::new(file);
let mut offset = 0;
loop {
let mut line = String::new();
let size = buf.read_line(&mut line)?;
if size == 0 {
break;
}
let data: Vec<(K, Option<V>)> = match serde_json::from_str(&line) {
Ok(data) => data,
Err(_e) => {
offset = offset + size;
continue;
}
};
for (key, value) in data {
match value {
Some(_value) => {
map.insert(
key,
Ref {
start: offset,
length: size,
p: PhantomData,
},
);
}
None => {
map.remove(&key);
}
}
}
offset = offset + size;
}
let file = buf.into_inner();
let result = JSStore(map, file, filename.as_ref().to_path_buf());
Ok(result)
}
}
impl<K: Eq + Ord + Serialize, V: Serialize> JSStore<K, V> {
pub fn insert<Key: Into<K>>(&mut self, k: Key, v: V) -> Result<(), std::io::Error> {
let k = k.into();
let mut file = &self.1;
file.seek(std::io::SeekFrom::End(0))?;
let offset = file.stream_position()?;
file.write_all(b"\n")?;
let buf = serde_json::to_vec(&[(&k, &v)])?;
file.write_all(&buf)?;
self.0.insert(
k,
Ref {
start: offset as usize,
length: buf.len() + 1,
p: PhantomData,
},
);
Ok(())
}
pub fn remove<Key: Into<K>>(&mut self, k: Key) -> Result<(), std::io::Error> {
let k = k.into();
let mut file = &self.1;
file.seek(std::io::SeekFrom::End(0))?;
file.write_all(b"\n")?;
let buf = serde_json::to_vec::<[(&K, Option<V>)]>(&[(&k, None)])?;
file.write_all(&buf)?;
self.0.remove(&k);
Ok(())
}
pub fn batch(&mut self) -> Batch<JSStore<K, V>, K, V> {
Batch::new(self)
}
pub fn dump<P: AsRef<Path>>(&self, filename: P) -> Result<Self, std::io::Error>
where
V: DeserializeOwned,
K: DeserializeOwned + Clone,
{
let mut store = Self::new(filename)?;
for (k, _v) in &self.0 {
let v = self.get(k.clone())?.unwrap();
store.insert(k.clone(), v)?;
}
Ok(store)
}
pub fn compress<P: AsRef<Path>>(&mut self, tempfilename: P) -> Result<(), std::io::Error>
where
V: DeserializeOwned,
K: DeserializeOwned + Clone,
{
let mut newstore = self.dump(tempfilename)?;
std::mem::swap(&mut newstore.0, &mut self.0);
std::mem::swap(&mut newstore.1, &mut self.1);
let tempfilename = newstore.2.clone();
drop(newstore);
let targetname = self.2.as_path();
std::fs::remove_file(&targetname)?;
std::fs::rename(&tempfilename, &targetname)?;
Ok(())
}
pub fn get<Key: Into<K>>(&self, key: Key) -> Result<Option<V>, std::io::Error>
where
V: DeserializeOwned,
K: DeserializeOwned,
{
let key = key.into();
let r = match self.0.get(&key) {
Some(r) => r,
None => {
return Ok(None);
}
};
let mut file = &self.1;
file.seek(std::io::SeekFrom::Start(r.start.try_into().unwrap()))?;
let mut buf = vec![0u8; r.length];
file.read_exact(&mut buf)?;
let s = unsafe { String::from_utf8_unchecked(buf) };
let data: Vec<(K, Option<V>)> = serde_json::from_str(&s).unwrap();
let data = data.into_iter().rev();
for (k, v) in data {
if k == key {
return Ok(v);
}
}
unreachable!()
}
pub fn contains_key<Key: Into<K>>(&self, key: Key) -> bool {
self.0.contains_key(&key.into())
}
pub fn iter<'a>(&'a self) -> Iter<'a, K, V> {
Iter(self.0.iter(), self)
}
pub fn keys(&self)->Keys<K, Ref<V>>{
self.0.keys()
}
}
pub struct Iter<'a, K, V>(
std::collections::btree_map::Iter<'a, K, Ref<V>>,
&'a JSStore<K, V>,
);
impl<'a, K, V> Iterator for Iter<'a, K, V>
where
K: Eq + Ord + Serialize + Clone + DeserializeOwned,
V: Serialize + DeserializeOwned,
{
type Item = (K, V);
fn next(&mut self) -> Option<Self::Item> {
self.0.next().and_then(|(k, _v)| {
let v = match self.1.get(k.to_owned()) {
Ok(Some(v)) => v,
_ => return None,
};
Some((k.to_owned(), v))
})
}
}
pub struct Batch<'a, B, K, V>(&'a mut B, Vec<(K, Option<V>)>);
impl<'a, B, K, V> Batch<'a, B, K, V>
where
B: BatchCommit<K, V>,
{
fn new(store: &'a mut B) -> Batch<'a, B, K, V> {
Batch(store, vec![])
}
pub fn insert<Key: Into<K>>(mut self, k: Key, v: V) -> Self {
let k = k.into();
self.1.push((k, Some(v)));
self
}
pub fn remove<Key: Into<K>>(mut self, k: Key) -> Self {
let k = k.into();
self.1.push((k, None));
self
}
pub fn commit(self) -> Result<(), std::io::Error>
where
K: Eq + Ord + Serialize,
V: Serialize,
{
self.0.commit(self.1)
}
}
pub trait BatchCommit<K, V> {
fn commit(&mut self, oplist: Vec<(K, Option<V>)>) -> Result<(), std::io::Error>;
}
impl<K, V> BatchCommit<K, V> for JSStore<K, V>
where
K: Eq + Ord + Serialize,
V: Serialize,
{
fn commit(&mut self, oplist: Vec<(K, Option<V>)>) -> Result<(), std::io::Error> {
let mut file = &self.1;
let v: Vec<(K, Option<V>)> = oplist;
file.seek(std::io::SeekFrom::End(0))?;
let offset = file.stream_position()?;
file.write_all(b"\n")?;
let buf = serde_json::to_vec(&v).unwrap();
file.write_all(&buf)?;
for (key, value) in v {
match value {
Some(_value) => {
self.0.insert(
key,
Ref {
start: offset as usize,
length: buf.len() + 1,
p: PhantomData,
},
);
}
None => {
self.0.remove(&key);
}
}
}
Ok(())
}
}
#[cfg(test)]
mod tests {
use crate::JSStore;
#[test]
fn test() {
let mut x: JSStore<String, i32> = JSStore::new("test.json").unwrap();
x.insert("x", 100).unwrap();
x.insert("z", 200).unwrap();
drop(x);
let mut x: JSStore<String, i32> = JSStore::new("test.json").unwrap();
assert!(x.get("x").unwrap() == Some(100));
x.batch()
.insert("x", 10)
.insert("y", 20)
.remove("z")
.commit()
.unwrap();
drop(x);
let mut x: JSStore<String, i32> = JSStore::new("test.json").unwrap();
assert!(x.get("z").unwrap() == None);
assert!(x.get("x").unwrap() == Some(10));
assert!(x.get("y").unwrap() == Some(20));
x.compress("tem.json").unwrap();
assert!(x.get("z").unwrap() == None);
assert!(x.get("x").unwrap() == Some(10));
assert!(x.get("y").unwrap() == Some(20));
assert!(x.contains_key("y"));
let n: Vec<_> = x.iter().collect();
assert!(n.len() == 2);
let x: Result<JSStore<String, i32>, _> = JSStore::new("test.json");
assert!(x.is_err());
}
}