walpurgis/src/ast_builder.rs

use rand::Rng;
use rand::seq::SliceRandom;
use rand_distr::Distribution;
use std::collections::HashMap;

use crate::params::Params;
use crate::params::a_lt;
use crate::ast::*;

#[derive(Default)]
struct StateTracking {
    int_recursion_depth: u32,
    real_recursion_depth: u32,
}

pub struct AstBuilder {
    params: Params,
    ast: GlobalBlock,
    name_counter: u64,
    typedef_map: HashMap<BaseType, Vec<String>>,
    typedef_vec: Vec<(String, BaseType)>,
    state: StateTracking,
    // Random generation
    rng: rand::rngs::ThreadRng,
    rng_float: FloatGenerator,
    rng_int: IntGenerator,
}

impl AstBuilder {
    pub fn from(params: Params) -> Self {
        let mut typedef_map = HashMap::new();
        let mut typedef_vec = Vec::new();

        typedef_map.insert(BaseType::Int, vec!["integer".to_string()]);
        typedef_map.insert(BaseType::Real, vec!["real".to_string()]);
        typedef_vec.push(("integer".to_string(), BaseType::Int));
        typedef_vec.push(("real".to_string(), BaseType::Real));

        Self {
            ast: GlobalBlock::default(),
            name_counter: 0,
            typedef_map,
            typedef_vec,
            state: StateTracking::default(),
            rng: rand::thread_rng(),
            rng_float: FloatGenerator::new(&params),
            rng_int: IntGenerator::new(&params),
            params,
        }
    }

    pub fn generate(&mut self) {
        loop {
            let mut p: f64 = self.rng.gen();

            if a_lt(&mut p, self.params.global_flow.end_generation) {
                break;
            } else if a_lt(&mut p, self.params.global_flow.gen_typedef) {
                let typedef = self.gen_typedef();
                self.ast.push(typedef);
            } else {
                let decl = self.gen_global_decl();
                self.ast.push(decl);
            }
        }
    }

    fn gen_global_decl(&mut self) -> Stat {
        let mut p: f64 = self.rng.gen();
        let ps = &self.params.statements.gen_global_decl;

        let t: BaseType;
        let assn: Expr;

        if a_lt(&mut p, ps.gen_integer) {
            t = BaseType::Int;
            assn = self.gen_integer();
        } else {
            t = BaseType::Real;
            assn = self.gen_real();
        }

        let v = self.gen_variable_quantified(t, Quantifier::Const);

        Stat::new_decl(
            DeclarationBuilder::default()
                .type_ident(self.get_ident_for_type(v.get_base()))
                .variable(v)
                .assn(Some(assn))
                .build()
                .unwrap()
        )
    }

    fn gen_typedef(&mut self) -> Stat {
        let new_name = self.gen_name();
        let (old_name, basetype) = self.typedef_vec
            .as_slice()
            .choose(&mut self.rng)
            .unwrap()
            .clone();

        self.typedef_map.get_mut(&basetype).unwrap().push(new_name.clone());
        self.typedef_vec.push((new_name.clone(), basetype));

        Stat::new_typedef(
            TypeDefBuilder::default()
                .old_name(old_name)
                .new_name(new_name)
                .basetype(basetype)
                .build()
                .unwrap()
        )
    }

    fn gen_variable_quantified(&mut self, t: BaseType, q: Quantifier) -> Variable {
        VariableBuilder::default()
            .type_(t)
            .name(self.gen_name())
            .quantifier(q)
            .build()
            .unwrap()
    }

    fn gen_integer(&mut self) -> Expr {
        let mut p: f64 = self.rng.gen();
        let mut ps = &self.params.types.gen_integer;

        self.state.int_recursion_depth += 1;

        if a_lt(&mut p, ps.get_instant) || self.state.int_recursion_depth >= ps.max_depth {
            self.state.int_recursion_depth -= 1;
            return Expr::new_literal(self.gen_literal(BaseType::Int));
        }

        let lhs = self.gen_integer();
        let rhs = self.gen_integer();
        self.state.int_recursion_depth -= 1;

        ps = &self.params.types.gen_integer;

        if a_lt(&mut p, ps.expr_add) {
            Expr::new_binary_op(BinaryOperator::add(lhs, rhs))
        } else if a_lt(&mut p, ps.expr_sub) {
            Expr::new_binary_op(BinaryOperator::subtract(lhs, rhs))
        } else if a_lt(&mut p, ps.expr_mul) {
            Expr::new_binary_op(BinaryOperator::multiply(lhs, rhs))
        } else if a_lt(&mut p, ps.expr_div) {
            Expr::new_binary_op(BinaryOperator::divide(lhs, rhs))
        } else {
            Expr::new_literal(self.gen_literal(BaseType::Int))  // TODO: this shouldn't be here
        }
    }

    fn gen_real(&mut self) -> Expr {
        let mut p: f64 = self.rng.gen();
        let mut ps = &self.params.types.gen_real;

        self.state.real_recursion_depth += 1;

        if a_lt(&mut p, ps.get_instant) || self.state.real_recursion_depth >= ps.max_depth {
            self.state.real_recursion_depth -= 1;
            return Expr::new_literal(self.gen_literal(BaseType::Real));
        } else if a_lt(&mut p, ps.gen_integer) {
            self.state.real_recursion_depth -= 1;
            return self.gen_integer();
        }

        let lhs = self.gen_real();
        let rhs = self.gen_real();
        self.state.real_recursion_depth -= 1;

        ps = &self.params.types.gen_real;

        if a_lt(&mut p, ps.expr_add) {
            Expr::new_binary_op(BinaryOperator::add(lhs, rhs))
        } else if a_lt(&mut p, ps.expr_sub) {
            Expr::new_binary_op(BinaryOperator::subtract(lhs, rhs))
        } else if a_lt(&mut p, ps.expr_mul) {
            Expr::new_binary_op(BinaryOperator::multiply(lhs, rhs))
        } else if a_lt(&mut p, ps.expr_div) {
            Expr::new_binary_op(BinaryOperator::divide(lhs, rhs))
        } else {
            Expr::new_literal(self.gen_literal(BaseType::Real))  // TODO: this shouldn't be here
        }
    }

    fn gen_literal(&mut self, t: BaseType) -> Literal {
        match t {
            BaseType::Int => {
                let i = self.rng_int.sample(&mut self.rng);
                Literal::Int(i)
            }
            BaseType::Real => {
                let f = self.rng_float.sample(&mut self.rng);
                Literal::Real(f)
            }
            BaseType::Never => panic!("Attempted to generate literal of type Never"),
            BaseType::Unset => panic!("Attempted to generate literal of type Unset"),
        }
    }

    fn get_ident_for_type(&mut self, t: BaseType) -> String {
        self.typedef_map[&t].choose(&mut self.rng).unwrap().clone()
    }

    fn gen_name(&mut self) -> String {
        let s = format!("_{}", self.name_counter);
        self.name_counter += 1;
        s
    }
}

impl ToString for AstBuilder {
    fn to_string(&self) -> String {
        let mut s = format!("// Generated by {} v{}\n",
            env!("CARGO_PKG_NAME"),
            env!("CARGO_PKG_VERSION"),
        );

        s.push_str(&self.ast.to_string());
        s
    }
}

/// Generates values from regions of interest for an integer
///
/// These are concentrated around i32::MIN, 0, and i32::MAX, though any i32 value is possible
struct IntGenerator {
    distro_pos: rand_distr::Normal<f32>,
    distro_zero: rand_distr::Normal<f32>,
    distro_neg: rand_distr::Normal<f32>,
}

impl IntGenerator {
    fn new(p: &Params) -> Self {
        Self {
            distro_pos: rand_distr::Normal::new(0.0, p.dev.int_gen_distro_pos_stddiv).unwrap(),
            distro_zero: rand_distr::Normal::new(0.0, 3.0).unwrap(),
            distro_neg: rand_distr::Normal::new(0.0, p.dev.int_gen_distro_neg_stddiv).unwrap(),
        }
    }
}

impl Distribution<i32> for IntGenerator {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> i32 {
        let p: f64 = rng.gen();

        if p < 0.3 {
            i32::MAX - (self.distro_pos.sample(rng).abs() as i32)
        } else if p < 0.7 {
            self.distro_zero.sample(rng) as i32
        } else {
            -1 * (i32::MAX - (self.distro_pos.sample(rng).abs() as i32))
        }
    }
}

/// Generates values from regions of interest for a float
///
/// These are concentrated around f32::MIN, 0, and f32::MAX, though any f32 value is possible
struct FloatGenerator {
    distro_pos: rand_distr::Normal<f32>,
    distro_zero: rand_distr::Normal<f32>,
    distro_neg: rand_distr::Normal<f32>,
}

impl FloatGenerator {
    fn new(p: &Params) -> Self {
        Self {
            distro_pos: rand_distr::Normal::new(1.0, p.dev.float_gen_distro_pos_stddiv).unwrap(),
            distro_zero: rand_distr::Normal::new(0.0, 3.0).unwrap(),
            distro_neg: rand_distr::Normal::new(1.0, p.dev.float_gen_distro_neg_stddiv).unwrap(),
        }
    }
}

impl Distribution<f32> for FloatGenerator {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> f32 {
        let p: f64 = rng.gen();

        if p < 0.33 {
            f32::MAX - self.distro_pos.sample(rng).abs()
        } else if p < 0.66 {
            self.distro_zero.sample(rng)
        } else {
            -1.0 * (f32::MAX - self.distro_neg.sample(rng).abs())
        }
    }
}