clz/ctz/popcnt

src/backend.rs

@@ -453,9 +453,9 @@ type Stack = Vec<StackValue>;
 
 pub enum MemoryAccessMode {
     /// This is slower than using `Unchecked` mode, but works in
-    /// any scenario, running on a system that can't index more
+    /// any scenario (the most important scenario being when we're
-    /// memory than the compiled Wasm can being the most important
+    /// running on a system that can't index much more memory than
-    /// one.
+    /// the Wasm).
     Checked,
     /// This means that checks are _not emitted by the compiler_!
     /// If you're using WebAssembly to run untrusted code, you
@@ -491,6 +491,35 @@ impl StackDepth {
     }
 }
 
+macro_rules! unop {
+    ($name:ident, $instr:ident, $reg_ty:ident, $typ:ty, $const_fallback:expr) => {
+        pub fn $name(&mut self) {
+            let val = self.pop();
+
+            let out_val = match val.location(&self.block_state.locals) {
+                ValueLocation::Immediate(imm) => Value::Immediate($const_fallback(imm as $typ) as _),
+                ValueLocation::Stack(offset) => {
+                    let offset = self.adjusted_offset(offset);
+                    let temp = self.block_state.regs.take_scratch_gpr();
+                    dynasm!(self.asm
+                        ; $instr $reg_ty(temp), [rsp + offset]
+                    );
+                    Value::Temp(temp)
+                }
+                ValueLocation::Reg(reg) => {
+                    let temp = self.block_state.regs.take_scratch_gpr();
+                    dynasm!(self.asm
+                        ; $instr $reg_ty(temp), $reg_ty(reg)
+                    );
+                    Value::Temp(temp)
+                }
+            };
+
+            self.push(out_val);
+        }
+    }
+}
+
 macro_rules! cmp_i32 {
     ($name:ident, $instr:ident, $reverse_instr:ident, $const_fallback:expr) => {
         pub fn $name(&mut self) {
@@ -1476,6 +1505,13 @@ impl Context<'_> {
         }
     }
 
+    unop!(i32_clz, lzcnt, Rd, u32, u32::leading_zeros);
+    unop!(i64_clz, lzcnt, Rq, u64, |a: u64| a.leading_zeros() as u64);
+    unop!(i32_ctz, tzcnt, Rd, u32, u32::trailing_zeros);
+    unop!(i64_ctz, tzcnt, Rq, u64, |a: u64| a.trailing_zeros() as u64);
+    unop!(i32_popcnt, popcnt, Rd, u32, u32::count_ones);
+    unop!(i64_popcnt, popcnt, Rq, u64, |a: u64| a.count_ones() as u64);
+
     // TODO: Use `lea` when the LHS operand isn't a temporary but both of the operands
     //       are in registers.
     commutative_binop_i32!(i32_add, add, |a, b| (a as i32).wrapping_add(b as i32));

src/function_body.rs

@@ -395,6 +395,9 @@ pub fn translate(
             Operator::I32Or => ctx.i32_or(),
             Operator::I32Xor => ctx.i32_xor(),
             Operator::I32Mul => ctx.i32_mul(),
+            Operator::I32Clz => ctx.i32_clz(),
+            Operator::I32Ctz => ctx.i32_ctz(),
+            Operator::I32Popcnt => ctx.i32_popcnt(),
             Operator::I64Eq => ctx.i64_eq(),
             Operator::I64Eqz => ctx.i64_eqz(),
             Operator::I64Ne => ctx.i64_neq(),
@@ -412,6 +415,9 @@ pub fn translate(
             Operator::I64Or => ctx.i64_or(),
             Operator::I64Xor => ctx.i64_xor(),
             Operator::I64Mul => ctx.i64_mul(),
+            Operator::I64Clz => ctx.i64_clz(),
+            Operator::I64Ctz => ctx.i64_ctz(),
+            Operator::I64Popcnt => ctx.i64_popcnt(),
             Operator::Drop => ctx.drop(),
             Operator::SetLocal { local_index } => ctx.set_local(local_index),
             Operator::GetLocal { local_index } => ctx.get_local(local_index),

src/module.rs

@@ -131,10 +131,9 @@ impl TranslatedModule {
             });
 
         let mem_size = self.memory.map(|m| m.limits.initial).unwrap_or(0) as usize;
-        let layout = Layout::new::<VmCtx>()
+        let (layout, _mem_offset) = Layout::new::<VmCtx>()
             .extend(Layout::array::<u8>(mem_size * WASM_PAGE_SIZE).unwrap())
-            .unwrap()
+            .unwrap();
-            .0;
 
         let ptr = unsafe { alloc::alloc_zeroed(layout) } as *mut VmCtx;
 

src/tests.rs

@@ -32,10 +32,8 @@ mod op32 {
 
                 lazy_static! {
                     static ref AS_PARAMS: ExecutableModule = translate_wat(&format!(
-                        "
+                        "(module (func (param i32) (param i32) (result i32)
-                        (module (func (param i32) (param i32) (result i32)
+                            (i32.{op} (get_local 0) (get_local 1))))",
-                            (i32.{op} (get_local 0) (get_local 1))))
-                    ",
                         op = OP
                     ));
                 }
@@ -82,6 +80,43 @@ mod op32 {
         };
     }
 
+    macro_rules! unop_test {
+        ($name:ident, $func:expr) => {
+            mod $name {
+                use super::{translate_wat, ExecutableModule};
+                use std::sync::Once;
+        
+                lazy_static! {
+                    static ref AS_PARAM: ExecutableModule = translate_wat(
+                        concat!("(module (func (param i32) (result i32)
+                            (i32.",stringify!($name)," (get_local 0))))"),
+                    );
+                }
+        
+                quickcheck! {
+                    fn as_param(a: u32) -> bool {
+                         AS_PARAM.execute_func::<(u32,), u32>(0, (a,)) == Ok($func(a))
+                    }
+        
+                    fn lit(a: u32) -> bool {
+                        let translated = translate_wat(&format!(concat!("
+                            (module (func (result i32)
+                                (i32.",stringify!($name)," (i32.const {val}))))
+                        "), val = a));
+                        static ONCE: Once = Once::new();
+                        ONCE.call_once(|| translated.disassemble());
+        
+                        translated.execute_func::<(), u32>(0, ()) == Ok($func(a))
+                    }
+                }
+            }
+        }
+    }
+
+    unop_test!(clz, u32::leading_zeros);
+    unop_test!(ctz, u32::trailing_zeros);
+    unop_test!(popcnt, u32::count_ones);
+
     binop_test!(add, i32::wrapping_add);
     binop_test!(sub, i32::wrapping_sub);
     binop_test!(and, std::ops::BitAnd::bitand);
@@ -157,6 +192,43 @@ mod op64 {
         };
     }
 
+    macro_rules! unop_test {
+        ($name:ident, $func:expr) => {
+            mod $name {
+                use super::{translate_wat, ExecutableModule};
+                use std::sync::Once;
+        
+                lazy_static! {
+                    static ref AS_PARAM: ExecutableModule = translate_wat(
+                        concat!("(module (func (param i64) (result i64)
+                            (i64.",stringify!($name)," (get_local 0))))"),
+                    );
+                }
+        
+                quickcheck! {
+                    fn as_param(a: u64) -> bool {
+                         AS_PARAM.execute_func::<(u64,), u64>(0, (a,)) == Ok($func(a))
+                    }
+        
+                    fn lit(a: u64) -> bool {
+                        let translated = translate_wat(&format!(concat!("
+                            (module (func (result i64)
+                                (i64.",stringify!($name)," (i64.const {val}))))
+                        "), val = a));
+                        static ONCE: Once = Once::new();
+                        ONCE.call_once(|| translated.disassemble());
+        
+                        translated.execute_func::<(), u64>(0, ()) == Ok($func(a))
+                    }
+                }
+            }
+        }
+    }
+
+    unop_test!(clz, |a: u64| a.leading_zeros() as _);
+    unop_test!(ctz, |a: u64| a.trailing_zeros() as _);
+    unop_test!(popcnt, |a: u64| a.count_ones() as _);
+
     binop_test!(add, i64::wrapping_add);
     binop_test!(sub, i64::wrapping_sub);
     binop_test!(and, std::ops::BitAnd::bitand);
@@ -698,7 +770,7 @@ fn wrong_index() {
     );
 }
 
-    fn iterative_fib_baseline(n: u32) -> u32 {
+fn iterative_fib_baseline(n: u32) -> u32 {
     let (mut a, mut b) = (1, 1);
 
     for _ in 0..n {
@@ -708,7 +780,7 @@ fn wrong_index() {
     }
 
     a
-    }
+}
 
 const FIBONACCI: &str = r#"
 (module
@@ -973,8 +1045,14 @@ fn call_indirect() {
 
     module.disassemble();
 
-    assert_eq!(module.execute_func::<(i32, i64), i64>(0, (1, 10)).unwrap(), 3628800);
+    assert_eq!(
-    assert_eq!(module.execute_func::<(i32, i64), i64>(0, (2, 10)).unwrap(), 89);
+        module.execute_func::<(i32, i64), i64>(0, (1, 10)).unwrap(),
+        3628800
+    );
+    assert_eq!(
+        module.execute_func::<(i32, i64), i64>(0, (2, 10)).unwrap(),
+        89
+    );
 }
 
 #[bench]