From 10ee9faba12366067767d61834a6aed77214ea0a Mon Sep 17 00:00:00 2001
From: Aumetra Weisman <aumetra@cryptolab.net>
Date: Tue, 8 Oct 2024 15:28:57 +0200
Subject: [PATCH] Add enum codegen

---
 Cargo.lock                                    |  1 +
 packages/cw-schema-codegen/Cargo.toml         |  3 ++
 packages/cw-schema-codegen/src/lib.rs         |  1 +
 packages/cw-schema-codegen/src/rust.rs        | 13 +++++
 .../templates/rust/enum.tpl.rs                |  5 ++
 packages/cw-schema-codegen/tests/rust_tpl.rs  | 53 +++++++++++++++++++
 .../snapshots/rust_tpl__complex_enum.snap     | 11 ++++
 .../tests/snapshots/rust_tpl__empty_enum.snap |  7 +++
 .../snapshots/rust_tpl__simple_enum.snap      | 11 ++++
 9 files changed, 105 insertions(+)
 create mode 100644 packages/cw-schema-codegen/src/lib.rs
 create mode 100644 packages/cw-schema-codegen/src/rust.rs
 create mode 100644 packages/cw-schema-codegen/tests/rust_tpl.rs
 create mode 100644 packages/cw-schema-codegen/tests/snapshots/rust_tpl__complex_enum.snap
 create mode 100644 packages/cw-schema-codegen/tests/snapshots/rust_tpl__empty_enum.snap
 create mode 100644 packages/cw-schema-codegen/tests/snapshots/rust_tpl__simple_enum.snap

diff --git a/Cargo.lock b/Cargo.lock
index 5ca227bbcf..c8c9a4112a 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -945,6 +945,7 @@ dependencies = [
  "cw-schema",
  "either",
  "frunk",
+ "insta",
  "log",
  "serde_json",
  "simple_logger",
diff --git a/packages/cw-schema-codegen/Cargo.toml b/packages/cw-schema-codegen/Cargo.toml
index 9d0daaaef8..4e70ae3826 100644
--- a/packages/cw-schema-codegen/Cargo.toml
+++ b/packages/cw-schema-codegen/Cargo.toml
@@ -16,3 +16,6 @@ log = "0.4.22"
 serde_json = "1.0.128"
 simple_logger = "5.0.0"
 strum = { version = "0.26.3", default-features = false, features = ["derive"] }
+
+[dev-dependencies]
+insta = "1.40.0"
diff --git a/packages/cw-schema-codegen/src/lib.rs b/packages/cw-schema-codegen/src/lib.rs
new file mode 100644
index 0000000000..0ad9e7d3d7
--- /dev/null
+++ b/packages/cw-schema-codegen/src/lib.rs
@@ -0,0 +1 @@
+pub mod rust;
diff --git a/packages/cw-schema-codegen/src/rust.rs b/packages/cw-schema-codegen/src/rust.rs
new file mode 100644
index 0000000000..3939704023
--- /dev/null
+++ b/packages/cw-schema-codegen/src/rust.rs
@@ -0,0 +1,13 @@
+use askama::Template;
+
+pub struct EnumVariantTemplate<'a> {
+    pub name: &'a str,
+    pub types: Option<&'a [&'a str]>,
+}
+
+#[derive(Template)]
+#[template(escape = "none", path = "rust/enum.tpl.rs")]
+pub struct EnumTemplate<'a> {
+    pub name: &'a str,
+    pub variants: &'a [EnumVariantTemplate<'a>],
+}
diff --git a/packages/cw-schema-codegen/templates/rust/enum.tpl.rs b/packages/cw-schema-codegen/templates/rust/enum.tpl.rs
index e69de29bb2..f140c84b3a 100644
--- a/packages/cw-schema-codegen/templates/rust/enum.tpl.rs
+++ b/packages/cw-schema-codegen/templates/rust/enum.tpl.rs
@@ -0,0 +1,5 @@
+pub enum {{ name }} {
+    {% for variant in variants %}
+        {{ variant.name }} {% if let Some(types) = variant.types %}({% for ty in types %} {{ ty }}, {% endfor %}) {% endif %},
+    {% endfor %}
+}
\ No newline at end of file
diff --git a/packages/cw-schema-codegen/tests/rust_tpl.rs b/packages/cw-schema-codegen/tests/rust_tpl.rs
new file mode 100644
index 0000000000..74062015b2
--- /dev/null
+++ b/packages/cw-schema-codegen/tests/rust_tpl.rs
@@ -0,0 +1,53 @@
+use askama::Template;
+use cw_schema_codegen::rust::{EnumTemplate, EnumVariantTemplate};
+
+#[test]
+fn simple_enum() {
+    let tpl = EnumTemplate {
+        name: "Simple",
+        variants: &[
+            EnumVariantTemplate {
+                name: "One",
+                types: None,
+            },
+            EnumVariantTemplate {
+                name: "Two",
+                types: None,
+            },
+        ],
+    };
+
+    let rendered = tpl.render().unwrap();
+    insta::assert_snapshot!(rendered);
+}
+
+#[test]
+fn complex_enum() {
+    let tpl = EnumTemplate {
+        name: "Complex",
+        variants: &[
+            EnumVariantTemplate {
+                name: "One",
+                types: Some(&["u64"]),
+            },
+            EnumVariantTemplate {
+                name: "Two",
+                types: Some(&["String", "u64"]),
+            },
+        ],
+    };
+
+    let rendered = tpl.render().unwrap();
+    insta::assert_snapshot!(rendered);
+}
+
+#[test]
+fn empty_enum() {
+    let tpl = EnumTemplate {
+        name: "Empty",
+        variants: &[],
+    };
+
+    let rendered = tpl.render().unwrap();
+    insta::assert_snapshot!(rendered);
+}
diff --git a/packages/cw-schema-codegen/tests/snapshots/rust_tpl__complex_enum.snap b/packages/cw-schema-codegen/tests/snapshots/rust_tpl__complex_enum.snap
new file mode 100644
index 0000000000..765052ffc2
--- /dev/null
+++ b/packages/cw-schema-codegen/tests/snapshots/rust_tpl__complex_enum.snap
@@ -0,0 +1,11 @@
+---
+source: packages/cw-schema-codegen/tests/rust_tpl.rs
+expression: rendered
+---
+pub enum Complex {
+    
+        One ( u64, ) ,
+    
+        Two ( String,  u64, ) ,
+    
+}
diff --git a/packages/cw-schema-codegen/tests/snapshots/rust_tpl__empty_enum.snap b/packages/cw-schema-codegen/tests/snapshots/rust_tpl__empty_enum.snap
new file mode 100644
index 0000000000..7819505dfd
--- /dev/null
+++ b/packages/cw-schema-codegen/tests/snapshots/rust_tpl__empty_enum.snap
@@ -0,0 +1,7 @@
+---
+source: packages/cw-schema-codegen/tests/rust_tpl.rs
+expression: rendered
+---
+pub enum Empty {
+    
+}
diff --git a/packages/cw-schema-codegen/tests/snapshots/rust_tpl__simple_enum.snap b/packages/cw-schema-codegen/tests/snapshots/rust_tpl__simple_enum.snap
new file mode 100644
index 0000000000..bd30912181
--- /dev/null
+++ b/packages/cw-schema-codegen/tests/snapshots/rust_tpl__simple_enum.snap
@@ -0,0 +1,11 @@
+---
+source: packages/cw-schema-codegen/tests/rust_tpl.rs
+expression: rendered
+---
+pub enum Simple {
+    
+        One ,
+    
+        Two ,
+    
+}