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
//! Recursive structure stored using a compact stack machine representation
//! Collapsed via stack machine evaluation.
//!
use crate::{
    functor::Functor,
    recursive::{Collapse, Expand},
    recursive_tree::{RecursiveTree, RecursiveTreeRef},
};

/// Used to mark structures that are expanded via depth first traversal and consumed via stack machine
/// This is a zero-size marker type and has the lowest memory cost (lower than boxed pointers)
/// at the cost of a slightly slower 'Collapse::collapse_layers' fn speed
///
/// NOTE: adds hard requirement, functor traversal order MUST be constant and arity must not change
#[derive(Debug, Clone, Copy)]
pub struct StackMarker;

impl<A, U, O: Functor<StackMarker, Unwrapped = A, To = U>> Expand<A, O>
    for RecursiveTree<U, StackMarker>
{
    fn expand_layers<F: Fn(A) -> O>(a: A, generate_layer: F) -> Self {
        let mut frontier = Vec::from([a]);
        let mut elems = vec![];

        // unfold to build a vec of elems while preserving topo order
        while let Some(seed) = frontier.pop() {
            let layer = generate_layer(seed);

            let mut topush = Vec::new();
            let layer = layer.fmap(|aa| {
                topush.push(aa);
                StackMarker
            });
            frontier.extend(topush.into_iter().rev());

            elems.push(layer);
        }

        elems.reverse();

        Self {
            elems,
            _underlying: std::marker::PhantomData,
        }
    }
}

impl<A, O, U: Functor<A, To = O, Unwrapped = StackMarker>> Collapse<A, O>
    for RecursiveTree<U, StackMarker>
{
    fn collapse_layers<F: FnMut(O) -> A>(self, mut fold_layer: F) -> A {
        let mut result_stack = Vec::new();

        for layer in self.elems.into_iter() {
            // each layer is only referenced once so just remove it, also we know it's there so unwrap is fine
            let layer = layer.fmap(|_| result_stack.pop().unwrap());

            result_stack.push(fold_layer(layer));
        }

        result_stack.pop().unwrap()
    }
}

impl<'a, A, O: 'a, U> Collapse<A, O> for RecursiveTreeRef<'a, U, StackMarker>
where
    &'a U: Functor<A, To = O, Unwrapped = StackMarker>,
{
    fn collapse_layers<F: FnMut(O) -> A>(self, mut fold_layer: F) -> A {
        let mut result_stack = Vec::with_capacity(32);

        for layer in self.elems.iter() {
            let layer = layer.fmap(|_| result_stack.pop().unwrap());

            result_stack.push(fold_layer(layer));
        }

        result_stack.pop().unwrap()
    }
}