nxpu_opt/

cse.rs

//! Common subexpression elimination pass.
//!
//! Walks expression arenas, hashes each expression by opcode and operand
//! indices, and rewrites duplicates to point to the first (canonical)
//! occurrence. Dead duplicates are then cleaned up by a subsequent DCE pass.

use std::collections::HashMap;
use std::hash::{DefaultHasher, Hash, Hasher};

use nxpu_ir::{Expression, Function, Handle, Module};

use crate::Pass;
use crate::dce::expression_operands;

/// Common subexpression elimination pass.
///
/// Detects structurally identical expressions within each function's
/// expression arena and rewrites all references to point to the canonical
/// (first) occurrence. Subsequent DCE removes the dead duplicates.
#[derive(Debug)]
pub struct CommonSubexprElimination;

impl Pass for CommonSubexprElimination {
    fn name(&self) -> &str {
        "cse"
    }

    fn run(&self, module: &mut Module) -> bool {
        let mut changed = false;
        // Run on global expressions.
        changed |= run_on_arena(&mut module.global_expressions);
        // Run on each function.
        for (_, func) in module.functions.iter_mut() {
            changed |= run_on_function(func);
        }
        // Run on entry point functions.
        for ep in &mut module.entry_points {
            changed |= run_on_function(&mut ep.function);
        }
        changed
    }
}

/// Run CSE on a single function's expression arena.
fn run_on_function(func: &mut Function) -> bool {
    run_on_arena(&mut func.expressions)
}

/// Run CSE on an expression arena.
///
/// Returns `true` if any expressions were rewritten.
fn run_on_arena(arena: &mut nxpu_ir::Arena<Expression>) -> bool {
    // Phase 1: compute canonical mapping.
    // Map hash → list of (handle, expression-discriminant) for collision resolution.
    let mut seen: HashMap<u64, Handle<Expression>> = HashMap::new();
    let mut remap: HashMap<Handle<Expression>, Handle<Expression>> = HashMap::new();

    let handles: Vec<Handle<Expression>> = arena.iter().map(|(h, _)| h).collect();

    for &handle in &handles {
        let hash = hash_expression(&arena[handle], &remap);
        if let Some(&canonical) = seen.get(&hash) {
            // Verify structural equality (not just hash match).
            if expressions_equal(&arena[canonical], &arena[handle], &remap) {
                remap.insert(handle, canonical);
                continue;
            }
        }
        seen.insert(hash, handle);
    }

    if remap.is_empty() {
        return false;
    }

    // Phase 2: rewrite operand references using the remap table.
    for &handle in &handles {
        if remap.contains_key(&handle) {
            // This expression is dead (will be cleaned by DCE).
            continue;
        }
        let rewritten = rewrite_operands(&arena[handle], &remap);
        if let Some(new_expr) = rewritten {
            arena[handle] = new_expr;
        }
    }

    true
}

/// Hash an expression by its discriminant and operand indices.
fn hash_expression(
    expr: &Expression,
    remap: &HashMap<Handle<Expression>, Handle<Expression>>,
) -> u64 {
    let mut hasher = DefaultHasher::new();
    std::mem::discriminant(expr).hash(&mut hasher);

    match expr {
        Expression::Literal(lit) => {
            // Hash the literal bytes for value equality.
            format!("{lit:?}").hash(&mut hasher);
        }
        Expression::Binary { op, left, right } => {
            std::mem::discriminant(op).hash(&mut hasher);
            resolve(left, remap).index().hash(&mut hasher);
            resolve(right, remap).index().hash(&mut hasher);
        }
        Expression::Unary { op, expr: inner } => {
            std::mem::discriminant(op).hash(&mut hasher);
            resolve(inner, remap).index().hash(&mut hasher);
        }
        Expression::Math {
            fun,
            arg,
            arg1,
            arg2,
            arg3,
        } => {
            std::mem::discriminant(fun).hash(&mut hasher);
            resolve(arg, remap).index().hash(&mut hasher);
            if let Some(a) = arg1 {
                resolve(a, remap).index().hash(&mut hasher);
            }
            if let Some(a) = arg2 {
                resolve(a, remap).index().hash(&mut hasher);
            }
            if let Some(a) = arg3 {
                resolve(a, remap).index().hash(&mut hasher);
            }
        }
        Expression::Load { pointer } => {
            resolve(pointer, remap).index().hash(&mut hasher);
        }
        Expression::Access { base, index } => {
            resolve(base, remap).index().hash(&mut hasher);
            resolve(index, remap).index().hash(&mut hasher);
        }
        Expression::AccessIndex { base, index } => {
            resolve(base, remap).index().hash(&mut hasher);
            index.hash(&mut hasher);
        }
        Expression::Compose { ty, components } => {
            ty.index().hash(&mut hasher);
            for c in components {
                resolve(c, remap).index().hash(&mut hasher);
            }
        }
        Expression::Select {
            condition,
            accept,
            reject,
        } => {
            resolve(condition, remap).index().hash(&mut hasher);
            resolve(accept, remap).index().hash(&mut hasher);
            resolve(reject, remap).index().hash(&mut hasher);
        }
        Expression::Splat { size, value } => {
            size.hash(&mut hasher);
            resolve(value, remap).index().hash(&mut hasher);
        }
        Expression::Swizzle {
            size,
            vector,
            pattern,
        } => {
            size.hash(&mut hasher);
            resolve(vector, remap).index().hash(&mut hasher);
            pattern.hash(&mut hasher);
        }
        Expression::As {
            expr: inner,
            kind,
            convert,
        } => {
            resolve(inner, remap).index().hash(&mut hasher);
            std::mem::discriminant(kind).hash(&mut hasher);
            convert.hash(&mut hasher);
        }
        Expression::ArrayLength(e) => {
            resolve(e, remap).index().hash(&mut hasher);
        }
        // Non-dedupable expressions: hash by handle index to keep them unique.
        Expression::FunctionArgument(i) => {
            i.hash(&mut hasher);
        }
        Expression::GlobalVariable(h) => {
            h.index().hash(&mut hasher);
        }
        Expression::LocalVariable(h) => {
            h.index().hash(&mut hasher);
        }
        Expression::CallResult(_) | Expression::AtomicResult { .. } | Expression::ZeroValue(_) => {
            // Side-effectful or unique — never merge.
            // Use a unique value to prevent collisions.
            std::ptr::from_ref(expr).hash(&mut hasher);
        }
    }

    hasher.finish()
}

/// Check structural equality of two expressions, resolving handles through the remap.
fn expressions_equal(
    a: &Expression,
    b: &Expression,
    remap: &HashMap<Handle<Expression>, Handle<Expression>>,
) -> bool {
    match (a, b) {
        (Expression::Literal(la), Expression::Literal(lb)) => {
            format!("{la:?}") == format!("{lb:?}")
        }
        (
            Expression::Binary {
                op: op_a,
                left: la,
                right: ra,
            },
            Expression::Binary {
                op: op_b,
                left: lb,
                right: rb,
            },
        ) => {
            op_a == op_b
                && resolve(la, remap) == resolve(lb, remap)
                && resolve(ra, remap) == resolve(rb, remap)
        }
        (Expression::Unary { op: op_a, expr: ea }, Expression::Unary { op: op_b, expr: eb }) => {
            op_a == op_b && resolve(ea, remap) == resolve(eb, remap)
        }
        (
            Expression::Math {
                fun: fa,
                arg: a0,
                arg1: a1,
                arg2: a2,
                arg3: a3,
            },
            Expression::Math {
                fun: fb,
                arg: b0,
                arg1: b1,
                arg2: b2,
                arg3: b3,
            },
        ) => {
            fa == fb
                && resolve(a0, remap) == resolve(b0, remap)
                && resolve_opt(a1.as_ref(), remap) == resolve_opt(b1.as_ref(), remap)
                && resolve_opt(a2.as_ref(), remap) == resolve_opt(b2.as_ref(), remap)
                && resolve_opt(a3.as_ref(), remap) == resolve_opt(b3.as_ref(), remap)
        }
        (Expression::Load { pointer: pa }, Expression::Load { pointer: pb }) => {
            resolve(pa, remap) == resolve(pb, remap)
        }
        (
            Expression::Compose {
                ty: ta,
                components: ca,
            },
            Expression::Compose {
                ty: tb,
                components: cb,
            },
        ) => {
            ta == tb
                && ca.len() == cb.len()
                && ca
                    .iter()
                    .zip(cb.iter())
                    .all(|(x, y)| resolve(x, remap) == resolve(y, remap))
        }
        (
            Expression::Access {
                base: ba,
                index: ia,
            },
            Expression::Access {
                base: bb,
                index: ib,
            },
        ) => resolve(ba, remap) == resolve(bb, remap) && resolve(ia, remap) == resolve(ib, remap),
        (
            Expression::AccessIndex {
                base: ba,
                index: ia,
            },
            Expression::AccessIndex {
                base: bb,
                index: ib,
            },
        ) => resolve(ba, remap) == resolve(bb, remap) && ia == ib,
        (
            Expression::Select {
                condition: ca,
                accept: aa,
                reject: ra,
            },
            Expression::Select {
                condition: cb,
                accept: ab,
                reject: rb,
            },
        ) => {
            resolve(ca, remap) == resolve(cb, remap)
                && resolve(aa, remap) == resolve(ab, remap)
                && resolve(ra, remap) == resolve(rb, remap)
        }
        (
            Expression::Splat {
                size: sa,
                value: va,
            },
            Expression::Splat {
                size: sb,
                value: vb,
            },
        ) => sa == sb && resolve(va, remap) == resolve(vb, remap),
        (
            Expression::Swizzle {
                size: sa,
                vector: va,
                pattern: pa,
            },
            Expression::Swizzle {
                size: sb,
                vector: vb,
                pattern: pb,
            },
        ) => sa == sb && resolve(va, remap) == resolve(vb, remap) && pa == pb,
        (
            Expression::As {
                expr: ea,
                kind: ka,
                convert: ca,
            },
            Expression::As {
                expr: eb,
                kind: kb,
                convert: cb,
            },
        ) => resolve(ea, remap) == resolve(eb, remap) && ka == kb && ca == cb,
        (Expression::ArrayLength(ea), Expression::ArrayLength(eb)) => {
            resolve(ea, remap) == resolve(eb, remap)
        }
        (Expression::FunctionArgument(ia), Expression::FunctionArgument(ib)) => ia == ib,
        (Expression::GlobalVariable(ha), Expression::GlobalVariable(hb)) => ha == hb,
        (Expression::LocalVariable(ha), Expression::LocalVariable(hb)) => ha == hb,
        (Expression::ZeroValue(ta), Expression::ZeroValue(tb)) => ta == tb,
        _ => false,
    }
}

/// Resolve a handle through the remap table.
fn resolve(
    h: &Handle<Expression>,
    remap: &HashMap<Handle<Expression>, Handle<Expression>>,
) -> Handle<Expression> {
    remap.get(h).copied().unwrap_or(*h)
}

fn resolve_opt(
    h: Option<&Handle<Expression>>,
    remap: &HashMap<Handle<Expression>, Handle<Expression>>,
) -> Option<Handle<Expression>> {
    h.map(|handle| resolve(handle, remap))
}

/// Rewrite operand references in an expression using the remap table.
/// Returns `Some(new_expr)` if any operand was remapped, `None` otherwise.
fn rewrite_operands(
    expr: &Expression,
    remap: &HashMap<Handle<Expression>, Handle<Expression>>,
) -> Option<Expression> {
    let operands = expression_operands(expr);
    let any_remapped = operands.iter().any(|h| remap.contains_key(h));
    if !any_remapped {
        return None;
    }

    Some(match expr {
        Expression::Binary { op, left, right } => Expression::Binary {
            op: *op,
            left: resolve(left, remap),
            right: resolve(right, remap),
        },
        Expression::Unary { op, expr: inner } => Expression::Unary {
            op: *op,
            expr: resolve(inner, remap),
        },
        Expression::Math {
            fun,
            arg,
            arg1,
            arg2,
            arg3,
        } => Expression::Math {
            fun: *fun,
            arg: resolve(arg, remap),
            arg1: resolve_opt(arg1.as_ref(), remap),
            arg2: resolve_opt(arg2.as_ref(), remap),
            arg3: resolve_opt(arg3.as_ref(), remap),
        },
        Expression::Load { pointer } => Expression::Load {
            pointer: resolve(pointer, remap),
        },
        Expression::Access { base, index } => Expression::Access {
            base: resolve(base, remap),
            index: resolve(index, remap),
        },
        Expression::AccessIndex { base, index } => Expression::AccessIndex {
            base: resolve(base, remap),
            index: *index,
        },
        Expression::Compose { ty, components } => Expression::Compose {
            ty: *ty,
            components: components.iter().map(|c| resolve(c, remap)).collect(),
        },
        Expression::Select {
            condition,
            accept,
            reject,
        } => Expression::Select {
            condition: resolve(condition, remap),
            accept: resolve(accept, remap),
            reject: resolve(reject, remap),
        },
        Expression::Splat { size, value } => Expression::Splat {
            size: *size,
            value: resolve(value, remap),
        },
        Expression::Swizzle {
            size,
            vector,
            pattern,
        } => Expression::Swizzle {
            size: *size,
            vector: resolve(vector, remap),
            pattern: *pattern,
        },
        Expression::As {
            expr: inner,
            kind,
            convert,
        } => Expression::As {
            expr: resolve(inner, remap),
            kind: *kind,
            convert: *convert,
        },
        Expression::ArrayLength(e) => Expression::ArrayLength(resolve(e, remap)),
        // These don't have expression operands to remap.
        other => other.clone(),
    })
}

#[cfg(test)]
mod tests {
    use super::*;
    use nxpu_ir::{BinaryOp, Literal, Range, Statement};

    #[test]
    fn cse_identical_literals() {
        let mut func = Function::new("test");
        let _a = func
            .expressions
            .append(Expression::Literal(Literal::F32(42.0)));
        let _b = func
            .expressions
            .append(Expression::Literal(Literal::F32(42.0)));
        let changed = run_on_function(&mut func);
        assert!(changed);
    }

    #[test]
    fn cse_identical_binary_ops() {
        let mut func = Function::new("test");
        let a = func
            .expressions
            .append(Expression::Literal(Literal::F32(1.0)));
        let b = func
            .expressions
            .append(Expression::Literal(Literal::F32(2.0)));
        let _add1 = func.expressions.append(Expression::Binary {
            op: BinaryOp::Add,
            left: a,
            right: b,
        });
        let _add2 = func.expressions.append(Expression::Binary {
            op: BinaryOp::Add,
            left: a,
            right: b,
        });
        let changed = run_on_function(&mut func);
        assert!(changed);
    }

    #[test]
    fn cse_different_ops_not_merged() {
        let mut func = Function::new("test");
        let a = func
            .expressions
            .append(Expression::Literal(Literal::F32(1.0)));
        let b = func
            .expressions
            .append(Expression::Literal(Literal::F32(2.0)));
        let _add = func.expressions.append(Expression::Binary {
            op: BinaryOp::Add,
            left: a,
            right: b,
        });
        let _mul = func.expressions.append(Expression::Binary {
            op: BinaryOp::Multiply,
            left: a,
            right: b,
        });
        let changed = run_on_function(&mut func);
        assert!(!changed);
    }

    #[test]
    fn cse_different_operands_not_merged() {
        let mut func = Function::new("test");
        let a = func
            .expressions
            .append(Expression::Literal(Literal::F32(1.0)));
        let b = func
            .expressions
            .append(Expression::Literal(Literal::F32(2.0)));
        let c = func
            .expressions
            .append(Expression::Literal(Literal::F32(3.0)));
        let _add1 = func.expressions.append(Expression::Binary {
            op: BinaryOp::Add,
            left: a,
            right: b,
        });
        let _add2 = func.expressions.append(Expression::Binary {
            op: BinaryOp::Add,
            left: a,
            right: c,
        });
        let changed = run_on_function(&mut func);
        assert!(!changed);
    }

    #[test]
    fn cse_transitive_dedup() {
        // x = a + b, y = a + b (dup of x)
        // z = x * c, w = y * c → after CSE y→x, so w = x * c = z (dup)
        let mut func = Function::new("test");
        let a = func
            .expressions
            .append(Expression::Literal(Literal::F32(1.0)));
        let b = func
            .expressions
            .append(Expression::Literal(Literal::F32(2.0)));
        let c = func
            .expressions
            .append(Expression::Literal(Literal::F32(3.0)));
        let x = func.expressions.append(Expression::Binary {
            op: BinaryOp::Add,
            left: a,
            right: b,
        });
        let y = func.expressions.append(Expression::Binary {
            op: BinaryOp::Add,
            left: a,
            right: b,
        });
        let _z = func.expressions.append(Expression::Binary {
            op: BinaryOp::Multiply,
            left: x,
            right: c,
        });
        let _w = func.expressions.append(Expression::Binary {
            op: BinaryOp::Multiply,
            left: y,
            right: c,
        });
        let changed = run_on_function(&mut func);
        assert!(changed);
    }

    #[test]
    fn cse_preserves_non_duplicates() {
        let mut func = Function::new("test");
        let _a = func
            .expressions
            .append(Expression::Literal(Literal::F32(1.0)));
        let _b = func
            .expressions
            .append(Expression::Literal(Literal::F32(2.0)));
        let _c = func
            .expressions
            .append(Expression::Literal(Literal::F32(3.0)));
        let changed = run_on_function(&mut func);
        assert!(!changed);
    }

    #[test]
    fn cse_empty_module() {
        let mut module = Module::default();
        let pass = CommonSubexprElimination;
        let changed = pass.run(&mut module);
        assert!(!changed);
    }

    #[test]
    fn cse_on_entry_point_functions() {
        use nxpu_ir::EntryPoint;

        let mut module = Module::default();
        let mut func = Function::new("ep");
        let a = func
            .expressions
            .append(Expression::Literal(Literal::F32(7.0)));
        let _b = func
            .expressions
            .append(Expression::Literal(Literal::F32(7.0)));
        func.body
            .push(Statement::Emit(Range::from_index_range(0..2)));
        func.body.push(Statement::Return { value: Some(a) });
        module.entry_points.push(EntryPoint {
            name: "main".into(),
            workgroup_size: [1, 1, 1],
            function: func,
        });
        let pass = CommonSubexprElimination;
        let changed = pass.run(&mut module);
        assert!(changed);
    }

    #[test]
    fn cse_on_global_expressions() {
        let mut module = Module::default();
        module
            .global_expressions
            .append(Expression::Literal(Literal::F32(99.0)));
        module
            .global_expressions
            .append(Expression::Literal(Literal::F32(99.0)));
        let pass = CommonSubexprElimination;
        let changed = pass.run(&mut module);
        assert!(changed);
    }

    #[test]
    fn cse_plus_dce_combined() {
        use crate::DeadCodeElimination;

        let mut module = Module::default();
        let mut func = Function::new("ep");
        let a = func
            .expressions
            .append(Expression::Literal(Literal::F32(1.0)));
        let b = func
            .expressions
            .append(Expression::Literal(Literal::F32(2.0)));
        let add1 = func.expressions.append(Expression::Binary {
            op: BinaryOp::Add,
            left: a,
            right: b,
        });
        let _add2 = func.expressions.append(Expression::Binary {
            op: BinaryOp::Add,
            left: a,
            right: b,
        });

        func.body
            .push(Statement::Emit(Range::from_index_range(0..4)));
        func.body.push(Statement::Return { value: Some(add1) });

        module.entry_points.push(nxpu_ir::EntryPoint {
            name: "main".into(),
            workgroup_size: [1, 1, 1],
            function: func,
        });

        let cse = CommonSubexprElimination;
        let cse_changed = cse.run(&mut module);
        assert!(cse_changed);

        let dce = DeadCodeElimination;
        let dce_changed = dce.run(&mut module);
        // DCE should be able to clean up the dead duplicate.
        // The dead emit range covering the duplicate may be removed.
        let _ = dce_changed;
    }

    #[test]
    fn cse_reduces_expression_count() {
        // Measure that CSE actually reduces the number of unique expressions
        // referenced by the IR (via remap count).
        let mut func = Function::new("test");
        let a = func
            .expressions
            .append(Expression::Literal(Literal::F32(1.0)));
        let b = func
            .expressions
            .append(Expression::Literal(Literal::F32(2.0)));
        // Create 5 identical Add(a,b) expressions.
        for _ in 0..5 {
            func.expressions.append(Expression::Binary {
                op: BinaryOp::Add,
                left: a,
                right: b,
            });
        }
        let expr_count_before = func.expressions.len();
        assert_eq!(expr_count_before, 7); // 2 literals + 5 adds

        let changed = run_on_function(&mut func);
        assert!(changed);

        // After CSE, 4 of the 5 Add expressions are remapped to the canonical one.
        // The arena size doesn't shrink (arena is append-only), but all duplicates
        // now point to the same canonical expression. Verify by counting unique
        // expression hashes after CSE rewrites.
        let mut unique = std::collections::HashSet::new();
        for (_, expr) in func.expressions.iter() {
            unique.insert(format!("{expr:?}"));
        }
        // All 5 Add(a,b) should still be in the arena, but 4 are dead duplicates
        // that subsequent DCE would remove. The key metric is that CSE reported
        // a change, meaning duplicates were detected and remapped.
        assert!(unique.len() <= expr_count_before);
    }
}