use crate::v11::{make_empty_map, make_map_with_root_and_bitwidth, Keyer, Map};
use cid::Cid;
use fvm_ipld_blockstore::Blockstore;
use fvm_ipld_hamt::{BytesKey, Error};
use serde::de::DeserializeOwned;
use serde::Serialize;
use serde::__private::PhantomData;
use std::collections::btree_map::Entry::{Occupied, Vacant};
use std::collections::BTreeMap;
pub struct MapMap<'a, BS, V, K1, K2> {
outer: Map<'a, BS, Cid>,
inner_bitwidth: u32,
cache: BTreeMap<Vec<u8>, Map<'a, BS, V>>,
key_types: PhantomData<(K1, K2)>,
}
impl<'a, BS, V, K1, K2> MapMap<'a, BS, V, K1, K2>
where
BS: Blockstore,
V: Serialize + DeserializeOwned + Clone + std::cmp::PartialEq,
K1: Keyer + std::fmt::Debug + std::fmt::Display,
K2: Keyer + std::fmt::Debug + std::fmt::Display,
{
pub fn new(bs: &'a BS, outer_bitwidth: u32, inner_bitwidth: u32) -> Self {
MapMap {
outer: make_empty_map(bs, outer_bitwidth),
inner_bitwidth,
cache: BTreeMap::<Vec<u8>, Map<BS, V>>::new(),
key_types: PhantomData,
}
}
pub fn from_root(
bs: &'a BS,
cid: &Cid,
outer_bitwidth: u32,
inner_bitwidth: u32,
) -> Result<Self, Error> {
Ok(MapMap {
outer: make_map_with_root_and_bitwidth(cid, bs, outer_bitwidth)?,
inner_bitwidth,
cache: BTreeMap::<Vec<u8>, Map<BS, V>>::new(),
key_types: PhantomData,
})
}
pub fn flush(&mut self) -> Result<Cid, Error> {
for (k, in_map) in self.cache.iter_mut() {
if in_map.is_empty() {
self.outer.delete(&BytesKey(k.to_vec()))?;
} else {
let new_in_root = in_map.flush()?;
self.outer.set(BytesKey(k.to_vec()), new_in_root)?;
}
}
self.outer.flush()
}
fn load_inner_map(&mut self, k: K1) -> Result<(bool, &mut Map<'a, BS, V>), Error> {
let in_map_thunk = || -> Result<(bool, Map<BS, V>), Error> {
match self.outer.get(&k.key())? {
Some(root) => Ok((
false,
make_map_with_root_and_bitwidth::<BS, V>(
root,
*self.outer.store(),
self.inner_bitwidth,
)?,
)),
None => Ok((
true,
make_empty_map(*self.outer.store(), self.inner_bitwidth),
)),
}
};
let raw_k = k.key().0;
match self.cache.entry(raw_k) {
Occupied(entry) => {
let in_map = entry.into_mut();
Ok((in_map.is_empty(), in_map))
}
Vacant(entry) => {
let (empty, in_map) = in_map_thunk()?;
Ok((empty, entry.insert(in_map)))
}
}
}
pub fn get(&mut self, outside_k: K1, inside_k: K2) -> Result<Option<&V>, Error> {
let (is_empty, in_map) = self.load_inner_map(outside_k)?;
if is_empty {
return Ok(None);
}
in_map.get(&inside_k.key())
}
pub fn for_each_outer<F>(&self, f: F) -> Result<(), Error>
where
F: FnMut(&BytesKey, &Cid) -> anyhow::Result<()>,
{
self.outer.for_each(f)
}
pub fn for_each<F>(&mut self, outside_k: K1, f: F) -> Result<(), Error>
where
F: FnMut(&BytesKey, &V) -> anyhow::Result<()>,
{
let (is_empty, in_map) = self.load_inner_map(outside_k)?;
if is_empty {
return Ok(());
}
in_map.for_each(f)
}
pub fn put(&mut self, outside_k: K1, inside_k: K2, value: V) -> Result<Option<V>, Error> {
let in_map = self.load_inner_map(outside_k)?.1;
in_map.set(inside_k.key(), value)
}
pub fn put_if_absent(&mut self, outside_k: K1, inside_k: K2, value: V) -> Result<bool, Error> {
let in_map = self.load_inner_map(outside_k)?.1;
in_map.set_if_absent(inside_k.key(), value)
}
pub fn put_many<I>(&mut self, outside_k: K1, values: I) -> Result<(), Error>
where
I: Iterator<Item = (K2, V)>,
{
let in_map = self.load_inner_map(outside_k)?.1;
for (k, v) in values {
in_map.set(k.key(), v)?;
}
Ok(())
}
pub fn remove(&mut self, outside_k: K1, inside_k: K2) -> Result<Option<V>, Error> {
let (is_empty, in_map) = self.load_inner_map(outside_k)?;
if is_empty {
return Ok(None);
}
in_map
.delete(&inside_k.key())
.map(|o: Option<(BytesKey, V)>| -> Option<V> { o.map(|p: (BytesKey, V)| -> V { p.1 }) })
}
}