Support lifting locals to SSA as well.

include/weval.h

@@ -246,6 +246,14 @@ void weval_write_stack(uint64_t* ptr, uint32_t index, uint64_t value)
  * effect never occurs). */
 uint64_t weval_pop_stack(uint64_t* ptr) WEVAL_WASM_IMPORT("pop.stack");
 
+/* Locals virtualization; locals are also flushed when the stack is
+ * flushed */
+
+uint64_t weval_read_local(uint64_t* ptr, uint32_t index)
+    WEVAL_WASM_IMPORT("read.local");
+void weval_write_local(uint64_t* ptr, uint32_t index, uint64_t value)
+    WEVAL_WASM_IMPORT("write.local");
+
 /* Debugging and stats intrinsics */
 
 void weval_trace_line(uint32_t line_number) WEVAL_WASM_IMPORT("trace.line");

lib/weval-stubs.wat

@@ -37,4 +37,7 @@
        unreachable)
  (func (export "write.stack") (param i32 i32 i64))
  (func (export "pop.stack") (param i32) (result i64)
-       unreachable))
+       unreachable)
+ (func (export "read.local") (param i32 i32) (result i64)
+       unreachable)
+ (func (export "write.local") (param i32 i32 i64)))

src/eval.rs

@@ -1601,6 +1601,71 @@ impl<'a> Evaluator<'a> {
                         ));
                         self.func.blocks[new_block].insts.push(store);
                     }
+
+                    for (_, (addr, data)) in std::mem::take(&mut state.flow.locals) {
+                        let addr = addr.value().unwrap();
+                        let data = data.value().unwrap();
+                        log::trace!("sync_stack: local addr {} data {}", addr, data);
+                        let args = self.func.arg_pool.double(addr, data);
+                        let store = self.func.add_value(ValueDef::Operator(
+                            Operator::I64Store {
+                                memory: MemoryArg {
+                                    align: 1,
+                                    offset: 0,
+                                    memory: self.image.main_heap().unwrap(),
+                                },
+                            },
+                            args,
+                            ListRef::default(),
+                        ));
+                        self.func.blocks[new_block].insts.push(store);
+                    }
+                    EvalResult::Elide
+                } else if Some(function_index) == self.intrinsics.read_local {
+                    let ptr = self.func.arg_pool[values][0];
+                    let idx = abs[1].as_const_u32().unwrap();
+                    match state.flow.locals.get(&idx) {
+                        None => {
+                            let args = self.func.arg_pool.single(ptr);
+                            let i64_ty = self.func.single_type_list(Type::I64);
+                            let load = self.func.add_value(ValueDef::Operator(
+                                Operator::I64Load {
+                                    memory: MemoryArg {
+                                        align: 1,
+                                        offset: 0,
+                                        memory: self.image.main_heap().unwrap(),
+                                    },
+                                },
+                                args,
+                                i64_ty,
+                            ));
+                            self.func.blocks[new_block].insts.push(load);
+                            EvalResult::Alias(AbstractValue::Runtime(None), load)
+                        }
+                        Some((_, RegValue::Value { data, abs, .. })) => {
+                            EvalResult::Alias(abs.clone(), *data)
+                        }
+                        _ => unreachable!(),
+                    }
+                } else if Some(function_index) == self.intrinsics.write_local {
+                    let ptr = self.func.arg_pool[values][0];
+                    let idx = abs[1].as_const_u32().unwrap();
+                    let data = self.func.arg_pool[values][2];
+                    state.flow.locals.insert(
+                        idx,
+                        (
+                            RegValue::Value {
+                                data: ptr,
+                                abs: abs[0].clone(),
+                                ty: Type::I32,
+                            },
+                            RegValue::Value {
+                                data,
+                                abs: abs[2].clone(),
+                                ty: Type::I64,
+                            },
+                        ),
+                    );
                     EvalResult::Elide
                 } else {
                     EvalResult::Unhandled
@@ -2191,6 +2256,10 @@ impl<'a> Evaluator<'a> {
                 handle_value(RegSlot::StackAddr(i as u32), addr)?;
                 handle_value(RegSlot::StackData(i as u32), data)?;
             }
+            for (&i, (addr, data)) in succ_state.locals.iter() {
+                handle_value(RegSlot::LocalAddr(i), addr)?;
+                handle_value(RegSlot::LocalData(i), data)?;
+            }
 
             for pred_idx in 0..self.func.blocks[block].preds.len() {
                 let pred = self.func.blocks[block].preds[pred_idx];
@@ -2202,6 +2271,8 @@ impl<'a> Evaluator<'a> {
                         RegSlot::Register(_) => pred_state.regs.get(&idx).as_ref().unwrap(),
                         RegSlot::StackAddr(i) => &pred_state.stack.get(i as usize).unwrap().0,
                         RegSlot::StackData(i) => &pred_state.stack.get(i as usize).unwrap().1,
+                        RegSlot::LocalAddr(i) => &pred_state.locals.get(&i).unwrap().0,
+                        RegSlot::LocalData(i) => &pred_state.locals.get(&i).unwrap().1,
                     };
                     let pred_val = pred_reg.value().unwrap();
                     self.func.blocks[pred]
@@ -2220,6 +2291,9 @@ impl<'a> Evaluator<'a> {
         // For each edge, look at known stack depth of pred and
         // succ. If succ's range is smaller, read regs from pred and
         // sync at end of pred.
+        //
+        // Also look at `locals` and find locals present in pred and
+        // not in some succ, and sync them.
         for (_, &block) in &self.block_map {
             if self.func.blocks[block].succs.is_empty() {
                 continue;
@@ -2257,6 +2331,42 @@ impl<'a> Evaluator<'a> {
                 ));
                 self.func.blocks[block].insts.push(store);
             }
+
+            let locals_to_sync = pred_state
+                .locals
+                .keys()
+                .filter(|key| {
+                    self.func.blocks[block]
+                        .succs
+                        .iter()
+                        .any(|succ| !self.state.block_entry[*succ].locals.contains_key(key))
+                })
+                .cloned()
+                .collect::<Vec<_>>();
+            for local in locals_to_sync {
+                let (addr, data) = pred_state.locals.get(&local).unwrap();
+                let addr = addr.value().unwrap();
+                let data = data.value().unwrap();
+                log::trace!(
+                    "spilling local {} back to real locals memory: addr {} data {}",
+                    local,
+                    addr,
+                    data
+                );
+                let args = self.func.arg_pool.double(addr, data);
+                let store = self.func.add_value(ValueDef::Operator(
+                    Operator::I64Store {
+                        memory: MemoryArg {
+                            align: 1,
+                            offset: 0,
+                            memory: self.image.main_heap().unwrap(),
+                        },
+                    },
+                    args,
+                    ListRef::default(),
+                ));
+                self.func.blocks[block].insts.push(store);
+            }
         }
     }
 

src/filter.rs

@@ -51,12 +51,14 @@ fn gen_replacement_bytecode(
         "write.global.0" => Ok(vec![wasm_encoder::Instruction::GlobalSet(global_base + 0)]),
         "write.global.1" => Ok(vec![wasm_encoder::Instruction::GlobalSet(global_base + 1)]),
         "write.global.2" => Ok(vec![wasm_encoder::Instruction::GlobalSet(global_base + 2)]),
-        
+
         // These can't be polyfilled so we rewrite them to
         // trap. They're only used in template-specialized variants
         // fed to weval requests.
         "read.reg" | "write.reg" | "push.stack" | "pop.stack" | "read.stack" | "write.stack"
-            | "sync.stack" => Ok(vec![wasm_encoder::Instruction::Unreachable]),
+        | "sync.stack" | "read.local" | "write.local" => {
+            Ok(vec![wasm_encoder::Instruction::Unreachable])
+        }
 
         // All other intrinsics have "pass through first arg" behavior
         // if they have a return value, and otherwise have no effect.

src/intrinsics.rs

@@ -26,6 +26,8 @@ pub struct Intrinsics {
     pub read_stack: Option<Func>,
     pub write_stack: Option<Func>,
     pub pop_stack: Option<Func>,
+    pub read_local: Option<Func>,
+    pub write_local: Option<Func>,
 }
 
 impl Intrinsics {
@@ -84,6 +86,18 @@ impl Intrinsics {
                 &[],
             ),
             pop_stack: find_imported_intrinsic(module, "pop.stack", &[Type::I32], &[Type::I64]),
+            read_local: find_imported_intrinsic(
+                module,
+                "read.local",
+                &[Type::I32, Type::I32],
+                &[Type::I64],
+            ),
+            write_local: find_imported_intrinsic(
+                module,
+                "write.local",
+                &[Type::I32, Type::I32, Type::I64],
+                &[],
+            ),
         }
     }
 }

src/state.rs

@@ -109,11 +109,16 @@ pub struct ProgPointState {
     ///
     /// Each entry is an (address, data) pair.
     pub stack: Vec<(RegValue, RegValue)>,
+    /// Virtualized locals, with (address, data) pairs for spilling
+    /// back to memory at sync points.
+    pub locals: BTreeMap<u32, (RegValue, RegValue)>,
 }
 
 #[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub enum RegSlot {
     Register(u32),
+    LocalAddr(u32),
+    LocalData(u32),
     StackData(u32),
     StackAddr(u32),
 }
@@ -288,6 +293,7 @@ impl ProgPointState {
             regs: BTreeMap::new(),
             globals,
             stack: vec![],
+            locals: BTreeMap::new(),
         }
     }
 
@@ -315,6 +321,13 @@ impl ProgPointState {
             this.1 = new_data;
         }
 
+        changed |= map_meet_with(
+            &mut self.locals,
+            &other.locals,
+            |(a0, a1), (b0, b1)| (RegValue::meet(a0, b0), RegValue::meet(a1, b1)),
+            None,
+        );
+
         changed
     }
 
@@ -337,6 +350,10 @@ impl ProgPointState {
             create_merge(addr);
             create_merge(data);
         }
+        for (addr, data) in self.locals.values_mut() {
+            create_merge(addr);
+            create_merge(data);
+        }
     }
 
     pub fn update_at_block_entry<C, GB: FnMut(&mut C, RegSlot, Type) -> Value>(
@@ -362,6 +379,10 @@ impl ProgPointState {
             handle_value(RegSlot::StackAddr(i as u32), addr);
             handle_value(RegSlot::StackData(i as u32), data);
         }
+        for (i, (addr, value)) in self.locals.iter_mut() {
+            handle_value(RegSlot::LocalAddr(*i), addr);
+            handle_value(RegSlot::LocalData(*i), value);
+        }
 
         Ok(())
     }