walpurgis/src/ast_builder.rs

use rand::Rng;
use rand_distr::Distribution;

use crate::params::Params;
use crate::ast::{
    Expr,
    Statement,
    Quantifier,
    BaseType,
    GlobalBlock,
    Block,
    Literal,
    Declaration,
    DeclarationBuilder,
    Variable,
    VariableBuilder,
    BinaryOperator,
};

pub struct AstBuilder {
    params: Params,
    ast: GlobalBlock,
    name_counter: u64,
    rng: rand::rngs::ThreadRng,
}

impl AstBuilder {
    pub fn from(params: Params) -> Self {
        Self {
            params,
            rng: rand::thread_rng(),
            name_counter: 0,
            ast: GlobalBlock::default(),
        }
    }

    pub fn generate(&mut self) {

        let mut p: f64;

        loop {
            p = self.rng.gen();

            if p < self.params.global_flow.end_generation {
                break;
            } else {
                let decl = self.gen_decl();
                self.ast.push(decl);
            }
        }
    }

    fn gen_decl(&mut self) -> Box<dyn Statement> {
        let t = BaseType::Int;
        let v = self.gen_variable_quantified(t, Quantifier::Const);

        Box::new(
            DeclarationBuilder::default()
                .variable(v)
                .assn(self.gen_literal(t))
                .build()
                .unwrap()
        )
    }

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

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


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

    //fn gen_expr(&mut self, type_: BaseType) -> Box<dyn Expr> {
    //}

    fn gen_literal(&mut self, t: BaseType) -> Literal {
        match t {
            BaseType::Int => {
                let r = rand_distr::Beta::new(0.5, 0.5).unwrap();
                let i: i32 = (r.sample(&mut self.rng) * (i32::MAX as f64)) as i32;
                Literal::Int(i)
            }
            BaseType::Real => Literal::Real(1.0),
            BaseType::Never => panic!("Attempted to generate literal of type Never"),
            BaseType::Unset => panic!("Attempted to generate literal of type Unset"),
        }
    }
}

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
    }
}