#![allow(clippy::upper_case_acronyms, dead_code, clippy::type_complexity)]
use banyan::{
index::{CompactSeq, Summarizable, VecSeq},
query::{AllQuery, AndQuery, OffsetRangeQuery, Query},
store::{BranchCache, MemStore, ReadOnlyStore},
StreamBuilder, Transaction, Tree, TreeTypes,
};
use futures::Future;
use libipld::{
cbor::DagCborCodec,
codec::{Codec, Decode, Encode},
Cid, DagCbor, Ipld,
};
use quickcheck::{Arbitrary, Gen, TestResult};
use range_collections::RangeSet;
use sha2::{Digest, Sha256};
use std::{
collections::HashSet,
convert::{TryFrom, TryInto},
fmt,
io::{Read, Seek, Write},
iter::FromIterator,
ops::Range,
};
#[allow(dead_code)]
pub type Forest = banyan::Forest<TT, MemStore<Sha256Digest>>;
#[allow(dead_code)]
pub type Txn = Transaction<TT, MemStore<Sha256Digest>, MemStore<Sha256Digest>>;
#[derive(Debug, Clone)]
pub struct TT;
#[derive(Debug, Clone, Copy, PartialEq, Eq, DagCbor)]
pub struct Key(pub u64);
#[derive(Debug, Clone, Copy, PartialEq, Eq, DagCbor)]
pub struct KeyRange(pub u64, pub u64);
impl KeyRange {
fn as_range_set(&self) -> RangeSet<u64> {
RangeSet::from(self.0..self.1.saturating_add(1))
}
}
#[derive(Debug, Clone)]
pub struct KeyQuery(pub RangeSet<u64>);
impl Query<TT> for KeyQuery {
fn containing(&self, _: u64, index: &banyan::index::LeafIndex<TT>, res: &mut [bool]) {
for (res, key) in res.iter_mut().zip(index.keys()) {
*res = *res && self.0.contains(&key.0);
}
}
fn intersecting(&self, _offset: u64, index: &banyan::index::BranchIndex<TT>, res: &mut [bool]) {
for (res, range) in res.iter_mut().zip(index.summaries()) {
*res = *res && !self.0.is_disjoint(&range.as_range_set());
}
}
}
#[derive(Debug, Clone, DagCbor)]
pub struct KeySeq(pub Vec<Key>);
impl CompactSeq for KeySeq {
type Item = Key;
fn get(&self, index: usize) -> Option<Key> {
self.0.get(index).cloned()
}
fn len(&self) -> usize {
self.0.len()
}
}
impl Summarizable<KeyRange> for KeySeq {
fn summarize(&self) -> KeyRange {
let min = self.0.iter().map(|x| x.0).min().unwrap();
let max = self.0.iter().map(|x| x.0).max().unwrap();
KeyRange(min, max)
}
}
impl Summarizable<KeyRange> for VecSeq<KeyRange> {
fn summarize(&self) -> KeyRange {
let min = self.as_ref().iter().map(|x| x.0).min().unwrap();
let max = self.as_ref().iter().map(|x| x.1).max().unwrap();
KeyRange(min, max)
}
}
impl FromIterator<Key> for KeySeq {
fn from_iter<T: IntoIterator<Item = Key>>(iter: T) -> Self {
KeySeq(iter.into_iter().collect())
}
}
impl TreeTypes for TT {
type Key = Key;
type KeySeq = KeySeq;
type Summary = KeyRange;
type SummarySeq = VecSeq<KeyRange>;
type Link = Sha256Digest;
}
impl Key {
fn combine(&mut self, rhs: &Key) {
self.0 |= rhs.0
}
}
impl Arbitrary for Key {
fn arbitrary(g: &mut Gen) -> Self {
Self(Arbitrary::arbitrary(g))
}
}
#[allow(dead_code)]
pub fn txn(store: MemStore<Sha256Digest>, cache_cap: usize) -> Txn {
let branch_cache = BranchCache::new(cache_cap);
Txn::new(Forest::new(store.clone(), branch_cache), store)
}
#[derive(Debug, Clone)]
pub struct TestFilter {
offset_range: Range<u64>,
keys: RangeSet<u64>,
}
impl TestFilter {
pub fn offset_range(offset_range: Range<u64>) -> Self {
Self {
offset_range,
keys: RangeSet::all(),
}
}
pub fn keys(values: impl IntoIterator<Item = u64>) -> Self {
let mut keys = RangeSet::empty();
for value in values {
keys |= RangeSet::from(value..value + 1);
}
Self {
offset_range: 0..u64::max_value(),
keys,
}
}
pub fn query(&self) -> impl Query<TT> + Clone {
AndQuery(
OffsetRangeQuery::from(self.offset_range.clone()),
KeyQuery(self.keys.clone()),
)
}
pub fn contains(&self, triple: &(u64, Key, u64)) -> bool {
let (offset, key, _) = triple;
self.offset_range.contains(offset) && self.keys.contains(&key.0)
}
}
impl Arbitrary for TestFilter {
fn arbitrary(g: &mut Gen) -> Self {
let offset_range = Arbitrary::arbitrary(g);
let points: Vec<u16> = Arbitrary::arbitrary(g);
let mut keys = RangeSet::empty();
for key in points {
let key = key as u64;
keys |= RangeSet::from(key..key + 1);
}
Self { offset_range, keys }
}
}
#[derive(Debug, Clone)]
pub enum TestTree {
Unpacked(UnpackedTestTree),
Packed(PackedTestTree),
}
impl TestTree {
pub fn packed(xs: Vec<(Key, u64)>) -> Self {
Self::Packed(PackedTestTree(xs))
}
pub fn unpacked(xss: Vec<Vec<(Key, u64)>>) -> Self {
Self::Unpacked(UnpackedTestTree(xss))
}
pub fn tree(self) -> anyhow::Result<(Tree<TT, u64>, Txn, Vec<(Key, u64)>)> {
let (b, txn, xs) = self.builder()?;
Ok((b.snapshot(), txn, xs))
}
pub fn builder(self) -> anyhow::Result<(StreamBuilder<TT, u64>, Txn, Vec<(Key, u64)>)> {
match self {
Self::Unpacked(x) => x.builder(),
Self::Packed(x) => x.builder(),
}
}
}
impl Arbitrary for TestTree {
fn arbitrary(g: &mut Gen) -> Self {
if Arbitrary::arbitrary(g) {
Self::Packed(Arbitrary::arbitrary(g))
} else {
Self::Unpacked(Arbitrary::arbitrary(g))
}
}
}
#[derive(Debug, Clone)]
pub struct PackedTestTree(Vec<(Key, u64)>);
impl PackedTestTree {
pub fn builder(self) -> anyhow::Result<(StreamBuilder<TT, u64>, Txn, Vec<(Key, u64)>)> {
let store = MemStore::new(usize::max_value(), Sha256Digest::digest);
let mut txn = txn(store, 1 << 20);
let mut builder = StreamBuilder::<TT, u64>::debug();
let xs = self.0.clone();
txn.extend(&mut builder, self.0)?;
txn.assert_invariants(&builder)?;
Ok((builder, txn, xs))
}
}
impl Arbitrary for PackedTestTree {
fn arbitrary(g: &mut Gen) -> Self {
Self(Arbitrary::arbitrary(g))
}
}
#[derive(Debug, Clone)]
pub struct UnpackedTestTree(Vec<Vec<(Key, u64)>>);
impl UnpackedTestTree {
pub fn builder(self) -> anyhow::Result<(StreamBuilder<TT, u64>, Txn, Vec<(Key, u64)>)> {
let store = MemStore::new(usize::max_value(), Sha256Digest::digest);
let mut txn = txn(store, 1 << 20);
let mut builder = StreamBuilder::<TT, u64>::debug();
let xs = self.0.iter().cloned().flatten().collect();
for xs in self.0 {
txn.extend_unpacked(&mut builder, xs)?;
}
txn.assert_invariants(&builder)?;
Ok((builder, txn, xs))
}
}
impl Arbitrary for UnpackedTestTree {
fn arbitrary(g: &mut Gen) -> Self {
let xs: Vec<(Key, u64)> = Arbitrary::arbitrary(g);
let mut cuts: Vec<usize> = Arbitrary::arbitrary(g);
if !xs.is_empty() {
for x in cuts.iter_mut() {
*x %= xs.len();
}
cuts.push(0);
cuts.push(xs.len());
cuts.sort_unstable();
cuts.dedup();
} else {
cuts.clear();
}
let res = cuts
.iter()
.zip(cuts.iter().skip(1))
.map(|(start, end)| xs[*start..*end].to_vec())
.collect::<Vec<_>>();
UnpackedTestTree(res)
}
}
pub fn links(
forest: &Forest,
tree: &Tree<TT, u64>,
links: &mut HashSet<Sha256Digest>,
) -> anyhow::Result<()> {
let link_opts = forest
.iter_index(tree, AllQuery)
.map(|x| x.map(|x| x.link().as_ref().cloned()))
.collect::<anyhow::Result<Vec<_>>>()?;
links.extend(link_opts.into_iter().flatten());
Ok(())
}
fn same_secrets<X, Y>(a: &Tree<TT, X>, b: &Tree<TT, Y>) -> bool {
match (a.secrets(), b.secrets()) {
(Some(sa), Some(sb)) => {
sa.index_key() == sb.index_key() || sa.value_key() == sb.value_key()
}
_ => {
true
}
}
}
type OffsetSet = RangeSet<u64, [u64; 16]>;
#[derive(DagCbor)]
struct IpldNode(u64, Vec<Cid>, Ipld);
impl IpldNode {
fn into_data(self) -> anyhow::Result<(u64, Vec<Cid>, Vec<u8>)> {
if let Ipld::Bytes(data) = self.2 {
Ok((self.0, self.1, data))
} else {
Err(anyhow::anyhow!("expected ipld bytes"))
}
}
}
fn block_range(forest: &Forest, hash: &Sha256Digest) -> anyhow::Result<Range<u64>> {
let blob = forest.store().get(hash)?;
let (offset, _, encrypted) = DagCborCodec.decode::<IpldNode>(&blob)?.into_data()?;
let len = encrypted.len() as u64;
let end_offset = offset
.checked_add(len)
.ok_or_else(|| anyhow::anyhow!("seek wraparound"))?;
Ok(offset..end_offset)
}
fn no_range_overlap(forest: &Forest, hashes: HashSet<Sha256Digest>) -> anyhow::Result<bool> {
let mut ranges = OffsetSet::empty();
let mut result = true;
for hash in hashes {
let range = OffsetSet::from(block_range(forest, &hash)?);
if !ranges.is_disjoint(&range) {
result = false;
}
ranges |= range;
}
Ok(result)
}
#[allow(dead_code)]
pub fn no_offset_overlap<'a>(
forest: &'a Forest,
trees: impl IntoIterator<Item = &'a Tree<TT, u64>>,
) -> anyhow::Result<bool> {
let mut trees = trees.into_iter().peekable();
if let Some(first) = trees.peek().cloned() {
let mut hashes = HashSet::new();
for tree in trees {
anyhow::ensure!(same_secrets(first, tree), "not the same secrets");
links(forest, tree, &mut hashes)?;
}
no_range_overlap(forest, hashes)
} else {
Ok(true)
}
}
#[allow(dead_code)]
pub async fn test<F: Future<Output = anyhow::Result<bool>>>(f: impl Fn() -> F) -> TestResult {
match f().await {
Ok(success) => TestResult::from_bool(success),
Err(cause) => TestResult::error(cause.to_string()),
}
}
#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq, PartialOrd, Ord)]
pub struct Sha256Digest([u8; 32]);
impl Decode<DagCborCodec> for Sha256Digest {
fn decode<R: Read + Seek>(c: DagCborCodec, r: &mut R) -> anyhow::Result<Self> {
Self::try_from(libipld::Cid::decode(c, r)?)
}
}
impl Encode<DagCborCodec> for Sha256Digest {
fn encode<W: Write>(&self, c: DagCborCodec, w: &mut W) -> anyhow::Result<()> {
libipld::Cid::encode(&Cid::from(*self), c, w)
}
}
impl Sha256Digest {
pub fn digest(data: &[u8]) -> Self {
let mut hasher = Sha256::new();
hasher.update(data);
let result = hasher.finalize();
Sha256Digest(result.try_into().unwrap())
}
pub fn read(data: &[u8]) -> anyhow::Result<Self> {
Ok(Self(data[0..32].try_into()?))
}
}
impl AsRef<[u8]> for Sha256Digest {
fn as_ref(&self) -> &[u8] {
&self.0
}
}
impl fmt::Display for Sha256Digest {
fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
write!(f, "{}", hex::encode(self.as_ref()))
}
}
impl From<Sha256Digest> for Cid {
fn from(value: Sha256Digest) -> Self {
let mh = multihash::Multihash::wrap(0x12, &value.0).unwrap();
Cid::new_v1(0x71, mh)
}
}
impl TryFrom<Cid> for Sha256Digest {
type Error = anyhow::Error;
fn try_from(value: Cid) -> Result<Self, Self::Error> {
anyhow::ensure!(value.codec() == 0x71, "Unexpected codec");
anyhow::ensure!(value.hash().code() == 0x12, "Unexpected hash algorithm");
let digest: [u8; 32] = value.hash().digest().try_into()?;
Ok(Self(digest))
}
}
pub(crate) trait IterExt<'a>
where
Self: Iterator + Sized + Send + 'a,
{
fn boxed(self) -> Box<dyn Iterator<Item = Self::Item> + Send + 'a> {
Box::new(self)
}
}
impl<'a, T: Iterator + Sized + Send + 'a> IterExt<'a> for T {}