Claude's declass improvements

compiler/src/transform/mod.rs

@@ -2,7 +2,8 @@ mod declass;
 
 use crate::{
     flatten::FlatFn,
-    libqjs::parse::{AtomRef, ClosureVar, ImmOperand, Instr, JsValue, opcodes::OpCode},
+    libqjs::parse::{ClosureVar, ImmOperand, Instr, JsValue, opcodes::OpCode},
+    model::{EXTRA_ARGS, FIELD_FN},
 };
 
 use declass::declass;
@@ -17,6 +18,12 @@ pub fn transform(flat_fns: Vec<FlatFn>, captures: &[Vec<ClosureVar>]) -> Vec<Fla
 
 fn transform_one(fn_id: usize, f: FlatFn, captures: &[Vec<ClosureVar>]) -> FlatFn {
     let declassed = declass(&f.code);
+
+    // Every function gains two parameters after its own, in this order. The
+    // indices are static, so neither needs a sentinel.
+    let this_idx = f.arg_count as u16;
+    let env_idx = f.arg_count as u16 + 1;
+
     let mut out_len = 0;
 
     for instr in &declassed {
@@ -40,7 +47,11 @@ fn transform_one(fn_id: usize, f: FlatFn, captures: &[Vec<ClosureVar>]) -> FlatF
             continue;
         }
 
-        if replace_capture_accs(instr, &captures[fn_id], &mut out) {
+        if replace_capture_accs(instr, &captures[fn_id], env_idx, &mut out) {
+            continue;
+        }
+
+        if replace_push_this(instr, this_idx, &mut out) {
             continue;
         }
 
@@ -51,17 +62,32 @@ fn transform_one(fn_id: usize, f: FlatFn, captures: &[Vec<ClosureVar>]) -> FlatF
         out.push(instr.clone());
     }
 
-    //debug_assert!(out.len() == out_len);
+    debug_assert!(out.len() == out_len);
 
     FlatFn {
         name: f.name,
-        arg_count: f.arg_count,
+        arg_count: f.arg_count + EXTRA_ARGS,
         var_count: f.var_count,
         cpool: f.cpool,
         code: out,
     }
 }
 
+/// `this` is an ordinary parameter now, so reading it is an ordinary arg read.
+fn replace_push_this(instr: &Instr, this_idx: u16, out: &mut Vec<Instr>) -> bool {
+    if instr.op != OpCode::PushThis {
+        return false;
+    }
+
+    out.push(Instr {
+        pc: instr.pc,
+        op: OpCode::GetArg,
+        imm: ImmOperand::Arg(this_idx),
+    });
+
+    true
+}
+
 struct FClosureInfo<'a> {
     cpool_idx: u32,
     child_captures: &'a [ClosureVar],
@@ -134,7 +160,7 @@ fn replace_fclosure(
     out.push(Instr {
         pc,
         op: OpCode::DefineField,
-        imm: ImmOperand::Atom(AtomRef::Predefined(0)),
+        imm: ImmOperand::Atom(FIELD_FN),
     });
 
     true
@@ -171,7 +197,12 @@ fn is_capture_acc(instr: &Instr) -> Option<CaptureAccess> {
     Some(acc)
 }
 
-fn replace_capture_accs(instr: &Instr, my_captures: &[ClosureVar], out: &mut Vec<Instr>) -> bool {
+fn replace_capture_accs(
+    instr: &Instr,
+    my_captures: &[ClosureVar],
+    env_idx: u16,
+    out: &mut Vec<Instr>,
+) -> bool {
     let acc = is_capture_acc(instr);
 
     let acc = match acc {
@@ -188,7 +219,7 @@ fn replace_capture_accs(instr: &Instr, my_captures: &[ClosureVar], out: &mut Vec
             out.push(Instr {
                 pc,
                 op: OpCode::GetArg,
-                imm: ImmOperand::Arg(u16::MAX),
+                imm: ImmOperand::Arg(env_idx),
             });
 
             out.push(Instr {
@@ -204,7 +235,7 @@ fn replace_capture_accs(instr: &Instr, my_captures: &[ClosureVar], out: &mut Vec
             out.push(Instr {
                 pc,
                 op: OpCode::GetArg,
-                imm: ImmOperand::Arg(u16::MAX),
+                imm: ImmOperand::Arg(env_idx),
             });
 
             out.push(Instr {