jsrt/compiler/src/infer.rs

#![allow(non_snake_case)]

use std::collections::HashMap;

use crate::{
    libqjs::parse::{
        Atom, AtomRef, ClosureVar, ImmOperand, JsString, JsValue, atom_names, opcodes::OpCode,
    },
    regify::{BasicBlock, RegFn, RegInstr},
};

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub enum FnKind {
    Id(usize),
    DefineField(AtomRef),
    PutField(AtomRef),
}

#[derive(Debug, Clone, Hash, PartialEq, Eq)]
pub enum Type {
    Unknown,
    Bool,
    Int,
    Float,
    Fn(FnKind),
    Any,
}

/// Type lattice join rules:
/// Unknown * T = T
///     Any * T = Any
///       T * T = T
///       T * U = Any
fn join_ty(a: &Type, b: &Type) -> Type {
    match (a, b) {
        (Type::Unknown, t) | (t, Type::Unknown) => t.clone(),
        (Type::Any, _) | (_, Type::Any) => Type::Any,
        (a, b) if a == b => a.clone(),
        _ => Type::Any,
    }
}

#[derive(Debug, Hash, PartialEq, Eq)]
pub struct InferRequest {
    fn_id: usize,
    arg_types: Vec<Type>,
}

#[derive(Debug)]
pub struct InferResult {
    val_types: Vec<Type>,
    ret_type: Type,
}

pub type TypeCache = HashMap<InferRequest, InferResult>;

pub fn infer_types(reg_fns: &[RegFn], atoms: &[Atom]) -> TypeCache {
    // Find main fn
    let main_fn_id = reg_fns
        .iter()
        .position(|f| atom_to_str(&f.name, atoms) == "main")
        .unwrap();

    let mut cache = TypeCache::with_capacity(reg_fns.len());
    let mut worklist = vec![];

    worklist.push(InferRequest {
        fn_id: main_fn_id,
        arg_types: vec![],
    });

    while let Some(req) = worklist.pop() {
        match walk_fn(&req, reg_fns, &mut cache) {
            Ok(res) => {
                cache.insert(req, res);
            }
            Err(new_req) => {
                worklist.push(new_req);
                worklist.push(req);
            }
        }
    }

    cache
}

fn walk_fn(
    req: &InferRequest,
    reg_fns: &[RegFn],
    cache: &mut TypeCache,
) -> Result<InferResult, InferRequest> {
    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];

    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)?;

        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
            | PushI16 | PushI32 => Type::Int,

            PushTrue | PushFalse => Type::Bool,

            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!(),
                },
                _ => panic!(),
            },

            // ─────────────────────────────────── Arithmetic ───────────────────────────────────
            Add | Sub | Mul | Div | Mod | Pow => {
                match (&val_types[instr.args[0]], &val_types[instr.args[1]]) {
                    (Type::Int, Type::Int) => Type::Int,
                    (Type::Float, Type::Float)
                    | (Type::Int, Type::Float)
                    | (Type::Float, Type::Int) => Type::Float,
                    _ => Type::Any,
                }
            }

            Neg | Plus | Inc | Dec => match val_types[instr.args[0]] {
                Type::Int => Type::Int,
                Type::Float => Type::Float,
                _ => Type::Any,
            },

            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()
            }

            SetLocUninitialized => Type::Unknown,

            // ────────────────────────────────────── Returns ───────────────────────────────────
            Return | ReturnAsync | ReturnUndef => val_types[instr.args[0]].clone(),

            _ => Type::Unknown,
        };

        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 {
            for &dst in &instr.dsts {
                val_types[dst] = join_ty(&val_types[dst], &ty);
            }
        }
    }

    Ok(())
}

pub fn atom_to_str(a: &AtomRef, atoms: &[Atom]) -> String {
    match a {
        AtomRef::TaggedInt(n) => n.to_string(),
        AtomRef::Predefined(i) => atom_names::PREDEFINED_ATOMS
            .get(*i as usize)
            .copied()
            .unwrap_or("<predef?>")
            .to_string(),
        AtomRef::FileLocal(i) => atoms
            .get(*i)
            .map(atom_text)
            .unwrap_or_else(|| format!("<file_atom {i}>")),
    }
}

fn atom_text(a: &Atom) -> String {
    match a {
        Atom::String(s) | Atom::GlobalSymbol(s) | Atom::Symbol(s) => js_string(s),
        Atom::Predefined(n) => format!("<predef {n}>"),
    }
}

fn js_string(s: &JsString) -> String {
    match s {
        JsString::Latin1(b) => b.iter().map(|&c| c as char).collect(),
        JsString::UTF16(b) => {
            let u: Vec<u16> = b
                .chunks_exact(2)
                .map(|c| u16::from_le_bytes([c[0], c[1]]))
                .collect();
            String::from_utf16_lossy(&u)
        }
    }
}