Absorb redundant atom-only SSA ternaries

crates/compiler/src/compiler.rs

@@ -450,10 +450,18 @@ impl Compiler {
         // Flattening may produce ternary expressions not in SSA form.
         self.do_pass::<SsaForming>(SsaFormingInput { rename_defs: false })?;
 
-        self.do_pass::<SsaConstPropagation>(())?;
+        self.do_pass::<SsaConstPropagation>(SsaConstPropagationInput::default())?;
 
         self.do_pass::<SsaForming>(SsaFormingInput { rename_defs: false })?;
 
+        // The final SSA pass can introduce local aliases around atom-fielded
+        // composites. Run a late cleanup before CSE/DCE, leaving direct
+        // composite field forwarding to the earlier pass.
+        self.do_pass::<SsaConstPropagation>(SsaConstPropagationInput {
+            forward_direct_composites: false,
+            propagate_constants: false,
+        })?;
+
         self.do_pass::<CommonSubexpressionEliminating>(())?;
 
         self.do_pass::<DeadCodeEliminating>(())?;

crates/passes/src/common_subexpression_elimination/ast.rs

@@ -25,7 +25,9 @@ impl AstReconstructor for CommonSubexpressionEliminatingVisitor<'_> {
     fn reconstruct_expression(&mut self, input: Expression, _additional: &()) -> (Expression, Self::AdditionalOutput) {
         // We simply forward every expression to `try_expr` rather than using the individual reconstruct
         // functions from the `AstReconstructor` trait.
-        (self.try_expr(input, None).expect("CSE Error: Error reconstructing expression").0, Default::default())
+        let mut input = input;
+        self.try_expr(&mut input, None);
+        (input, Default::default())
     }
 
     fn reconstruct_block(&mut self, mut block: Block) -> (Block, Self::AdditionalOutput) {
@@ -38,23 +40,20 @@ impl AstReconstructor for CommonSubexpressionEliminatingVisitor<'_> {
     fn reconstruct_definition(&mut self, mut input: DefinitionStatement) -> (Statement, Self::AdditionalOutput) {
         match input.place {
             DefinitionPlace::Single(place) => {
-                let (value, definition_not_needed) = self
-                    .try_expr(input.value, Some(place.name))
-                    .expect("CSE Error: Error reconstructing single definition");
-
-                if definition_not_needed {
-                    // We don't need this definition - everywhere its variable is referred to, we'll map it to some other
-                    // Path.
-                    (Statement::dummy(), Default::default())
+                if let Some(definition_not_needed) = self.try_expr(&mut input.value, Some(place.name)) {
+                    if definition_not_needed {
+                        // We don't need this definition - everywhere its variable is referred to, we'll map it to some other
+                        // Path.
+                        (Statement::dummy(), Default::default())
+                    } else {
+                        (input.into(), Default::default())
+                    }
                 } else {
-                    input.value = value;
                     (input.into(), Default::default())
                 }
             }
             DefinitionPlace::Multiple(_) => {
-                let (value, _) =
-                    self.try_expr(input.value, None).expect("CSE Error: Error reconstructing multiple definitions");
-                input.value = value;
+                self.try_expr(&mut input.value, None);
                 (input.into(), Default::default())
             }
         }

crates/passes/src/common_subexpression_elimination/visitor.rs

@@ -134,11 +134,15 @@ impl CommonSubexpressionEliminatingVisitor<'_> {
     ///
     /// - `place` If this expression is the right hand side of a definition, `place` is the left hand side,
     ///
-    /// Returns (transformed expression, place_not_needed). `place_not_needed` is true iff it has been mapped to
+    /// Mutates `expression` in place while probing it. `None` means the whole
+    /// expression is not CSE-trackable, but any rewritten atom operands must be
+    /// preserved so deleted aliases are not reintroduced.
+    ///
+    /// Returns `place_not_needed`, which is true iff `place` has been mapped to
     /// another path, and thus its definition is no longer needed.
-    pub fn try_expr(&mut self, mut expression: Expression, place: Option<Symbol>) -> Option<(Expression, bool)> {
+    pub fn try_expr(&mut self, expression: &mut Expression, place: Option<Symbol>) -> Option<bool> {
         let span = expression.span();
-        let expr: Expr = match &mut expression {
+        let expr: Expr = match expression {
             Expression::ArrayAccess(array_access) => {
                 let array = self.try_atom(&mut array_access.array)?;
                 let index = self.try_atom(&mut array_access.index)?;
@@ -202,25 +206,25 @@ impl CommonSubexpressionEliminatingVisitor<'_> {
                 for arg in &mut intrinsic.arguments {
                     self.try_atom(arg)?;
                 }
-                return Some((expression, false));
+                return Some(false);
             }
 
             Expression::Call(call) => {
                 // Don't worry about the const expressions.
                 for arg in &mut call.arguments {
                     self.try_atom(arg)?;
                 }
-                return Some((expression, false));
+                return Some(false);
             }
 
             Expression::Cast(cast) => {
                 self.try_atom(&mut cast.expression)?;
-                return Some((expression, false));
+                return Some(false);
             }
 
             Expression::MemberAccess(member_access) => {
                 self.try_atom(&mut member_access.inner)?;
-                return Some((expression, false));
+                return Some(false);
             }
 
             Expression::Composite(composite_expression) => {
@@ -229,7 +233,7 @@ impl CommonSubexpressionEliminatingVisitor<'_> {
                         self.try_atom(expr)?;
                     }
                 }
-                return Some((expression, false));
+                return Some(false);
             }
 
             Expression::Tuple(tuple_expression) => {
@@ -238,9 +242,12 @@ impl CommonSubexpressionEliminatingVisitor<'_> {
                 tuple_expression.elements = tuple_expression
                     .elements
                     .drain(..)
-                    .map(|expr| self.try_expr(expr, None).map(|x| x.0))
-                    .collect::<Option<Vec<_>>>()?;
-                return Some((expression, false));
+                    .map(|mut expr| {
+                        self.try_expr(&mut expr, None);
+                        expr
+                    })
+                    .collect::<Vec<_>>();
+                return Some(false);
             }
 
             Expression::TupleAccess(_) => panic!("Tuple access expressions should not exist in this pass."),
@@ -258,11 +265,11 @@ impl CommonSubexpressionEliminatingVisitor<'_> {
                     }
                     DynamicOpKind::Read { .. } => {}
                 }
-                return Some((expression, false));
+                return Some(false);
             }
 
             Expression::Async(_) | Expression::Err(_) | Expression::Unit(_) => {
-                return Some((expression, false));
+                return Some(false);
             }
         };
 
@@ -278,9 +285,11 @@ impl CommonSubexpressionEliminatingVisitor<'_> {
                     // We were defining a new variable, whose right hand side is already defined, so map
                     // this variable to the previous variable.
                     self.scopes.last_mut().unwrap().expressions.insert(Atom::Path(vec![place]).into(), name);
-                    return Some((Path::from(identifier).to_local().into(), true));
+                    *expression = Path::from(identifier).to_local().into();
+                    return Some(true);
                 }
-                return Some((Path::from(identifier).to_local().into(), false));
+                *expression = Path::from(identifier).to_local().into();
+                return Some(false);
             }
         }
 
@@ -289,6 +298,6 @@ impl CommonSubexpressionEliminatingVisitor<'_> {
             self.scopes.last_mut().unwrap().expressions.insert(expr, place);
         }
 
-        Some((expression, false))
+        Some(false)
     }
 }

crates/passes/src/peephole_optimization/mod.rs

@@ -22,6 +22,7 @@
 //! - **Identity operation folding**: eliminates `add x 0`, `mul x 1`, `or x false`, etc.
 //! - **Trivial assert elimination**: removes `assert.eq <lit> <lit>` when both sides are equal,
 //!   preserving the dummy assert required for empty closures/finalizes.
+//! - **Ternary absorption**: collapses redundant nested selects over the same condition.
 //! - **Dead register elimination**: removes pure instructions whose destination register is never read.
 //! - **Consecutive cast folding**: merges adjacent casts through the same intermediate register.
 //! - **Register renumbering**: compacts register indices after instruction removal.
@@ -44,6 +45,7 @@ impl Pass for PeepholeOptimizing {
         input.for_each_statement_list(|stmts, num_inputs| {
             fold_identity_operations(stmts);
             eliminate_trivial_asserts(stmts);
+            fold_redundant_ternaries(stmts);
             eliminate_dead_registers(stmts);
             fold_consecutive_casts(stmts);
             renumber_registers(stmts, num_inputs);
@@ -94,16 +96,37 @@ fn is_true(expr: &AleoExpr) -> bool {
     matches!(expr, AleoExpr::Bool(true))
 }
 
-/// Returns true if the expression is a literal (not a register or composite expression).
+/// Returns true if the expression is a generated-Aleo literal.
 fn is_literal(expr: &AleoExpr) -> bool {
-    !matches!(
-        expr,
+    match expr {
+        AleoExpr::Address(_)
+        | AleoExpr::Identifier(_)
+        | AleoExpr::Bool(_)
+        | AleoExpr::Field(_)
+        | AleoExpr::Group(_)
+        | AleoExpr::Signature(_)
+        | AleoExpr::Scalar(_)
+        | AleoExpr::String(_)
+        | AleoExpr::U8(_)
+        | AleoExpr::U16(_)
+        | AleoExpr::U32(_)
+        | AleoExpr::U64(_)
+        | AleoExpr::U128(_)
+        | AleoExpr::I8(_)
+        | AleoExpr::I16(_)
+        | AleoExpr::I32(_)
+        | AleoExpr::I64(_)
+        | AleoExpr::I128(_) => true,
         AleoExpr::Reg(_)
-            | AleoExpr::Tuple(_)
-            | AleoExpr::ArrayAccess(_, _)
-            | AleoExpr::MemberAccess(_, _)
-            | AleoExpr::RawName(_)
-    )
+        | AleoExpr::Tuple(_)
+        | AleoExpr::ArrayAccess(_, _)
+        | AleoExpr::MemberAccess(_, _)
+        | AleoExpr::RawName(_) => false,
+    }
+}
+
+fn is_ternary_absorption_operand(expr: &AleoExpr) -> bool {
+    matches!(expr, AleoExpr::Reg(_)) || is_literal(expr)
 }
 
 // Identity operation folding
@@ -431,6 +454,39 @@ fn eliminate_trivial_asserts(stmts: &mut Vec<AleoStmt>) {
     }
 }
 
+// Ternary absorption
+
+fn fold_redundant_ternaries(stmts: &mut [AleoStmt]) {
+    let mut ternaries: HashMap<AleoReg, (AleoExpr, AleoExpr, AleoExpr)> = HashMap::new();
+
+    for stmt in stmts {
+        if let AleoStmt::Ternary(condition, if_true, if_false, dest) = stmt {
+            if let AleoExpr::Reg(inner) = if_true.clone()
+                && let Some((inner_condition, inner_true, inner_false)) = ternaries.get(&inner)
+                && condition == inner_condition
+                && if_false == inner_false
+            {
+                *if_true = inner_true.clone();
+            }
+
+            if let AleoExpr::Reg(inner) = if_false.clone()
+                && let Some((inner_condition, inner_true, inner_false)) = ternaries.get(&inner)
+                && condition == inner_condition
+                && if_true == inner_true
+            {
+                *if_false = inner_false.clone();
+            }
+
+            if is_ternary_absorption_operand(condition)
+                && is_ternary_absorption_operand(if_true)
+                && is_ternary_absorption_operand(if_false)
+            {
+                ternaries.insert(dest.clone(), (condition.clone(), if_true.clone(), if_false.clone()));
+            }
+        }
+    }
+}
+
 // Dead register elimination
 
 /// Collect all registers that are read (used as operands) in a statement.
@@ -1092,3 +1148,55 @@ fn renumber_registers(stmts: &mut [AleoStmt], num_inputs: usize) {
         }
     }
 }
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    fn reg(index: u64) -> AleoExpr {
+        AleoExpr::Reg(AleoReg::R(index))
+    }
+
+    fn reg_dest(index: u64) -> AleoReg {
+        AleoReg::R(index)
+    }
+
+    #[test]
+    fn redundant_ternary_absorption_uses_atom_operands() {
+        let mut stmts = vec![
+            AleoStmt::Ternary(reg(0), reg(1), reg(2), reg_dest(3)),
+            AleoStmt::Ternary(reg(0), reg(3), reg(2), reg_dest(4)),
+        ];
+
+        fold_redundant_ternaries(&mut stmts);
+
+        assert_eq!(stmts[1], AleoStmt::Ternary(reg(0), reg(1), reg(2), reg_dest(4)));
+    }
+
+    #[test]
+    fn redundant_ternary_absorption_uses_literal_operands() {
+        let one = AleoExpr::U8(1);
+        let two = AleoExpr::U8(2);
+        let mut stmts = vec![
+            AleoStmt::Ternary(reg(0), one.clone(), two.clone(), reg_dest(3)),
+            AleoStmt::Ternary(reg(0), reg(3), two.clone(), reg_dest(4)),
+        ];
+
+        fold_redundant_ternaries(&mut stmts);
+
+        assert_eq!(stmts[1], AleoStmt::Ternary(reg(0), one, two, reg_dest(4)));
+    }
+
+    #[test]
+    fn redundant_ternary_absorption_ignores_non_atom_operands() {
+        let member = AleoExpr::MemberAccess(Box::new(reg(1)), "field".to_string());
+        let mut stmts = vec![
+            AleoStmt::Ternary(reg(0), member, reg(2), reg_dest(3)),
+            AleoStmt::Ternary(reg(0), reg(3), reg(2), reg_dest(4)),
+        ];
+
+        fold_redundant_ternaries(&mut stmts);
+
+        assert_eq!(stmts[1], AleoStmt::Ternary(reg(0), reg(3), reg(2), reg_dest(4)));
+    }
+}