Improve regify/infer

compiler/src/infer.rs

@@ -3,10 +3,9 @@
 use std::collections::HashMap;
 
 use crate::{
-    libqjs::parse::{
+    infer::Type::ReturnOf,
-        Atom, AtomRef, ClosureVar, ImmOperand, JsString, JsValue, atom_names, opcodes::OpCode,
+    libqjs::parse::{Atom, AtomRef, ImmOperand, JsString, JsValue, atom_names, opcodes::OpCode},
-    },
+    regify::RegFn,
-    regify::{BasicBlock, RegFn, RegInstr},
 };
 
 #[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() {
+        if instr.dsts.is_empty() && !is_ret {
-        let block = &f.blocks[block_id];
+            continue;
-        walk_block(&block, &req.arg_types, f, &mut val_types, &mut ret_type)?;
+        }
 
-        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
+            Dup | SetLoc => val_types[instr.args[0]].clone(),
-            | PutLocCheckInit | SetLoc0 | SetLoc1 | SetLoc2 | SetLoc3 | SetLoc | SetLoc8 => {
+
-                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);
+            ret_type = join_ty(&ret_type, &ty);
-        } else if !is_loc_read {
+        } 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 {

compiler/src/regify/basic_block.rs

@@ -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 is_terminal_instr(instr) && i + 1 < code.len() {
-            if i + 1 < code.len() {
+            blocks.push(BasicBlock::new(i + 1));
-                blocks.push(BasicBlock::new(i + 1));
-            }
         }
     }
 

compiler/src/regify/mod.rs

@@ -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(
+            walk_block(block, &flat_fn.code, &mut stack, &mut next_val, &mut out);
-                &blocks[block_idx],
-                &flat_fn.code,
-                &mut stack,
-                &mut next_val,
-                &mut out,
-            );
 
-            for &succ_idx in &blocks[block_idx].succ {
+            for succ_idx in block.succ.clone() {
-                if entry_stacks[succ_idx].is_none() {
+                let succ_stack = &mut entry_stacks[succ_idx];
-                    entry_stacks[succ_idx] = Some(stack.clone());
+
-                    worklist.push(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
+        // Locals are not renamed, local read/writes need special handling
-        let local_access = match instr.op {
+        // Handle read of a local value
-            OpCode::GetLoc0 | OpCode::PutLoc0 | OpCode::SetLoc0 => Some(0),
+        match is_local_read(instr) {
-            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 {
             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
+                continue;
-            None if instr.op == OpCode::GetLoc0Loc1 => {
-                for k in [0, 1] {
-                    stack.push(k);
-                    out_instr.dsts.push(k);
-                }
             }
+            None if instr.op == OpCode::SetLocUninitialized => {
+                out_instr.op = OpCode::Nop;
+                out_instr.imm = ImmOperand::None;
+                out[instr_idx] = out_instr;
+                continue;
+            }
+            None => (),
+        }
 
-            None => {
+        // Handle this instr's pushes
-                for _ in 0..n_push {
+        for _ in 0..n_push {
-                    let v = fresh_val();
+            let v = fresh_val();
-                    stack.push(v);
+            stack.push(v);
-                    out_instr.dsts.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()
 }