Improve infer and transforms

compiler/src/transform/mod.rs

@@ -0,0 +1,243 @@
+mod declass;
+
+use crate::{
+    flatten::FlatFn,
+    libqjs::parse::{AtomRef, ClosureVar, ImmOperand, Instr, JsValue, opcodes::OpCode},
+};
+
+use declass::declass;
+
+pub fn transform(flat_fns: Vec<FlatFn>, captures: &[Vec<ClosureVar>]) -> Vec<FlatFn> {
+    flat_fns
+        .into_iter()
+        .enumerate()
+        .map(|(i, f)| transform_one(i, f, captures))
+        .collect()
+}
+
+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 &declassed {
+        if let Some(info) = is_fclosure(instr, &f, captures) {
+            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 &declassed {
+        if replace_fclosure(instr, &f, captures, &mut out) {
+            continue;
+        }
+
+        if replace_capture_accs(instr, &captures[fn_id], &mut out) {
+            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,
+        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 {
+        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: ImmOperand::None,
+    });
+
+    for c in info.child_captures {
+        out.push(Instr {
+            pc,
+            op: OpCode::GetLoc,
+            imm: ImmOperand::Loc(c.var_idx as u16),
+        });
+
+        out.push(Instr {
+            pc,
+            op: OpCode::DefineField,
+            imm: ImmOperand::Atom(c.name),
+        });
+    }
+
+    out.push(Instr {
+        pc,
+        op: OpCode::PushConst,
+        imm: ImmOperand::Const(info.cpool_idx),
+    });
+
+    out.push(Instr {
+        pc,
+        op: OpCode::DefineField,
+        imm: 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 {
+            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 {
+            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: ImmOperand::Arg(u16::MAX),
+            });
+
+            out.push(Instr {
+                pc,
+                op: OpCode::GetField,
+                imm: ImmOperand::Atom(atom),
+            });
+        }
+
+        CaptureAccess::Write(write_var) => {
+            let atom = my_captures[write_var as usize].name;
+
+            out.push(Instr {
+                pc,
+                op: OpCode::GetArg,
+                imm: ImmOperand::Arg(u16::MAX),
+            });
+
+            out.push(Instr {
+                pc,
+                op: OpCode::PutField,
+                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
+}