clvm_fuzzing/
make_tree.rs

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