From 71ffb5980e4f8462f5624c30d857834d0020cb8c Mon Sep 17 00:00:00 2001
From: Akemi Izuko <akemi@noway.moe>
Date: Thu, 16 Nov 2023 21:47:27 -0700
Subject: [PATCH] AstBuilder: custom int distribution

---
 params.toml        |  2 ++
 src/ast_builder.rs | 40 ++++++++++++++++++++++++++++++++++++----
 src/params.rs      |  2 ++
 3 files changed, 40 insertions(+), 4 deletions(-)

diff --git a/params.toml b/params.toml
index 3a2c575..24c02fc 100644
--- a/params.toml
+++ b/params.toml
@@ -95,3 +95,5 @@
 [dev]
     float_gen_distro_pos_stddiv = 56713724e30  # DO NOT expand with zeros
     float_gen_distro_neg_stddiv = 56713724e30
+    int_gen_distro_pos_stddiv = 214748364.0
+    int_gen_distro_neg_stddiv = 214748364.0
diff --git a/src/ast_builder.rs b/src/ast_builder.rs
index 17db0bf..8174717 100644
--- a/src/ast_builder.rs
+++ b/src/ast_builder.rs
@@ -23,6 +23,7 @@ pub struct AstBuilder {
     name_counter: u64,
     rng: rand::rngs::ThreadRng,
     rng_float: FloatGenerator,
+    rng_int: IntGenerator,
 }
 
 impl AstBuilder {
@@ -32,6 +33,7 @@ impl AstBuilder {
             name_counter: 0,
             rng: rand::thread_rng(),
             rng_float: FloatGenerator::new(&params),
+            rng_int: IntGenerator::new(&params),
             params,
         }
     }
@@ -62,7 +64,6 @@ impl AstBuilder {
             BaseType::Real
         };
 
-
         let v = self.gen_variable_quantified(t, Quantifier::Const);
 
         Box::new(
@@ -105,8 +106,7 @@ impl AstBuilder {
     fn gen_literal(&mut self, t: BaseType) -> Literal {
         match t {
             BaseType::Int => {
-                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;
+                let i = self.rng_int.sample(&mut self.rng);
                 Literal::Int(i)
             }
             BaseType::Real => {
@@ -131,6 +131,38 @@ impl ToString for AstBuilder {
     }
 }
 
+/// 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.33 {
+            i32::MAX - (self.distro_pos.sample(rng).abs() as i32)
+        } else if p < 0.66 {
+            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
@@ -142,7 +174,7 @@ struct FloatGenerator {
 
 impl FloatGenerator {
     fn new(p: &Params) -> Self {
-        FloatGenerator {
+        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(),
diff --git a/src/params.rs b/src/params.rs
index 271b410..54d0d66 100644
--- a/src/params.rs
+++ b/src/params.rs
@@ -154,6 +154,8 @@ toml_struct!{
     struct Dev {
         float_gen_distro_pos_stddiv: f32,
         float_gen_distro_neg_stddiv: f32,
+        int_gen_distro_pos_stddiv: f32,
+        int_gen_distro_neg_stddiv: f32,
     }
 }