add traits, passes, and utils from more_waffle

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
-name = "waffle"
+name = "portal-pc-waffle"
-version = "0.0.27"
+version = "0.0.27+portal"
 description = "Wasm Analysis Framework For Lightweight Experiments"
 authors = ["Chris Fallin <chris@cfallin.org>"]
 license = "Apache-2.0 WITH LLVM-exception"
@@ -29,3 +29,6 @@ wasm-smith = { version = "0.202", optional = true }
 [features]
 default = []
 fuzzing = ["libfuzzer-sys", "wasm-smith"]
+
+[lib]
+name = "waffle"

src/ir/module.rs

@@ -205,6 +205,40 @@ impl<'a> Module<'a> {
         }
     }
 
+    pub fn take_per_func_body<F: FnMut(&mut Self,&mut FunctionBody)>(&mut self, mut f: F) {
+        for func_decl in self.funcs.iter().collect::<Vec<_>>() {
+            let mut x = std::mem::take(&mut self.funcs[func_decl]);
+            if let Some(body) = x.body_mut(){
+                f(self,body);
+            }
+            self.funcs[func_decl] = x;
+        }
+    }
+
+    pub fn try_per_func_body<F: FnMut(&mut FunctionBody) -> Result<(),E>,E>(&mut self, mut f: F) -> Result<(),E>{
+        for func_decl in self.funcs.values_mut() {
+            if let Some(body) = func_decl.body_mut() {
+                f(body)?;
+            }
+        }
+        Ok(())
+    }
+
+    pub fn try_take_per_func_body<F: FnMut(&mut Self,&mut FunctionBody) -> Result<(),E>,E>(&mut self, mut f: F) -> Result<(),E>{
+        for func_decl in self.funcs.iter().collect::<Vec<_>>() {
+            let mut x = std::mem::take(&mut self.funcs[func_decl]);
+            let mut y = None;
+            if let Some(body) = x.body_mut(){
+                y = Some(f(self,body));
+            }
+            self.funcs[func_decl] = x;
+            if let Some(z) = y{
+                z?;
+            }
+        }
+        Ok(())
+    }
+
     pub fn expand_func<'b>(&'b mut self, id: Func) -> Result<&'b mut FuncDecl<'a>> {
         if let FuncDecl::Lazy(..) = self.funcs[id] {
             // End the borrow. This is cheap (a slice copy).

src/lib.rs

@@ -23,6 +23,7 @@ pub use ops::{Ieee32, Ieee64, MemoryArg, Operator};
 
 mod interp;
 pub use interp::*;
+pub mod util;
 
 pub use passes::basic_opt::OptOptions;
 

src/op_traits.rs

@@ -2151,3 +2151,122 @@ pub fn op_rematerialize(op: &Operator) -> bool {
         _ => false,
     }
 }
+
+pub fn rewrite_mem<V,E>(
+    o: &mut Operator,
+    v: &mut [V],
+    mut go: impl FnMut(&mut crate::Memory, Option<&mut V>) -> Result<(),E>,
+) -> Result<(),E>{
+    match o {
+        Operator::I32Load { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::I64Load { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::F32Load { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::F64Load { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::I32Load8S { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::I32Load8U { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::I32Load16S { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::I32Load16U { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::I64Load8S { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::I64Load8U { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::I64Load16S { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::I64Load16U { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::I64Load32S { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::I64Load32U { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::I32Store { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::I64Store { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::F32Store { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::F64Store { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::I32Store8 { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::I32Store16 { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::I64Store8 { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::I64Store16 { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::I64Store32 { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::MemorySize { mem } => go(mem, None),
+        Operator::MemoryGrow { mem } => go(mem, None),
+        Operator::MemoryCopy { dst_mem, src_mem } => {
+            go(dst_mem, Some(&mut v[0]))?;
+            go(src_mem, Some(&mut v[1]))?;
+            Ok(())
+        }
+        Operator::MemoryFill { mem } => go(mem, Some(&mut v[0])),
+
+        Operator::V128Load { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Load8x8S { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Load8x8U { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Load16x4S { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Load16x4U { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Load32x2S { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Load32x2U { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Load8Splat { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Load16Splat { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Load32Splat { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Load64Splat { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Load32Zero { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Load64Zero { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Store { memory } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Load8Lane { memory, .. } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Load16Lane { memory, .. } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Load32Lane { memory, .. } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Load64Lane { memory, .. } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Store8Lane { memory, .. } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Store16Lane { memory, .. } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Store32Lane { memory, .. } => go(&mut memory.memory, Some(&mut v[0])),
+        Operator::V128Store64Lane { memory, .. } => go(&mut memory.memory, Some(&mut v[0])),
+        _ => Ok(())
+    }
+}
+pub fn mem_count(o: &Operator) -> usize {
+    match o {
+        Operator::I32Load { memory } => 1,
+        Operator::I64Load { memory } => 1,
+        Operator::F32Load { memory } => 1,
+        Operator::F64Load { memory } => 1,
+        Operator::I32Load8S { memory } => 1,
+        Operator::I32Load8U { memory } => 1,
+        Operator::I32Load16S { memory } => 1,
+        Operator::I32Load16U { memory } => 1,
+        Operator::I64Load8S { memory } => 1,
+        Operator::I64Load8U { memory } => 1,
+        Operator::I64Load16S { memory } => 1,
+        Operator::I64Load16U { memory } => 1,
+        Operator::I64Load32S { memory } => 1,
+        Operator::I64Load32U { memory } => 1,
+        Operator::I32Store { memory } => 1,
+        Operator::I64Store { memory } => 1,
+        Operator::F32Store { memory } => 1,
+        Operator::F64Store { memory } => 1,
+        Operator::I32Store8 { memory } => 1,
+        Operator::I32Store16 { memory } => 1,
+        Operator::I64Store8 { memory } => 1,
+        Operator::I64Store16 { memory } => 1,
+        Operator::I64Store32 { memory } => 1,
+        Operator::MemorySize { mem } => 1,
+        Operator::MemoryGrow { mem } => 1,
+        Operator::MemoryCopy { dst_mem, src_mem } => 2,
+        Operator::MemoryFill { mem } => 1,
+
+        Operator::V128Load { memory } => 1,
+        Operator::V128Load8x8S { memory } => 1,
+        Operator::V128Load8x8U { memory } => 1,
+        Operator::V128Load16x4S { memory } => 1,
+        Operator::V128Load16x4U { memory } => 1,
+        Operator::V128Load32x2S { memory } => 1,
+        Operator::V128Load32x2U { memory } => 1,
+        Operator::V128Load8Splat { memory } => 1,
+        Operator::V128Load16Splat { memory } => 1,
+        Operator::V128Load32Splat { memory } => 1,
+        Operator::V128Load64Splat { memory } => 1,
+        Operator::V128Load32Zero { memory } => 1,
+        Operator::V128Load64Zero { memory } => 1,
+        Operator::V128Store { memory } => 1,
+        Operator::V128Load8Lane { memory, .. } => 1,
+        Operator::V128Load16Lane { memory, .. } => 1,
+        Operator::V128Load32Lane { memory, .. } => 1,
+        Operator::V128Load64Lane { memory, .. } => 1,
+        Operator::V128Store8Lane { memory, .. } => 1,
+        Operator::V128Store16Lane { memory, .. } => 1,
+        Operator::V128Store32Lane { memory, .. } => 1,
+        Operator::V128Store64Lane { memory, .. } => 1,
+        _ => 0,
+    }
+}

src/passes.rs

@@ -7,3 +7,5 @@ pub mod maxssa;
 pub mod resolve_aliases;
 pub mod ssa;
 pub mod trace;
+pub mod mem_fusing;
+pub mod unmem;

src/passes/mem_fusing.rs

@@ -0,0 +1,194 @@
+use std::{collections::BTreeMap, convert::Infallible, iter::empty};
+
+use anyhow::Context;
+// use libc::name_t;
+
+// use crate::append_before;
+use crate::{
+    entity::{EntityRef, EntityVec},
+    ExportKind, Func, FuncDecl, FunctionBody, Memory, Module, Operator, Type, ValueDef,
+};
+// use itertools::Itertools;
+
+pub struct Fuse {
+    pub resolve: Func,
+    pub grow: Func,
+    pub size: Func,
+}
+pub fn get_exports(m: &Module) -> BTreeMap<String, ExportKind> {
+    let mut b = BTreeMap::new();
+    for e in m.exports.iter() {
+        let e = e.clone();
+        b.insert(e.name, e.kind);
+    }
+    return b;
+}
+pub fn finalize(m: &mut Module) {
+    let mem = m.memories[Memory::new(0)].clone();
+    m.memories = EntityVec::default();
+    m.memories.push(mem);
+}
+impl Fuse {
+    pub fn new(m: &Module) -> Option<Self> {
+        let e = get_exports(m);
+        let Some(ExportKind::Func(a)) = e.get("sk%resolve") else {
+            return None;
+        };
+        let a = *a;
+        let Some(ExportKind::Func(b)) = e.get("sk%grow") else {
+            return None;
+        };
+        let b = *b;
+        let Some(ExportKind::Func(c)) = e.get("sk%size") else {
+            return None;
+        };
+        let c = *c;
+        return Some(Fuse {
+            resolve: a,
+            grow: b,
+            size: c,
+        });
+    }
+    pub fn finalize(self, m: &mut Module) {
+        let mem = m.memories[Memory::new(0)].clone();
+        let l = m.memories.len() - 1;
+        m.memories = EntityVec::default();
+        m.memories.push(mem);
+        let v = vec![self.resolve, self.grow, self.size];
+        let mut new = vec![];
+        for f in v.clone() {
+            let n = m.funcs[f].clone();
+            let s = n.sig();
+            let name = n.name().to_owned();
+            // let n = m.funcs.push(n);
+            let mut b = FunctionBody::new(&m, s);
+            let mut p = b.blocks[b.entry]
+                .params
+                .iter()
+                .map(|a| a.1)
+                .collect::<Vec<_>>();
+            let vz = b.arg_pool.from_iter(empty());
+            let tz = b.type_pool.from_iter(empty());
+            let ti = b.type_pool.from_iter(vec![Type::I32].into_iter());
+            let i = b.add_value(ValueDef::Operator(
+                Operator::I32Const { value: l as u32 },
+                vz,
+                ti,
+            ));
+            b.append_to_block(b.entry, i);
+            let i = b.arg_pool.from_iter(vec![p[p.len() - 1], i].into_iter());
+            let i = b.add_value(ValueDef::Operator(Operator::I32Add, i, ti));
+            let l = p.len();
+            p[l - 1] = i;
+            b.append_to_block(b.entry, i);
+            b.set_terminator(b.entry, crate::Terminator::ReturnCall { func: f, args: p });
+            let n = FuncDecl::Body(s, name, b);
+            let n = m.funcs.push(n);
+            new.push(n);
+        }
+        for x in m.exports.iter_mut() {
+            let ExportKind::Func(xf) = &mut x.kind else {
+                continue;
+            };
+            for (o, n) in v.iter().zip(new.iter()) {
+                if xf == o {
+                    *xf = *n
+                }
+            }
+        }
+    }
+    pub fn process(&self, f: &mut FunctionBody) {
+        let vz = f.arg_pool.from_iter(empty());
+        let tz = f.type_pool.from_iter(empty());
+        let ti = f.type_pool.from_iter(vec![Type::I32].into_iter());
+        let mut ka = BTreeMap::new();
+        for k in f.blocks.iter().collect::<Vec<_>>() {
+            // eprintln!("dbg: {v}");
+            // let k = f.value_blocks[v];
+            for v in std::mem::take(&mut f.blocks[k].insts) {
+                let mut w = f.values[v].clone();
+                // let vi = v;
+                if let ValueDef::Operator(a, b, c) = &mut w {
+                    let mut bp = f.arg_pool[*b].to_vec();
+                    match a.clone() {
+                        Operator::MemorySize { mem } => {
+                            if mem.index() != 0 {
+                                let ia = f.add_value(ValueDef::Operator(
+                                    Operator::I32Const {
+                                        value: mem.index() as u32,
+                                    },
+                                    vz,
+                                    ti,
+                                ));
+                                // append_before(f, ia, vi, k);
+                                f.append_to_block(k, ia);
+                                *a = Operator::Call {
+                                    function_index: self.size,
+                                };
+                                bp.push(ia);
+                            }
+                        }
+                        Operator::MemoryGrow { mem } => {
+                            if mem.index() != 0 {
+                                let ia = f.add_value(ValueDef::Operator(
+                                    Operator::I32Const {
+                                        value: mem.index() as u32,
+                                    },
+                                    vz,
+                                    ti,
+                                ));
+                                // append_before(f, ia, vi, k);
+                                f.append_to_block(k, ia);
+                                *a = Operator::Call {
+                                    function_index: self.grow,
+                                };
+                                bp.push(ia);
+                            }
+                        }
+                        _ => crate::op_traits::rewrite_mem(a, &mut bp, |m, v| {
+                            if m.index() != 0 {
+                                let ia = f.add_value(ValueDef::Operator(
+                                    Operator::I32Const {
+                                        value: m.index() as u32,
+                                    },
+                                    vz,
+                                    ti,
+                                ));
+                                f.append_to_block(k, ia); 
+                                // append_before(f, ia, vi, k);
+                                if let Some(v) = v {
+                                    let w = f.arg_pool.from_iter(vec![*v, ia].into_iter());
+                                    let x = f.add_value(ValueDef::Operator(
+                                        Operator::Call {
+                                            function_index: self.resolve,
+                                        },
+                                        w,
+                                        ti,
+                                    ));
+                                    f.append_to_block(k, x);
+                                    // crate::append_before(f, x, vi, k);
+                                    *v = x;
+                                }
+                                *m = Memory::new(0);
+                            }
+                            Ok::<(),Infallible>(())
+                        }).unwrap(),
+                    }
+                    *b = *ka
+                        .entry(bp.clone())
+                        .or_insert_with(|| f.arg_pool.from_iter(bp.into_iter()));
+                }
+                f.values[v] = w;
+                f.append_to_block(k, v);
+            }
+        }
+    }
+}
+pub fn fuse(m: &mut Module) -> anyhow::Result<()> {
+    let f = Fuse::new(m).context("in getting the fuse funcs")?;
+    crate::passes::unmem::metafuse_all(m, &mut crate::passes::unmem::All {});
+    // crate::passes::splice::splice_module(m)?;
+    m.per_func_body(|b| f.process(b));
+    f.finalize(m);
+    return Ok(());
+}

src/passes/unmem.rs

@@ -0,0 +1,152 @@
+use crate::util::{add_start, new_sig};
+use std::{collections::BTreeMap, iter::empty};
+use crate::{
+    Block, ExportKind, Func, FuncDecl, FunctionBody, Import, ImportKind, Memory, MemoryArg,
+    MemoryData, Module, Operator, Signature, SignatureData, Table, Type, Value, ValueDef,
+};
+
+// use super::reload::{ImportedCfg, ImportedMemoriesToFunc};
+
+pub fn fuse_iter(m: &mut Module, x: impl Iterator<Item = (usize, u8)>, mem: Memory) {
+    let null = new_sig(
+        m,
+        SignatureData {
+            params: vec![],
+            returns: vec![],
+        },
+    );
+    let mut b = FunctionBody::new(m, null);
+    let vz = b.arg_pool.from_iter(empty());
+    let tz = b.type_pool.from_iter(empty());
+    let ti = b.type_pool.from_iter(vec![Type::I32].into_iter());
+    let ia = b.add_value(ValueDef::Operator(Operator::I32Const { value: 0 }, vz, ti));
+    b.append_to_block(b.entry, ia);
+    for (a, c) in x {
+        // let ia = b.add_value(ValueDef::Operator(Operator::I32Const { value: a as u32 }, vz, ti));
+        // b.append_to_block(b.entry, ia);
+        let ic = b.add_value(ValueDef::Operator(
+            Operator::I32Const { value: c as u32 },
+            vz,
+            ti,
+        ));
+        b.append_to_block(b.entry, ic);
+        let ia = b.add_value(ValueDef::Operator(
+            Operator::I32Const { value: a as u32 },
+            vz,
+            ti,
+        ));
+        b.append_to_block(b.entry, ia);
+        let vs = b.arg_pool.from_iter(vec![ia, ic].into_iter());
+        let j = b.add_value(ValueDef::Operator(
+            Operator::I32Store8 {
+                memory: MemoryArg {
+                    align: 0,
+                    offset: 0,
+                    memory: mem,
+                },
+            },
+            vs,
+            tz,
+        ));
+        b.append_to_block(b.entry, j);
+    }
+    b.set_terminator(b.entry, crate::Terminator::Return { values: vec![] });
+    let f = m.funcs.push(FuncDecl::Body(null, format!("z"), b));
+    add_start(m, f);
+}
+pub fn metafuse_iter(m: &mut Module, x: &[(usize, u8)], mem: Memory) {
+    for w in x.chunks(4096) {
+        fuse_iter(m, w.iter().map(|a| *a), mem);
+    }
+}
+pub fn metafuse(m: &mut Module, mem: Memory, dat: MemoryData) {
+    let null = new_sig(
+        m,
+        SignatureData {
+            params: vec![],
+            returns: vec![],
+        },
+    );
+    let mut v = vec![];
+    for s in dat.segments.iter() {
+        v.extend(s.data.iter().enumerate().map(|(a, b)| (a + s.offset, *b)));
+    }
+    metafuse_iter(m, &v, mem);
+    let mut b = FunctionBody::new(m, null);
+    let vz = b.arg_pool.from_iter(empty());
+    let tz = b.type_pool.from_iter(empty());
+    let ti = b.type_pool.from_iter(vec![Type::I32].into_iter());
+    let ia = b.add_value(ValueDef::Operator(
+        Operator::I32Const {
+            value: dat.initial_pages as u32,
+        },
+        vz,
+        ti,
+    ));
+    b.append_to_block(b.entry, ia);
+    let vs = b.arg_pool.from_iter(vec![ia].into_iter());
+    let ib = b.add_value(ValueDef::Operator(
+        Operator::MemoryGrow { mem: mem },
+        vs,
+        tz,
+    ));
+    b.append_to_block(b.entry, ib);
+    b.set_terminator(b.entry, crate::Terminator::Return { values: vec![] });
+    let f = m.funcs.push(FuncDecl::Body(null, format!("z"), b));
+    add_start(m, f);
+}
+pub fn metafuse_all(m: &mut Module, cfg: &mut impl Cfg) {
+    let mut b = BTreeMap::new();
+    for mem in m
+        .memories
+        .entries_mut()
+        .collect::<Vec<_>>()
+        .into_iter()
+        .rev()
+    {
+        let memory64 = mem.1.memory64;
+        b.insert(
+            mem.0,
+            std::mem::replace(
+                mem.1,
+                MemoryData {
+                    initial_pages: 0,
+                    maximum_pages: None,
+                    segments: vec![],
+                    memory64,
+                },
+            ),
+        );
+    }
+    for mem in m.memories.iter().collect::<Vec<_>>() {
+        if !cfg.unmemmable(m, mem) {
+            m.memories[mem] = b.remove(&mem).unwrap();
+        }
+    }
+    for (c, d) in b.into_iter() {
+        metafuse(m, c, d);
+    }
+}
+pub trait Cfg {
+    fn unmemmable(&mut self, module: &mut Module, mem: Memory) -> bool;
+}
+pub struct All {}
+impl Cfg for All {
+    fn unmemmable(&mut self, module: &mut Module, mem: Memory) -> bool {
+        true
+    }
+}
+
+pub struct ImportsOnly{}
+impl Cfg for ImportsOnly{
+    fn unmemmable(&mut self, module: &mut Module, mem: Memory) -> bool {
+        for i in module.imports.iter().map(|a| a.clone()).collect::<Vec<_>>() {
+            // if self.cfg.do_lower(&i.module, &i.name) {
+                if i.kind == ImportKind::Memory(mem) {
+                    return true;
+                }
+            // }
+        }
+        false
+    }
+}

src/util.rs

@@ -0,0 +1,39 @@
+use crate::*;
+pub fn add_start(m: &mut Module, tf: Func) {
+    let s = SignatureData {
+        params: vec![],
+        returns: vec![],
+    };
+    let s = new_sig(m, s);
+    let mut f = FunctionBody::new(&m, s);
+    let vz = f.arg_pool.from_iter(std::iter::empty());
+    let t = m.funcs[tf].sig();
+    let t = m.signatures[t].clone().returns;
+    let tz = f.type_pool.from_iter(t.into_iter());
+    let v = f.add_value(ValueDef::Operator(
+        Operator::Call { function_index: tf },
+        vz,
+        tz,
+    ));
+    f.append_to_block(f.entry, v);
+    f.set_terminator(
+        f.entry,
+        match m.start_func {
+            Some(a) => Terminator::ReturnCall {
+                func: a,
+                args: vec![],
+            },
+            None => Terminator::Return { values: vec![] },
+        },
+    );
+    let f = m.funcs.push(FuncDecl::Body(s, format!("start"), f));
+    m.start_func = Some(f);
+}
+pub fn new_sig(m: &mut Module, s: SignatureData) -> Signature {
+    for (a, b) in m.signatures.entries() {
+        if *b == s {
+            return a;
+        }
+    }
+    return m.signatures.push(s);
+}