From e0535acc917dc2c4f9dfd89f16bcbcb65f6ad3fb Mon Sep 17 00:00:00 2001
From: Akemi Izuko <akemi@noway.moe>
Date: Fri, 17 Nov 2023 15:46:13 -0700
Subject: [PATCH] AstBuilder: use typedef in decls

---
 src/ast/mod.rs       |  2 +-
 src/ast/statement.rs |  3 ++-
 src/ast_builder.rs   | 13 +++++++++----
 3 files changed, 12 insertions(+), 6 deletions(-)

diff --git a/src/ast/mod.rs b/src/ast/mod.rs
index 6912899..6d898dd 100644
--- a/src/ast/mod.rs
+++ b/src/ast/mod.rs
@@ -44,7 +44,7 @@ impl ToString for Quantifier {
     }
 }
 
-#[derive(Copy, Clone, Eq)]
+#[derive(Copy, Clone, Hash, Eq)]
 pub enum BaseType {
     Int,
     Real,
diff --git a/src/ast/statement.rs b/src/ast/statement.rs
index 2ce8a9e..b0bdb47 100644
--- a/src/ast/statement.rs
+++ b/src/ast/statement.rs
@@ -84,6 +84,7 @@ impl ToString for Block {
 
 #[derive(Default, Builder)]
 pub struct Declaration {
+    type_ident: String,
     variable: Variable,
     assn: Option<Expr>,
 }
@@ -94,7 +95,7 @@ impl ToString for Declaration {
 
         s.push_str(&self.variable.get_quantifier().to_string());
         s.push(' ');
-        s.push_str(&self.variable.get_base().to_string());
+        s.push_str(&self.type_ident);
         s.push(' ');
         s.push_str(&self.variable.to_string());
 
diff --git a/src/ast_builder.rs b/src/ast_builder.rs
index bd683bb..eed232e 100644
--- a/src/ast_builder.rs
+++ b/src/ast_builder.rs
@@ -10,7 +10,7 @@ pub struct AstBuilder {
     params: Params,
     ast: GlobalBlock,
     name_counter: u64,
-    typedef_map: HashMap<String, BaseType>,
+    typedef_map: HashMap<BaseType, Vec<String>>,
     typedef_vec: Vec<(String, BaseType)>,
     // Random generation
     rng: rand::rngs::ThreadRng,
@@ -23,8 +23,8 @@ impl AstBuilder {
         let mut typedef_map = HashMap::new();
         let mut typedef_vec = Vec::new();
 
-        typedef_map.insert("integer".to_string(), BaseType::Int);
-        typedef_map.insert("real".to_string(), BaseType::Real);
+        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));
 
@@ -80,6 +80,7 @@ impl AstBuilder {
 
         Stat::new_decl(
             DeclarationBuilder::default()
+                .type_ident(self.get_ident_for_type(v.get_base()))
                 .variable(v)
                 .assn(Some(assn))
                 .build()
@@ -95,7 +96,7 @@ impl AstBuilder {
             .unwrap()
             .clone();
 
-        self.typedef_map.insert(new_name.clone(), basetype);
+        self.typedef_map.get_mut(&basetype).unwrap().push(new_name.clone());
         self.typedef_vec.push((new_name.clone(), basetype));
 
         Stat::new_typedef(
@@ -140,6 +141,10 @@ impl AstBuilder {
         }
     }
 
+    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;