Claude's declass improvements

compiler/src/infer.rs

@@ -5,7 +5,8 @@ use std::collections::HashMap;
 use crate::{
     infer::Type::ReturnOf,
     libqjs::parse::{Atom, AtomRef, ImmOperand, JsString, JsValue, atom_names, opcodes::OpCode},
-    regify::RegFn,
+    model::{EXTRA_ARGS, FIELD_FN, FIELD_PROTO, MAX_PROTO_DEPTH},
+    regify::{RegFn, RegInstr},
 };
 
 #[derive(Debug, Clone, Hash, PartialEq, Eq)]
@@ -41,19 +42,48 @@ fn join_ty(a: &Type, b: &Type) -> Type {
     }
 }
 
+/// Result of looking a field up on one object, ignoring its prototype.
+enum Lookup {
+    Found(Type),
+    Missing,
+    Any,
+}
+
+/// Look `field` up on `ty`. A miss on the object itself follows the prototype
+/// link and retries, which is how a method defined on a prototype is found
+/// from an instance.
 fn get_field(ty: &Type, field: AtomRef, cache: &mut TypeCache) -> Type {
     let mut current = ty.clone();
 
+    for _ in 0..MAX_PROTO_DEPTH {
+        match own_field(&current, field, cache) {
+            Lookup::Found(t) => return t,
+            Lookup::Any => return Type::Any,
+            Lookup::Missing => match own_field(&current, FIELD_PROTO, cache) {
+                Lookup::Found(proto) => current = proto,
+                _ => return Type::Undefined,
+            },
+        }
+    }
+
+    // The chain loops, or is deeper than anything real. Stop guessing.
+    Type::Any
+}
+
+/// Walk one object's own field chain, without following its prototype.
+fn own_field(ty: &Type, field: AtomRef, cache: &mut TypeCache) -> Lookup {
+    let mut current = ty.clone();
+
     loop {
         match current {
-            Type::Any => return Type::Any,
+            Type::Any => return Lookup::Any,
 
             Type::ReturnOf((fn_kind, args)) => match fn_kind {
                 FnKind::DefineField(atom) => {
-                    assert!(args.len() == 2);
+                    debug_assert!(args.len() == 2);
 
                     if atom == field {
-                        return args[0].clone();
+                        return Lookup::Found(args[0].clone());
                     }
 
                     current = args[1].clone();
@@ -65,16 +95,74 @@ fn get_field(ty: &Type, field: AtomRef, cache: &mut TypeCache) -> Type {
                         arg_types: args.clone(),
                     }) {
                         Some(info) => current = info.ret_type.clone(),
-                        None => panic!(),
+                        // Not inferred yet, so the shape is unknown rather
+                        // than impossible.
+                        None => return Lookup::Any,
                     }
                 }
             },
 
-            _ => return Type::Undefined,
+            _ => return Lookup::Missing,
         }
     }
 }
 
+/// A call frame, in the order the operands were pushed:
+///
+///     [this,] callee, arg0 .. arg_n
+///
+/// `instr.args` is in pop order, so the callee is last for a plain call, and
+/// second to last for a method call, which carries a receiver beneath it. A
+/// plain call has no receiver slot, so its `this` is undefined.
+struct CallFrame {
+    /// The closure object being called. It is also the callee's environment.
+    callee: Type,
+    this: Type,
+    /// The actual arguments, back in source order.
+    args: Vec<Type>,
+}
+
+fn call_frame(instr: &RegInstr, val_types: &[Type]) -> CallFrame {
+    let is_method = matches!(instr.op, OpCode::CallMethod | OpCode::TailCallMethod);
+    let n = instr.args.len();
+
+    let (callee, this, rest) = if is_method {
+        (
+            val_types[instr.args[n - 2]].clone(),
+            val_types[instr.args[n - 1]].clone(),
+            &instr.args[..n - 2],
+        )
+    } else {
+        (
+            val_types[instr.args[n - 1]].clone(),
+            Type::Undefined,
+            &instr.args[..n - 1],
+        )
+    };
+
+    CallFrame {
+        callee,
+        this,
+        args: rest.iter().rev().map(|a| val_types[*a].clone()).collect(),
+    }
+}
+
+/// The parameters a lowered function sees: its own, then `this`, then the
+/// environment. The actuals are padded or truncated to the declared arity, so
+/// the last two always land at a fixed index whatever the caller passed.
+fn arg_types_for(callee: &RegFn, frame: &CallFrame) -> Vec<Type> {
+    let declared = callee.arg_count.saturating_sub(EXTRA_ARGS);
+    let mut out = Vec::with_capacity(declared + EXTRA_ARGS);
+
+    for i in 0..declared {
+        out.push(frame.args.get(i).cloned().unwrap_or(Type::Undefined));
+    }
+
+    out.push(frame.this.clone());
+    out.push(frame.callee.clone());
+    out
+}
+
 #[derive(Debug, Hash, PartialEq, Eq)]
 pub struct InferRequest {
     pub fn_id: usize,
@@ -98,9 +186,11 @@ pub fn infer_types(reg_fns: &[RegFn], atoms: &[Atom]) -> TypeCache {
 
     let mut cache = TypeCache::with_capacity(reg_fns.len());
 
+    // main is called by nobody, so every parameter it has — including this
+    // and its environment — is undefined.
     let req = InferRequest {
         fn_id: main_fn_id,
-        arg_types: vec![],
+        arg_types: vec![Type::Undefined; reg_fns[main_fn_id].arg_count],
     };
 
     let res = walk_fn(&req, reg_fns, &mut cache);
@@ -175,18 +265,11 @@ fn walk_fn(req: &InferRequest, reg_fns: &[RegFn], cache: &mut TypeCache) -> Infe
                 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) => {
+                    ImmOperand::Arg(k) => req
-                        let k = if k == u16::MAX {
+                        .arg_types
-                            req.arg_types.len() as u16 - 1
+                        .get(k as usize)
-                        } else {
+                        .cloned()
-                            k
+                        .unwrap_or(Type::Undefined),
-                        };
-
-                        req.arg_types
-                            .get(k as usize)
-                            .cloned()
-                            .unwrap_or(Type::Undefined)
-                    }
                     _ => unreachable!(),
                 },
 
@@ -214,37 +297,30 @@ fn walk_fn(req: &InferRequest, reg_fns: &[RegFn], cache: &mut TypeCache) -> Infe
                 }
 
                 // ──────────────────────────────── Function Calls ──────────────────────────────────
-                Call | Call0 | Call1 | Call2 | Call3 | TailCall => {
+                Call | Call0 | Call1 | Call2 | Call3 | TailCall | CallMethod | TailCallMethod => {
-                    let closure = instr.args.last().unwrap();
+                    let frame = call_frame(instr, &val_types);
-                    let closure = &val_types[*closure];
-                    let fn_id = get_field(closure, AtomRef::Predefined(0), cache);
 
-                    let fn_id = match fn_id {
+                    match get_field(&frame.callee, FIELD_FN, cache) {
-                        Type::Fn(id) => id,
+                        Type::Fn(fn_id) => {
-                        _ => unreachable!(),
+                            let req = InferRequest {
-                    };
+                                fn_id,
+                                arg_types: arg_types_for(&reg_fns[fn_id], &frame),
+                            };
 
-                    let arg_tys = instr
+                            match cache.get(&req) {
-                        .args
+                                Some(res) => res.ret_type.clone(),
-                        .iter()
+                                None => {
-                        .map(|a| val_types[*a].clone())
+                                    let res = walk_fn(&req, reg_fns, cache);
-                        .collect::<Vec<Type>>();
+                                    let ret_ty = res.ret_type.clone();
-
+                                    cache.insert(req, res);
-                    let req = InferRequest {
+                                    ret_ty
-                        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
                         }
+
+                        // A callee we cannot resolve to one function: a
+                        // builtin, or a value we lost track of.
+                        _ => Type::Any,
                     }
                 }
 
@@ -257,6 +333,13 @@ fn walk_fn(req: &InferRequest, reg_fns: &[RegFn], cache: &mut TypeCache) -> Infe
 
             if is_ret {
                 ret_type = join_ty(&ret_type, &ty);
+            } else if instr.op == DefineField {
+                // A define states its result completely, from its operands.
+                // Joining it with the object's previous type would only
+                // destroy the chain that a write-back is building up.
+                for &dst in &instr.dsts {
+                    val_types[dst] = ty.clone();
+                }
             } else {
                 for &dst in &instr.dsts {
                     val_types[dst] = join_ty(&val_types[dst], &ty);

compiler/src/model.rs

@@ -0,0 +1,24 @@
+//! Reserved field names of the lowered object model.
+//!
+//! Every lowered object is a plain bag of fields. Two names are reserved for
+//! the compiler. Both are quickjs internal atoms, spelled `<...>`, which
+//! JavaScript source cannot produce, so neither can collide with a user
+//! property.
+
+use crate::libqjs::parse::AtomRef;
+
+/// The function held by a closure object. Every closure carries one, and a
+/// call reads it to find the function behind the object it is calling.
+pub const FIELD_FN: AtomRef = AtomRef::Predefined(0); // "<null>"
+
+/// An instance's prototype object. A field lookup that misses on an object
+/// itself follows this link and retries.
+pub const FIELD_PROTO: AtomRef = AtomRef::Predefined(124); // "<home_object>"
+
+/// How many prototype links one lookup follows before giving up and
+/// answering `Any`. Guards against a chain that loops.
+pub const MAX_PROTO_DEPTH: usize = 16;
+
+/// Parameters a lowered function gains beyond its source arity: `this`,
+/// then the environment. They are the last two, in that order.
+pub const EXTRA_ARGS: usize = 2;

compiler/src/regify/mod.rs

@@ -167,6 +167,17 @@ fn walk_block(
             None => (),
         }
 
+        // A field write is a define whose result lands back in the object's own
+        // value, so later reads of that object see the field. This is what lets
+        // a chain of writes to one local build a single object type, and it is
+        // why no later stage needs to know PutField at all.
+        if instr.op == OpCode::PutField {
+            out_instr.op = OpCode::DefineField;
+            out_instr.dsts.push(out_instr.args[1]);
+            out[instr_idx] = out_instr;
+            continue;
+        }
+
         // Handle this instr's pushes
         for _ in 0..n_push {
             let v = fresh_val();

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 {