From dfabbaf6f54f14aba14f41512a6c28c38c42a870 Mon Sep 17 00:00:00 2001
From: Akemi Izuko <akemi@noway.moe>
Date: Thu, 16 Nov 2023 21:24:05 -0700
Subject: [PATCH] AstBuilder: custom float distribution

---
 src/ast.rs         |  1 -
 src/ast_builder.rs | 50 ++++++++++++++++++++++++++++++++++++++++------
 2 files changed, 44 insertions(+), 7 deletions(-)

diff --git a/src/ast.rs b/src/ast.rs
index 5f45bc8..87bab55 100644
--- a/src/ast.rs
+++ b/src/ast.rs
@@ -74,7 +74,6 @@ impl ToString for GlobalBlock {
 
         for stat in &self.statements {
             s.push_str(&stat.to_string());
-            s.push(' ');
         }
 
         s.push('\n');
diff --git a/src/ast_builder.rs b/src/ast_builder.rs
index aea0ac4..b017fec 100644
--- a/src/ast_builder.rs
+++ b/src/ast_builder.rs
@@ -22,15 +22,17 @@ pub struct AstBuilder {
     ast: GlobalBlock,
     name_counter: u64,
     rng: rand::rngs::ThreadRng,
+    rng_float: FloatGenerator,
 }
 
 impl AstBuilder {
     pub fn from(params: Params) -> Self {
         Self {
-            params,
-            rng: rand::thread_rng(),
-            name_counter: 0,
             ast: GlobalBlock::default(),
+            name_counter: 0,
+            rng: rand::thread_rng(),
+            rng_float: FloatGenerator::new(&params),
+            params,
         }
     }
 
@@ -51,7 +53,7 @@ impl AstBuilder {
     }
 
     fn gen_decl(&mut self) -> Box<dyn Statement> {
-        let t = BaseType::Int;
+        let t = BaseType::Real;
         let v = self.gen_variable_quantified(t, Quantifier::Const);
 
         Box::new(
@@ -94,11 +96,14 @@ impl AstBuilder {
     fn gen_literal(&mut self, t: BaseType) -> Literal {
         match t {
             BaseType::Int => {
-                let r = rand_distr::Beta::new(0.5, 0.5).unwrap();
+                let r = rand_distr::Beta::new(0.01, 0.01).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::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"),
         }
@@ -116,3 +121,36 @@ impl ToString for AstBuilder {
         s
     }
 }
+
+/// 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 {
+        FloatGenerator {
+            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())
+        }
+    }
+}