//! Localification: a simple form of register allocation that picks
//! locations for SSA values in Wasm locals.

use crate::backend::treeify::Trees;
use crate::cfg::CFGInfo;
use crate::entity::{EntityVec, PerEntity};
use crate::ir::{Block, FunctionBody, Local, Type, Value, ValueDef};
use smallvec::{smallvec, SmallVec};
use std::collections::{HashMap, HashSet};

#[derive(Clone, Debug, Default)]
pub struct Localifier {
    pub values: PerEntity<Value, SmallVec<[Local; 2]>>,
    pub locals: EntityVec<Local, Type>,
}

impl Localifier {
    pub fn compute(body: &FunctionBody, cfg: &CFGInfo, trees: &Trees) -> Self {
        Context::new(body, cfg, trees).compute()
    }
}

struct Context<'a> {
    body: &'a FunctionBody,
    cfg: &'a CFGInfo,
    trees: &'a Trees,
    results: Localifier,

    // Affinities (blockparam value to input arg value).
    affinities: HashMap<Value, SmallVec<[Value; 4]>>,

    // Livein to each block from dominators.
    livein_values: PerEntity<Block, HashSet<Value>>,
    livein_locals: PerEntity<Block, HashSet<Local>>,
}

impl<'a> Context<'a> {
    fn new(body: &'a FunctionBody, cfg: &'a CFGInfo, trees: &'a Trees) -> Self {
        let mut results = Localifier::default();
        // Create locals for function args.
        for &(ty, value) in &body.blocks[body.entry].params {
            let param_local = results.locals.push(ty);
            results.values[value] = smallvec![param_local];
        }
        // Compute affinities.
        let mut affinities = HashMap::default();
        for block in body.blocks.values() {
            block.terminator.visit_targets(|target| {
                for (&arg, &param) in target
                    .args
                    .iter()
                    .zip(body.blocks[target.block].params.iter().map(|(_, val)| val))
                {
                    affinities
                        .entry(arg)
                        .or_insert_with(|| smallvec![])
                        .push(param);
                }
            });
        }

        Self {
            body,
            cfg,
            trees,
            results,
            affinities,
            livein_values: PerEntity::default(),
            livein_locals: PerEntity::default(),
        }
    }

    fn compute(mut self) -> Localifier {
        // Create domtree preorder for traversal (we iterate in
        // reverse preorder below).
        let mut order = vec![];
        order.push(self.body.entry);
        let mut i = 0;
        while i < order.len() {
            for child in self.cfg.dom_children(order[i]) {
                order.push(child);
            }
            i += 1;
        }

        for &block in order.iter().rev() {
            self.process(block);
        }

        debug_assert!(self.livein_values[self.body.entry].is_empty());
        debug_assert!(self.livein_locals[self.body.entry].is_empty());

        self.results
    }

    fn process(&mut self, block: Block) {
        let mut live_values = HashSet::new();
        let mut live_locals = HashSet::new();

        // Collect liveins of all dominated blocks; this is our initial live-set.
        for child in self.cfg.dom_children(block) {
            for &livein_value in &self.livein_values[child] {
                log::trace!(
                    "localify: block {} gets livein value {} from block {}",
                    block,
                    livein_value,
                    child
                );
                live_values.insert(livein_value);
            }
            for &livein_local in &self.livein_locals[child] {
                live_locals.insert(livein_local);
            }
        }

        log::trace!(
            "localify: process block {}: liveout values {:?} locals {:?}",
            block,
            live_values,
            live_locals
        );

        // For each use/def in reverse order, update live-set; on last
        // use (first observed), allocate a local.
        fn handle_use(
            body: &FunctionBody,
            cfg: &CFGInfo,
            block: Block,
            u: Value,
            live_values: &mut HashSet<Value>,
            live_locals: &mut HashSet<Local>,
            results: &mut Localifier,
            affinities: &HashMap<Value, SmallVec<[Value; 4]>>,
        ) {
            let u = body.resolve_alias(u);
            if live_values.insert(u) {
                // If there is already an allocation (e.g. for a
                // function parameter), return.
                if !results.values[u].is_empty() {
                    return;
                }

                // Need to create an allocation.
                let def = &body.values[u];
                match def {
                    &ValueDef::Alias(_value) => {
                        unreachable!();
                    }
                    &ValueDef::PickOutput(value, idx, _) => {
                        handle_use(
                            body,
                            cfg,
                            block,
                            value,
                            live_values,
                            live_locals,
                            results,
                            affinities,
                        );
                        results.values[u] = smallvec![results.values[value][idx]];
                    }
                    &ValueDef::BlockParam(..) | &ValueDef::Operator(..) => {
                        // Uses from dominating blocks may be live
                        // over a loop backedge, so for simplicity,
                        // let's not reuse any locals for uses of
                        // non-local defs.
                        let can_reuse = cfg.def_block[u] == block;

                        let locals = def
                            .tys()
                            .iter()
                            .map(|&ty| {
                                let reused = if can_reuse {
                                    // Try to find a local of the right type that is not live.
                                    let affinities =
                                        affinities.get(&u).map(|v| &v[..]).unwrap_or(&[]);
                                    let mut try_list = affinities
                                        .iter()
                                        .filter_map(|&aff_val| {
                                            let local = *results.values[aff_val].get(0)?;
                                            Some((local, results.locals[local]))
                                        })
                                        .chain(
                                            results
                                                .locals
                                                .entries()
                                                .map(|(local, &ty)| (local, ty)),
                                        );

                                    try_list
                                        .find(|&(local, local_ty)| {
                                            local_ty == ty && live_locals.insert(local)
                                        })
                                        .map(|(local, _)| local)
                                } else {
                                    None
                                };

                                reused.unwrap_or_else(|| {
                                    let local = results.locals.push(ty);
                                    live_locals.insert(local);
                                    local
                                })
                            })
                            .collect::<SmallVec<_>>();
                        results.values[u] = locals;
                    }
                    &ValueDef::Placeholder(_) | &ValueDef::None => unreachable!(),
                }
            }
        }

        fn handle_def(
            body: &FunctionBody,
            d: Value,
            live_values: &mut HashSet<Value>,
            live_locals: &mut HashSet<Local>,
            results: &Localifier,
        ) {
            if let ValueDef::Alias(..) = &body.values[d] {
                return;
            }

            if live_values.remove(&d) {
                for &local in &results.values[d] {
                    live_locals.remove(&local);
                }
            }
        }

        self.body.blocks[block].terminator.visit_uses(|u| {
            handle_use(
                self.body,
                self.cfg,
                block,
                u,
                &mut live_values,
                &mut live_locals,
                &mut self.results,
                &self.affinities,
            )
        });

        for &inst in self.body.blocks[block].insts.iter().rev() {
            handle_def(
                self.body,
                inst,
                &mut live_values,
                &mut live_locals,
                &self.results,
            );
            self.body.values[inst].visit_uses(|u| {
                // If treeified, then don't process use.
                if !self.trees.owner.contains_key(&u) {
                    handle_use(
                        self.body,
                        self.cfg,
                        block,
                        u,
                        &mut live_values,
                        &mut live_locals,
                        &mut self.results,
                        &self.affinities,
                    )
                }
            });
        }

        for &(_, param) in &self.body.blocks[block].params {
            handle_def(
                self.body,
                param,
                &mut live_values,
                &mut live_locals,
                &self.results,
            );
        }

        log::trace!(
            "localify: process block {}: livein values {:?} locals {:?}",
            block,
            live_values,
            live_locals
        );
        self.livein_locals[block] = live_locals;
        self.livein_values[block] = live_values;
    }
}