ngr/src/syntax/arbitrary.rs

use crate::syntax::ast::{ConstantType, Expression, Name, Program, Statement, Value};
use crate::syntax::location::Location;
use proptest::sample::select;
use proptest::{
    prelude::{Arbitrary, BoxedStrategy, Strategy},
    strategy::{Just, Union},
};
use std::collections::HashMap;
use std::ops::Range;

const VALID_VARIABLE_NAMES: &str = r"[a-z][a-zA-Z0-9_]*";

impl ConstantType {
    fn get_operators(&self) -> &'static [(&'static str, usize)] {
        match self {
            ConstantType::I8 | ConstantType::I16 | ConstantType::I32 | ConstantType::I64 => {
                &[("+", 2), ("-", 1), ("-", 2), ("*", 2), ("/", 2)]
            }
            ConstantType::U8 | ConstantType::U16 | ConstantType::U32 | ConstantType::U64 => {
                &[("+", 2), ("-", 2), ("*", 2), ("/", 2)]
            }
        }
    }
}

#[derive(Clone)]
pub struct GenerationEnvironment {
    allow_inference: bool,
    block_length: Range<usize>,
    bindings: HashMap<Name, ConstantType>,
    return_type: ConstantType,
}

impl Default for GenerationEnvironment {
    fn default() -> Self {
        GenerationEnvironment {
            allow_inference: true,
            block_length: 2..10,
            bindings: HashMap::new(),
            return_type: ConstantType::U64,
        }
    }
}

impl GenerationEnvironment {
    pub fn new(allow_inference: bool) -> Self {
        GenerationEnvironment {
            allow_inference,
            ..Default::default()
        }
    }
}

impl Arbitrary for Program {
    type Parameters = GenerationEnvironment;
    type Strategy = BoxedStrategy<Self>;

    fn arbitrary_with(genenv: Self::Parameters) -> Self::Strategy {
        proptest::collection::vec(
            ProgramStatementInfo::arbitrary(),
            genenv.block_length.clone(),
        )
        .prop_flat_map(move |mut items| {
            let mut statements = Vec::new();
            let mut genenv = genenv.clone();

            for psi in items.drain(..) {
                if genenv.bindings.is_empty() || psi.should_be_binding {
                    genenv.return_type = psi.binding_type;
                    let expr = Expression::arbitrary_with(genenv.clone());
                    genenv.bindings.insert(psi.name.clone(), psi.binding_type);
                    statements.push(
                        expr.prop_map(move |expr| {
                            Statement::Binding(Location::manufactured(), psi.name.clone(), expr)
                        })
                        .boxed(),
                    );
                } else {
                    let printers = genenv.bindings.keys().map(|n| {
                        Just(Statement::Print(
                            Location::manufactured(),
                            Name::manufactured(n),
                        ))
                    });
                    statements.push(Union::new(printers).boxed());
                }
            }

            statements
                .prop_map(|statements| Program { statements })
                .boxed()
        })
        .boxed()
    }
}

impl Arbitrary for Name {
    type Parameters = ();
    type Strategy = BoxedStrategy<Self>;

    fn arbitrary_with(_: Self::Parameters) -> Self::Strategy {
        VALID_VARIABLE_NAMES.prop_map(Name::manufactured).boxed()
    }
}

#[derive(Debug)]
struct ProgramStatementInfo {
    should_be_binding: bool,
    name: Name,
    binding_type: ConstantType,
}

impl Arbitrary for ProgramStatementInfo {
    type Parameters = ();
    type Strategy = BoxedStrategy<Self>;

    fn arbitrary_with(_args: Self::Parameters) -> Self::Strategy {
        (
            Union::new(vec![Just(true), Just(true), Just(false)]),
            Name::arbitrary(),
            ConstantType::arbitrary(),
        )
            .prop_map(
                |(should_be_binding, name, binding_type)| ProgramStatementInfo {
                    should_be_binding,
                    name,
                    binding_type,
                },
            )
            .boxed()
    }
}

impl Arbitrary for Expression {
    type Parameters = GenerationEnvironment;
    type Strategy = BoxedStrategy<Self>;

    fn arbitrary_with(genenv: Self::Parameters) -> Self::Strategy {
        // Value(Location, Value). These are the easiest variations to create, because we can always
        // create one.
        let value_strategy = Value::arbitrary_with(genenv.clone())
            .prop_map(|x| Expression::Value(Location::manufactured(), x))
            .boxed();

        // Reference(Location, String), These are slightly trickier, because we can end up in a situation
        // where either no variables are defined, or where none of the defined variables have a type we
        // can work with. So what we're going to do is combine this one with the previous one as a "leaf
        // strategy" -- our non-recursive items -- if we can, or just set that to be the value strategy
        // if we can't actually create an references.
        let mut bound_variables_of_type = genenv
            .bindings
            .iter()
            .filter(|(_, v)| genenv.return_type == **v)
            .map(|(n, _)| n)
            .collect::<Vec<_>>();
        let leaf_strategy = if bound_variables_of_type.is_empty() {
            value_strategy
        } else {
            let mut strats = bound_variables_of_type
                .drain(..)
                .map(|x| {
                    Just(Expression::Reference(
                        Location::manufactured(),
                        x.name.clone(),
                    ))
                    .boxed()
                })
                .collect::<Vec<_>>();
            strats.push(value_strategy);
            Union::new(strats).boxed()
        };

        // now we generate our recursive types, given our leaf strategy
        leaf_strategy
            .prop_recursive(3, 10, 2, move |strat| {
                (
                    select(genenv.return_type.get_operators()),
                    strat.clone(),
                    strat,
                )
                    .prop_map(|((oper, count), left, right)| {
                        let mut args = vec![left, right];
                        while args.len() > count {
                            args.pop();
                        }
                        Expression::Primitive(Location::manufactured(), oper.to_string(), args)
                    })
            })
            .boxed()
    }
}

impl Arbitrary for Value {
    type Parameters = GenerationEnvironment;
    type Strategy = BoxedStrategy<Self>;

    fn arbitrary_with(genenv: Self::Parameters) -> Self::Strategy {
        let printed_base_strategy = Union::new([
            Just(None::<u8>),
            Just(Some(2)),
            Just(Some(8)),
            Just(Some(10)),
            Just(Some(16)),
        ]);
        let value_strategy = u64::arbitrary();

        (printed_base_strategy, bool::arbitrary(), value_strategy)
            .prop_map(move |(base, declare_type, value)| {
                let converted_value = match genenv.return_type {
                    ConstantType::I8 => value % (i8::MAX as u64),
                    ConstantType::U8 => value % (u8::MAX as u64),
                    ConstantType::I16 => value % (i16::MAX as u64),
                    ConstantType::U16 => value % (u16::MAX as u64),
                    ConstantType::I32 => value % (i32::MAX as u64),
                    ConstantType::U32 => value % (u32::MAX as u64),
                    ConstantType::I64 => value % (i64::MAX as u64),
                    ConstantType::U64 => value,
                };
                let ty = if declare_type || !genenv.allow_inference {
                    Some(genenv.return_type)
                } else {
                    None
                };
                Value::Number(base, ty, converted_value)
            })
            .boxed()
    }
}

impl Arbitrary for ConstantType {
    type Parameters = ();
    type Strategy = BoxedStrategy<Self>;

    fn arbitrary_with(_: Self::Parameters) -> Self::Strategy {
        Union::new([
            Just(ConstantType::I8),
            Just(ConstantType::I16),
            Just(ConstantType::I32),
            Just(ConstantType::I64),
            Just(ConstantType::U8),
            Just(ConstantType::U16),
            Just(ConstantType::U32),
            Just(ConstantType::U64),
        ])
        .boxed()
    }
}