Improve infer and transforms

compiler/src/infer.rs

@@ -12,7 +12,6 @@ use crate::{
 pub enum FnKind {
     Id(usize),
     DefineField(AtomRef),
-    PutField(AtomRef),
 }
 
 #[derive(Debug, Clone, Hash, PartialEq, Eq)]
@@ -23,7 +22,7 @@ pub enum Type {
     Int,
     Float,
     Empty,
-    Fn(FnKind),
+    Fn(usize),
     ReturnOf((FnKind, Vec<Type>)),
     Any,
 }
@@ -42,16 +41,50 @@ fn join_ty(a: &Type, b: &Type) -> Type {
     }
 }
 
+fn get_field(ty: &Type, field: AtomRef, cache: &mut TypeCache) -> Type {
+    let mut current = ty.clone();
+
+    loop {
+        match current {
+            Type::Any => return Type::Any,
+
+            Type::ReturnOf((fn_kind, args)) => match fn_kind {
+                FnKind::DefineField(atom) => {
+                    assert!(args.len() == 2);
+
+                    if atom == field {
+                        return args[0].clone();
+                    }
+
+                    current = args[1].clone();
+                }
+
+                FnKind::Id(id) => {
+                    match cache.get(&InferRequest {
+                        fn_id: id,
+                        arg_types: args.clone(),
+                    }) {
+                        Some(info) => current = info.ret_type.clone(),
+                        None => panic!(),
+                    }
+                }
+            },
+
+            _ => return Type::Undefined,
+        }
+    }
+}
+
 #[derive(Debug, Hash, PartialEq, Eq)]
 pub struct InferRequest {
-    fn_id: usize,
-    arg_types: Vec<Type>,
+    pub fn_id: usize,
+    pub arg_types: Vec<Type>,
 }
 
 #[derive(Debug)]
 pub struct InferResult {
-    val_types: Vec<Type>,
-    ret_type: Type,
+    pub val_types: Vec<Type>,
+    pub ret_type: Type,
 }
 
 pub type TypeCache = HashMap<InferRequest, InferResult>;
@@ -64,139 +97,181 @@ pub fn infer_types(reg_fns: &[RegFn], atoms: &[Atom]) -> TypeCache {
         .unwrap();
 
     let mut cache = TypeCache::with_capacity(reg_fns.len());
-    let mut worklist = vec![];
 
-    worklist.push(InferRequest {
+    let req = 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);
-            }
-        }
-    }
+    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,
-) -> Result<InferResult, InferRequest> {
+fn walk_fn(req: &InferRequest, reg_fns: &[RegFn], cache: &mut TypeCache) -> InferResult {
     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;
 
-    for instr in &f.code {
-        let is_ret = matches!(instr.op, Return | ReturnAsync | ReturnUndef);
+    for block in &f.blocks {
+        let code = &f.code[block.start..block.end];
+        let code = code.iter().chain(&block.postfix);
 
-        if instr.dsts.is_empty() && !is_ret {
-            continue;
-        }
+        for instr in code {
+            let is_ret = matches!(instr.op, Return | ReturnAsync | ReturnUndef);
 
-        let ty = match instr.op {
-            // ──────────────────────────────────── Literals ────────────────────────────────────
-            PushMinus1 | Push0 | Push1 | Push2 | Push3 | Push4 | Push5 | Push6 | Push7 | PushI8
-            | PushI16 | PushI32 => Type::Int,
+            if instr.dsts.is_empty() && !is_ret {
+                continue;
+            }
 
-            PushTrue | PushFalse => Type::Bool,
+            // TODO: Some of these ops may not every show up by this point
+            let ty = match instr.op {
+                // ──────────────────────────────────── Literals ────────────────────────────────────
+                PushMinus1 | Push0 | Push1 | Push2 | Push3 | Push4 | Push5 | Push6 | Push7
+                | PushI8 | PushI16 | PushI32 => Type::Int,
 
-            // 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)),
+                PushTrue | PushFalse => Type::Bool,
+
+                PushConst | PushConst8 => 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(*id),
+                        _ => unreachable!(),
+                    },
                     _ => unreachable!(),
                 },
-                _ => unreachable!(),
-            },
 
-            // ─────────────────────────────────── 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,
+                // ─────────────────────────────────── Arithmetic ───────────────────────────────────
+                Add | Sub | Mul | 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,
+                },
+                Div => Type::Float,
 
-            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,
 
-            Lt | Lte | Gt | Gte | Eq | Neq | StrictEq | StrictNeq | Not => Type::Bool,
+                // ────────────────────────────────── Value Movers ──────────────────────────────────
+                Dup | SetLoc | SetName => val_types[instr.args[0]].clone(),
+                Undefined => Type::Undefined,
 
-            // ────────────────────────────────── Value Movers ──────────────────────────────────
-            Dup | SetLoc => val_types[instr.args[0]].clone(),
+                // ──────────────────────────────────── Arguments ───────────────────────────────────
+                GetArg0 => req.arg_types.first().cloned().unwrap_or(Type::Undefined),
+                GetArg1 => req.arg_types.get(1).cloned().unwrap_or(Type::Undefined),
+                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) => {
+                        let k = if k == u16::MAX {
+                            req.arg_types.len() as u16 - 1
+                        } else {
+                            k
+                        };
 
-            // ──────────────────────────────────── 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,
+                        req.arg_types
+                            .get(k as usize)
+                            .cloned()
+                            .unwrap_or(Type::Undefined)
+                    }
                     _ => unreachable!(),
-                };
+                },
 
-                let arg0_ty = val_types[instr.args[0]].clone();
-                let arg1_ty = val_types[instr.args[1]].clone();
+                // ──────────────────────────────────── Objects ─────────────────────────────────────
+                Object => Type::Empty,
+                DefineField => {
+                    let field = match instr.imm {
+                        ImmOperand::Atom(atom) => atom,
+                        _ => unreachable!(),
+                    };
 
-                ReturnOf((FnKind::DefineField(atom), vec![arg0_ty, arg1_ty]))
-            }
+                    let arg0_ty = val_types[instr.args[0]].clone();
+                    let arg1_ty = val_types[instr.args[1]].clone();
 
-            // ────────────────────────────────────── Returns ───────────────────────────────────
-            Return | ReturnAsync => val_types[instr.args[0]].clone(),
-            ReturnUndef => Type::Empty,
+                    ReturnOf((FnKind::DefineField(field), vec![arg0_ty, arg1_ty]))
+                }
+                GetField => {
+                    let field = match instr.imm {
+                        ImmOperand::Atom(atom) => atom,
+                        _ => unreachable!(),
+                    };
 
-            _ => unreachable!(),
-        };
+                    let ty = val_types[instr.args[0]].clone();
+                    get_field(&ty, field, cache)
+                }
 
-        if is_ret {
-            ret_type = join_ty(&ret_type, &ty);
-        } else {
-            for &dst in &instr.dsts {
-                val_types[dst] = join_ty(&val_types[dst], &ty);
+                // ──────────────────────────────── Function Calls ──────────────────────────────────
+                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(),
+                ReturnUndef => Type::Empty,
+
+                _ => unreachable!("{instr:?}"),
+            };
+
+            if is_ret {
+                ret_type = join_ty(&ret_type, &ty);
+            } else {
+                for &dst in &instr.dsts {
+                    val_types[dst] = join_ty(&val_types[dst], &ty);
+                }
             }
         }
     }
 
-    Ok(InferResult {
+    InferResult {
         val_types,
         ret_type,
-    })
+    }
 }
 
-pub fn atom_to_str(a: &AtomRef, atoms: &[Atom]) -> String {
+fn atom_to_str(a: &AtomRef, atoms: &[Atom]) -> String {
     match a {
         AtomRef::TaggedInt(n) => n.to_string(),
         AtomRef::Predefined(i) => atom_names::PREDEFINED_ATOMS

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
-        } else if is_capture_acc(instr).is_some() {
-            out_len += 2;
-        } else {
-            out_len += 1;
+            out_len += 1 + (2 * info.child_captures.len()) + 2; // object; (push; define_field)*; push_fn_const; define_field
+            continue;
         }
+
+        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
+}