Improve handling of generics

src/boilerplate.rs

@@ -126,9 +126,10 @@ pub fn init_value_for_type(ctx: &Context, ty: String) -> String {
         return "[]".to_string();
     }
 
-    if ctx.structs.contains_key(&ty) {
+    let typ = ty.split('[').next().unwrap();
+    if ctx.structs.contains_key(typ) {
         // Type is a struct, initialize fields recursively
-        let fields = ctx.structs.get(&ty).unwrap();
+        let fields = ctx.structs.get(typ).unwrap();
         let struct_value = fields
             .iter()
             .map(|field| {
@@ -139,17 +140,17 @@ pub fn init_value_for_type(ctx: &Context, ty: String) -> String {
                 )
             })
             .collect_vec()
-            .join(",");
+            .join(", ");
         return format!("{{ {} }}", struct_value);
     }
 
-    match ty.as_str() {
+    match typ {
         "str" => "\"\"".to_string(),
         "int" => "0".to_string(),
         "Addr" => "\"s1\"".to_string(),
         _ => {
             eprintln!("No init value for type: {ty}");
-            "<missing-type>".to_string()
+            "\"<missing-value>\"".to_string()
         }
     }
 }

src/lib.rs

@@ -161,6 +161,7 @@ fn translate_items(tcx: TyCtxt, crate_name: &str, items: Vec<&rustc_hir::Item>)
         record_fields: vec![],
         struct_fields: vec![],
         pat_fields: vec![],
+        generics: vec![],
         current_item_name: "".to_string(),
     };
 

src/state_extraction.rs

@@ -1,7 +1,9 @@
 use itertools::Itertools;
+use regex::Regex;
 use rustc_span::{symbol::Ident, Symbol};
 
 use crate::{translate::Translatable, types::Context};
+extern crate regex;
 
 // State extraction is the only place where we need to deal with
 // rustc_middle::ty::Ty. This is the type given by typeck, which is the only way
@@ -11,19 +13,30 @@ use crate::{translate::Translatable, types::Context};
 impl Translatable for rustc_middle::ty::Ty<'_> {
     fn translate(&self, ctx: &mut Context) -> String {
         // FIXME: This should be quite unstable, but I couldn't figure out how to navigate `ty` here
+
         let var_name = &format!("{:#?}", self);
-        let name = var_name
-            .split("::")
-            .last()
-            .unwrap()
+        let name_core = var_name
+            .clone()
             .split(' ')
             .last()
             .unwrap()
             .split(')')
             .next()
-            .unwrap();
+            .unwrap()
+            .to_string();
+
+        // The name_core can look like namespace::namespace::type<namespace::namespace::type_arg>
+        // The following regex replacement has the goal of transforming it into type<type_arg>
+        let re = Regex::new(r"(\w+::)*(\w+)(<(\w+::)*(\w+)>)?").unwrap();
+        let name = re.replace(name_core.as_str(), |caps: &regex::Captures| {
+            if let Some(m) = caps.get(4) {
+                format!("{}[{}]", &caps[2], m.as_str())
+            } else {
+                caps[2].to_string()
+            }
+        });
 
-        Ident::with_dummy_span(Symbol::intern(name)).translate(ctx)
+        Ident::with_dummy_span(Symbol::intern(name.to_string().as_str())).translate(ctx)
     }
 }
 impl Translatable for rustc_middle::ty::GenericArg<'_> {

src/translate.rs

@@ -29,13 +29,15 @@ pub fn missing_translation<T: Translatable + Debug>(item: T, descr: &str) -> Str
 }
 
 impl Translatable for Ident {
-    fn translate(&self, _ctx: &mut Context) -> String {
+    fn translate(&self, ctx: &mut Context) -> String {
         const QUINT_KEYWORDS: [&str; 31] = [
             "Set", "List", "Rec", "Tup", "not", "and", "or", "iff", "implies", "all", "any", "if",
             "else", "module", "import", "from", "export", "as", "const", "var", "def", "val",
             "pure", "nondet", "action", "temporal", "assume", "type", "Bool", "Int", "Nat",
         ];
         let name_with_underscore = format!("{}_", self);
+        let lowercase_name = self.as_str().to_lowercase();
+
         let translated_str = match self.as_str() {
             "Vec" => "List",
             "String" => "str",
@@ -49,6 +51,7 @@ impl Translatable for Ident {
             "deps" => "state",
             "_deps" => "state",
             s if QUINT_KEYWORDS.contains(&s) => name_with_underscore.as_str(),
+            s if ctx.generics.contains(&s.to_string()) => lowercase_name.as_str(),
             s => s,
         };
 
@@ -496,6 +499,34 @@ impl Translatable for Function<'_> {
     }
 }
 
+impl Translatable for rustc_hir::GenericParam<'_> {
+    fn translate(&self, ctx: &mut Context) -> String {
+        match self.name {
+            rustc_hir::ParamName::Plain(ident) => {
+                let name = ident.translate(ctx);
+                if name.starts_with('\'') {
+                    // Lifetime, ignore it
+                    return "".to_string();
+                }
+                ctx.generics.push(name.clone());
+                name.to_lowercase()
+            }
+            _ => "".to_string(),
+        }
+    }
+}
+
+impl Translatable for rustc_hir::Generics<'_> {
+    fn translate(&self, ctx: &mut Context) -> String {
+        let type_args = translate_list(self.params, ctx, ", ");
+        if type_args.is_empty() {
+            "".to_string()
+        } else {
+            format!("[{type_args}]")
+        }
+    }
+}
+
 impl Translatable for rustc_hir::Item<'_> {
     fn translate(&self, ctx: &mut Context) -> String {
         let name = self.ident.as_str();
@@ -599,24 +630,31 @@ impl Translatable for rustc_hir::Item<'_> {
   }}"
                 )
             }
-            rustc_hir::ItemKind::Struct(variant_data, _generics) => {
+            rustc_hir::ItemKind::Struct(variant_data, generics) => {
+                let type_args = generics.translate(ctx);
                 let translated_fields = variant_data.translate(ctx);
                 let fields = ctx.struct_fields.clone();
                 ctx.struct_fields.clear();
+                ctx.generics.clear();
 
                 ctx.structs.insert(name.to_string(), fields);
 
-                format!("  type {name} = {translated_fields}")
+                format!("  type {name}{type_args} = {translated_fields}")
             }
 
-            rustc_hir::ItemKind::Enum(enum_def, _generics) => {
-                format!("  type {name} =\n{}", enum_def.translate(ctx))
+            rustc_hir::ItemKind::Enum(enum_def, generics) => {
+                let type_args = generics.translate(ctx);
+                let enum_variants = enum_def.translate(ctx);
+                ctx.generics.clear();
+                format!("  type {name}{type_args} =\n{enum_variants}",)
             }
-            rustc_hir::ItemKind::TyAlias(ty, _generics) => {
+            rustc_hir::ItemKind::TyAlias(ty, generics) => {
+                let type_args = generics.translate(ctx);
                 let translated_type = ty.translate(ctx);
                 ctx.type_aliases
                     .insert(name.to_string(), translated_type.clone());
-                format!("  type {name} = {translated_type}")
+                ctx.generics.clear();
+                format!("  type {name}{type_args} = {translated_type}")
             }
             // FIXME: We should translate this (at least Use) into imports.
             // This is only necessary for crates that import things from other crates

src/types.rs

@@ -46,6 +46,7 @@ pub struct Context<'tcx, 'c> {
     pub record_fields: Vec<String>,
     pub struct_fields: Vec<Field>,
     pub pat_fields: Vec<String>,
+    pub generics: Vec<String>,
     pub current_item_name: String,
 }
 

tests/snapshots/integration_tests__ctf03.snap

@@ -113,8 +113,8 @@ module flash_loan {
   }
 
   pure val init_contract_state = {
-    config: <missing-type>,
-    flash_loan: <missing-type>
+    config: "<missing-value>",
+    flash_loan: "<missing-value>"
   }
 
   action execute_step = all {
@@ -145,7 +145,7 @@ module flash_loan {
     val funds = [{ denom: denom, amount: amount }]
     val info = { sender: sender, funds: funds }
 
-    pure val message: InstantiateMsg = <missing-type>
+    pure val message: InstantiateMsg = "<missing-value>"
     pure val r = instantiate(init_contract_state, { block: { time: 0, height: 1 } }, info, message)
 
     all {
@@ -255,7 +255,7 @@ module mock_arb {
   }
 
   pure val init_contract_state = {
-    config: <missing-type>
+    config: "<missing-value>"
   }
 
   action execute_step = all {
@@ -282,7 +282,7 @@ module mock_arb {
     val funds = [{ denom: denom, amount: amount }]
     val info = { sender: sender, funds: funds }
 
-    pure val message: InstantiateMsg = <missing-type>
+    pure val message: InstantiateMsg = "<missing-value>"
     pure val r = instantiate(init_contract_state, { block: { time: 0, height: 1 } }, info, message)
 
     all {
@@ -390,7 +390,7 @@ module proxy {
   }
 
   pure val init_contract_state = {
-    config: <missing-type>
+    config: "<missing-value>"
   }
 
   action execute_step = all {
@@ -417,7 +417,7 @@ module proxy {
     val funds = [{ denom: denom, amount: amount }]
     val info = { sender: sender, funds: funds }
 
-    pure val message: InstantiateMsg = <missing-type>
+    pure val message: InstantiateMsg = "<missing-value>"
     pure val r = instantiate(init_contract_state, { block: { time: 0, height: 1 } }, info, message)
 
     all {

tests/snapshots/integration_tests__ctf05.snap

@@ -138,7 +138,7 @@ module oaksecurity_cosmwasm_ctf_05 {
   }
 
   pure val init_contract_state = {
-    state: { current_owner: "s1",proposed_owner: <missing-type> },
+    state: { current_owner: "s1", proposed_owner: "<missing-value>" },
     balances: Map()
   }
 

tests/snapshots/integration_tests__ctf06.snap

@@ -110,8 +110,8 @@ module oaksecurity_cosmwasm_ctf_06 {
   }
 
   pure val init_contract_state = {
-    config: { voting_window: 0,voting_token: "s1",owner: "s1" },
-    proposal: { proposer: "s1",timestamp: 0 }
+    config: { voting_window: 0, voting_token: "s1", owner: "s1" },
+    proposal: { proposer: "s1", timestamp: 0 }
   }
 
   action execute_step = all {
@@ -141,7 +141,7 @@ module oaksecurity_cosmwasm_ctf_06 {
     val funds = [{ denom: denom, amount: amount }]
     val info = { sender: sender, funds: funds }
 
-    pure val message: InstantiateMsg = { token: "",owner: "",window: 0 }
+    pure val message: InstantiateMsg = { token: "", owner: "", window: 0 }
     pure val r = instantiate(init_contract_state, { block: { time: 0, height: 1 } }, info, message)
 
     all {

tests/snapshots/integration_tests__ctf07.snap

@@ -142,7 +142,7 @@ module oaksecurity_cosmwasm_ctf_07 {
     val funds = [{ denom: denom, amount: amount }]
     val info = { sender: sender, funds: funds }
 
-    pure val message: InstantiateMsg = { owner: "",threshold: 0 }
+    pure val message: InstantiateMsg = { owner: "", threshold: 0 }
     pure val r = instantiate(init_contract_state, { block: { time: 0, height: 1 } }, info, message)
 
     all {

tests/snapshots/integration_tests__ctf08.snap

@@ -157,7 +157,7 @@ module oaksecurity_cosmwasm_ctf_08 {
     config: { nft_contract: "s1" },
     sales: Map(),
     trades: Map(),
-    operations: { n_trades: 0,n_sales: 0 }
+    operations: { n_trades: 0, n_sales: 0 }
   }
 
   action execute_step = all {

tests/snapshots/integration_tests__ctf09.snap

@@ -109,7 +109,7 @@ module oaksecurity_cosmwasm_ctf_09 {
   }
 
   pure val init_contract_state = {
-    state: { owner: "s1",total_staked: 0,global_index: 0 },
+    state: { owner: "s1", total_staked: 0, global_index: 0 },
     users: Map()
   }
 

tests/snapshots/integration_tests__ctf10.snap

@@ -68,7 +68,7 @@ module oaksecurity_cosmwasm_ctf_10 {
   }
 
   pure val init_contract_state = {
-    config: { nft_contract: "s1",mint_per_user: 0,total_tokens: 0 },
+    config: { nft_contract: "s1", mint_per_user: 0, total_tokens: 0 },
     whitelist: { users: [] }
   }
 
@@ -95,7 +95,7 @@ module oaksecurity_cosmwasm_ctf_10 {
     val funds = [{ denom: denom, amount: amount }]
     val info = { sender: sender, funds: funds }
 
-    pure val message: InstantiateMsg = { cw721_code_id: 0,mint_per_user: 0,whitelisted_users: [] }
+    pure val message: InstantiateMsg = { cw721_code_id: 0, mint_per_user: 0, whitelisted_users: [] }
     pure val r = instantiate(init_contract_state, { block: { time: 0, height: 1 } }, info, message)
 
     all {