From 5477c7ddbebbee462785171bf7b2d0681a7e147a Mon Sep 17 00:00:00 2001
From: Nicholas Ward <npward@berkeley.edu>
Date: Wed, 15 Feb 2023 11:20:09 -0800
Subject: [PATCH] fixes

---
 evm/src/cpu/kernel/asm/bignum/addmul.asm | 20 +++++------
 evm/src/cpu/kernel/tests/bignum.rs       | 46 ++++++++++++++++++++++--
 2 files changed, 53 insertions(+), 13 deletions(-)

diff --git a/evm/src/cpu/kernel/asm/bignum/addmul.asm b/evm/src/cpu/kernel/asm/bignum/addmul.asm
index 3ead50c5..701a2afa 100644
--- a/evm/src/cpu/kernel/asm/bignum/addmul.asm
+++ b/evm/src/cpu/kernel/asm/bignum/addmul.asm
@@ -50,20 +50,20 @@ addmul_loop:
     ADD
     // stack: prod_hi' = prod_hi + prod_lo_carry, prod_lo', carry, i, a_cur_loc, b_cur_loc, val, retdest
     DUP3
-    // stack: carry, prod_lo', prod_hi', carry, i, a_cur_loc, b_cur_loc, val, retdest
-    DUP2
-    // stack: prod_lo', carry, prod_lo', prod_hi', carry, i, a_cur_loc, b_cur_loc, val, retdest
+    // stack: carry, prod_hi', prod_lo', carry, i, a_cur_loc, b_cur_loc, val, retdest
+    DUP3
+    // stack: prod_lo', carry, prod_hi', prod_lo', carry, i, a_cur_loc, b_cur_loc, val, retdest
     ADD
     %shl_const(128)
     %shr_const(128)
-    // stack: to_write = (prod_lo' + carry) % 2^128, prod_lo', prod_hi', carry, i, a_cur_loc, b_cur_loc, val, retdest
-    SWAP1
-    // stack: prod_lo', to_write, prod_hi', carry, i, a_cur_loc, b_cur_loc, val, retdest
-    DUP2
-    // stack: to_write, prod_lo', to_write, prod_hi', carry, i, a_cur_loc, b_cur_loc, val, retdest
+    // stack: to_write = (prod_lo' + carry) % 2^128, prod_hi', prod_lo', carry, i, a_cur_loc, b_cur_loc, val, retdest
+    SWAP2
+    // stack: prod_lo', prod_hi', to_write, carry, i, a_cur_loc, b_cur_loc, val, retdest
+    DUP3
+    // stack: to_write, prod_lo', prod_hi', to_write, carry, i, a_cur_loc, b_cur_loc, val, retdest
     LT
-    // stack: carry_new = to_write < prod_lo', to_write, prod_hi', carry, i, a_cur_loc, b_cur_loc, val, retdest
-    %stack (cn, tw, ph, c) -> (cn, ph, tw)
+    // stack: carry_new = to_write < prod_lo', prod_hi', to_write, carry, i, a_cur_loc, b_cur_loc, val, retdest
+    %stack (vals: 3, c) -> (vals)
     // stack: carry_new, prod_hi', to_write, i, a_cur_loc, b_cur_loc, val, retdest
     ADD
     // stack: carry = carry_new' + prod_hi', to_write, i, a_cur_loc, b_cur_loc, val, retdest
diff --git a/evm/src/cpu/kernel/tests/bignum.rs b/evm/src/cpu/kernel/tests/bignum.rs
index 49d28342..6f0edb4c 100644
--- a/evm/src/cpu/kernel/tests/bignum.rs
+++ b/evm/src/cpu/kernel/tests/bignum.rs
@@ -3,6 +3,7 @@ use ethereum_types::U256;
 use itertools::Itertools;
 use num::{BigUint, Signed};
 use num_bigint::RandBigInt;
+use rand::Rng;
 
 use crate::cpu::kernel::aggregator::KERNEL;
 use crate::cpu::kernel::interpreter::Interpreter;
@@ -21,7 +22,7 @@ fn pack_bignums(biguints: &[BigUint], length: usize) -> Vec<U256> {
 
 fn gen_bignum(bit_size: usize) -> BigUint {
     let mut rng = rand::thread_rng();
-    rng.gen_bigint(bit_size as u64).abs().to_biguint().unwrap()
+    rng.gen_biguint(bit_size as u64)
 }
 
 fn bignum_len(a: &BigUint) -> usize {
@@ -31,8 +32,8 @@ fn bignum_len(a: &BigUint) -> usize {
 fn gen_two_bignums_ordered(bit_size: usize) -> (BigUint, BigUint) {
     let mut rng = rand::thread_rng();
     let (a, b) = {
-        let a = rng.gen_bigint(bit_size as u64).abs().to_biguint().unwrap();
-        let b = rng.gen_bigint(bit_size as u64).abs().to_biguint().unwrap();
+        let a = rng.gen_biguint(bit_size as u64);
+        let b = rng.gen_biguint(bit_size as u64);
         (a.clone().max(b.clone()), a.min(b))
     };
 
@@ -176,6 +177,45 @@ fn test_add_bignum() -> Result<()> {
     Ok(())
 }
 
+#[test]
+fn test_addmul_bignum() -> Result<()> {
+    let mut rng = rand::thread_rng();
+    let (a, b, length, memory) = prepare_two_bignums(1000);
+    let val: u64 = 1;
+    let val_u256 = U256::from(val);
+
+    // Determine expected result.
+    let result = a + b * BigUint::from(val);
+    let expected_result: Vec<U256> = biguint_to_mem_vec(result);
+
+    let a_start_loc = 0.into();
+    let b_start_loc = length;
+
+    // Prepare stack.
+    let retdest = 0xDEADBEEFu32.into();
+    let mut initial_stack: Vec<U256> = vec![length, a_start_loc, b_start_loc, val_u256, retdest];
+    initial_stack.reverse();
+
+    // Prepare interpreter.
+    let addmul_bignum = KERNEL.global_labels["addmul_bignum"];
+    let mut interpreter = Interpreter::new_with_kernel(addmul_bignum, initial_stack);
+    interpreter.set_kernel_general_memory(memory);
+
+    // Run add function.
+    interpreter.run()?;
+
+    // Determine actual result.
+    let new_memory = interpreter.get_kernel_general_memory();
+    let actual_result: Vec<_> = new_memory[..expected_result.len()].into();
+
+    dbg!(interpreter.stack());
+
+    // Compare.
+    assert_eq!(actual_result, expected_result);
+
+    Ok(())
+}
+
 #[test]
 fn test_mul_bignum() -> Result<()> {
     let (a, b, length, memory) = prepare_two_bignums(1000);