Improve regify/infer

This commit is contained in:
UditDey
2026-09-12 12:32:23 +00:00
parent 465d9095bb
commit d026165e7a
3 changed files with 176 additions and 123 deletions
+54 -53
View File
@@ -3,10 +3,9 @@
use std::collections::HashMap;
use crate::{
libqjs::parse::{
Atom, AtomRef, ClosureVar, ImmOperand, JsString, JsValue, atom_names, opcodes::OpCode,
},
regify::{BasicBlock, RegFn, RegInstr},
infer::Type::ReturnOf,
libqjs::parse::{Atom, AtomRef, ImmOperand, JsString, JsValue, atom_names, opcodes::OpCode},
regify::RegFn,
};
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
@@ -19,10 +18,13 @@ pub enum FnKind {
#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub enum Type {
Unknown,
Undefined,
Bool,
Int,
Float,
Empty,
Fn(FnKind),
ReturnOf((FnKind, Vec<Type>)),
Any,
}
@@ -89,35 +91,19 @@ fn walk_fn(
reg_fns: &[RegFn],
cache: &mut TypeCache,
) -> Result<InferResult, InferRequest> {
use OpCode::*;
let f = &reg_fns[req.fn_id];
let mut val_types = vec![Type::Unknown; f.val_count];
let mut ret_type = Type::Unknown;
let mut worklist = vec![0];
for instr in &f.code {
let is_ret = matches!(instr.op, Return | ReturnAsync | ReturnUndef);
while let Some(block_id) = worklist.pop() {
let block = &f.blocks[block_id];
walk_block(&block, &req.arg_types, f, &mut val_types, &mut ret_type)?;
if instr.dsts.is_empty() && !is_ret {
continue;
}
worklist.extend(block.succ.iter());
}
Ok(InferResult {
val_types,
ret_type,
})
}
fn walk_block(
block: &BasicBlock,
arg_types: &[Type],
f: &RegFn,
val_types: &mut [Type],
ret_type: &mut Type,
) -> Result<(), InferRequest> {
use OpCode::*;
for instr in &f.code[block.start..block.end] {
let ty = match instr.op {
// ──────────────────────────────────── Literals ────────────────────────────────────
PushMinus1 | Push0 | Push1 | Push2 | Push3 | Push4 | Push5 | Push6 | Push7 | PushI8
@@ -125,15 +111,16 @@ fn walk_block(
PushTrue | PushFalse => Type::Bool,
// TODO: Fclosure will disappear soon
PushConst | PushConst8 | Fclosure | Fclosure8 => match instr.imm {
ImmOperand::Const(k) => match f.cpool.get(k as usize) {
Some(JsValue::Int(_)) => Type::Int,
Some(JsValue::Float(_)) => Type::Float,
Some(JsValue::BoolTrue) | Some(JsValue::BoolFalse) => Type::Bool,
Some(&JsValue::FnById(id)) => Type::Fn(FnKind::Id(id)),
_ => panic!(),
_ => unreachable!(),
},
_ => panic!(),
_ => unreachable!(),
},
// ─────────────────────────────────── Arithmetic ───────────────────────────────────
@@ -156,43 +143,57 @@ fn walk_block(
Lt | Lte | Gt | Gte | Eq | Neq | StrictEq | StrictNeq | Not => Type::Bool,
// ────────────────────────────────── Value Movers ──────────────────────────────────
Dup | PutLoc0 | PutLoc1 | PutLoc2 | PutLoc3 | PutLoc | PutLoc8 | PutLocCheck
| PutLocCheckInit | SetLoc0 | SetLoc1 | SetLoc2 | SetLoc3 | SetLoc | SetLoc8 => {
val_types[instr.args[0]].clone()
Dup | SetLoc => val_types[instr.args[0]].clone(),
// ──────────────────────────────────── Arguments ───────────────────────────────────
OpCode::GetArg0 => req.arg_types.first().cloned().unwrap_or(Type::Undefined),
OpCode::GetArg1 => req.arg_types.get(1).cloned().unwrap_or(Type::Undefined),
OpCode::GetArg2 => req.arg_types.get(2).cloned().unwrap_or(Type::Undefined),
OpCode::GetArg3 => req.arg_types.get(3).cloned().unwrap_or(Type::Undefined),
OpCode::GetArg => match instr.imm {
ImmOperand::Arg(k) => req
.arg_types
.get(k as usize)
.cloned()
.unwrap_or(Type::Undefined),
_ => unreachable!(),
},
// ──────────────────────────────────── Objects ─────────────────────────────────────
OpCode::Object => Type::Empty,
OpCode::DefineField => {
let atom = match instr.imm {
ImmOperand::Atom(atom) => atom,
_ => unreachable!(),
};
let arg0_ty = val_types[instr.args[0]].clone();
let arg1_ty = val_types[instr.args[1]].clone();
ReturnOf((FnKind::DefineField(atom), vec![arg0_ty, arg1_ty]))
}
SetLocUninitialized => Type::Unknown,
// ────────────────────────────────────── Returns ───────────────────────────────────
Return | ReturnAsync | ReturnUndef => val_types[instr.args[0]].clone(),
Return | ReturnAsync => val_types[instr.args[0]].clone(),
ReturnUndef => Type::Empty,
_ => Type::Unknown,
_ => unreachable!(),
};
let is_ret = matches!(instr.op, Return | ReturnAsync | ReturnUndef);
let is_loc_read = matches!(
instr.op,
OpCode::GetLoc0
| OpCode::GetLoc1
| OpCode::GetLoc2
| OpCode::GetLoc3
| OpCode::GetLoc0Loc1
| OpCode::GetLoc
| OpCode::GetLoc8
| OpCode::GetLocCheck
);
// TODO: is_loc_read logic is ungraceful
if is_ret {
*ret_type = join_ty(ret_type, &ty);
} else if !is_loc_read {
ret_type = join_ty(&ret_type, &ty);
} else {
for &dst in &instr.dsts {
val_types[dst] = join_ty(&val_types[dst], &ty);
}
}
}
Ok(())
Ok(InferResult {
val_types,
ret_type,
})
}
pub fn atom_to_str(a: &AtomRef, atoms: &[Atom]) -> String {
+5 -4
View File
@@ -1,3 +1,4 @@
use super::RegInstr;
use crate::libqjs::parse::{ImmOperand, Instr, opcodes::OpCode};
type BlockId = usize;
@@ -7,6 +8,7 @@ pub struct BasicBlock {
pub start: usize,
pub end: usize,
pub succ: Vec<BlockId>,
pub postfix: Vec<RegInstr>,
}
impl BasicBlock {
@@ -15,6 +17,7 @@ impl BasicBlock {
start,
end: 0,
succ: vec![],
postfix: vec![],
}
}
}
@@ -35,10 +38,8 @@ pub fn find_basic_blocks(code: &[Instr]) -> Vec<BasicBlock> {
}
// If terminal, next instr (if exists) is a start
if is_terminal_instr(instr) {
if i + 1 < code.len() {
blocks.push(BasicBlock::new(i + 1));
}
if is_terminal_instr(instr) && i + 1 < code.len() {
blocks.push(BasicBlock::new(i + 1));
}
}
+117 -66
View File
@@ -50,7 +50,7 @@ pub struct RegFn {
impl RegFn {
fn from_flat(flat_fn: FlatFn) -> Self {
let blocks = find_basic_blocks(&flat_fn.code);
let mut blocks = find_basic_blocks(&flat_fn.code);
let mut out = vec![RegInstr::nop(); flat_fn.code.len()];
let mut next_val = flat_fn.var_count; // Reserve values for locals
@@ -62,20 +62,20 @@ impl RegFn {
worklist.push(0);
while let Some(block_idx) = worklist.pop() {
let block = &mut blocks[block_idx];
let mut stack: Vec<usize> = entry_stacks[block_idx].clone().unwrap_or_default();
walk_block(
&blocks[block_idx],
&flat_fn.code,
&mut stack,
&mut next_val,
&mut out,
);
walk_block(block, &flat_fn.code, &mut stack, &mut next_val, &mut out);
for &succ_idx in &blocks[block_idx].succ {
if entry_stacks[succ_idx].is_none() {
entry_stacks[succ_idx] = Some(stack.clone());
worklist.push(succ_idx);
for succ_idx in block.succ.clone() {
let succ_stack = &mut entry_stacks[succ_idx];
match succ_stack {
Some(succ_stack) => add_block_postfix(block, &stack, succ_stack),
None => {
*succ_stack = Some(stack.clone());
worklist.push(succ_idx);
}
}
}
}
@@ -115,21 +115,8 @@ fn walk_block(
args: vec![],
};
// Some instrs have extra pops not reflected in info.n_pop
let extra_pop = match instr.op {
OpCode::Call0 => 0,
OpCode::Call1 => 1,
OpCode::Call2 => 2,
OpCode::Call3 => 3,
_ => match instr.imm_operands {
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_pop;
let n_pop = info.n_pop as usize + extra_pops(instr);
let n_push = info.n_push as usize;
// Handle this instr's pops
@@ -138,61 +125,125 @@ fn walk_block(
out_instr.args.push(v);
}
// Locals are not renamed, needs special handling
let local_access = match instr.op {
OpCode::GetLoc0 | OpCode::PutLoc0 | OpCode::SetLoc0 => Some(0),
OpCode::GetLoc1 | OpCode::PutLoc1 | OpCode::SetLoc1 => Some(1),
OpCode::GetLoc2 | OpCode::PutLoc2 | OpCode::SetLoc2 => Some(2),
OpCode::GetLoc3 | OpCode::PutLoc3 | OpCode::SetLoc3 => Some(3),
OpCode::GetLoc
| OpCode::GetLoc8
| OpCode::GetLocCheck
| OpCode::PutLoc
| OpCode::PutLoc8
| OpCode::PutLocCheck
| OpCode::PutLocCheckInit
| OpCode::SetLoc
| OpCode::SetLoc8
| OpCode::SetLocUninitialized => match instr.imm_operands {
ImmOperand::Loc(k) => Some(k as usize),
_ => None,
},
_ => None,
};
// Handle this instr's pushes, handling local accesses separately
match local_access {
// Locals are not renamed, local read/writes need special handling
// Handle read of a local value
match is_local_read(instr) {
Some(loc) => {
stack.push(loc);
out_instr.op = OpCode::Nop;
out_instr.imm = ImmOperand::None;
out[instr_idx] = out_instr;
continue;
}
None if instr.op == OpCode::GetLoc0Loc1 => {
stack.push(0);
stack.push(1);
out_instr.op = OpCode::Nop;
out_instr.imm = ImmOperand::None;
out[instr_idx] = out_instr;
continue;
}
None => (),
}
// Handle write of a local value
match is_local_write(instr) {
Some(loc) => {
out_instr.op = OpCode::SetLoc;
out_instr.imm = ImmOperand::None;
out_instr.dsts.push(loc);
out[instr_idx] = out_instr;
if n_push == 1 {
stack.push(loc);
}
}
// The one op that reads two locals at once
None if instr.op == OpCode::GetLoc0Loc1 => {
for k in [0, 1] {
stack.push(k);
out_instr.dsts.push(k);
}
continue;
}
None if instr.op == OpCode::SetLocUninitialized => {
out_instr.op = OpCode::Nop;
out_instr.imm = ImmOperand::None;
out[instr_idx] = out_instr;
continue;
}
None => (),
}
None => {
for _ in 0..n_push {
let v = fresh_val();
stack.push(v);
out_instr.dsts.push(v);
}
}
// Handle this instr's pushes
for _ in 0..n_push {
let v = fresh_val();
stack.push(v);
out_instr.dsts.push(v);
}
out[instr_idx] = out_instr;
}
}
fn extra_pops(instr: &Instr) -> usize {
match instr.op {
OpCode::Call0 => 0,
OpCode::Call1 => 1,
OpCode::Call2 => 2,
OpCode::Call3 => 3,
_ => match instr.imm_operands {
ImmOperand::Npop(npop) => npop as usize,
ImmOperand::NpopU16 { npop, .. } => npop as usize,
_ => 0,
},
}
}
fn is_local_read(instr: &Instr) -> Option<usize> {
match instr.op {
OpCode::GetLoc0 => Some(0),
OpCode::GetLoc1 => Some(1),
OpCode::GetLoc2 => Some(2),
OpCode::GetLoc3 => Some(3),
OpCode::GetLoc | OpCode::GetLoc8 | OpCode::GetLocCheck => match instr.imm_operands {
ImmOperand::Loc(k) => Some(k as usize),
_ => unreachable!(),
},
_ => None,
}
}
fn is_local_write(instr: &Instr) -> Option<usize> {
match instr.op {
OpCode::PutLoc0 | OpCode::SetLoc0 => Some(0),
OpCode::PutLoc1 | OpCode::SetLoc1 => Some(1),
OpCode::PutLoc2 | OpCode::SetLoc2 => Some(2),
OpCode::PutLoc3 | OpCode::SetLoc3 => Some(3),
OpCode::PutLoc
| OpCode::PutLoc8
| OpCode::PutLocCheck
| OpCode::PutLocCheckInit
| OpCode::SetLoc
| OpCode::SetLoc8 => match instr.imm_operands {
ImmOperand::Loc(k) => Some(k as usize),
_ => unreachable!(),
},
_ => None,
}
}
fn add_block_postfix(block: &mut BasicBlock, stack: &[usize], succ_stack: &[usize]) {
let overlap = stack.iter().rev().zip(succ_stack.iter().rev());
for (val, succ_val) in overlap {
block.postfix.push(RegInstr {
dsts: vec![*succ_val],
op: OpCode::SetLoc,
imm: ImmOperand::None,
args: vec![*val],
});
}
}
pub fn regify(flat_fns: Vec<FlatFn>) -> Vec<RegFn> {
flat_fns.into_iter().map(RegFn::from_flat).collect()
}