Do over/underflow checks for addition (SafeMath)

Related to #153

analyzer/src/namespace/types.rs

@@ -12,6 +12,7 @@ use std::num::{
 
 use crate::FunctionAttributes;
 use num_bigint::BigInt;
+use strum::IntoStaticStr;
 
 pub fn u256_max() -> BigInt {
     BigInt::from(2).pow(256) - 1
@@ -125,7 +126,7 @@ pub enum Base {
     Address,
 }
 
-#[derive(Clone, Debug, PartialEq, PartialOrd, Ord, Eq)]
+#[derive(Clone, Debug, PartialEq, PartialOrd, Ord, Eq, IntoStaticStr)]
 pub enum Integer {
     U256,
     U128,

compiler/src/yul/mappers/assignments.rs

@@ -128,7 +128,7 @@ mod tests {
     #[rstest(
         assignment,
         expected_yul,
-        case("foo = 1 + 2", "$foo := add(1, 2)"),
+        case("foo = 1 + 2", "$foo := checked_add_u256(1, 2)"),
         case("foo = 1 - 2", "$foo := sub(1, 2)"),
         case("foo = 1 * 2", "$foo := mul(1, 2)"),
         case("foo = 1 / 2", "$foo := div(1, 2)"),

compiler/src/yul/mappers/expressions.rs

@@ -11,8 +11,10 @@ use fe_analyzer::builtins::{
     GlobalMethod,
 };
 use fe_analyzer::namespace::types::{
+    Base,
     FeSized,
     FixedSize,
+    Integer,
     Type,
 };
 use fe_analyzer::{
@@ -239,7 +241,45 @@ pub fn expr_bin_operation(
             .typ;
 
         return match op.node {
-            fe::BinOperator::Add => Ok(expression! { add([yul_left], [yul_right]) }),
+            fe::BinOperator::Add => match typ {
+                Type::Base(Base::Numeric(Integer::I256)) => {
+                    Ok(expression! { checked_add_i256([yul_left], [yul_right]) })
+                }
+                Type::Base(Base::Numeric(Integer::I128)) => {
+                    Ok(expression! { checked_add_i128([yul_left], [yul_right]) })
+                }
+                Type::Base(Base::Numeric(Integer::I64)) => {
+                    Ok(expression! { checked_add_i64([yul_left], [yul_right]) })
+                }
+                Type::Base(Base::Numeric(Integer::I32)) => {
+                    Ok(expression! { checked_add_i32([yul_left], [yul_right]) })
+                }
+                Type::Base(Base::Numeric(Integer::I16)) => {
+                    Ok(expression! { checked_add_i16([yul_left], [yul_right]) })
+                }
+                Type::Base(Base::Numeric(Integer::I8)) => {
+                    Ok(expression! { checked_add_i8([yul_left], [yul_right]) })
+                }
+                Type::Base(Base::Numeric(Integer::U256)) => {
+                    Ok(expression! { checked_add_u256([yul_left], [yul_right]) })
+                }
+                Type::Base(Base::Numeric(Integer::U128)) => {
+                    Ok(expression! { checked_add_u128([yul_left], [yul_right]) })
+                }
+                Type::Base(Base::Numeric(Integer::U64)) => {
+                    Ok(expression! { checked_add_u64([yul_left], [yul_right]) })
+                }
+                Type::Base(Base::Numeric(Integer::U32)) => {
+                    Ok(expression! { checked_add_u32([yul_left], [yul_right]) })
+                }
+                Type::Base(Base::Numeric(Integer::U16)) => {
+                    Ok(expression! { checked_add_u16([yul_left], [yul_right]) })
+                }
+                Type::Base(Base::Numeric(Integer::U8)) => {
+                    Ok(expression! { checked_add_u8([yul_left], [yul_right]) })
+                }
+                _ => unimplemented!("Addition for non-numeric types not yet supported"),
+            },
             fe::BinOperator::Sub => Ok(expression! { sub([yul_left], [yul_right]) }),
             fe::BinOperator::Mult => Ok(expression! { mul([yul_left], [yul_right]) }),
             fe::BinOperator::Div => match typ.is_signed_integer() {
@@ -651,7 +691,7 @@ mod tests {
     #[rstest(
         expression,
         expected_yul,
-        case("1 + 2", "add(1, 2)"),
+        case("1 + 2", "checked_add_u256(1, 2)"),
         case("1 - 2", "sub(1, 2)"),
         case("1 * 2", "mul(1, 2)"),
         case("1 / 2", "div(1, 2)"),

compiler/src/yul/runtime/functions/math.rs

@@ -0,0 +1,128 @@
+use fe_analyzer::namespace::types::Integer;
+use yultsur::*;
+
+fn checked_add_unsigned(size: Integer, max_value: yul::Expression) -> yul::Statement {
+    if size.is_signed() {
+        panic!("Expected unsigned integer")
+    }
+    let size: &str = size.into();
+    let fn_name = identifier! {(format!("checked_add_{}", size.to_lowercase()))};
+    function_definition! {
+        function [fn_name](val1, val2) -> sum {
+            // overflow, if val1 > (max_value - val2)
+            (if (gt(val1, (sub([max_value], val2)))) { (revert(0, 0)) })
+            (sum := add(val1, val2))
+        }
+    }
+}
+
+/// Add two u256 numbers. Revert if result overflows.
+pub fn checked_add_u256() -> yul::Statement {
+    checked_add_unsigned(
+        Integer::U256,
+        literal_expression! {0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff},
+    )
+}
+
+/// Add two u128 numbers. Revert if result overflows.
+pub fn checked_add_u128() -> yul::Statement {
+    checked_add_unsigned(
+        Integer::U128,
+        literal_expression! {0xffffffffffffffffffffffffffffffff},
+    )
+}
+
+/// Add two u64 numbers. Revert if result overflows.
+pub fn checked_add_u64() -> yul::Statement {
+    checked_add_unsigned(Integer::U64, literal_expression! {0xffffffffffffffff})
+}
+
+/// Add two u32 numbers. Revert if result overflows.
+pub fn checked_add_u32() -> yul::Statement {
+    checked_add_unsigned(Integer::U32, literal_expression! {0xffffffff})
+}
+
+/// Add two u16 numbers. Revert if result overflows.
+pub fn checked_add_u16() -> yul::Statement {
+    checked_add_unsigned(Integer::U16, literal_expression! {0xffff})
+}
+
+/// Add two u8 numbers. Revert if result overflows.
+pub fn checked_add_u8() -> yul::Statement {
+    checked_add_unsigned(Integer::U8, literal_expression! {0xff})
+}
+
+fn checked_add_signed(
+    size: Integer,
+    min_value: yul::Expression,
+    max_value: yul::Expression,
+) -> yul::Statement {
+    if !size.is_signed() {
+        panic!("Expected signed integer")
+    }
+    let size: &str = size.into();
+    let fn_name = identifier! {(format!("checked_add_{}", size.to_lowercase()))};
+    function_definition! {
+        function [fn_name](val1, val2) -> sum {
+            // overflow, if val1 >= 0 and val2 > (max_value - val1)
+            (if (and((iszero((slt(val1, 0)))), (sgt(val2, (sub([max_value], val1)))))) { (revert(0, 0)) })
+            // underflow, if val1 < 0 and val2 < (min_val - val1)
+            (if (and((slt(val1, 0)), (slt(val2, (sub([min_value], val1)))))) { (revert(0, 0)) })
+            (sum := add(val1, val2))
+        }
+    }
+}
+
+/// Add two i256 numbers. Revert if result over- or underflows.
+pub fn checked_add_i256() -> yul::Statement {
+    checked_add_signed(
+        Integer::I256,
+        literal_expression! {0x8000000000000000000000000000000000000000000000000000000000000000},
+        literal_expression! {0x7fffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff},
+    )
+}
+
+/// Add two i128 numbers. Revert if result over- or underflows.
+pub fn checked_add_i128() -> yul::Statement {
+    checked_add_signed(
+        Integer::I128,
+        literal_expression! {0xffffffffffffffffffffffffffffffff80000000000000000000000000000000},
+        literal_expression! {0x7fffffffffffffffffffffffffffffff},
+    )
+}
+
+/// Add two i64 numbers. Revert if result over- or underflows.
+pub fn checked_add_i64() -> yul::Statement {
+    checked_add_signed(
+        Integer::I64,
+        literal_expression! {0xffffffffffffffffffffffffffffffffffffffffffffffff8000000000000000},
+        literal_expression! {0x7fffffffffffffff},
+    )
+}
+
+/// Add two i32 numbers. Revert if result over- or underflows.
+pub fn checked_add_i32() -> yul::Statement {
+    checked_add_signed(
+        Integer::I32,
+        literal_expression! {0xffffffffffffffffffffffffffffffffffffffffffffffffffffffff80000000},
+        literal_expression! {0x7fffffff},
+    )
+}
+
+/// Add two i16 numbers. Revert if result over- or underflows.
+pub fn checked_add_i16() -> yul::Statement {
+    checked_add_signed(
+        Integer::I16,
+        literal_expression! {0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff8000},
+        literal_expression! {0x7fff},
+    )
+}
+
+/// Add two i8 numbers. Revert if result over- or underflows.
+pub fn checked_add_i8() -> yul::Statement {
+    checked_add_signed(
+        Integer::I8,
+        literal_expression! {0xffffffffffffffffffffffffffffffffffffffffffffffffffffffffffffff80},
+        literal_expression! {0x7f},
+    )
+}

compiler/src/yul/runtime/functions/mod.rs

@@ -4,6 +4,7 @@ use yultsur::*;
 pub mod abi;
 pub mod contracts;
 pub mod data;
+pub mod math;
 pub mod structs;
 
 /// Returns all functions that should be available during runtime.
@@ -35,5 +36,17 @@ pub fn std() -> Vec<yul::Statement> {
         abi::pack(AbiDecodeLocation::Memory),
         contracts::create2(),
         contracts::create(),
+        math::checked_add_u256(),
+        math::checked_add_u128(),
+        math::checked_add_u64(),
+        math::checked_add_u32(),
+        math::checked_add_u16(),
+        math::checked_add_u8(),
+        math::checked_add_i256(),
+        math::checked_add_i128(),
+        math::checked_add_i64(),
+        math::checked_add_i32(),
+        math::checked_add_i16(),
+        math::checked_add_i8(),
     ]
 }

compiler/tests/evm_contracts.rs

@@ -584,78 +584,7 @@ fn test_numeric_sizes() {
     with_executor(&|mut executor| {
         let harness = deploy_contract(&mut executor, "numeric_sizes.fe", "Foo", &[]);
 
-        struct SizeConfig {
-            size: usize,
-            u_min: ethabi::Token,
-            i_min: ethabi::Token,
-            u_max: ethabi::Token,
-            i_max: ethabi::Token,
-        }
-
-        let zero = uint_token(0);
-        let u64_max = ethabi::Token::Uint(U256::from(2).pow(U256::from(64)) - 1);
-        let i64_min = ethabi::Token::Int(get_2s_complement_for_negative(
-            U256::from(2).pow(U256::from(63)),
-        ));
-
-        let u128_max = ethabi::Token::Uint(U256::from(2).pow(U256::from(128)) - 1);
-        let i128_max = ethabi::Token::Int(U256::from(2).pow(U256::from(127)) - 1);
-        let i128_min = ethabi::Token::Int(get_2s_complement_for_negative(
-            U256::from(2).pow(U256::from(127)),
-        ));
-
-        let u256_max = ethabi::Token::Uint(U256::MAX);
-        let i256_max = ethabi::Token::Int(U256::from(2).pow(U256::from(255)) - 1);
-        let i256_min = ethabi::Token::Int(get_2s_complement_for_negative(
-            U256::from(2).pow(U256::from(255)),
-        ));
-
-        let sizes = [
-            SizeConfig {
-                size: 8,
-                u_min: zero.clone(),
-                i_min: int_token(-128),
-                u_max: uint_token(255),
-                i_max: int_token(127),
-            },
-            SizeConfig {
-                size: 16,
-                u_min: zero.clone(),
-                i_min: int_token(-32768),
-                u_max: uint_token(65535),
-                i_max: int_token(32767),
-            },
-            SizeConfig {
-                size: 32,
-                u_min: zero.clone(),
-                i_min: int_token(-2147483648),
-                u_max: uint_token(4294967295),
-                i_max: int_token(2147483647),
-            },
-            SizeConfig {
-                size: 64,
-                u_min: zero.clone(),
-                i_min: i64_min.clone(),
-                u_max: u64_max.clone(),
-                i_max: int_token(9223372036854775807),
-            },
-            SizeConfig {
-                size: 128,
-                u_min: zero.clone(),
-                i_min: i128_min.clone(),
-                u_max: u128_max.clone(),
-                i_max: i128_max.clone(),
-            },
-            SizeConfig {
-                size: 256,
-                u_min: zero.clone(),
-                i_min: i256_min.clone(),
-                u_max: u256_max.clone(),
-                i_max: i256_max.clone(),
-            },
-        ];
-
-        for config in sizes.iter() {
+        for config in NumericAbiTokenBounds::get_all().iter() {
             harness.test_function(
                 &mut executor,
                 &format!("get_u{}_min", config.size),
@@ -707,6 +636,65 @@ fn sized_vals_in_sto() {
     });
 }
 
+#[test]
+fn checked_arithmetic() {
+    with_executor(&|mut executor| {
+        let harness = deploy_contract(
+            &mut executor,
+            "checked_arithmetic.fe",
+            "CheckedArithmetic",
+            &[],
+        );
+
+        for config in NumericAbiTokenBounds::get_all().iter() {
+            // unsigned: max_value + 1 fails
+            harness.test_function_reverts(
+                &mut executor,
+                &format!("add_u{}", config.size),
+                &[config.u_max.clone(), uint_token(1)],
+            );
+
+            // unsigned: max_value + 0 works
+            harness.test_function(
+                &mut executor,
+                &format!("add_u{}", config.size),
+                &[config.u_max.clone(), uint_token(0)],
+                Some(&config.u_max),
+            );
+
+            // signed: max_value + 1 fails
+            harness.test_function_reverts(
+                &mut executor,
+                &format!("add_i{}", config.size),
+                &[config.i_max.clone(), int_token(1)],
+            );
+
+            // signed: max_value + 0 works
+            harness.test_function(
+                &mut executor,
+                &format!("add_i{}", config.size),
+                &[config.i_max.clone(), int_token(0)],
+                Some(&config.i_max),
+            );
+
+            // signed: min_value + -1 fails
+            harness.test_function_reverts(
+                &mut executor,
+                &format!("add_i{}", config.size),
+                &[config.i_min.clone(), int_token(-1)],
+            );
+
+            // signed: min_value + 0 works
+            harness.test_function(
+                &mut executor,
+                &format!("add_i{}", config.size),
+                &[config.i_min.clone(), int_token(0)],
+                Some(&config.i_min),
+            );
+        }
+    });
+}
+
 #[test]
 fn structs() {
     with_executor(&|mut executor| {