Rewrite with working async

compiler/src/regify/mod.rs

@@ -10,7 +10,7 @@ pub use basic_block::BasicBlock;
 use basic_block::find_basic_blocks;
 use walk::walk_block;
 
-#[derive(Debug, Clone, Copy)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub enum Operand {
     Local(usize),
     Arg(usize),
@@ -46,34 +46,119 @@ pub struct RegFn {
     pub blocks: Vec<BasicBlock>,
 }
 
-impl RegFn {
-    fn from_flat(f: FlatFn) -> Self {
-        let blocks = find_basic_blocks(&f.code);
-        let mut out = vec![RegInstr::nop(); f.code.len()];
+pub fn regify(flat_fns: Vec<FlatFn>) -> Vec<RegFn> {
+    flat_fns.into_iter().map(RegFn::from_flat).collect()
+}
 
-        // The stack each block recieves at entry from its predecessors
-        let mut entry_stacks: Vec<Option<Vec<Operand>>> = vec![None; blocks.len()];
+impl RegFn {
+    /// The locals and args that a closure shares: the locations of its `define_field_ref`s
+    /// (see transform/closure.rs)
+    pub fn captured_slots(&self) -> Vec<Operand> {
+        let refs = self.code.iter().filter(|i| i.op == OpCode::DefineFieldRef);
+        let mut slots: Vec<Operand> = refs
+            .map(|i| i.args[0])
+            .filter(|slot| !matches!(slot, Operand::Reg(_)))
+            .collect();
+
+        slots.sort_unstable();
+        slots.dedup();
+        slots
+    }
+
+    /// Joins: the stack that a block gets at entry is new regs, made when the walk first gets
+    /// to the block. Each predecessor copies its own stack into them, in its postfix. So every
+    /// block has one entry stack, whichever way control comes in (this does the job of phi
+    /// nodes in SSA)
+    ///
+    /// A block with two successors copies into the entry regs of both, before its jump. The
+    /// copies for the edge that is not taken do no harm: a block's entry regs are only read
+    /// in that block, and every way into it writes them first
+    fn from_flat(f: FlatFn) -> Self {
+        let mut blocks = find_basic_blocks(&f.code);
+        let mut out = vec![RegInstr::nop(); f.code.len()];
         let mut next_reg = 0;
 
-        // Walk blocks in CFG order, starting from block 0
-        let mut worklist = vec![];
+        let mut entry_stacks: Vec<Option<Vec<Operand>>> = vec![None; blocks.len()];
         entry_stacks[0] = Some(vec![]);
-        worklist.push(0);
 
-        while let Some(block_idx) = worklist.pop() {
-            let mut stack = entry_stacks[block_idx].clone().unwrap_or_default();
+        // Walk blocks in CFG order, starting from block 0. A block that the walk never
+        // gets to is dead code, and stays `nop`
+        let mut worklist = vec![0];
 
-            walk_block(
-                &blocks[block_idx],
-                &f.code,
-                &mut stack,
-                &mut next_reg,
-                &mut out,
+        while let Some(b) = worklist.pop() {
+            let mut stack = entry_stacks[b].clone().unwrap();
+            walk_block(&blocks[b], &f.code, &mut stack, &mut next_reg, &mut out);
+
+            for s in blocks[b].succ.clone() {
+                let entry = entry_stacks[s].get_or_insert_with(|| {
+                    worklist.push(s);
+                    stack.iter().map(|_| fresh_reg(&mut next_reg)).collect()
+                });
+
+                assert_eq!(
+                    entry.len(),
+                    stack.len(),
+                    "a join of stacks of different depths"
+                );
+
+                let copies = stack
+                    .iter()
+                    .zip(entry.iter())
+                    .filter(|(src, dst)| src != dst);
+
+                for (&src, &dst) in copies {
+                    blocks[b].postfix.push(RegInstr {
+                        dsts: vec![dst],
+                        op: OpCode::SetLoc,
+                        imm: ImmOperand::None,
+                        args: vec![src],
+                    });
+                }
+            }
+
+            // The copies run one after the other, so no copy may write a reg that a later
+            // copy reads. A block's stack only holds its own entry regs and the regs it made,
+            // so that could only be a block that jumps to itself and moves its stack around
+            let postfix = &blocks[b].postfix;
+            assert!(
+                postfix
+                    .iter()
+                    .all(|c| postfix.iter().all(|d| c.args[0] != d.dsts[0])),
+                "the copies of a join overlap"
             );
-
-            // TODO: handle sucessors
         }
 
-        todo!()
+        mark_puts(&mut out);
+
+        Self {
+            arg_count: f.arg_count,
+            local_count: f.local_count,
+            reg_count: next_reg,
+            cpool: f.cpool,
+            code: out,
+            blocks,
+        }
+    }
+}
+
+fn fresh_reg(next_reg: &mut usize) -> Operand {
+    let r = *next_reg;
+    *next_reg += 1;
+    Operand::Reg(r)
+}
+
+/// A define_field into an object that no instr of this fn made (a class field on `this`, a
+/// static field on the class) is a write like any other: a put_field. So a define_field is
+/// always a write into an object before it goes anywhere, and its field is never absent
+fn mark_puts(code: &mut [RegInstr]) {
+    let made: Vec<Operand> = (code.iter())
+        .filter(|i| matches!(i.op, OpCode::Object | OpCode::Fclosure | OpCode::Fclosure8))
+        .map(|i| i.dsts[0])
+        .collect();
+
+    for instr in code {
+        if instr.op == OpCode::DefineField && !made.contains(&instr.args[1]) {
+            instr.op = OpCode::PutField;
+        }
     }
 }

compiler/src/regify/walk.rs

@@ -42,11 +42,16 @@ pub fn walk_block(
             continue;
         }
 
-        if try_handle_misc(instr, stack, out_instr) {
+        if try_handle_shuffle(instr, stack, out_instr) {
             continue;
         }
 
         handle_normal_case(instr, stack, &mut fresh_reg, out_instr);
+
+        if instr.op == OpCode::DefineFieldRef {
+            let prev = std::mem::replace(&mut out[instr_idx - 1], RegInstr::nop());
+            resolve_location(&prev, &mut out[instr_idx]);
+        }
     }
 }
 
@@ -163,61 +168,53 @@ fn try_handle_named_write(instr: &Instr, stack: &mut Vec<Operand>, out: &mut Reg
     true
 }
 
-fn try_handle_misc(instr: &Instr, stack: &mut Vec<Operand>, out_instr: &mut RegInstr) -> bool {
-    use OpCode::*;
-
-    *out_instr = match instr.op {
-        // Duplicate stack top, out_instr becomes `nop`
-        Dup => {
-            let top = stack[stack.len() - 1].clone();
-            stack.push(top);
-
-            RegInstr {
-                dsts: vec![],
-                op: Nop,
-                imm: ImmOperand::None,
-                args: vec![],
-            }
-        }
-
-        // Duplicate the top 2 elems in stack, out_instr becomes `nop`
-        Dup2 => {
-            let a = stack[stack.len() - 1].clone();
-            let b = stack[stack.len() - 2].clone();
-
-            stack.push(b);
-            stack.push(a);
-
-            RegInstr {
-                dsts: vec![],
-                op: Nop,
-                imm: ImmOperand::None,
-                args: vec![],
-            }
-        }
-
-        // Swap top 2 elems in stack, out_instr becomes `nop`
-        Swap => {
-            let a = stack.pop().unwrap();
-            let b = stack.pop().unwrap();
-
-            stack.push(a);
-            stack.push(b);
-
-            RegInstr {
-                dsts: vec![],
-                op: Nop,
-                imm: ImmOperand::None,
-                args: vec![],
-            }
-        }
-
-        _ => return false,
+/// Handles an instr that only moves stack values around: it pops `n` values, and pushes some
+/// of them again. No value is made, so `out_instr` becomes `nop`
+fn try_handle_shuffle(instr: &Instr, stack: &mut Vec<Operand>, out_instr: &mut RegInstr) -> bool {
+    let Some((n, pushes)) = shuffle(instr.op) else {
+        return false;
     };
 
+    let popped = stack.split_off(stack.len() - n);
+    stack.extend(pushes.iter().map(|&i| popped[i]));
+
+    *out_instr = RegInstr::nop();
     true
 }
 
+/// The shuffles: how many values an instr pops, and which of them it pushes, in stack order
+/// (bottom first). The numbers index the popped values, bottom first
+fn shuffle(op: OpCode) -> Option<(usize, &'static [usize])> {
+    use OpCode::*;
+
+    Some(match op {
+        Drop => (1, &[]),                 // a ->
+        Nip => (2, &[1]),                 // a b -> b
+        Nip1 => (3, &[1, 2]),             // a b c -> b c
+        Dup => (1, &[0, 0]),              // a -> a a
+        Dup1 => (2, &[0, 0, 1]),          // a b -> a a b
+        Dup2 => (2, &[0, 1, 0, 1]),       // a b -> a b a b
+        Dup3 => (3, &[0, 1, 2, 0, 1, 2]), // a b c -> a b c a b c
+        Insert2 => (2, &[1, 0, 1]),       // obj a -> a obj a
+        Insert3 => (3, &[2, 0, 1, 2]),    // obj prop a -> a obj prop a
+        Insert4 => (4, &[3, 0, 1, 2, 3]), // this obj prop a -> a this obj prop a
+        Perm3 => (3, &[1, 0, 2]),         // obj a b -> a obj b
+        Perm4 => (4, &[2, 0, 1, 3]),      // obj prop a b -> a obj prop b
+        Perm5 => (5, &[3, 0, 1, 2, 4]),   // this obj prop a b -> a this obj prop b
+        Swap => (2, &[1, 0]),             // a b -> b a
+        Swap2 => (4, &[2, 3, 0, 1]),      // a b c d -> c d a b
+        Rot3l => (3, &[1, 2, 0]),         // x a b -> a b x
+        Rot3r => (3, &[2, 0, 1]),         // a b x -> x a b
+        Rot4l => (4, &[1, 2, 3, 0]),      // x a b c -> a b c x
+        Rot5l => (5, &[1, 2, 3, 4, 0]),   // x a b c d -> a b c d x
+        // Binds [[HomeObject]], which only `super` reads. `super` is not supported
+        SetHomeObject => (2, &[0, 1]),
+        // Makes the key a property key. A key is kept as it is (see infer)
+        ToPropkey2 => (2, &[0, 1]),
+        _ => return None,
+    })
+}
+
 /// Handle a normal instruction with simulated stack pushes/pops
 fn handle_normal_case(
     instr: &Instr,
@@ -225,6 +222,8 @@ fn handle_normal_case(
     mut fresh_reg: impl FnMut() -> Operand,
     out_instr: &mut RegInstr,
 ) {
+    use OpCode::*;
+
     // op and imm carried as-is
     *out_instr = RegInstr {
         dsts: vec![],
@@ -244,12 +243,22 @@ fn handle_normal_case(
     let n_pop = info.n_pop as usize + extra_pops;
     let n_push = info.n_push as usize;
 
-    // Simulate this instr's pops
+    // Simulate this instr's pops, top first
     for _ in 0..n_pop {
         let x = stack.pop().unwrap();
         out_instr.args.push(x);
     }
 
+    if matches!(instr.op, Call | CallMethod | CallConstructor) {
+        out_instr.args = call_operands(instr.op, std::mem::take(&mut out_instr.args));
+    }
+
+    // These leave the object they wrote on the stack: the same value, not a new one
+    if matches!(instr.op, DefineField | DefineFieldRef) {
+        stack.push(out_instr.args[1]);
+        return;
+    }
+
     // Simulate this instr's pushes
     for _ in 0..n_push {
         let x = fresh_reg();
@@ -257,3 +266,57 @@ fn handle_normal_case(
         out_instr.dsts.push(x);
     }
 }
+
+/// Put a call's popped operands in the order of the callee's params: the callee (which is its
+/// env), `this`, then the args in source order
+///
+/// `call_method` has `this` at the bottom of its frame, below the callee. `call_constructor`
+/// has the callee at the bottom and `this` above it (see transform/new.rs). `call` has `this`
+/// on top, where transform pushed it
+fn call_operands(op: OpCode, mut popped: Vec<Operand>) -> Vec<Operand> {
+    let (callee, this) = match op {
+        OpCode::CallMethod => {
+            let this = popped.pop().unwrap();
+            (popped.pop().unwrap(), this)
+        }
+        OpCode::CallConstructor => {
+            let callee = popped.pop().unwrap();
+            (callee, popped.pop().unwrap())
+        }
+        _ => {
+            let callee = popped.pop().unwrap();
+            (callee, popped.remove(0))
+        }
+    };
+
+    // What remains are the args, last arg first
+    popped.reverse();
+
+    [callee, this].into_iter().chain(popped).collect()
+}
+
+/// `define_field_ref x` defines `x` as the location that its value was read from (see
+/// transform/closure.rs), and the instr before it is that read. Its value arg becomes the
+/// location:
+/// - a local or arg: the field is that binding
+/// - the reg of an object `o`, for a read `get_field x` of `o`: the field is `o.x`, which is
+///   a location itself (a capture that the owner gives to its own child)
+///
+/// The read itself is then not used, so it becomes `nop`
+fn resolve_location(prev: &RegInstr, out_instr: &mut RegInstr) {
+    let value = out_instr.args[0];
+    assert_eq!(
+        prev.dsts,
+        [value],
+        "define_field_ref must follow the read of its value"
+    );
+
+    out_instr.args[0] = match prev.op {
+        OpCode::SetLoc => prev.args[0],
+        OpCode::GetField => {
+            assert!(prev.imm.atom() == out_instr.imm.atom());
+            prev.args[0]
+        }
+        op => panic!("define_field_ref after {op:?}, which reads no location"),
+    };
+}