simplify_enums: cache folded expressions

This commit is contained in:
Jacob Lifshay 2026-06-01 21:39:26 -07:00
parent 6902aea3a6
commit b1116c4a1a
Signed by: programmerjake
SSH key fingerprint: SHA256:HnFTLGpSm4Q4Fj502oCFisjZSoakwEuTsJJMSke63RQ

View file

@ -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> {