From b1116c4a1a88f68f3a5612f80c38fbfb7ba6f558 Mon Sep 17 00:00:00 2001
From: Jacob Lifshay <programmerjake@gmail.com>
Date: Mon, 1 Jun 2026 21:39:26 -0700
Subject: [PATCH] simplify_enums: cache folded expressions

---
 .../src/module/transform/simplify_enums.rs    | 37 ++++++++++++++-----
 1 file changed, 27 insertions(+), 10 deletions(-)

diff --git a/crates/fayalite/src/module/transform/simplify_enums.rs b/crates/fayalite/src/module/transform/simplify_enums.rs
index 5f136f2..c52939c 100644
--- a/crates/fayalite/src/module/transform/simplify_enums.rs
+++ b/crates/fayalite/src/module/transform/simplify_enums.rs
@@ -96,6 +96,7 @@ enum EnumTypeState {
 
 struct ModuleState {
     module_name: NameId,
+    expr_cache: HashMap<ExprEnum, ExprEnum>,
 }
 
 impl ModuleState {
@@ -675,6 +676,7 @@ impl Folder for State {
     fn fold_module<T: BundleType>(&mut self, v: Module<T>) -> Result<Module<T>, Self::Error> {
         self.module_state_stack.push(ModuleState {
             module_name: v.name_id(),
+            expr_cache: HashMap::default(),
         });
         let retval = Fold::default_fold(v, self);
         self.module_state_stack.pop();
@@ -682,30 +684,39 @@ impl Folder for State {
     }
 
     fn fold_expr_enum(&mut self, op: ExprEnum) -> Result<ExprEnum, Self::Error> {
-        match op {
+        if let Some(folded_op) = self
+            .module_state_stack
+            .last()
+            .expect("known to be in module")
+            .expr_cache
+            .get(&op)
+        {
+            return Ok(*folded_op);
+        }
+        let folded_op = match op {
             ExprEnum::EnumLiteral(op) => {
                 let folded_variant_value = op.variant_value().map(|v| v.fold(self)).transpose()?;
-                Ok(*Expr::expr_enum(self.handle_enum_literal(
+                *Expr::expr_enum(self.handle_enum_literal(
                     op.ty(),
                     op.variant_index(),
                     folded_variant_value,
-                )?))
+                )?)
             }
             ExprEnum::VariantAccess(op) => {
                 let folded_base_expr = Expr::canonical(op.base()).fold(self)?;
-                Ok(*Expr::expr_enum(self.handle_variant_access(
+                *Expr::expr_enum(self.handle_variant_access(
                     op.base().ty(),
                     folded_base_expr,
                     op.variant_index(),
-                )?))
+                )?)
             }
-            ExprEnum::MemPort(mem_port) => Ok(
+            ExprEnum::MemPort(mem_port) => {
                 if let Some(&wire) = self.replacement_mem_ports.get(&mem_port) {
                     ExprEnum::Wire(wire)
                 } else {
                     ExprEnum::MemPort(mem_port.fold(self)?)
-                },
-            ),
+                }
+            }
             ExprEnum::UIntLiteral(_)
             | ExprEnum::SIntLiteral(_)
             | ExprEnum::BoolLiteral(_)
@@ -813,8 +824,14 @@ impl Folder for State {
             | ExprEnum::Wire(_)
             | ExprEnum::Reg(_)
             | ExprEnum::RegSync(_)
-            | ExprEnum::RegAsync(_) => op.default_fold(self),
-        }
+            | ExprEnum::RegAsync(_) => op.default_fold(self)?,
+        };
+        self.module_state_stack
+            .last_mut()
+            .expect("known to be in module")
+            .expr_cache
+            .insert(op, folded_op);
+        Ok(folded_op)
     }
 
     fn fold_block(&mut self, block: Block) -> Result<Block, Self::Error> {