New declosure impl

compiler/src/declosure.rs

@@ -0,0 +1,204 @@
+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> {
+    flat_fns
+        .into_iter()
+        .enumerate()
+        .map(|(i, f)| declosure_one(i, f, captures))
+        .collect()
+}
+
+fn declosure_one(fn_id: usize, f: FlatFn, captures: &[Vec<ClosureVar>]) -> FlatFn {
+    let mut out_len = 0;
+
+    for instr in &f.code {
+        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;
+        }
+    }
+
+    let mut out = Vec::with_capacity(out_len);
+
+    for instr in &f.code {
+        if replace_fclosure(instr, &f, captures, &mut out) {
+            continue;
+        }
+
+        if replace_capture_accs(instr, &captures[fn_id], &mut out) {
+            continue;
+        }
+
+        out.push(instr.clone());
+    }
+
+    FlatFn {
+        name: f.name,
+        arg_count: f.arg_count,
+        var_count: f.var_count,
+        cpool: f.cpool,
+        code: out,
+    }
+}
+
+struct FClosureInfo<'a> {
+    cpool_idx: u32,
+    child_captures: &'a [ClosureVar],
+}
+
+fn is_fclosure<'a>(
+    instr: &Instr,
+    f: &FlatFn,
+    captures: &'a [Vec<ClosureVar>],
+) -> Option<FClosureInfo<'a>> {
+    if instr.op != OpCode::Fclosure8 {
+        return None;
+    }
+
+    let cpool_idx = match instr.imm_operands {
+        ImmOperand::Const(idx) => idx,
+        _ => unreachable!(),
+    };
+
+    let child_fn_id = match f.cpool[cpool_idx as usize] {
+        JsValue::FnById(id) => id,
+        _ => unreachable!(),
+    };
+
+    Some(FClosureInfo {
+        cpool_idx,
+        child_captures: &captures[child_fn_id],
+    })
+}
+
+fn replace_fclosure(
+    instr: &Instr,
+    f: &FlatFn,
+    captures: &[Vec<ClosureVar>],
+    out: &mut Vec<Instr>,
+) -> bool {
+    let info = match is_fclosure(instr, f, captures) {
+        Some(info) => info,
+        None => return false,
+    };
+
+    let pc = instr.pc;
+
+    out.push(Instr {
+        pc,
+        op: OpCode::Object,
+        imm_operands: ImmOperand::None,
+    });
+
+    for c in info.child_captures {
+        out.push(Instr {
+            pc,
+            op: OpCode::GetLoc,
+            imm_operands: ImmOperand::Loc(c.var_idx as u16),
+        });
+
+        out.push(Instr {
+            pc,
+            op: OpCode::DefineField,
+            imm_operands: ImmOperand::Atom(c.name),
+        });
+    }
+
+    out.push(Instr {
+        pc,
+        op: OpCode::PushConst,
+        imm_operands: ImmOperand::Const(info.cpool_idx),
+    });
+
+    out.push(Instr {
+        pc,
+        op: OpCode::DefineField,
+        imm_operands: ImmOperand::Atom(AtomRef::Predefined(0)),
+    });
+
+    true
+}
+
+enum CaptureAccess {
+    Read(u16),
+    Write(u16),
+}
+
+fn is_capture_acc(instr: &Instr) -> Option<CaptureAccess> {
+    let acc = match instr.op {
+        OpCode::GetVarRef0 => CaptureAccess::Read(0),
+        OpCode::GetVarRef1 => CaptureAccess::Read(1),
+        OpCode::GetVarRef2 => CaptureAccess::Read(2),
+        OpCode::GetVarRef3 => CaptureAccess::Read(3),
+        OpCode::GetVarRef | OpCode::GetVarRefCheck => match instr.imm_operands {
+            ImmOperand::VarRef(var) => CaptureAccess::Read(var),
+            _ => unreachable!(),
+        },
+
+        OpCode::PutVarRef0 => CaptureAccess::Write(0),
+        OpCode::PutVarRef1 => CaptureAccess::Write(1),
+        OpCode::PutVarRef2 => CaptureAccess::Write(2),
+        OpCode::PutVarRef3 => CaptureAccess::Write(3),
+        OpCode::PutVarRef | OpCode::PutVarRefCheck => match instr.imm_operands {
+            ImmOperand::VarRef(var) => CaptureAccess::Write(var),
+            _ => unreachable!(),
+        },
+
+        _ => return None,
+    };
+
+    Some(acc)
+}
+
+fn replace_capture_accs(instr: &Instr, my_captures: &[ClosureVar], out: &mut Vec<Instr>) -> bool {
+    let acc = is_capture_acc(instr);
+
+    let acc = match acc {
+        Some(acc) => acc,
+        None => return false,
+    };
+
+    let pc = instr.pc;
+
+    match acc {
+        CaptureAccess::Read(read_var) => {
+            let atom = my_captures[read_var as usize].name;
+
+            out.push(Instr {
+                pc,
+                op: OpCode::GetArg,
+                imm_operands: ImmOperand::Arg(u16::MAX),
+            });
+
+            out.push(Instr {
+                pc,
+                op: OpCode::GetField,
+                imm_operands: ImmOperand::Atom(atom),
+            });
+        }
+
+        CaptureAccess::Write(write_var) => {
+            let atom = my_captures[write_var as usize].name;
+
+            out.push(Instr {
+                pc,
+                op: OpCode::GetArg,
+                imm_operands: ImmOperand::Arg(u16::MAX),
+            });
+
+            out.push(Instr {
+                pc,
+                op: OpCode::PutField,
+                imm_operands: ImmOperand::Atom(atom),
+            });
+        }
+    }
+
+    true
+}

compiler/src/flatten.rs

@@ -6,17 +6,17 @@ pub struct FlatFn {
     pub arg_count: usize,
     pub var_count: usize,
     pub cpool: Vec<JsValue>,
-    pub captures: Vec<ClosureVar>,
     pub code: Vec<Instr>,
 }
 
-pub fn flatten(bin: QuickJSBinary) -> (Vec<FlatFn>, Vec<Atom>) {
+pub fn flatten(bin: QuickJSBinary) -> (Vec<FlatFn>, Vec<Vec<ClosureVar>>, Vec<Atom>) {
     let mut fns = vec![];
-    walk(&mut fns, bin.root_mod.root_fn);
+    let mut captures = vec![];
-    (fns, bin.atoms)
+    walk(bin.root_mod.root_fn, &mut fns, &mut captures);
+    (fns, captures, bin.atoms)
 }
 
-fn walk(flat_fns: &mut Vec<FlatFn>, func: Function) -> usize {
+fn walk(func: Function, flat_fns: &mut Vec<FlatFn>, captures: &mut Vec<Vec<ClosureVar>>) -> usize {
     let id = flat_fns.len();
 
     flat_fns.push(FlatFn {
@@ -24,15 +24,16 @@ fn walk(flat_fns: &mut Vec<FlatFn>, func: Function) -> usize {
         arg_count: func.arg_count,
         var_count: func.var_count,
         cpool: vec![],
-        captures: func.closure_vars,
         code: func.byte_code,
     });
 
+    captures.push(func.closure_vars);
+
     let cpool = func
         .cpool
         .into_iter()
         .map(|c| match c {
-            JsValue::Function(child) => JsValue::FnById(walk(flat_fns, child)),
+            JsValue::Function(child) => JsValue::FnById(walk(child, flat_fns, captures)),
             other => other,
         })
         .collect();

compiler/src/libqjs/parse.rs

@@ -514,6 +514,7 @@ pub enum ImmOperand {
     },
 }
 
+#[derive(Clone)]
 pub struct Instr {
     pub pc: usize,
     pub op: opcodes::OpCode,

compiler/src/main.rs

@@ -1,3 +1,4 @@
+mod declosure;
 mod flatten;
 mod infer;
 mod libqjs;
@@ -6,6 +7,7 @@ mod regify;
 use std::fs::File;
 use std::io::Read;
 
+use declosure::declosure;
 use flatten::flatten;
 use infer::infer_types;
 use regify::regify;
@@ -26,12 +28,8 @@ fn main() {
     let binary = libqjs::compile(&src, filename);
 
     // Process it
-    let (flat_fns, atoms) = flatten(binary);
+    let (flat_fns, captures, atoms) = flatten(binary);
+    let flat_fns = declosure(flat_fns, &captures);
     let reg_fns = regify(flat_fns);
-    println!("{reg_fns:#?}");
-    println!("{atoms:#?}");
-
     let type_cache = infer_types(&reg_fns, &atoms);
-
-    println!("{type_cache:#?}");
 }

compiler/src/regify/mod.rs

@@ -43,7 +43,6 @@ pub struct RegFn {
     pub arg_count: usize,
     pub val_count: usize,
     pub cpool: Vec<JsValue>,
-    pub captures: Vec<ClosureVar>,
     pub code: Vec<RegInstr>,
     pub blocks: Vec<BasicBlock>,
 }
@@ -85,7 +84,6 @@ impl RegFn {
             arg_count: flat_fn.arg_count,
             val_count: next_val,
             cpool: flat_fn.cpool,
-            captures: flat_fn.captures,
             code: out,
             blocks,
         }