feat: add Config associated type to Database trait

This commit is contained in:
Magamedrasul Ibragimov 2023-02-10 14:52:45 +04:00
parent 69c9a5dd03
commit ee50bb7318
5 changed files with 69 additions and 27 deletions

View File

@ -4,13 +4,16 @@ use std::collections::HashMap;
/// Trait that must be implemented for a Database
pub trait Database {
/// Config for database. Default is necessary for a default() pmtree function
type Config: Default;
/// Creates new instance of db
fn new(dbpath: &str) -> Result<Self>
fn new(config: Self::Config) -> Result<Self>
where
Self: Sized;
/// Loades existing db (existence check required)
fn load(dbpath: &str) -> Result<Self>
fn load(config: Self::Config) -> Result<Self>
where
Self: Sized;

View File

@ -45,13 +45,13 @@ where
{
/// Creates tree with specified depth and default "pmtree_db" dbpath.
pub fn default(depth: usize) -> Result<Self> {
Self::new(depth, "pmtree_db")
Self::new(depth, D::Config::default())
}
/// Creates new `MerkleTree` and store it to the specified path/db
pub fn new(depth: usize, dbpath: &str) -> Result<Self> {
pub fn new(depth: usize, db_config: D::Config) -> Result<Self> {
// Create new db instance
let mut db = D::new(dbpath)?;
let mut db = D::new(db_config)?;
// Insert depth val into db
let depth_val = depth.to_be_bytes().to_vec();
@ -85,9 +85,9 @@ where
}
/// Loads existing Merkle Tree from the specified path/db
pub fn load(dbpath: &str) -> Result<Self> {
pub fn load(db_config: D::Config) -> Result<Self> {
// Load existing db instance
let db = D::load(dbpath)?;
let db = D::load(db_config)?;
// Load root
let root = H::deserialize(db.get(Key(0, 0).into())?.unwrap());

View File

@ -6,12 +6,17 @@ use tiny_keccak::{Hasher as _, Keccak};
struct MemoryDB(HashMap<DBKey, Value>);
struct MyKeccak(Keccak);
#[derive(Default)]
struct MemoryDBConfig;
impl Database for MemoryDB {
fn new(_dbpath: &str) -> Result<Self> {
type Config = MemoryDBConfig;
fn new(_db_config: MemoryDBConfig) -> Result<Self> {
Ok(MemoryDB(HashMap::new()))
}
fn load(_dbpath: &str) -> Result<Self> {
fn load(_db_config: MemoryDBConfig) -> Result<Self> {
Err(Box::new(Error("Cannot load in-memory DB".to_string())))
}
@ -60,7 +65,7 @@ impl Hasher for MyKeccak {
#[test]
fn insert_delete() -> Result<()> {
let mut mt = MerkleTree::<MemoryDB, MyKeccak>::new(2, "abacaba")?;
let mut mt = MerkleTree::<MemoryDB, MyKeccak>::new(2, MemoryDBConfig)?;
assert_eq!(mt.capacity(), 4);
assert_eq!(mt.depth(), 2);
@ -106,7 +111,7 @@ fn insert_delete() -> Result<()> {
#[test]
fn batch_insertions() -> Result<()> {
let mut mt = MerkleTree::<MemoryDB, MyKeccak>::new(2, "abacaba")?;
let mut mt = MerkleTree::<MemoryDB, MyKeccak>::new(2, MemoryDBConfig)?;
let leaves = [
hex!("0000000000000000000000000000000000000000000000000000000000000001"),

View File

@ -43,12 +43,17 @@ impl Hasher for PoseidonHash {
}
}
#[derive(Default)]
struct MemoryDBConfig;
impl Database for MemoryDB {
fn new(_dbpath: &str) -> Result<Self> {
type Config = MemoryDBConfig;
fn new(_db_config: MemoryDBConfig) -> Result<Self> {
Ok(MemoryDB(HashMap::new()))
}
fn load(_dbpath: &str) -> Result<Self> {
fn load(_db_config: MemoryDBConfig) -> Result<Self> {
Err(Box::new(Error("Cannot load in-memory DB".to_string())))
}
@ -69,9 +74,16 @@ impl Database for MemoryDB {
}
}
#[derive(Default)]
struct SledConfig {
path: String,
}
impl Database for MySled {
fn new(dbpath: &str) -> Result<Self> {
let db = sled::open(dbpath).unwrap();
type Config = SledConfig;
fn new(db_config: SledConfig) -> Result<Self> {
let db = sled::open(db_config.path).unwrap();
if db.was_recovered() {
return Err(Box::new(Error("Database exists, try load()!".to_string())));
}
@ -79,11 +91,11 @@ impl Database for MySled {
Ok(MySled(db))
}
fn load(dbpath: &str) -> Result<Self> {
let db = sled::open(dbpath).unwrap();
fn load(db_config: SledConfig) -> Result<Self> {
let db = sled::open(&db_config.path).unwrap();
if !db.was_recovered() {
fs::remove_dir_all(dbpath).expect("Error removing db");
fs::remove_dir_all(&db_config.path).expect("Error removing db");
return Err(Box::new(Error(
"Trying to load non-existing database!".to_string(),
)));
@ -119,7 +131,7 @@ impl Database for MySled {
#[test]
fn poseidon_memory() -> Result<()> {
let mut mt = MerkleTree::<MemoryDB, PoseidonHash>::new(TEST_TREE_HEIGHT, "abacaba")?;
let mut mt = MerkleTree::<MemoryDB, PoseidonHash>::new(TEST_TREE_HEIGHT, MemoryDBConfig)?;
let leaf_index = 3;
@ -243,7 +255,12 @@ fn poseidon_memory() -> Result<()> {
#[test]
fn poseidon_sled() -> Result<()> {
let mut mt = MerkleTree::<MySled, PoseidonHash>::new(TEST_TREE_HEIGHT, "abacaba")?;
let mut mt = MerkleTree::<MySled, PoseidonHash>::new(
TEST_TREE_HEIGHT,
SledConfig {
path: String::from("abacaba"),
},
)?;
let leaf_index = 3;

View File

@ -7,9 +7,16 @@ use tiny_keccak::{Hasher as _, Keccak};
struct MyKeccak(Keccak);
struct MySled(sled::Db);
#[derive(Default)]
struct SledConfig {
path: String,
}
impl Database for MySled {
fn new(dbpath: &str) -> Result<Self> {
let db = sled::open(dbpath).unwrap();
type Config = SledConfig;
fn new(db_config: SledConfig) -> Result<Self> {
let db = sled::open(db_config.path).unwrap();
if db.was_recovered() {
return Err(Box::new(Error("Database exists, try load()!".to_string())));
}
@ -17,11 +24,11 @@ impl Database for MySled {
Ok(MySled(db))
}
fn load(dbpath: &str) -> Result<Self> {
let db = sled::open(dbpath).unwrap();
fn load(db_config: SledConfig) -> Result<Self> {
let db = sled::open(&db_config.path).unwrap();
if !db.was_recovered() {
fs::remove_dir_all(dbpath).expect("Error removing db");
fs::remove_dir_all(&db_config.path).expect("Error removing db");
return Err(Box::new(Error(
"Trying to load non-existing database!".to_string(),
)));
@ -83,7 +90,12 @@ impl Hasher for MyKeccak {
#[test]
fn insert_delete() -> Result<()> {
let mut mt = MerkleTree::<MySled, MyKeccak>::new(2, "abacabas")?;
let mut mt = MerkleTree::<MySled, MyKeccak>::new(
2,
SledConfig {
path: String::from("abacabas"),
},
)?;
assert_eq!(mt.capacity(), 4);
assert_eq!(mt.depth(), 2);
@ -131,7 +143,12 @@ fn insert_delete() -> Result<()> {
#[test]
fn batch_insertions() -> Result<()> {
let mut mt = MerkleTree::<MySled, MyKeccak>::new(2, "abacabasa")?;
let mut mt = MerkleTree::<MySled, MyKeccak>::new(
2,
SledConfig {
path: String::from("abacabasa"),
},
)?;
let leaves = [
hex!("0000000000000000000000000000000000000000000000000000000000000001"),