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

//! Generates a zoo of terms form a [`Signature`]. For each function symbol in the signature
//! a closed term is generated and added to the zoo.

use libafl::bolts::rands::Rand;

use crate::{
    algebra::{
        atoms::Function,
        signature::{FunctionDefinition, Signature},
        Matcher, Term,
    },
    fuzzer::mutations::util::Choosable,
};

const MAX_DEPTH: u16 = 8; // how deep terms we allow max
const MAX_TRIES: u16 = 100; // How often we want to try to generate before stopping

pub struct TermZoo<M: Matcher> {
    terms: Vec<Term<M>>,
}

impl<M: Matcher> TermZoo<M> {
    pub fn generate<R: Rand>(signature: &Signature, rand: &mut R) -> Self {
        let terms = signature
            .functions
            .iter()
            .filter_map(|def| {
                let mut counter = MAX_TRIES;

                loop {
                    if counter == 0 {
                        break None;
                    }

                    counter -= 1;

                    if let Some(term) = Self::generate_term(signature, def, MAX_DEPTH, rand) {
                        break Some(term);
                    }
                }
            })
            .collect::<Vec<_>>();

        Self { terms }
    }

    fn generate_term<R: Rand>(
        signature: &Signature,
        (shape, dynamic_fn): &FunctionDefinition,
        depth: u16,
        rand: &mut R,
    ) -> Option<Term<M>> {
        if depth == 0 {
            // Reached max depth
            return None;
        }

        let required_types = &shape.argument_types;

        let mut subterms = Vec::with_capacity(required_types.len());

        for typ in required_types {
            if let Some(possibilities) = signature.functions_by_typ.get(typ) {
                if let Some(possibility) = possibilities.choose(rand) {
                    if let Some(subterm) =
                        Self::generate_term(signature, possibility, depth - 1, rand)
                    {
                        subterms.push(subterm)
                    } else {
                        // Max depth reached
                        return None;
                    }
                } else {
                    // No possibilities available
                    return None;
                }
            } else {
                // Type not available
                return None;
            }
        }

        Some(Term::Application(
            Function::new(shape.clone(), dynamic_fn.clone()),
            subterms,
        ))
    }

    pub fn choose_filtered<P, R: Rand>(&self, filter: P, rand: &mut R) -> Option<&Term<M>>
    where
        P: FnMut(&&Term<M>) -> bool,
    {
        self.terms.choose_filtered(filter, rand)
    }

    pub fn terms(&self) -> &[Term<M>] {
        &self.terms
    }
}