Improve infer and transforms

compiler/src/infer.rs

@@ -22,7 +22,7 @@ pub enum Type {
     Int,
     Float,
     Empty,
-    Fn(FnKind),
+    Fn(usize),
     ReturnOf((FnKind, Vec<Type>)),
     Any,
 }
@@ -42,30 +42,36 @@ fn join_ty(a: &Type, b: &Type) -> Type {
 }
 
 fn get_field(ty: &Type, field: AtomRef, cache: &mut TypeCache) -> Type {
-    match ty {
+    let mut current = ty.clone();
-        Type::ReturnOf((fn_kind, args)) => match fn_kind {
-            FnKind::Id(id) => {
-                let ret_ty = cache
-                    .get(&InferRequest {
-                        fn_id: *id,
-                        arg_types: args.clone(),
-                    })
-                    .unwrap()
-                    .ret_type
-                    .clone();
 
-                get_field(&ret_ty, field, cache)
+    loop {
-            }
+        match current {
-            FnKind::DefineField(atom) => {
+            Type::Any => return Type::Any,
-                if *atom == field {
+
-                    args[0].clone()
+            Type::ReturnOf((fn_kind, args)) => match fn_kind {
-                } else {
+                FnKind::DefineField(atom) => {
-                    Type::Undefined
+                    assert!(args.len() == 2);
+
+                    if atom == field {
+                        return args[0].clone();
+                    }
+
+                    current = args[1].clone();
                 }
-            }
+
-        },
+                FnKind::Id(id) => {
-        Type::Any => Type::Any,
+                    match cache.get(&InferRequest {
-        _ => Type::Undefined,
+                        fn_id: id,
+                        arg_types: args.clone(),
+                    }) {
+                        Some(info) => current = info.ret_type.clone(),
+                        None => panic!(),
+                    }
+                }
+            },
+
+            _ => return Type::Undefined,
+        }
     }
 }
 
@@ -97,12 +103,13 @@ pub fn infer_types(reg_fns: &[RegFn], atoms: &[Atom]) -> TypeCache {
         arg_types: vec![],
     };
 
-    walk_fn(req, reg_fns, &mut cache);
+    let res = walk_fn(&req, reg_fns, &mut cache);
+    cache.insert(req, res);
 
     cache
 }
 
-fn walk_fn(req: InferRequest, reg_fns: &[RegFn], cache: &mut TypeCache) {
+fn walk_fn(req: &InferRequest, reg_fns: &[RegFn], cache: &mut TypeCache) -> InferResult {
     use OpCode::*;
 
     let f = &reg_fns[req.fn_id];
@@ -133,7 +140,7 @@ fn walk_fn(req: InferRequest, reg_fns: &[RegFn], cache: &mut TypeCache) {
                         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)),
+                        Some(JsValue::FnById(id)) => Type::Fn(*id),
                         _ => unreachable!(),
                     },
                     _ => unreachable!(),
@@ -160,6 +167,7 @@ fn walk_fn(req: InferRequest, reg_fns: &[RegFn], cache: &mut TypeCache) {
 
                 // ────────────────────────────────── Value Movers ──────────────────────────────────
                 Dup | SetLoc | SetName => val_types[instr.args[0]].clone(),
+                Undefined => Type::Undefined,
 
                 // ──────────────────────────────────── Arguments ───────────────────────────────────
                 GetArg0 => req.arg_types.first().cloned().unwrap_or(Type::Undefined),
@@ -167,11 +175,18 @@ fn walk_fn(req: InferRequest, reg_fns: &[RegFn], cache: &mut TypeCache) {
                 GetArg2 => req.arg_types.get(2).cloned().unwrap_or(Type::Undefined),
                 GetArg3 => req.arg_types.get(3).cloned().unwrap_or(Type::Undefined),
                 GetArg => match instr.imm {
-                    ImmOperand::Arg(k) => req
+                    ImmOperand::Arg(k) => {
-                        .arg_types
+                        let k = if k == u16::MAX {
-                        .get(k as usize)
+                            req.arg_types.len() as u16 - 1
-                        .cloned()
+                        } else {
-                        .unwrap_or(Type::Undefined),
+                            k
+                        };
+
+                        req.arg_types
+                            .get(k as usize)
+                            .cloned()
+                            .unwrap_or(Type::Undefined)
+                    }
                     _ => unreachable!(),
                 },
 
@@ -199,7 +214,39 @@ fn walk_fn(req: InferRequest, reg_fns: &[RegFn], cache: &mut TypeCache) {
                 }
 
                 // ──────────────────────────────── Function Calls ──────────────────────────────────
-                Call | Call0 | Call1 | Call2 | Call3 | TailCall => todo!(),
+                Call | Call0 | Call1 | Call2 | Call3 | TailCall => {
+                    let closure = instr.args.last().unwrap();
+                    let closure = &val_types[*closure];
+                    let fn_id = get_field(closure, AtomRef::Predefined(0), cache);
+
+                    let fn_id = match fn_id {
+                        Type::Fn(id) => id,
+                        _ => unreachable!(),
+                    };
+
+                    let arg_tys = instr
+                        .args
+                        .iter()
+                        .map(|a| val_types[*a].clone())
+                        .collect::<Vec<Type>>();
+
+                    let req = InferRequest {
+                        fn_id,
+                        arg_types: arg_tys,
+                    };
+
+                    let res = cache.get(&req);
+
+                    match res {
+                        Some(res) => res.ret_type.clone(),
+                        None => {
+                            let res = walk_fn(&req, reg_fns, cache);
+                            let ret_ty = res.ret_type.clone();
+                            cache.insert(req, res);
+                            ret_ty
+                        }
+                    }
+                }
 
                 // ────────────────────────────────────── Returns ───────────────────────────────────
                 Return | ReturnAsync => val_types[instr.args[0]].clone(),
@@ -218,13 +265,10 @@ fn walk_fn(req: InferRequest, reg_fns: &[RegFn], cache: &mut TypeCache) {
         }
     }
 
-    cache.insert(
+    InferResult {
-        req,
+        val_types,
-        InferResult {
+        ret_type,
-            val_types,
+    }
-            ret_type,
-        },
-    );
 }
 
 fn atom_to_str(a: &AtomRef, atoms: &[Atom]) -> String {

compiler/src/libqjs/mod.rs

@@ -42,12 +42,10 @@ fn get_exception(ctx: *mut ffi::JSContext) -> String {
         let err = ffi::JS_GetException(ctx);
         let text = ffi::JS_ToCStringLen2(ctx, std::ptr::null_mut(), err, false);
 
-        let message = if text.is_null() {
+        if text.is_null() {
             "quickjs raised an exception it could not describe".to_string()
         } else {
             CStr::from_ptr(text).to_string_lossy().into_owned()
-        };
+        }
-
-        message
     }
 }

compiler/src/libqjs/parse.rs

@@ -518,13 +518,13 @@ pub enum ImmOperand {
 pub struct Instr {
     pub pc: usize,
     pub op: opcodes::OpCode,
-    pub imm_operands: ImmOperand,
+    pub imm: ImmOperand,
 }
 
 impl Debug for Instr {
     fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
         write!(f, "{:04} {}", self.pc, self.op.name())?;
-        match &self.imm_operands {
+        match &self.imm {
             ImmOperand::None => Ok(()),
             other => write!(f, " {other:?}"),
         }
@@ -637,7 +637,7 @@ fn parse_bytecode(code: &[u8]) -> IResult<&[u8], Vec<Instr>> {
         out.push(Instr {
             pc,
             op: info.op,
-            imm_operands: operands,
+            imm: operands,
         });
 
         pc += size;

compiler/src/main.rs

@@ -1,16 +1,17 @@
-mod declosure;
 mod flatten;
 mod infer;
 mod libqjs;
 mod regify;
+mod transform;
 
 use std::fs::File;
 use std::io::Read;
 
-use declosure::declosure;
 use flatten::flatten;
 use infer::infer_types;
+use libqjs::compile;
 use regify::regify;
+use transform::transform;
 
 fn main() {
     // Load source file
@@ -25,11 +26,11 @@ fn main() {
     file.read_to_end(&mut src).unwrap();
 
     // Compile to bytecode
-    let binary = libqjs::compile(&src, filename);
+    let binary = compile(&src, filename);
 
     // Process it
     let (flat_fns, captures, atoms) = flatten(binary);
-    let flat_fns = declosure(flat_fns, &captures);
+    let flat_fns = transform(flat_fns, &captures);
     let reg_fns = regify(flat_fns);
     let type_cache = infer_types(&reg_fns, &atoms);
 }

compiler/src/regify/basic_block.rs

@@ -1,3 +1,5 @@
+#![allow(non_snake_case)]
+
 use super::RegInstr;
 use crate::libqjs::parse::{ImmOperand, Instr, opcodes::OpCode};
 
@@ -90,7 +92,7 @@ fn is_jump_instr(instr: &Instr, code: &[Instr]) -> Option<usize> {
         return None;
     }
 
-    let target = match instr.imm_operands {
+    let target = match instr.imm {
         ImmOperand::Label { target } => code.iter().position(|x| x.pc == target).unwrap(),
         _ => unreachable!(),
     };

compiler/src/regify/mod.rs

@@ -2,7 +2,7 @@ mod basic_block;
 
 use crate::{
     flatten::FlatFn,
-    libqjs::parse::{AtomRef, ClosureVar, ImmOperand, Instr, JsValue, opcodes::OpCode},
+    libqjs::parse::{AtomRef, ImmOperand, Instr, JsValue, opcodes::OpCode},
 };
 
 pub use basic_block::BasicBlock;
@@ -109,7 +109,7 @@ fn walk_block(
         let mut out_instr = RegInstr {
             dsts: vec![],
             op: instr.op,
-            imm: instr.imm_operands,
+            imm: instr.imm,
             args: vec![],
         };
 
@@ -184,7 +184,7 @@ fn extra_pops(instr: &Instr) -> usize {
         OpCode::Call1 => 1,
         OpCode::Call2 => 2,
         OpCode::Call3 => 3,
-        _ => match instr.imm_operands {
+        _ => match instr.imm {
             ImmOperand::Npop(npop) => npop as usize,
             ImmOperand::NpopU16 { npop, .. } => npop as usize,
             _ => 0,
@@ -199,7 +199,7 @@ fn is_local_read(instr: &Instr) -> Option<usize> {
         OpCode::GetLoc2 => Some(2),
         OpCode::GetLoc3 => Some(3),
 
-        OpCode::GetLoc | OpCode::GetLoc8 | OpCode::GetLocCheck => match instr.imm_operands {
+        OpCode::GetLoc | OpCode::GetLoc8 | OpCode::GetLocCheck => match instr.imm {
             ImmOperand::Loc(k) => Some(k as usize),
             _ => unreachable!(),
         },
@@ -220,7 +220,7 @@ fn is_local_write(instr: &Instr) -> Option<usize> {
         | OpCode::PutLocCheck
         | OpCode::PutLocCheckInit
         | OpCode::SetLoc
-        | OpCode::SetLoc8 => match instr.imm_operands {
+        | OpCode::SetLoc8 => match instr.imm {
             ImmOperand::Loc(k) => Some(k as usize),
             _ => unreachable!(),
         },

compiler/src/transform/declass.rs

@@ -0,0 +1,141 @@
+use crate::libqjs::parse::{AtomRef, ImmOperand, Instr, opcodes::OpCode};
+
+#[derive(Debug)]
+struct MethodSequence {
+    fn_cpool_idx: u32,
+    name: AtomRef,
+}
+
+#[derive(Debug)]
+pub struct ClassSequence {
+    start_pc: usize,
+    ctor_fn_cpool_idx: u32,
+    methods: Vec<MethodSequence>,
+}
+
+impl ClassSequence {
+    pub fn emit_prototype_form(&self, out: &mut Vec<Instr>) {
+        out.push(Instr {
+            pc: self.start_pc,
+            op: OpCode::Fclosure8,
+            imm: ImmOperand::Const(self.ctor_fn_cpool_idx),
+        });
+
+        for m in &self.methods {
+            out.push(Instr {
+                pc: self.start_pc,
+                op: OpCode::GetField,
+                imm: ImmOperand::Atom(AtomRef::Predefined(63)), // .prototype
+            });
+            out.push(Instr {
+                pc: self.start_pc,
+                op: OpCode::Fclosure8,
+                imm: ImmOperand::Const(m.fn_cpool_idx),
+            });
+            out.push(Instr {
+                pc: self.start_pc,
+                op: OpCode::DefineField,
+                imm: ImmOperand::Atom(m.name),
+            });
+        }
+    }
+}
+
+pub fn declass(mut code: &[Instr]) -> Vec<Instr> {
+    let mut out = Vec::with_capacity(code.len());
+
+    while !code.is_empty() {
+        match match_class_seq(code) {
+            Ok((rest, seq)) => {
+                seq.emit_prototype_form(&mut out);
+                code = rest;
+            }
+            Err(rest) => {
+                let consumed = code.len() - rest.len(); // Handle partial consumption
+                let consumed = consumed.max(1);
+                out.extend_from_slice(&code[..consumed]);
+                code = &code[consumed..];
+            }
+        }
+    }
+
+    out
+}
+
+// ─────────────────────────── Class Sequence Grammar ─────────────────────────────────────
+// push_const8 #ctor ; define_class @Name                       <-- Class decl header
+// ( fclosure8 #method ; define_method @meth )*                 <-- Methods
+// ( fclosure8 #fields ; set_home_object ; put_loc ; drop )?    <-- Field initializer stuff
+//
+// match fns return Ok(advanced_instrs) or Err(original_instrs)
+
+fn match_class_seq(ins: &[Instr]) -> PResult<'_, ClassSequence> {
+    let pc = ins[0].pc;
+
+    let (ins, push_const) = match_instr(ins, |i| i.op == OpCode::PushConst8)?;
+
+    let is_derived_class = matches!(push_const.imm, ImmOperand::AtomU8 { arg, .. } if arg & 1 == 1);
+    assert!(!is_derived_class);
+
+    let (ins, _) = match_instr(ins, |i| i.op == OpCode::DefineClass)?;
+    let (ins, methods) = match_many(ins, match_method_seq);
+
+    let ctor_fn_id = match push_const.imm {
+        ImmOperand::Const(id) => id,
+        _ => unreachable!(),
+    };
+
+    let seq = ClassSequence {
+        start_pc: pc,
+        ctor_fn_cpool_idx: ctor_fn_id,
+        methods,
+    };
+
+    Ok((ins, seq))
+}
+
+// TODO: accessors are not handled
+fn match_method_seq(ins: &[Instr]) -> PResult<'_, MethodSequence> {
+    let (ins, fclosure) = match_instr(ins, |i| i.op == OpCode::Fclosure8)?;
+    let (ins, define_method) = match_instr(ins, |i| i.op == OpCode::DefineMethod)?;
+
+    let fn_id = match fclosure.imm {
+        ImmOperand::Const(id) => id,
+        _ => unreachable!(),
+    };
+
+    let name = match define_method.imm {
+        ImmOperand::AtomU8 { atom, .. } => atom,
+        _ => unreachable!(),
+    };
+
+    Ok((
+        ins,
+        MethodSequence {
+            fn_cpool_idx: fn_id,
+            name,
+        },
+    ))
+}
+
+// ───────────────────────────── Matcher Primitives ───────────────────────────────────────
+
+type PResult<'a, T> = Result<(&'a [Instr], T), &'a [Instr]>;
+
+fn match_instr(i: &[Instr], f: impl Fn(&Instr) -> bool) -> PResult<'_, &Instr> {
+    match i.first() {
+        Some(instr) if f(instr) => Ok((&i[1..], instr)),
+        _ => Err(i),
+    }
+}
+
+fn match_many<T>(mut i: &[Instr], f: impl Fn(&[Instr]) -> PResult<T>) -> (&[Instr], Vec<T>) {
+    let mut res = vec![];
+
+    while let Ok((rest, r)) = f(i) {
+        res.push(r);
+        i = rest;
+    }
+
+    (i, res)
+}

compiler/src/transform/mod.rs

@@ -1,32 +1,41 @@
+mod declass;
+
 use crate::{
     flatten::FlatFn,
     libqjs::parse::{AtomRef, ClosureVar, ImmOperand, Instr, JsValue, opcodes::OpCode},
 };
 
-pub fn declosure(flat_fns: Vec<FlatFn>, captures: &[Vec<ClosureVar>]) -> Vec<FlatFn> {
+use declass::declass;
+
+pub fn transform(flat_fns: Vec<FlatFn>, captures: &[Vec<ClosureVar>]) -> Vec<FlatFn> {
     flat_fns
         .into_iter()
         .enumerate()
-        .map(|(i, f)| declosure_one(i, f, captures))
+        .map(|(i, f)| transform_one(i, f, captures))
         .collect()
 }
 
-fn declosure_one(fn_id: usize, f: FlatFn, captures: &[Vec<ClosureVar>]) -> FlatFn {
+fn transform_one(fn_id: usize, f: FlatFn, captures: &[Vec<ClosureVar>]) -> FlatFn {
+    let declassed = declass(&f.code);
     let mut out_len = 0;
 
-    for instr in &f.code {
+    for instr in &declassed {
         if let Some(info) = is_fclosure(instr, &f, captures) {
-            out_len += 1 + (2 * info.child_captures.len()) + 1 + 1; // object; (push; define_field)*; push_fn_const; define_field
+            out_len += 1 + (2 * info.child_captures.len()) + 2; // object; (push; define_field)*; push_fn_const; define_field
-        } else if is_capture_acc(instr).is_some() {
+            continue;
-            out_len += 2;
-        } else {
-            out_len += 1;
         }
+
+        if is_capture_acc(instr).is_some() || is_tail_call(instr) {
+            out_len += 2;
+            continue;
+        }
+
+        out_len += 1;
     }
 
     let mut out = Vec::with_capacity(out_len);
 
-    for instr in &f.code {
+    for instr in &declassed {
         if replace_fclosure(instr, &f, captures, &mut out) {
             continue;
         }
@@ -35,9 +44,15 @@ fn declosure_one(fn_id: usize, f: FlatFn, captures: &[Vec<ClosureVar>]) -> FlatF
             continue;
         }
 
+        if replace_tail_call(instr, &mut out) {
+            continue;
+        }
+
         out.push(instr.clone());
     }
 
+    //debug_assert!(out.len() == out_len);
+
     FlatFn {
         name: f.name,
         arg_count: f.arg_count,
@@ -61,7 +76,7 @@ fn is_fclosure<'a>(
         return None;
     }
 
-    let cpool_idx = match instr.imm_operands {
+    let cpool_idx = match instr.imm {
         ImmOperand::Const(idx) => idx,
         _ => unreachable!(),
     };
@@ -93,33 +108,33 @@ fn replace_fclosure(
     out.push(Instr {
         pc,
         op: OpCode::Object,
-        imm_operands: ImmOperand::None,
+        imm: ImmOperand::None,
     });
 
     for c in info.child_captures {
         out.push(Instr {
             pc,
             op: OpCode::GetLoc,
-            imm_operands: ImmOperand::Loc(c.var_idx as u16),
+            imm: ImmOperand::Loc(c.var_idx as u16),
         });
 
         out.push(Instr {
             pc,
             op: OpCode::DefineField,
-            imm_operands: ImmOperand::Atom(c.name),
+            imm: ImmOperand::Atom(c.name),
         });
     }
 
     out.push(Instr {
         pc,
         op: OpCode::PushConst,
-        imm_operands: ImmOperand::Const(info.cpool_idx),
+        imm: ImmOperand::Const(info.cpool_idx),
     });
 
     out.push(Instr {
         pc,
         op: OpCode::DefineField,
-        imm_operands: ImmOperand::Atom(AtomRef::Predefined(0)),
+        imm: ImmOperand::Atom(AtomRef::Predefined(0)),
     });
 
     true
@@ -136,7 +151,7 @@ fn is_capture_acc(instr: &Instr) -> Option<CaptureAccess> {
         OpCode::GetVarRef1 => CaptureAccess::Read(1),
         OpCode::GetVarRef2 => CaptureAccess::Read(2),
         OpCode::GetVarRef3 => CaptureAccess::Read(3),
-        OpCode::GetVarRef | OpCode::GetVarRefCheck => match instr.imm_operands {
+        OpCode::GetVarRef | OpCode::GetVarRefCheck => match instr.imm {
             ImmOperand::VarRef(var) => CaptureAccess::Read(var),
             _ => unreachable!(),
         },
@@ -145,7 +160,7 @@ fn is_capture_acc(instr: &Instr) -> Option<CaptureAccess> {
         OpCode::PutVarRef1 => CaptureAccess::Write(1),
         OpCode::PutVarRef2 => CaptureAccess::Write(2),
         OpCode::PutVarRef3 => CaptureAccess::Write(3),
-        OpCode::PutVarRef | OpCode::PutVarRefCheck => match instr.imm_operands {
+        OpCode::PutVarRef | OpCode::PutVarRefCheck => match instr.imm {
             ImmOperand::VarRef(var) => CaptureAccess::Write(var),
             _ => unreachable!(),
         },
@@ -173,13 +188,13 @@ fn replace_capture_accs(instr: &Instr, my_captures: &[ClosureVar], out: &mut Vec
             out.push(Instr {
                 pc,
                 op: OpCode::GetArg,
-                imm_operands: ImmOperand::Arg(u16::MAX),
+                imm: ImmOperand::Arg(u16::MAX),
             });
 
             out.push(Instr {
                 pc,
                 op: OpCode::GetField,
-                imm_operands: ImmOperand::Atom(atom),
+                imm: ImmOperand::Atom(atom),
             });
         }
 
@@ -189,16 +204,40 @@ fn replace_capture_accs(instr: &Instr, my_captures: &[ClosureVar], out: &mut Vec
             out.push(Instr {
                 pc,
                 op: OpCode::GetArg,
-                imm_operands: ImmOperand::Arg(u16::MAX),
+                imm: ImmOperand::Arg(u16::MAX),
             });
 
             out.push(Instr {
                 pc,
                 op: OpCode::PutField,
-                imm_operands: ImmOperand::Atom(atom),
+                imm: ImmOperand::Atom(atom),
             });
         }
     }
 
     true
 }
+
+fn is_tail_call(instr: &Instr) -> bool {
+    instr.op == OpCode::TailCall
+}
+
+fn replace_tail_call(instr: &Instr, out: &mut Vec<Instr>) -> bool {
+    if !is_tail_call(instr) {
+        return false;
+    }
+
+    out.push(Instr {
+        pc: instr.pc,
+        op: OpCode::Call,
+        imm: instr.imm,
+    });
+
+    out.push(Instr {
+        pc: instr.pc,
+        op: OpCode::Return,
+        imm: ImmOperand::None,
+    });
+
+    true
+}