1
  2
  3
  4
  5
  6
  7
  8
  9
 10
 11
 12
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148

// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.
//
// Copyright (c) DUSK NETWORK. All rights reserved.

use alloc::vec::Vec;

#[cfg(feature = "persistence")]
use arbitrary::Arbitrary;

use canonical::{Canon, CanonError, EncodeToVec, Id, Source};
use canonical_derive::Canon;

use crate::link::Link;
use crate::{Annotation, Compound};

const TAG_EMPTY: u8 = 0;
const TAG_LEAF: u8 = 1;
const TAG_LINK: u8 = 2;

/// A generic annotation
#[derive(Clone, Canon, Debug, PartialEq)]
#[cfg_attr(feature = "persistence", derive(Arbitrary))]
pub struct GenericAnnotation(Vec<u8>);

/// A generic leaf
#[derive(Clone, Canon, Debug, PartialEq)]
#[cfg_attr(feature = "persistence", derive(Arbitrary))]
pub struct GenericLeaf(Vec<u8>);

impl GenericLeaf {
    pub(crate) fn new<C: Canon>(c: &C) -> Self {
        let vec = c.encode_to_vec();
        let res = GenericLeaf(vec);
        res
    }

    /// Cast the generic leaf to a concrete type
    pub fn cast<T: Canon>(&self) -> Result<T, CanonError> {
        T::decode(&mut Source::new(&self.0))
    }
}

impl GenericAnnotation {
    pub(crate) fn new<A: Canon>(a: &A) -> Self {
        GenericAnnotation(a.encode_to_vec())
    }

    /// Cast the generic leaf to a concrete type
    pub fn cast<A: Canon>(&self) -> Result<A, CanonError> {
        A::decode(&mut Source::new(&self.0))
    }
}

/// A generic child of a collection
#[derive(Clone, Debug, PartialEq)]
#[cfg_attr(feature = "persistence", derive(Arbitrary))]
pub enum GenericChild {
    /// Child is empty
    Empty,
    /// Child is a leaf    
    Leaf(GenericLeaf),
    /// Child is a link        
    Link(Id, GenericAnnotation),
}

impl Canon for GenericChild {
    fn encode(&self, sink: &mut canonical::Sink) {
        match self {
            Self::Empty => TAG_EMPTY.encode(sink),
            Self::Leaf(leaf) => {
                TAG_LEAF.encode(sink);
                leaf.encode(sink)
            }
            Self::Link(id, annotation) => {
                TAG_LINK.encode(sink);
                id.encode(sink);
                annotation.encode(sink);
            }
        }
    }

    fn decode(source: &mut canonical::Source) -> Result<Self, CanonError> {
        match u8::decode(source)? {
            TAG_EMPTY => Ok(GenericChild::Empty),
            TAG_LEAF => Ok(GenericChild::Leaf(GenericLeaf::decode(source)?)),
            TAG_LINK => {
                let id = Id::decode(source)?;
                let anno = GenericAnnotation::decode(source)?;
                Ok(GenericChild::Link(id, anno))
            }
            _ => Err(CanonError::InvalidEncoding),
        }
    }

    fn encoded_len(&self) -> usize {
        const TAG_LEN: usize = 1;
        match self {
            Self::Empty => TAG_LEN,
            Self::Leaf(leaf) => TAG_LEN + leaf.encoded_len(),
            Self::Link(id, anno) => {
                TAG_LEN + id.encoded_len() + anno.encoded_len()
            }
        }
    }
}

/// The generic tree structure, this is a generic version of any Compound tree,
/// which has had it's leaves and annotations replaced with generic variants of
/// prefixed lengths, so that the tree structure can still be followed even if
/// you don't know the concrete associated and generic types of the Compound
/// structure that was persisted
#[derive(Default, Clone, Canon, Debug, PartialEq)]
#[cfg_attr(feature = "persistence", derive(Arbitrary))]
pub struct GenericTree(Vec<GenericChild>);

impl GenericTree {
    pub(crate) fn new() -> Self {
        GenericTree(vec![])
    }

    pub(crate) fn push_empty(&mut self) {
        self.0.push(GenericChild::Empty)
    }

    pub(crate) fn push_leaf<L: Canon>(&mut self, leaf: &L) {
        let leaf = GenericLeaf::new(leaf);
        let child = GenericChild::Leaf(leaf);
        self.0.push(child)
    }

    pub(crate) fn push_link<C, A>(&mut self, link: &Link<C, A>)
    where
        C: Compound<A>,
        C::Leaf: Canon,
        A: Annotation<C::Leaf>,
    {
        let id = link.id();
        let anno = GenericAnnotation::new(&*link.annotation());
        self.0.push(GenericChild::Link(id, anno));
    }

    /// Provides an iterator over the generic children of the node
    pub fn children(&self) -> &[GenericChild] {
        &self.0
    }
}