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

// 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 crate::branch::Branch;
use crate::branch_mut::BranchMut;
use crate::compound::{Child, Compound, MutableLeaves};

use core::marker::PhantomData;

use ranno::Annotation;

/// The return value from a closure to [`walk`] the tree.
///
/// Determines how the [`Branch`] is constructed
///
/// [`walk`]: Walker::walk
pub enum Step {
    /// The correct leaf was found!
    Found(usize),
    /// Traverse the branch deeper
    Into(usize),
    /// Advance search
    Advance,
    /// Abort search
    Abort,
}

/// The argument given to a [`Walker`] to traverse through nodes.
pub struct Walk<'a, C, A> {
    index: usize,
    compound: &'a C,
    _marker: PhantomData<A>,
}

impl<'a, C, A> Walk<'a, C, A> {
    pub(crate) fn new(compound: &'a C, index: usize) -> Self {
        Walk {
            index,
            compound,
            _marker: PhantomData,
        }
    }
}

impl<'a, C, A> Walk<'a, C, A>
where
    C: Compound<A>,
{
    /// Returns the child at specific index relative to the branch index
    pub fn child(&self, index: usize) -> Child<'a, C, A> {
        self.compound.child(index + self.index)
    }
}

/// The trait used to construct a [`Branch`] or to iterate through a tree.
pub trait Walker<C, A> {
    /// Walk the tree node, returning the appropriate [`Step`]
    fn walk(&mut self, walk: Walk<C, A>) -> Step;
}

/// Walker that visits all leaves
pub struct AllLeaves;

impl<C, A> Walker<C, A> for AllLeaves
where
    C: Compound<A>,
{
    fn walk(&mut self, walk: Walk<C, A>) -> Step {
        for i in 0.. {
            match walk.child(i) {
                Child::Leaf(_) => return Step::Found(i),
                Child::Node(_) => return Step::Into(i),
                Child::Empty => (),
                Child::EndOfNode => return Step::Advance,
            }
        }
        unreachable!()
    }
}

/// Trait that provides a [`first`] and [`first_mut`] method to any
/// [`Compound`].
///
/// [`first`]: First::first
/// [`first_mut`]: First::first_mut
pub trait First<A>: Sized + Compound<A> {
    /// Construct a [`Branch`] pointing to the first element, if not empty
    fn first(&self) -> Option<Branch<Self, A>>;

    /// Construct a [`BranchMut`] pointing to the first element, if not empty
    fn first_mut(&mut self) -> Option<BranchMut<Self, A>>
    where
        Self: MutableLeaves;
}

impl<C, A> First<A> for C
where
    C: Compound<A>,
    A: Annotation<C>,
{
    fn first(&self) -> Option<Branch<Self, A>> {
        Branch::walk(self, AllLeaves)
    }

    fn first_mut(&mut self) -> Option<BranchMut<Self, A>>
    where
        C: MutableLeaves,
    {
        BranchMut::walk(self, AllLeaves)
    }
}