use std::collections::{HashMap, HashSet};
use std::fmt;

use anyhow::{anyhow, Result};

use petgraph::algo::coloring::dsatur_coloring;
use petgraph::graph::{NodeIndex, UnGraph};

use crate::config::{Config, SLICE_CAP, SPILL_AREA};
use crate::mir::{Instruction, Mir};
use crate::serializer::MirWriter;
use crate::symbol::Loc;
use crate::utils::Reg;

#[derive(Clone, Debug)]
struct Vertex {
    start: u32,
    end: u32,
    reg: Reg,
}

// #[derive(Clone)]
pub struct ColoringAllocator {
    pub code: MirWriter, // the revised mir
    regs: Vec<Reg>,
    locs: HashMap<Loc, Reg>,
    loads: HashSet<Loc>,
    count_statics: u32, // number of statis registers
    graph: UnGraph<Vertex, ()>,
    config: Config,
}

impl fmt::Debug for ColoringAllocator {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (i, ins) in self.code.iter().enumerate() {
            writeln!(f, "{:05}\t{:?}", i, ins)?;
        }

        writeln!(f, "...................")?;
        writeln!(f, "{:#?}", self.graph)?;

        Ok(())
    }
}

impl ColoringAllocator {
    pub fn new(config: Config) -> ColoringAllocator {
        let count_scratch = config.count_scratch();

        ColoringAllocator {
            code: MirWriter::new(),
            regs: vec![Reg::Ret; count_scratch as usize],
            count_statics: 0,
            graph: UnGraph::new_undirected(),
            locs: HashMap::new(),
            loads: HashSet::new(),
            config,
        }
    }

    pub fn optimize(&mut self, mir: &mut Mir) -> Result<()> {
        // create single-static-assignment form
        self.create(mir);

        // add edges to the graph based on which register
        // pairs overlap in the same region
        self.add_edges();

        // allocate registers using a coloring algorithm and
        // replace the static registers with the corresponding
        // logical (colored) registers.
        // Note that no error is raised if color() does not
        // work. Instead, it fails silently.

        if self.color().is_ok() {
            // checking is_ok because color may fail silently.
            mir.code = std::mem::take(&mut self.code);
        }
        Ok(())
    }

    fn push(&mut self, ins: Instruction) {
        self.code.push(&ins);
    }

    fn ip(&self) -> u32 {
        self.code.ip as u32
    }

    // resets cache on call boundaries to account for cobbled registers
    fn reset(&mut self) {
        self.regs = vec![Reg::Ret; self.config.count_scratch() as usize];
        self.locs.clear();
    }

    fn overlap(&self, idx1: NodeIndex, idx2: NodeIndex) -> bool {
        let v1 = &self.graph[idx1];
        let v2 = &self.graph[idx2];
        u32::min(v1.end, v2.end) > u32::max(v1.start, v2.start)
    }

    fn add_edges(&mut self) {
        let n = self.graph.node_count();

        for i1 in 0..n {
            let idx1 = NodeIndex::new(i1);
            for i2 in i1 + 1..n {
                let idx2 = NodeIndex::new(i2);

                if self.overlap(idx1, idx2) {
                    self.graph.add_edge(idx1, idx2, ());
                }
            }
        }
    }

    fn create_static(&mut self) -> Reg {
        let s = Reg::Static(self.count_statics);
        self.count_statics += 1;
        self.graph.add_node(Vertex {
            start: self.ip(),
            end: 0,
            reg: Reg::Ret,
        });
        s
    }

    // consumes a logical register, update the end interval forthe corresponding
    // static register and then returns the static register.
    // note that the cache (self.regs) is not invalidated and may
    // be used afterward
    fn consume(&mut self, src: Reg) -> Reg {
        if let Reg::Gen(r) = src {
            let s = self.regs[r as usize];
            if let Reg::Static(k) = s {
                self.graph[NodeIndex::new(k as usize)].end = self.ip();
            };
            s
        } else {
            src
        }
    }

    // converts a destination logical register to a static one
    fn produce(&mut self, dst: Reg) -> Reg {
        if let Reg::Gen(r) = dst {
            let s = self.create_static();
            self.regs[r as usize] = s;
            s
        } else {
            dst
        }
    }

    // helper function to ease producing dst and consuming s1.
    fn unary_op(&mut self, dst: Reg, s1: Reg) -> (Reg, Reg) {
        // note than RHS is consumed before producing LHS
        let s1 = self.consume(s1);
        let dst = self.produce(dst);
        (dst, s1)
    }

    // helper function to ease producing dst and consuming s1 and s2.
    fn binary_op(&mut self, dst: Reg, s1: Reg, s2: Reg) -> (Reg, Reg, Reg) {
        // note than RHS's are consumed before producing LHS
        let s1 = self.consume(s1);
        let s2 = self.consume(s2);
        let dst = self.produce(dst);
        (dst, s1, s2)
    }

    // helper function to ease producing dst and consuming s1, s2, and s3.
    fn ternary_op(&mut self, dst: Reg, s1: Reg, s2: Reg, s3: Reg) -> (Reg, Reg, Reg, Reg) {
        // note than RHS's are consumed before producing LHS
        let s1 = self.consume(s1);
        let s2 = self.consume(s2);
        let s3 = self.consume(s3);
        let dst = self.produce(dst);
        (dst, s1, s2, s3)
    }

    fn load(&mut self, dst: Reg, loc: Loc) {
        if let Reg::Gen(r) = dst {
            // if the desired location is already in a static reg, we just
            // use this reg
            if let Some(dst) = self.locs.get(&loc) {
                self.regs[r as usize] = *dst;
                return;
            } else {
                let dst = self.create_static();
                self.regs[r as usize] = dst;
                self.locs.insert(loc, dst);
                self.push(Instruction::Load { dst, loc });
            };
        } else {
            self.push(Instruction::Load { dst, loc });
        }

        self.loads.insert(loc);
    }

    fn save(&mut self, src: Reg, loc: Loc) {
        let src = self.consume(src);
        self.push(Instruction::Save { src, loc });

        if let Reg::Static(..) = src {
            self.locs.insert(loc, src);
        }
    }

    pub fn create(&mut self, mir: &Mir) {
        for ins in mir.code.iter() {
            match ins {
                Instruction::Nop => self.push(Instruction::Nop),
                Instruction::End => self.push(Instruction::End),
                Instruction::Uni { op, dst, s1 } => {
                    let (dst, s1) = self.unary_op(dst, s1);
                    self.push(Instruction::Uni { op, dst, s1 });
                }
                Instruction::Bi { op, dst, s1, s2 } => {
                    let (dst, s1, s2) = self.binary_op(dst, s1, s2);
                    self.push(Instruction::Bi { op, dst, s1, s2 });
                }
                Instruction::LoadConst { dst, idx } => {
                    let dst = self.produce(dst);
                    self.push(Instruction::LoadConst { dst, idx });
                }
                Instruction::Load { dst, loc } => {
                    self.load(dst, loc);
                }
                Instruction::Save { src, loc } => {
                    self.save(src, loc);
                }
                Instruction::LoadComplex { xd, yd, loc } => {
                    self.load(xd, loc);
                    self.load(yd, loc.imag());
                }
                Instruction::SaveComplex { xs, ys, loc } => {
                    self.save(xs, loc);
                    self.save(ys, loc.imag());
                }
                Instruction::Mov { dst, s1 } => {
                    let (dst, s1) = self.unary_op(dst, s1);
                    self.push(Instruction::Mov { dst, s1 });
                }
                Instruction::Fused { op, dst, a, b, c } => {
                    let (dst, a, b, c) = self.ternary_op(dst, a, b, c);
                    self.push(Instruction::Fused { op, dst, a, b, c });
                }
                Instruction::IfElse {
                    dst,
                    true_val,
                    false_val,
                    cond,
                } => {
                    let (dst, true_val, false_val) = self.binary_op(dst, true_val, false_val);
                    self.loads.insert(cond);
                    self.push(Instruction::IfElse {
                        dst,
                        true_val,
                        false_val,
                        cond,
                    });
                }
                Instruction::Branch { .. } => {
                    if let Instruction::Branch { label } = ins.clone() {
                        self.push(Instruction::Branch { label });
                    }
                }
                Instruction::BranchIf { .. } => {
                    if let Instruction::BranchIf {
                        cond,
                        label,
                        is_else,
                    } = ins.clone()
                    {
                        let cond = self.consume(cond);
                        self.push(Instruction::BranchIf {
                            cond,
                            label,
                            is_else,
                        });
                    }
                }
                Instruction::Call { .. } | Instruction::Label { .. } => {
                    self.push(ins.clone());
                    self.reset();
                }
                Instruction::LoadMath { op, dst, s1, loc } => {
                    let s1 = self.consume(s1);
                    let dst = self.produce(dst);
                    self.push(Instruction::LoadMath { op, dst, s1, loc });
                    self.loads.insert(loc);
                }
                Instruction::LoadConstMath { op, dst, s1, idx } => {
                    let s1 = self.consume(s1);
                    let dst = self.produce(dst);
                    self.push(Instruction::LoadConstMath { op, dst, s1, idx });
                }
                Instruction::ComplexBi {
                    op,
                    xd,
                    yd,
                    x1,
                    y1,
                    x2,
                    y2,
                } => {
                    let x1 = self.consume(x1);
                    let y1 = self.consume(y1);
                    let x2 = self.consume(x2);
                    let y2 = self.consume(y2);
                    let xd = self.produce(xd);
                    let yd = self.produce(yd);
                    self.push(Instruction::ComplexBi {
                        op,
                        xd,
                        yd,
                        x1,
                        y1,
                        x2,
                        y2,
                    });
                }
            }
        }
    }

    // converts a static register to the corresponding logical register
    // calculated by the coloring algorithm
    fn alloc(&self, dst: Reg) -> Reg {
        if let Reg::Static(s) = dst {
            let idx = NodeIndex::new(s as usize);
            self.graph[idx].reg
        } else {
            dst
        }
    }

    fn color(&mut self) -> Result<()> {
        let (coloring, count_colors) = dsatur_coloring(&self.graph);

        if count_colors > self.config.count_scratch() as usize {
            return Err(anyhow!(
                "Graph coloring exceeds the limit: {:?}",
                count_colors
            ));
        }

        for (idx, r) in coloring.iter() {
            self.graph[*idx].reg = Reg::Gen(*r as u8);
        }

        let code = std::mem::take(&mut self.code);

        // replace all static regs with the corresponding logical ones
        for ins in code.iter() {
            match ins {
                Instruction::Nop => self.push(Instruction::Nop),
                Instruction::End => self.push(Instruction::End),
                Instruction::Uni { op, dst, s1 } => self.push(Instruction::Uni {
                    op,
                    dst: self.alloc(dst),
                    s1: self.alloc(s1),
                }),
                Instruction::Bi { op, dst, s1, s2 } => self.push(Instruction::Bi {
                    op,
                    dst: self.alloc(dst),
                    s1: self.alloc(s1),
                    s2: self.alloc(s2),
                }),
                Instruction::LoadConst { dst, idx } => self.push(Instruction::LoadConst {
                    dst: self.alloc(dst),
                    idx,
                }),
                Instruction::Load { dst, loc } => self.push(Instruction::Load {
                    dst: self.alloc(dst),
                    loc,
                }),
                Instruction::Save { src, loc } => {
                    if !matches!(loc, Loc::Stack(..)) || self.loads.contains(&loc) {
                        self.push(Instruction::Save {
                            src: self.alloc(src),
                            loc,
                        })
                    }
                }
                Instruction::LoadComplex { xd, yd, loc } => self.push(Instruction::LoadComplex {
                    xd: self.alloc(xd),
                    yd: self.alloc(yd),
                    loc,
                }),
                Instruction::SaveComplex { xs, ys, loc } => {
                    if !matches!(loc, Loc::Stack(..)) || self.loads.contains(&loc) {
                        self.push(Instruction::SaveComplex {
                            xs: self.alloc(xs),
                            ys: self.alloc(ys),
                            loc,
                        })
                    }
                }
                Instruction::Mov { dst, s1 } => self.push(Instruction::Mov {
                    dst: self.alloc(dst),
                    s1: self.alloc(s1),
                }),
                Instruction::Fused { op, dst, a, b, c } => {
                    self.push(Instruction::Fused {
                        op,
                        dst: self.alloc(dst),
                        a: self.alloc(a),
                        b: self.alloc(b),
                        c: self.alloc(c),
                    });
                }
                Instruction::IfElse {
                    dst,
                    true_val,
                    false_val,
                    cond,
                } => self.push(Instruction::IfElse {
                    dst: self.alloc(dst),
                    true_val: self.alloc(true_val),
                    false_val: self.alloc(false_val),
                    cond,
                }),
                Instruction::Branch { .. } => {
                    if let Instruction::Branch { label } = ins.clone() {
                        self.push(Instruction::Branch { label });
                    }
                }
                Instruction::BranchIf { .. } => {
                    if let Instruction::BranchIf {
                        cond,
                        label,
                        is_else,
                    } = ins.clone()
                    {
                        let cond = self.alloc(cond);
                        self.push(Instruction::BranchIf {
                            cond,
                            label,
                            is_else,
                        });
                    }
                }
                Instruction::Call { .. } | Instruction::Label { .. } => {
                    self.push(ins.clone());
                }
                Instruction::LoadMath { op, dst, s1, loc } => {
                    self.push(Instruction::LoadMath {
                        op,
                        dst: self.alloc(dst),
                        s1: self.alloc(s1),
                        loc,
                    });
                }
                Instruction::LoadConstMath { op, dst, s1, idx } => {
                    self.push(Instruction::LoadConstMath {
                        op,
                        dst: self.alloc(dst),
                        s1: self.alloc(s1),
                        idx,
                    });
                }
                Instruction::ComplexBi {
                    op,
                    xd,
                    yd,
                    x1,
                    y1,
                    x2,
                    y2,
                } => {
                    self.push(Instruction::ComplexBi {
                        op,
                        xd: self.alloc(xd),
                        yd: self.alloc(yd),
                        x1: self.alloc(x1),
                        y1: self.alloc(y1),
                        x2: self.alloc(x2),
                        y2: self.alloc(y2),
                    });
                }
            }
        }

        Ok(())
    }
}

/************************* GreedyAllocator ***************************/

#[derive(Debug, Clone)]
struct Static {
    reg: Reg,
    end: usize,
}

#[derive(Debug, Clone)]
struct Alloc {
    owners: HashSet<usize>,
    loc: Option<Loc>,
}

impl Alloc {
    fn new() -> Alloc {
        Alloc {
            owners: HashSet::new(),
            loc: None,
        }
    }
}

// #[derive(Clone)]
pub struct GreedyAllocator {
    pub code: MirWriter,       // the revised mir
    regs: Vec<Option<usize>>,  // map of logical registers to static ones
    locs: HashMap<Loc, usize>, // map between locs and statics
    count_statics: usize,      // number of statis registers
    statics: Vec<Static>,      // the list of static registers
    allocs: Vec<Alloc>,        // allocation for logical registers
    config: Config,
    count_regs: usize,
}

impl fmt::Debug for GreedyAllocator {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        for (i, ins) in self.code.iter().enumerate() {
            writeln!(f, "{:05}\t{:?}", i, ins)?;
        }

        writeln!(f, "...................")?;

        for (i, s) in self.statics.iter().enumerate() {
            writeln!(f, "σ{} := ({:?})", i, &s)?;
        }

        Ok(())
    }
}

impl GreedyAllocator {
    pub fn new(config: Config, count_regs: usize) -> GreedyAllocator {
        GreedyAllocator {
            code: MirWriter::new(),
            regs: vec![None; count_regs],
            count_statics: 0,
            locs: HashMap::new(),
            statics: Vec::new(),
            allocs: vec![Alloc::new(); count_regs],
            config,
            count_regs,
        }
    }

    pub fn optimize(&mut self, mir: &mut Mir) -> Result<()> {
        // create single-static-assignment form
        self.create(mir);

        // allocate registers using a greedy algorithm and
        // replace the static registers with the corresponding
        // logical (colored) registers
        self.color()?;
        self.contract()?;
        // contract the code by removing unnecessary instructions
        mir.code = std::mem::take(&mut self.code);

        Ok(())
    }

    fn push(&mut self, ins: Instruction) {
        self.code.push(&ins);
    }

    // reset for the first pass (logical -> static pass)
    fn reset_regs(&mut self) {
        self.regs = vec![None; self.count_regs];
    }

    // reset for the second pass (static -> logical pass)
    fn reset_allocs(&mut self) {
        self.allocs = vec![Alloc::new(); self.count_regs];
    }

    fn create_static(&mut self, ip: usize, r: Reg) -> usize {
        let idx = self.count_statics;
        self.count_statics += 1;
        self.statics.push(Static { reg: r, end: ip });
        idx
    }

    // consumes a logical register, update the end interval for the corresponding
    // static register and then returns the static register.
    fn consume(&mut self, ip: usize, src: Reg) -> Reg {
        if let Reg::Gen(r) = src {
            let s = self.regs[r as usize]
                .unwrap_or_else(|| panic!("reg {:?} not found (line {})", src, ip));
            self.statics[s].end = ip;
            Reg::Static(s as u32)
        } else {
            src
        }
    }

    // converts a destination logical register to a static one
    fn produce(&mut self, ip: usize, dst: Reg) -> Reg {
        if let Reg::Gen(r) = dst {
            let s = self.create_static(ip, dst);
            self.regs[r as usize] = Some(s);
            Reg::Static(s as u32)
        } else {
            dst
        }
    }

    // helper function to ease producing dst and consuming s1.
    fn unary_op(&mut self, ip: usize, dst: Reg, s1: Reg) -> (Reg, Reg) {
        // note than RHS is consumed before producing LHS
        let s1 = self.consume(ip, s1);
        let dst = self.produce(ip, dst);
        (dst, s1)
    }

    // helper function to ease producing dst and consuming s1 and s2.
    fn binary_op(&mut self, ip: usize, dst: Reg, s1: Reg, s2: Reg) -> (Reg, Reg, Reg) {
        // note than RHS's are consumed before producing LHS
        let s1 = self.consume(ip, s1);
        let s2 = self.consume(ip, s2);
        let dst = self.produce(ip, dst);
        (dst, s1, s2)
    }

    // helper function to ease producing dst and consuming s1, s2, and s3.
    fn ternary_op(
        &mut self,
        ip: usize,
        dst: Reg,
        s1: Reg,
        s2: Reg,
        s3: Reg,
    ) -> (Reg, Reg, Reg, Reg) {
        // note than RHS's are consumed before producing LHS
        let s1 = self.consume(ip, s1);
        let s2 = self.consume(ip, s2);
        let s3 = self.consume(ip, s3);
        let dst = self.produce(ip, dst);
        (dst, s1, s2, s3)
    }

    // The first pass.
    // converts logical to static (SSA-form) registers
    pub fn create(&mut self, mir: &Mir) {
        for (ip, ins) in mir.code.iter().enumerate() {
            match ins {
                Instruction::Nop => self.push(Instruction::Nop),
                Instruction::End => self.push(Instruction::End),
                Instruction::Uni { op, dst, s1 } => {
                    let (dst, s1) = self.unary_op(ip, dst, s1);
                    self.push(Instruction::Uni { op, dst, s1 });
                }
                Instruction::Bi { op, dst, s1, s2 } => {
                    let (dst, s1, s2) = self.binary_op(ip, dst, s1, s2);
                    self.push(Instruction::Bi { op, dst, s1, s2 });
                }
                Instruction::LoadConst { dst, idx } => {
                    let dst = self.produce(ip, dst);
                    self.push(Instruction::LoadConst { dst, idx });
                }
                Instruction::Load { dst, loc } => {
                    let dst = self.produce(ip, dst);
                    self.push(Instruction::Load { dst, loc });
                }
                Instruction::Save { src, loc } => {
                    let src = self.consume(ip, src);
                    self.push(Instruction::Save { src, loc });
                }
                Instruction::LoadComplex { xd, yd, loc } => {
                    let xd = self.produce(ip, xd);
                    let yd = self.produce(ip, yd);
                    self.push(Instruction::LoadComplex { xd, yd, loc });
                }
                Instruction::SaveComplex { xs, ys, loc } => {
                    let xs = self.consume(ip, xs);
                    let ys = self.consume(ip, ys);
                    self.push(Instruction::SaveComplex { xs, ys, loc });
                }
                Instruction::Mov { dst, s1 } => {
                    let (dst, s1) = self.unary_op(ip, dst, s1);
                    self.push(Instruction::Mov { dst, s1 });
                }
                Instruction::Fused { op, dst, a, b, c } => {
                    let (dst, a, b, c) = self.ternary_op(ip, dst, a, b, c);
                    self.push(Instruction::Fused { op, dst, a, b, c });
                }
                Instruction::IfElse {
                    dst,
                    true_val,
                    false_val,
                    cond,
                } => {
                    let (dst, true_val, false_val) = self.binary_op(ip, dst, true_val, false_val);
                    self.push(Instruction::IfElse {
                        dst,
                        true_val,
                        false_val,
                        cond,
                    });
                }
                Instruction::Branch { .. } => {
                    if let Instruction::Branch { label } = ins.clone() {
                        self.push(Instruction::Branch { label });
                    }
                }
                Instruction::BranchIf { .. } => {
                    if let Instruction::BranchIf {
                        cond,
                        label,
                        is_else,
                    } = ins.clone()
                    {
                        let cond = self.consume(ip, cond);
                        self.push(Instruction::BranchIf {
                            cond,
                            label,
                            is_else,
                        });
                    }
                }
                Instruction::Call { .. } | Instruction::Label { .. } => {
                    self.push(ins.clone());
                    self.reset_regs();
                }
                Instruction::LoadMath { op, dst, s1, loc } => {
                    let s1 = self.consume(ip, s1);
                    let dst = self.produce(ip, dst);
                    self.push(Instruction::LoadMath { op, dst, s1, loc });
                }
                Instruction::LoadConstMath { op, dst, s1, idx } => {
                    let s1 = self.consume(ip, s1);
                    let dst = self.produce(ip, dst);
                    self.push(Instruction::LoadConstMath { op, dst, s1, idx });
                }
                Instruction::ComplexBi {
                    op,
                    xd,
                    yd,
                    x1,
                    y1,
                    x2,
                    y2,
                } => {
                    let x1 = self.consume(ip, x1);
                    let y1 = self.consume(ip, y1);
                    let x2 = self.consume(ip, x2);
                    let y2 = self.consume(ip, y2);
                    let xd = self.produce(ip, xd);
                    let yd = self.produce(ip, yd);
                    self.push(Instruction::ComplexBi {
                        op,
                        xd,
                        yd,
                        x1,
                        y1,
                        x2,
                        y2,
                    });
                }
            }
        }
    }

    // returns the logical register corresponding to the static register
    // dst back to the pool.
    fn deallocate(&mut self, ip: usize, dst: Reg) -> Reg {
        if let Reg::Static(s) = dst {
            let s = s as usize;
            let reg = self.statics[s].reg;
            if self.statics[s].end == ip {
                if let Reg::Gen(r) = reg {
                    self.allocs[r as usize].owners.remove(&s);
                }
            }
            reg
        } else {
            dst
        }
    }

    fn assign(&mut self, r: usize, s: usize, loc: Option<Loc>) -> Reg {
        self.allocs[r].owners.insert(s);
        self.allocs[r].loc = loc;
        let reg = Reg::Gen(r as u8);
        self.statics[s].reg = reg;
        reg
    }

    // allocates a new logical register from the pool and assigns it to
    // the static register dst, optionally with a location.
    fn allocate(&mut self, dst: Reg, loc: Option<Loc>) -> (Reg, bool) {
        if let Reg::Static(s) = dst {
            let s = s as usize;

            if loc.is_some() {
                if let Some(r) = self.allocs.iter().position(|x| x.loc == loc) {
                    return (self.assign(r, s, loc), true);
                }
            }

            if let Some(r) = self
                .allocs
                .iter()
                .position(|x| x.owners.is_empty() && x.loc.is_none())
            {
                return (self.assign(r, s, loc), false);
            }

            if let Some(r) = self.allocs.iter().position(|x| x.owners.is_empty()) {
                return (self.assign(r, s, loc), false);
            }

            panic!("register pool is empty");
        } else {
            (dst, false)
        }
    }

    // The second pass.
    // Converts static to logical registers.
    fn color(&mut self) -> Result<()> {
        self.reset_regs();
        let code = std::mem::take(&mut self.code);

        // replace all static regs with the corresponding logical ones
        for (ip, ins) in code.iter().enumerate() {
            match ins {
                Instruction::Nop => self.push(Instruction::Nop),
                Instruction::End => self.push(Instruction::End),
                Instruction::Uni { op, dst, s1 } => {
                    let s1 = self.deallocate(ip, s1);
                    let (dst, _) = self.allocate(dst, None);
                    self.push(Instruction::Uni { op, dst, s1 })
                }
                Instruction::Bi { op, dst, s1, s2 } => {
                    if self.config.is_sse() {
                        let (dst, _) = self.allocate(dst, None);
                        let s1 = self.deallocate(ip, s1);
                        let s2 = self.deallocate(ip, s2);
                        self.push(Instruction::Bi { op, dst, s1, s2 })
                    } else {
                        let s1 = self.deallocate(ip, s1);
                        let s2 = self.deallocate(ip, s2);
                        let (dst, _) = self.allocate(dst, None);
                        self.push(Instruction::Bi { op, dst, s1, s2 })
                    }
                }
                Instruction::LoadConst { dst, idx } => {
                    let (dst, _) = self.allocate(dst, None);
                    self.push(Instruction::LoadConst { dst, idx })
                }
                Instruction::Load { dst, loc } => {
                    let (dst, moved) = self.allocate(dst, Some(loc));
                    if !moved {
                        self.push(Instruction::Load { dst, loc });
                        self.locs.insert(loc, 0);
                    }
                }
                Instruction::Save { src, loc } => {
                    let src = self.deallocate(ip, src);
                    self.push(Instruction::Save { src, loc });

                    // this for loop is added due to a bug discovered while
                    // compiling Symbolica expressions (e.g., x^3 + y^3).
                    // A loc should be in only one register (the wrong registers
                    // were used). Therefore, a save invalidates all previous
                    // assignments to that loc.
                    for a in self.allocs.iter_mut() {
                        if Some(loc) == a.loc {
                            a.loc = None;
                        }
                    }

                    if let Reg::Gen(r) = src {
                        self.allocs[r as usize].loc = Some(loc);
                    }
                }
                Instruction::LoadComplex { xd, yd, loc } => {
                    let (xd, moved) = self.allocate(xd, Some(loc));
                    let (yd, _) = self.allocate(yd, Some(loc.imag()));
                    if !moved {
                        self.push(Instruction::LoadComplex { xd, yd, loc });
                        self.locs.insert(loc, 0);
                        self.locs.insert(loc.imag(), 0);
                    }
                }
                Instruction::SaveComplex { xs, ys, loc } => {
                    let xs = self.deallocate(ip, xs);
                    let ys = self.deallocate(ip, ys);
                    self.push(Instruction::SaveComplex { xs, ys, loc });

                    // this for loop is added due to a bug discovered while
                    // compiling Symbolica expressions (e.g., x^3 + y^3).
                    // A loc should be in only one register (the wrong registers
                    // were used). Therefore, a save invalidates all previous
                    // assignments to that loc.
                    for a in self.allocs.iter_mut() {
                        if Some(loc) == a.loc {
                            a.loc = None;
                        }
                        if Some(loc.imag()) == a.loc {
                            a.loc = None;
                        }
                    }

                    if let Reg::Gen(r) = xs {
                        self.allocs[r as usize].loc = Some(loc);
                    }
                    if let Reg::Gen(r) = ys {
                        self.allocs[r as usize].loc = Some(loc.imag());
                    }
                }
                Instruction::Mov { dst, s1 } => {
                    let s1 = self.deallocate(ip, s1);

                    let loc = if let Reg::Gen(r) = s1 {
                        self.allocs[r as usize].loc
                    } else {
                        None
                    };

                    let (dst, _) = self.allocate(dst, loc);
                    self.push(Instruction::Mov { dst, s1 })
                }
                Instruction::Fused { op, dst, a, b, c } => {
                    let a = self.deallocate(ip, a);
                    let b = self.deallocate(ip, b);
                    let c = self.deallocate(ip, c);
                    let (dst, _) = self.allocate(dst, None);
                    self.push(Instruction::Fused { op, dst, a, b, c });
                }
                Instruction::IfElse {
                    dst,
                    true_val,
                    false_val,
                    cond,
                } => {
                    if self.config.is_sse() {
                        let (dst, _) = self.allocate(dst, None);
                        let true_val = self.deallocate(ip, true_val);
                        let false_val = self.deallocate(ip, false_val);
                        self.push(Instruction::IfElse {
                            dst,
                            true_val,
                            false_val,
                            cond,
                        })
                    } else {
                        let true_val = self.deallocate(ip, true_val);
                        let false_val = self.deallocate(ip, false_val);
                        let (dst, _) = self.allocate(dst, None);
                        self.push(Instruction::IfElse {
                            dst,
                            true_val,
                            false_val,
                            cond,
                        })
                    }
                    self.locs.insert(cond, 0);
                }
                Instruction::Branch { .. } => {
                    if let Instruction::Branch { label } = ins.clone() {
                        self.push(Instruction::Branch { label });
                    }
                    self.reset_allocs(); // needed?
                }
                Instruction::BranchIf { .. } => {
                    if let Instruction::BranchIf {
                        cond,
                        label,
                        is_else,
                    } = ins.clone()
                    {
                        let cond = self.deallocate(ip, cond);
                        self.push(Instruction::BranchIf {
                            cond,
                            label,
                            is_else,
                        });
                        self.reset_allocs();
                    }
                }
                Instruction::Call { .. } | Instruction::Label { .. } => {
                    self.push(ins.clone());
                    self.reset_allocs();
                }
                Instruction::LoadMath { op, dst, s1, loc } => {
                    let s1 = self.deallocate(ip, s1);
                    let (dst, _) = self.allocate(dst, None);
                    self.push(Instruction::LoadMath { op, dst, s1, loc });
                    self.locs.insert(loc, 0);
                }
                Instruction::LoadConstMath { op, dst, s1, idx } => {
                    let s1 = self.deallocate(ip, s1);
                    let (dst, _) = self.allocate(dst, None);
                    self.push(Instruction::LoadConstMath { op, dst, s1, idx });
                }
                Instruction::ComplexBi {
                    op,
                    xd,
                    yd,
                    x1,
                    y1,
                    x2,
                    y2,
                } => {
                    if self.config.is_sse() {
                        let (xd, _) = self.allocate(xd, None);
                        let (yd, _) = self.allocate(yd, None);
                        let x1 = self.deallocate(ip, x1);
                        let y1 = self.deallocate(ip, y1);
                        let x2 = self.deallocate(ip, x2);
                        let y2 = self.deallocate(ip, y2);
                        self.push(Instruction::ComplexBi {
                            op,
                            xd,
                            yd,
                            x1,
                            y1,
                            x2,
                            y2,
                        })
                    } else {
                        let x1 = self.deallocate(ip, x1);
                        let y1 = self.deallocate(ip, y1);
                        let x2 = self.deallocate(ip, x2);
                        let y2 = self.deallocate(ip, y2);
                        let (xd, _) = self.allocate(xd, None);
                        let (yd, _) = self.allocate(yd, None);
                        self.push(Instruction::ComplexBi {
                            op,
                            xd,
                            yd,
                            x1,
                            y1,
                            x2,
                            y2,
                        })
                    }
                }
            }
        }

        Ok(())
    }

    // The third pass.
    // Removes unnessasary instructions.
    fn contract(&mut self) -> Result<()> {
        let code = std::mem::take(&mut self.code);
        let fixed = if self.config.is_complex() {
            (2 * SLICE_CAP + SPILL_AREA) as u32
        } else {
            (SLICE_CAP + SPILL_AREA) as u32
        };

        for ins in code.iter() {
            match ins {
                // This rule is commented out to prevent eliding Args during
                // external calls. If needed, we can restore this rule as long
                // as Args are added to self.locs.
                Instruction::Save { src, loc } => {
                    let keep = if let Loc::Stack(idx) = loc {
                        idx < fixed || self.locs.contains_key(&loc)
                    } else {
                        true
                    };

                    if keep {
                        self.push(Instruction::Save { src, loc })
                    }
                }
                Instruction::SaveComplex { xs, ys, loc } => {
                    let keep = if let Loc::Stack(idx) = loc {
                        idx < fixed
                            || (self.locs.contains_key(&loc) && self.locs.contains_key(&loc.imag()))
                    } else {
                        true
                    };

                    if keep {
                        self.push(Instruction::SaveComplex { xs, ys, loc })
                    }
                }
                Instruction::Mov { dst, s1 } => {
                    if dst != s1 {
                        self.push(Instruction::Mov { dst, s1 });
                    }
                }
                Instruction::Nop => {}
                _ => self.push(ins),
            }
        }

        Ok(())
    }
}
