module ssa

// Parallel-SSA worker support, ported from v2's module.v and adapted to v3's
// raw-block-id operand model: branch/phi/switch operands that name blocks are
// remapped by the block offset, while ordinary value operands are remapped by
// the value offset.

// FuncSSAData holds the SSA a worker produced for a single function.
pub struct FuncSSAData {
pub:
    func_idx int       // index into the main module's funcs[]
    blocks   []BlockID // worker-local block ids
    params   []ValueID // worker-local param value ids
}

// is_block_operand reports whether operand `idx` of this instruction names a
// basic block (a raw block id) rather than an SSA value.
pub fn (i &Instruction) is_block_operand(idx int) bool {
    return match i.op {
        .jmp { idx == 0 }
        .br { idx == 1 || idx == 2 }
        // switch_: cond, default_blk, [case_val, blk]... -> blocks at odd indices.
        .switch_ { idx % 2 == 1 }
        // phi: [val, blk]... -> blocks at odd indices.
        .phi { idx % 2 == 1 }
        else { false }
    }
}

// new_worker_module creates a lightweight Module for parallel SSA building. The
// worker gets a deep copy of the type store and is seeded with the main module's
// values/instrs/blocks so ids produced in earlier (serial) phases stay valid.
pub fn (mut m Module) new_worker_module() &Module {
    mut wf := []Function{cap: m.funcs.len}
    for f in m.funcs {
        wf << f
    }
    mut wg := []GlobalVar{cap: m.globals.len}
    for g in m.globals {
        wg << g
    }
    mut wts := TypeStore{
        cache: map[string]TypeID{}
    }
    wts.types = []Type{cap: m.type_store.types.len}
    for t in m.type_store.types {
        wts.types << t
    }
    for k, v in m.type_store.cache {
        wts.cache[k] = v
    }

    mut w := &Module{
        name:              m.name
        target:            m.target
        type_store:        wts
        funcs:             wf
        globals:           wg
        c_struct_names:    map[int]string{}
        c_typedef_structs: map[int]bool{}
        const_cache:       m.const_cache.clone()
    }
    w.values = []Value{cap: m.values.len + 1024}
    for v in m.values {
        w.values << v
    }
    w.instrs = []Instruction{cap: m.instrs.len + 1024}
    for instr in m.instrs {
        w.instrs << instr
    }
    w.blocks = []BasicBlock{cap: m.blocks.len + 256}
    for blk in m.blocks {
        w.blocks << blk
    }
    return w
}

// worker_type_cache_key supports worker type cache key handling for ssa.
fn worker_type_cache_key(t Type, type_remap []TypeID) string {
    return match t.kind {
        .int_t {
            if t.is_unsigned {
                'u${t.width}'
            } else {
                'i${t.width}'
            }
        }
        .float_t {
            'f${t.width}'
        }
        .ptr_t {
            remapped := if int(t.elem_type) < type_remap.len {
                type_remap[int(t.elem_type)]
            } else {
                t.elem_type
            }
            'p${remapped}'
        }
        .array_t {
            remapped := if int(t.elem_type) < type_remap.len {
                type_remap[int(t.elem_type)]
            } else {
                t.elem_type
            }
            'a${remapped}_${t.len}'
        }
        else {
            ''
        }
    }
}

// merge_worker_module folds a worker's freshly-built SSA (everything beyond the
// seed lengths) back into the main module, remapping type/value/instr/block/func
// ids. func_data carries each built function's worker-local blocks/params.
pub fn (mut m Module) merge_worker_module(w &Module, func_data []FuncSSAData, seed_values int, seed_instrs int, seed_blocks int, seed_types int, seed_funcs int) {
    // 1. Remap worker types >= seed_types into the main type store.
    mut type_remap := []TypeID{len: w.type_store.types.len, init: 0}
    for ti := 0; ti < seed_types && ti < type_remap.len; ti++ {
        type_remap[ti] = ti
    }
    for ti := seed_types; ti < w.type_store.types.len; ti++ {
        wt := w.type_store.types[ti]
        cache_key := worker_type_cache_key(wt, type_remap)
        if cache_key.len > 0 {
            if existing := m.type_store.cache[cache_key] {
                if existing > 0 {
                    type_remap[ti] = existing
                    continue
                }
            }
        }
        remap_t := fn (id TypeID, type_remap []TypeID, seed_types int) TypeID {
            if int(id) >= seed_types && int(id) < type_remap.len {
                return type_remap[int(id)]
            }
            return id
        }
        mut new_fields := []TypeID{cap: wt.fields.len}
        for f in wt.fields {
            new_fields << remap_t(f, type_remap, seed_types)
        }
        mut new_params := []TypeID{cap: wt.params.len}
        for p in wt.params {
            new_params << remap_t(p, type_remap, seed_types)
        }
        new_type := Type{
            kind:        wt.kind
            width:       wt.width
            is_unsigned: wt.is_unsigned
            elem_type:   remap_t(wt.elem_type, type_remap, seed_types)
            len:         wt.len
            fields:      new_fields
            field_names: wt.field_names
            params:      new_params
            ret_type:    remap_t(wt.ret_type, type_remap, seed_types)
            is_c_struct: wt.is_c_struct
            is_union:    wt.is_union
        }
        new_id := m.type_store.register(new_type)
        if cache_key.len > 0 {
            m.type_store.cache[cache_key] = new_id
        }
        type_remap[ti] = new_id
    }

    value_off := m.values.len - seed_values
    instr_off := m.instrs.len - seed_instrs
    block_off := m.blocks.len - seed_blocks

    remap_val := fn (id ValueID, seed_values int, value_off int) ValueID {
        return if id >= seed_values { id + value_off } else { id }
    }
    remap_blk := fn (id BlockID, seed_blocks int, block_off int) BlockID {
        return if id >= seed_blocks { id + block_off } else { id }
    }
    remap_typ := fn (id TypeID, seed_types int, type_remap []TypeID) TypeID {
        return if int(id) >= seed_types && int(id) < type_remap.len {
            type_remap[int(id)]
        } else {
            id
        }
    }

    // 2. Append worker values beyond the seed.
    for wi := seed_values; wi < w.values.len; wi++ {
        wv := w.values[wi]
        mut new_index := wv.index
        if wv.kind == .instruction {
            new_index += instr_off
        } else if wv.kind == .basic_block {
            new_index += block_off
        }
        mut new_uses := []ValueID{cap: wv.uses.len}
        for u in wv.uses {
            new_uses << remap_val(u, seed_values, value_off)
        }
        m.values << Value{
            id:    wi + value_off
            kind:  wv.kind
            typ:   remap_typ(wv.typ, seed_types, type_remap)
            name:  wv.name
            index: new_index
            uses:  new_uses
        }
    }

    // 3. Append worker instructions beyond the seed, remapping operands by kind.
    for ii := seed_instrs; ii < w.instrs.len; ii++ {
        instr := w.instrs[ii]
        mut new_ops := []ValueID{cap: instr.operands.len}
        for oi in 0 .. instr.operands.len {
            op := instr.operands[oi]
            if instr.is_block_operand(oi) {
                new_ops << ValueID(remap_blk(BlockID(op), seed_blocks, block_off))
            } else {
                new_ops << remap_val(op, seed_values, value_off)
            }
        }
        m.instrs << Instruction{
            op:         instr.op
            operands:   new_ops
            block:      remap_blk(instr.block, seed_blocks, block_off)
            typ:        remap_typ(instr.typ, seed_types, type_remap)
            pos:        instr.pos
            atomic_ord: instr.atomic_ord
            inline:     instr.inline
        }
    }

    // 4. Append worker blocks beyond the seed.
    for bi := seed_blocks; bi < w.blocks.len; bi++ {
        blk := w.blocks[bi]
        mut new_instrs := []ValueID{cap: blk.instrs.len}
        for vi in blk.instrs {
            new_instrs << remap_val(vi, seed_values, value_off)
        }
        mut new_preds := []BlockID{cap: blk.preds.len}
        for p in blk.preds {
            new_preds << remap_blk(p, seed_blocks, block_off)
        }
        mut new_succs := []BlockID{cap: blk.succs.len}
        for s in blk.succs {
            new_succs << remap_blk(s, seed_blocks, block_off)
        }
        m.blocks << BasicBlock{
            id:     blk.id + block_off
            val_id: remap_val(blk.val_id, seed_values, value_off)
            name:   blk.name
            parent: blk.parent
            instrs: new_instrs
            preds:  new_preds
            succs:  new_succs
        }
    }

    // 5. Patch the built functions' blocks/params in the main module.
    for fd in func_data {
        if fd.func_idx >= m.funcs.len {
            continue
        }
        mut f := m.funcs[fd.func_idx]
        if f.blocks.len > 0 {
            continue // already merged by an earlier worker
        }
        mut new_blocks := []BlockID{cap: fd.blocks.len}
        for blk_id in fd.blocks {
            new_blocks << remap_blk(blk_id, seed_blocks, block_off)
        }
        f.blocks = new_blocks
        mut new_params := []ValueID{cap: fd.params.len}
        for p in fd.params {
            new_params << remap_val(p, seed_values, value_off)
        }
        f.params = new_params
        m.funcs[fd.func_idx] = f
    }

    // 6. Append any brand-new functions a worker created (runtime stubs, etc.).
    for fi := seed_funcs; fi < w.funcs.len; fi++ {
        wfunc := w.funcs[fi]
        mut already := false
        for existing in m.funcs {
            if existing.name == wfunc.name {
                already = true
                break
            }
        }
        if already {
            continue
        }
        mut new_blocks := []BlockID{cap: wfunc.blocks.len}
        for blk_id in wfunc.blocks {
            new_blocks << remap_blk(blk_id, seed_blocks, block_off)
        }
        mut new_params := []ValueID{cap: wfunc.params.len}
        for p in wfunc.params {
            new_params << remap_val(p, seed_values, value_off)
        }
        m.funcs << Function{
            id:           m.funcs.len
            name:         wfunc.name
            typ:          remap_typ(wfunc.typ, seed_types, type_remap)
            blocks:       new_blocks
            params:       new_params
            is_c_extern:  wfunc.is_c_extern
            is_prototype: wfunc.is_prototype
            linkage:      wfunc.linkage
            call_conv:    wfunc.call_conv
        }
    }
}