/*
 * This file was initially derived from the files
 * `js/src/jit/BacktrackingAllocator.h` and
 * `js/src/jit/BacktrackingAllocator.cpp` in Mozilla Firefox, and was
 * originally licensed under the Mozilla Public License 2.0. We
 * subsequently relicensed it to Apache-2.0 WITH LLVM-exception (see
 * https://github.com/bytecodealliance/regalloc2/issues/7).
 *
 * Since the initial port, the design has been substantially evolved
 * and optimized.
 */

//! Backtracking register allocator. See doc/DESIGN.md for details of
//! its design.

use crate::ssa::validate_ssa;
use crate::{Function, MachineEnv, PReg, RegAllocError, RegClass, VecExt};
pub(crate) mod data_structures;
pub use data_structures::Ctx;
pub use data_structures::Stats;
use data_structures::*;
pub(crate) mod reg_traversal;
use reg_traversal::*;
pub(crate) mod requirement;
use requirement::*;
pub(crate) mod redundant_moves;
use redundant_moves::*;
pub(crate) mod liveranges;
use liveranges::*;
pub(crate) mod merge;
pub(crate) mod process;
use process::*;
use smallvec::smallvec;
pub(crate) mod dump;
pub(crate) mod moves;
pub(crate) mod spill;

impl<'a, F: Function> Env<'a, F> {
    pub(crate) fn new(func: &'a F, env: &'a MachineEnv, ctx: &'a mut Ctx) -> Self {
        let ninstrs = func.num_insts();
        let nblocks = func.num_blocks();

        ctx.liveins.preallocate(nblocks);
        ctx.liveouts.preallocate(nblocks);
        ctx.blockparam_ins.clear();
        ctx.blockparam_outs.clear();
        ctx.ranges.preallocate(4 * ninstrs);
        ctx.bundles.preallocate(ninstrs);
        ctx.spillsets.preallocate(ninstrs);
        ctx.vregs.preallocate(ninstrs);
        for preg in ctx.pregs.iter_mut() {
            preg.is_stack = false;
            preg.allocations.btree.clear();
        }
        ctx.allocation_queue.heap.clear();
        ctx.spilled_bundles.clear();
        ctx.scratch_spillset_pool
            .extend(ctx.spillslots.drain(..).map(|mut s| {
                s.ranges.btree.clear();
                s.ranges
            }));
        ctx.slots_by_class = core::array::from_fn(|_| SpillSlotList::default());
        ctx.extra_spillslots_by_class = core::array::from_fn(|_| smallvec![]);
        ctx.preferred_victim_by_class = [PReg::invalid(); 3];
        ctx.multi_fixed_reg_fixups.clear();
        ctx.allocated_bundle_count = 0;
        ctx.debug_annotations.clear();
        ctx.scratch_bump
            .get_mut()
            .expect("we dropped all refs to this")
            .reset();

        ctx.output.allocs.preallocate(4 * ninstrs);
        ctx.output.inst_alloc_offsets.clear();
        ctx.output.num_spillslots = 0;
        ctx.output.debug_locations.clear();
        ctx.output.edits.clear();
        ctx.output.stats = Stats::default();

        Self { func, env, ctx }
    }

    pub(crate) fn init(&mut self) -> Result<(), RegAllocError> {
        self.create_pregs_and_vregs();
        self.compute_liveness()?;
        self.build_liveranges();
        self.fixup_multi_fixed_vregs();
        self.merge_vreg_bundles();
        self.queue_bundles();
        if trace_enabled!() {
            self.dump_state();
        }
        Ok(())
    }

    pub(crate) fn run(&mut self) -> Result<Edits, RegAllocError> {
        self.process_bundles()?;
        self.try_allocating_regs_for_spilled_bundles();
        self.allocate_spillslots();
        let moves = self.apply_allocations_and_insert_moves();
        Ok(self.resolve_inserted_moves(moves))
    }
}

pub fn run<F: Function>(
    func: &F,
    mach_env: &MachineEnv,
    ctx: &mut Ctx,
    enable_annotations: bool,
    enable_ssa_checker: bool,
) -> Result<(), RegAllocError> {
    ctx.cfginfo.init(func, &mut ctx.cfginfo_ctx)?;

    if enable_ssa_checker {
        validate_ssa(func, &ctx.cfginfo)?;
    }

    ctx.annotations_enabled = enable_annotations;
    let mut env = Env::new(func, mach_env, ctx);
    env.init()?;

    let mut edits = env.run()?;

    if enable_annotations {
        env.dump_results();
    }

    ctx.output.edits.extend(edits.drain_edits());

    Ok(())
}