clvm_fuzzing/
make_tree.rs

1use arbitrary::{Arbitrary, Result, Unstructured};
2use clvmr::{Allocator, NodePtr};
3
4enum Op {
5    Pair(bool),
6    SubTree,
7}
8
9#[derive(Arbitrary)]
10enum NodeType {
11    Pair,
12    Bytes,
13    U8,
14    U16,
15    U32,
16    Previous,
17}
18
19#[derive(Debug)]
20pub struct ArbitraryClvmTree<const MAX_NODES: i64 = 600_000, const REUSE_NODES: bool = true> {
21    pub allocator: Allocator,
22    pub tree: NodePtr,
23    pub num_nodes: u32,
24}
25
26impl<'a, const MAX_NODES: i64, const REUSE_NODES: bool> Arbitrary<'a>
27    for ArbitraryClvmTree<MAX_NODES, REUSE_NODES>
28{
29    fn arbitrary(u: &mut Unstructured<'a>) -> Result<Self> {
30        let mut a = Allocator::new();
31        let (tree, num_nodes) =
32            make_tree_limits(&mut a, u, MAX_NODES, REUSE_NODES).expect("make_tree");
33        Ok(Self {
34            allocator: a,
35            tree,
36            num_nodes,
37        })
38    }
39}
40
41pub fn make_tree(a: &mut Allocator, unstructured: &mut Unstructured<'_>) -> (NodePtr, u32) {
42    make_tree_limits(a, unstructured, 600_000, true).expect("out of memory")
43}
44
45/// returns an arbitrary CLVM tree structure and the number of (unique) nodes
46/// it's made up of. That's both pairs and atoms.
47pub fn make_tree_limits(
48    a: &mut Allocator,
49    unstructured: &mut Unstructured<'_>,
50    mut max_nodes: i64,
51    reuse_nodes: bool,
52) -> anyhow::Result<(NodePtr, u32)> {
53    let mut previous_nodes = Vec::new();
54    let mut value_stack = Vec::new();
55    let mut op_stack = vec![Op::SubTree];
56    // the number of Op::SubTree items on the op_stack
57    let mut sub_trees: i64 = 1;
58    let mut counter = 0;
59
60    while let Some(op) = op_stack.pop() {
61        match op {
62            Op::Pair(swap) => {
63                let left = value_stack.pop().expect("internal error, empty stack");
64                let right = value_stack.pop().expect("internal error, empty stack");
65                let pair = if swap {
66                    a.new_pair(left, right)?
67                } else {
68                    a.new_pair(right, left)?
69                };
70                counter += 1;
71                value_stack.push(pair);
72                previous_nodes.push(pair);
73            }
74            Op::SubTree => {
75                sub_trees -= 1;
76                if unstructured.is_empty() {
77                    value_stack.push(NodePtr::NIL);
78                    continue;
79                }
80                match unstructured.arbitrary::<NodeType>() {
81                    Err(..) => value_stack.push(NodePtr::NIL),
82                    Ok(NodeType::Pair) => {
83                        if sub_trees > unstructured.len() as i64 || max_nodes <= 0 {
84                            // there isn't much entropy left, don't grow the
85                            // tree anymore
86                            value_stack.push(if reuse_nodes {
87                                *unstructured
88                                    .choose(&previous_nodes)
89                                    .unwrap_or(&NodePtr::NIL)
90                            } else {
91                                NodePtr::NIL
92                            });
93                        } else {
94                            // swap left and right arbitrarily, to avoid
95                            // having a bias because we build the tree depth
96                            // first, until we run out of entropy
97                            let swap = unstructured.arbitrary::<bool>() == Ok(true);
98                            op_stack.push(Op::Pair(swap));
99                            op_stack.push(Op::SubTree);
100                            op_stack.push(Op::SubTree);
101                            sub_trees += 2;
102                            max_nodes -= 2;
103                        }
104                    }
105                    Ok(NodeType::Bytes) => {
106                        counter += 1;
107                        value_stack.push(match unstructured.arbitrary::<Vec<u8>>() {
108                            Err(..) => NodePtr::NIL,
109                            Ok(val) => {
110                                let node = a.new_atom(&val)?;
111                                previous_nodes.push(node);
112                                node
113                            }
114                        });
115                    }
116                    Ok(NodeType::U8) => {
117                        counter += 1;
118                        value_stack.push(match unstructured.arbitrary::<u8>() {
119                            Err(..) => NodePtr::NIL,
120                            Ok(val) => a.new_small_number(val.into())?,
121                        });
122                    }
123                    Ok(NodeType::U16) => {
124                        counter += 1;
125                        value_stack.push(match unstructured.arbitrary::<u16>() {
126                            Err(..) => NodePtr::NIL,
127                            Ok(val) => a.new_small_number(val.into())?,
128                        });
129                    }
130                    Ok(NodeType::U32) => {
131                        counter += 1;
132                        value_stack.push(match unstructured.arbitrary::<u32>() {
133                            Err(..) => NodePtr::NIL,
134                            Ok(val) => a.new_number(val.into())?,
135                        });
136                    }
137                    Ok(NodeType::Previous) => {
138                        value_stack.push(if reuse_nodes {
139                            *unstructured
140                                .choose(&previous_nodes)
141                                .unwrap_or(&NodePtr::NIL)
142                        } else {
143                            NodePtr::NIL
144                        });
145                    }
146                }
147            }
148        }
149    }
150    assert_eq!(value_stack.len(), 1);
151    Ok((value_stack.remove(0), counter))
152}
153
154pub fn make_list(a: &mut Allocator, unstructured: &mut Unstructured<'_>) -> NodePtr {
155    let mut ret = NodePtr::NIL;
156
157    let length = unstructured.arbitrary::<u8>().unwrap_or(0);
158    let nil_terminated = unstructured.arbitrary::<bool>().unwrap_or(false);
159
160    for _ in 0..length {
161        let value = match unstructured
162            .arbitrary::<NodeType>()
163            .unwrap_or(NodeType::Previous)
164        {
165            NodeType::U8 => a
166                .new_small_number(unstructured.arbitrary::<u8>().unwrap_or(0).into())
167                .unwrap(),
168
169            NodeType::U16 => a
170                .new_small_number(unstructured.arbitrary::<u16>().unwrap_or(0).into())
171                .unwrap(),
172
173            NodeType::U32 => a
174                .new_number(unstructured.arbitrary::<u32>().unwrap_or(0).into())
175                .unwrap(),
176
177            NodeType::Bytes => a
178                .new_atom(&unstructured.arbitrary::<Vec<u8>>().unwrap_or_default())
179                .unwrap(),
180
181            NodeType::Pair => {
182                let left = NodePtr::NIL;
183                let right = NodePtr::NIL;
184                a.new_pair(left, right).unwrap()
185            }
186
187            NodeType::Previous => NodePtr::NIL,
188        };
189
190        ret = if nil_terminated {
191            a.new_pair(value, ret).unwrap()
192        } else {
193            value
194        };
195    }
196
197    ret
198}
clvm_fuzzing/make_tree.rs

clvm_fuzzing/
make_tree.rs
