jsrt/compiler/src/regify/walk.rs

//! This module converts one basic block's instrs from stack to register form using a simulated stack
//!
//! The output reg form instrs are structured as:
//! `dsts = op(imm, args)`
//!
//! Every instr's pushes allocates new register operands which are pushed on the stack and form the
//! output instr's `dsts`, the instr's pops pop operands from the stack and form the output instr's
//! `args`
//!
//! Each instr's `op` and `imm` is preserved except for some special cases
//!
//! `set_loc` is henceforth used to move data from one operand (local, arg, reg) to another

#![allow(non_snake_case)]

use super::{BasicBlock, Operand, RegInstr};
use crate::libqjs::parse::{ImmOperand, Instr, opcodes::OpCode};

/// Walk a basic block's instructions, converting each instr to register form
pub fn walk_block(
    block: &BasicBlock,
    code: &[Instr],
    stack: &mut Vec<Operand>,
    next_reg: &mut usize,
    out: &mut [RegInstr],
) {
    let mut fresh_reg = || {
        let r = *next_reg;
        *next_reg += 1;
        Operand::Reg(r)
    };

    for instr_idx in block.start..block.end {
        let instr = &code[instr_idx];
        let out_instr = &mut out[instr_idx];

        if try_handle_named_read(instr, stack, &mut fresh_reg, out_instr) {
            continue;
        }

        if try_handle_named_write(instr, stack, out_instr) {
            continue;
        }

        if try_handle_misc(instr, stack, out_instr) {
            continue;
        }

        handle_normal_case(instr, stack, &mut fresh_reg, out_instr);
    }
}

/// Handles a "named read" i.e. a read of a local or arg (and not a stack value)
///
/// We simulate pushing the operand to stack, and replace the instr with `SetLoc`, which
/// is our catch-all for a data-move instr
fn try_handle_named_read(
    instr: &Instr,
    stack: &mut Vec<Operand>,
    mut fresh_reg: impl FnMut() -> Operand,
    out: &mut RegInstr,
) -> bool {
    use OpCode::*;

    let arg = match instr.op {
        GetLoc0 => Operand::Local(0),
        GetLoc1 => Operand::Local(1),
        GetLoc2 => Operand::Local(2),
        GetLoc3 => Operand::Local(3),

        GetLoc | GetLoc8 | GetLocCheck => match instr.imm {
            ImmOperand::Loc(k) => Operand::Local(k as usize),
            _ => unreachable!(),
        },

        GetArg0 => Operand::Arg(0),
        GetArg1 => Operand::Arg(1),
        GetArg2 => Operand::Arg(2),
        GetArg3 => Operand::Arg(3),

        GetArg => match instr.imm {
            ImmOperand::Arg(k) => Operand::Arg(k as usize),
            _ => unreachable!(),
        },

        _ => return false,
    };

    let dst = fresh_reg();

    // Converted SetLoc, reading from arg, writing to dst
    *out = RegInstr {
        dsts: vec![dst],
        op: SetLoc,
        imm: ImmOperand::None,
        args: vec![arg],
    };

    // Simulate stack push
    stack.push(dst);
    true
}

/// Handles a "named write" i.e. a write of a local or arg (and not a stack value)
///
/// We simulate popping the operand off stack, and replace the instr with `SetLoc`, which
/// is our catch-all for a data-move instr
///
/// A distinction needs to be made between `Set*` and `Put*` since `Set*` instrs leave
/// the value-to-write on the stack and `Put*` pops it off
fn try_handle_named_write(instr: &Instr, stack: &mut Vec<Operand>, out: &mut RegInstr) -> bool {
    use OpCode::*;

    let (dst, preserved) = match instr.op {
        PutLoc0 => (Operand::Local(0), false),
        PutLoc1 => (Operand::Local(1), false),
        PutLoc2 => (Operand::Local(2), false),
        PutLoc3 => (Operand::Local(3), false),

        PutLoc | PutLoc8 | PutLocCheck | PutLocCheckInit => match instr.imm {
            ImmOperand::Loc(k) => (Operand::Local(k as usize), false),
            _ => unreachable!(),
        },

        PutArg0 => (Operand::Arg(0), false),
        PutArg1 => (Operand::Arg(1), false),
        PutArg2 => (Operand::Arg(2), false),
        PutArg3 => (Operand::Arg(3), false),

        SetLoc0 => (Operand::Local(0), true),
        SetLoc1 => (Operand::Local(1), true),
        SetLoc2 => (Operand::Local(2), true),
        SetLoc3 => (Operand::Local(3), true),

        SetLoc | SetLoc8 => match instr.imm {
            ImmOperand::Loc(k) => (Operand::Local(k as usize), true),
            _ => unreachable!(),
        },

        SetArg0 => (Operand::Arg(0), true),
        SetArg1 => (Operand::Arg(1), true),
        SetArg2 => (Operand::Arg(2), true),
        SetArg3 => (Operand::Arg(3), true),

        _ => return false,
    };

    // Simulate stack pop or peek, depending on instr
    let arg = if preserved {
        stack.last().unwrap().clone()
    } else {
        stack.pop().unwrap()
    };

    // Converted SetLoc, reading from arg, writing to dst
    *out = RegInstr {
        dsts: vec![dst],
        op: SetLoc,
        imm: ImmOperand::None,
        args: vec![arg],
    };

    true
}

fn try_handle_misc(instr: &Instr, stack: &mut Vec<Operand>, out_instr: &mut RegInstr) -> bool {
    use OpCode::*;

    *out_instr = match instr.op {
        // Duplicate stack top, out_instr becomes `nop`
        Dup => {
            let top = stack[stack.len() - 1].clone();
            stack.push(top);

            RegInstr {
                dsts: vec![],
                op: Nop,
                imm: ImmOperand::None,
                args: vec![],
            }
        }

        // Duplicate the top 2 elems in stack, out_instr becomes `nop`
        Dup2 => {
            let a = stack[stack.len() - 1].clone();
            let b = stack[stack.len() - 2].clone();

            stack.push(b);
            stack.push(a);

            RegInstr {
                dsts: vec![],
                op: Nop,
                imm: ImmOperand::None,
                args: vec![],
            }
        }

        // Swap top 2 elems in stack, out_instr becomes `nop`
        Swap => {
            let a = stack.pop().unwrap();
            let b = stack.pop().unwrap();

            stack.push(a);
            stack.push(b);

            RegInstr {
                dsts: vec![],
                op: Nop,
                imm: ImmOperand::None,
                args: vec![],
            }
        }

        _ => return false,
    };

    true
}

/// Handle a normal instruction with simulated stack pushes/pops
fn handle_normal_case(
    instr: &Instr,
    stack: &mut Vec<Operand>,
    mut fresh_reg: impl FnMut() -> Operand,
    out_instr: &mut RegInstr,
) {
    // op and imm carried as-is
    *out_instr = RegInstr {
        dsts: vec![],
        op: instr.op,
        imm: instr.imm,
        args: vec![],
    };

    // Some instrs have extra pops not reflected in info.n_pop
    let extra_pops = match instr.imm {
        ImmOperand::Npop(npop) => npop as usize,
        ImmOperand::NpopU16 { npop, .. } => npop as usize,
        _ => 0,
    };

    let info = instr.op.info();
    let n_pop = info.n_pop as usize + extra_pops;
    let n_push = info.n_push as usize;

    // Simulate this instr's pops
    for _ in 0..n_pop {
        let x = stack.pop().unwrap();
        out_instr.args.push(x);
    }

    // Simulate this instr's pushes
    for _ in 0..n_push {
        let x = fresh_reg();
        stack.push(x);
        out_instr.dsts.push(x);
    }
}