From fd6bef4cce9d290c10b3999bf77d34d5a45763f1 Mon Sep 17 00:00:00 2001
From: Joey Kraut <108701651+joeykraut@users.noreply.github.com>
Date: Sat, 25 Jan 2025 19:41:34 -0800
Subject: [PATCH] codegen: poseidon-codegen: Define codegen utils (#7)

---
 .gitignore                           |  1 +
 codegen/poseidon-codegen/Cargo.toml  | 13 ++++
 codegen/poseidon-codegen/src/main.rs | 98 ++++++++++++++++++++++++++++
 3 files changed, 112 insertions(+)
 create mode 100644 codegen/poseidon-codegen/Cargo.toml
 create mode 100644 codegen/poseidon-codegen/src/main.rs

diff --git a/.gitignore b/.gitignore
index 478b9a7..ec67977 100644
--- a/.gitignore
+++ b/.gitignore
@@ -4,6 +4,7 @@
 # Compiler files
 cache/
 out/
+**/target
 
 # Ignores development broadcast logs
 !/broadcast
diff --git a/codegen/poseidon-codegen/Cargo.toml b/codegen/poseidon-codegen/Cargo.toml
new file mode 100644
index 0000000..f0ac933
--- /dev/null
+++ b/codegen/poseidon-codegen/Cargo.toml
@@ -0,0 +1,13 @@
+[package]
+name = "poseidon-codegen"
+version = "0.1.0"
+edition = "2021"
+
+[dependencies]
+renegade-crypto = { git = "https://github.com/renegade-fi/renegade.git" }
+renegade-constants = { package = "constants", git = "https://github.com/renegade-fi/renegade.git", default-features = false }
+
+anyhow = "1.0"
+clap = { version = "4.5.1", features = ["derive"] }
+itertools = "0.14"
+num-bigint = "0.4"
diff --git a/codegen/poseidon-codegen/src/main.rs b/codegen/poseidon-codegen/src/main.rs
new file mode 100644
index 0000000..af21c16
--- /dev/null
+++ b/codegen/poseidon-codegen/src/main.rs
@@ -0,0 +1,98 @@
+use anyhow::Result;
+use clap::Parser;
+use itertools::Itertools;
+use num_bigint::BigUint;
+use renegade_constants::Scalar;
+use renegade_crypto::fields::scalar_to_biguint;
+use renegade_crypto::hash::{FULL_ROUND_CONSTANTS, PARTIAL_ROUND_CONSTANTS, R_F, R_P};
+use std::fs;
+use std::path::PathBuf;
+
+#[derive(Parser, Debug)]
+#[command(author, version, about, long_about = None)]
+struct Args {
+    /// Output file path for the generated Huff code
+    #[arg(short, long)]
+    output: PathBuf,
+}
+
+// --- Code Generation --- //
+
+/// Generates the Huff source code as a string
+///
+/// Maintains the invariant that the state is always the first three words of the stack
+fn generate_huff_code() -> String {
+    let mut code = String::new();
+    const HALF_FULL: usize = R_F / 2;
+    let partial_rounds = R_P;
+
+    for i in 0..HALF_FULL {
+        code.push_str(&generate_full_round(i));
+    }
+
+    for i in 0..partial_rounds {
+        code.push_str(&generate_partial_round(i));
+    }
+
+    for i in HALF_FULL..R_F {
+        code.push_str(&generate_full_round(i));
+    }
+
+    code
+}
+
+/// Generate a full round of the hash function
+fn generate_full_round(i: usize) -> String {
+    let mut code = String::new();
+    let round_constants = FULL_ROUND_CONSTANTS[i];
+    let round_constants_str = round_constants
+        .iter()
+        .map(|fp| Scalar::new(*fp))
+        .map(scalar_to_hex)
+        .collect_vec();
+
+    code.push_str(&format!("// --- Full Round {i} --- //\n"));
+    code.push_str(&format!(
+        "EXTERNAL_ROUND({}, {}, {})\n\n",
+        round_constants_str[0], round_constants_str[1], round_constants_str[2]
+    ));
+    code
+}
+
+/// Generate a partial round of the hash function
+fn generate_partial_round(i: usize) -> String {
+    let mut code = String::new();
+
+    let round_constant = PARTIAL_ROUND_CONSTANTS[i];
+    let round_constant_str = scalar_to_hex(Scalar::new(round_constant));
+
+    code.push_str(&format!("// --- Partial Round {i} --- //\n"));
+    code.push_str(&format!("INTERNAL_ROUND({round_constant_str})\n\n"));
+    code
+}
+
+// --- Helpers --- //
+
+/// Convert a `Scalar` element to a big-endian hex string
+fn scalar_to_hex(scalar: Scalar) -> String {
+    let biguint = scalar_to_biguint(&scalar);
+    format!("{biguint:#x}")
+}
+
+// --- Main --- //
+
+fn main() -> Result<()> {
+    let args = Args::parse();
+
+    // Generate the Huff code
+    let huff_code = generate_huff_code();
+
+    // Write to the output file
+    fs::write(&args.output, huff_code)?;
+    println!(
+        "Successfully generated Huff code at: {}",
+        args.output.display()
+    );
+
+    Ok(())
+}