// Copyright (c) 2026 Alexander Medvednikov. All rights reserved.
// Use of this source code is governed by an MIT license
// that can be found in the LICENSE file.

module optimize

import v2.ssa

fn dead_code_elimination(mut m ssa.Module) bool {
    mut any_changed := false
    mut changed := true
    for changed {
        changed = false
        for fi in 0 .. m.funcs.len {
            // Dead store elimination: find stores to local allocas that are never read
            dead_stores := find_dead_stores(m, &m.funcs[fi])

            for blk_id in m.funcs[fi].blocks {
                // First pass: find dead instructions (don't modify yet)
                mut dead_instrs := []int{}
                for val_id in m.blocks[blk_id].instrs {
                    val := m.values[val_id]
                    if val.kind == .instruction {
                        instr := m.instrs[val.index]
                        // Check if this is a dead store
                        if instr.op == .store && val_id in dead_stores {
                            dead_instrs << val_id
                            continue
                        }
                        side_effects := instr.op in [.store, .call, .call_indirect, .call_sret,
                            .ret, .br, .jmp, .switch_, .unreachable, .assign, .fence, .cmpxchg,
                            .atomicrmw, .go_call, .spawn_call]
                        if !side_effects && val.uses.len == 0 {
                            dead_instrs << val_id
                        }
                    }
                }

                // Only rebuild the instruction list if we found dead instructions
                if dead_instrs.len > 0 {
                    // Remove uses from dead instructions
                    for val_id in dead_instrs {
                        instr := m.instrs[m.values[val_id].index]
                        for op_id in instr.operands {
                            remove_use(mut m, op_id, val_id)
                        }
                    }

                    // Build set for O(1) lookup
                    mut dead_set := map[int]bool{}
                    for d in dead_instrs {
                        dead_set[d] = true
                    }

                    // Filter out dead instructions
                    mut new_instrs := []int{cap: m.blocks[blk_id].instrs.len}
                    for val_id in m.blocks[blk_id].instrs {
                        if !dead_set[val_id] {
                            new_instrs << val_id
                        }
                    }
                    // Avoid m.blocks[X].instrs = ... -- chained field assign broken in ARM64 self-hosted
                    mut dce_blk := m.blocks[blk_id]
                    dce_blk.instrs = new_instrs
                    m.blocks[blk_id] = dce_blk
                    changed = true
                    any_changed = true
                }
            }
        }
    }
    return any_changed
}

fn remove_use(mut m ssa.Module, val_id int, user_id int) {
    if val_id >= m.values.len {
        return
    }
    mut val := &m.values[val_id]
    // Remove all occurrences (handles instructions that use same value multiple times)
    // Iterate in reverse to safely delete while iterating
    for i := val.uses.len - 1; i >= 0; i-- {
        if val.uses[i] == user_id {
            val.uses.delete(i)
        }
    }
}

// Find stores to local allocas that are never read (dead stores).
// A store is dead if:
// 1. It stores to a local alloca
// 2. The alloca is ONLY used by stores (no loads, GEPs, calls, or other uses)
//
// We must be conservative: if the alloca pointer escapes (used by GEP, passed to
// a function, etc.), we cannot eliminate stores because the value might be read
// through the escaped pointer.
fn find_dead_stores(m &ssa.Module, func &ssa.Function) map[int]bool {
    mut dead_stores := map[int]bool{}

    // Find all local allocas
    mut allocas := map[int]bool{}
    mut alloca_stores := map[int][]int{} // alloca_id -> list of store val_ids
    mut alloca_escapes := map[int]bool{} // alloca_id -> true if pointer escapes

    for blk_id in func.blocks {
        for val_id in m.blocks[blk_id].instrs {
            val := m.values[val_id]
            if val.kind != .instruction {
                continue
            }
            instr := m.instrs[val.index]

            if instr.op == .alloca || instr.op == .heap_alloc {
                allocas[val_id] = true
            }
        }
    }

    // Check how each alloca is used
    for alloca_id, _ in allocas {
        alloca_val := m.values[alloca_id]
        for use_id in alloca_val.uses {
            use_val := m.values[use_id]
            if use_val.kind != .instruction {
                // Non-instruction use - alloca escapes
                alloca_escapes[alloca_id] = true
                continue
            }
            use_instr := m.instrs[use_val.index]
            match use_instr.op {
                .store {
                    // Check if alloca is the destination (operand[1]), not the value being stored
                    if use_instr.operands.len > 1 && use_instr.operands[1] == alloca_id {
                        alloca_stores[alloca_id] << use_id
                    } else {
                        // Alloca is being stored somewhere - it escapes
                        alloca_escapes[alloca_id] = true
                    }
                }
                .load {
                    // Alloca is being read - stores are not dead
                    alloca_escapes[alloca_id] = true
                }
                else {
                    // Any other use (GEP, call, etc.) - alloca escapes
                    alloca_escapes[alloca_id] = true
                }
            }
        }
    }

    // Mark stores to non-escaping allocas as dead
    for alloca_id, store_ids in alloca_stores {
        if !alloca_escapes[alloca_id] {
            for i := 0; i < store_ids.len; i++ {
                store_id := store_ids[i]
                dead_stores[store_id] = true
            }
        }
    }

    return dead_stores
}