1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113

//! # Anchor Safe Math
//!
//! `anchor_safe_math` is a collection of helper numeric operation functions that removes the 
//! verbosity of checking for overflow, underflow and division by zero errors.
//! 
//! # Examples
//!
//! ```
//! use anchor_lang::prelude::;
//! use anchor_safe_math::{SafeMath};
//! 
//! #[program]
//! pub mod example {
//!   use super::*;
//! 
//!   pub fn instruction(ctx: Context<Instruction>, amount: u64) -> ProgramResult {
//!     let state = &mut ctx.accounts.state;
//! 
//!     // You can apply any of the following operations
//!     state.total_amount = state.total_amount.safe_add(amount)?;
//!     state.total_amount = state.total_amount.safe_sub(amount)?;
//!     state.total_amount = state.total_amount.safe_mul(amount)?;
//!     state.total_amount = state.total_amount.safe_div(amount)?;
//!     state.total_amount = state.total_amount.safe_pow(8_u32)?;
//!   }
//! }
//! 
//! #[derive(Accounts)]
//! pub struct Instruction<'info> {
//!   ...
//! }
//! ```
use anchor_lang::prelude::{
  error,
  error_code,
  ProgramError,
};

use std::{
  result::Result as StdResult
};

#[error_code]
/// Errors that can be triggered by executing one of the supported numeric operations
pub enum ErrorCode {
  #[msg("overflow")]
  Overflow,
  #[msg("underflow")]
  Underflow,
  #[msg("division by zero")]
  DivisionByZero,
}

/// Defines a set of safe math operations that return a `ProgramError` which is expected in an anchor instruction execution.
pub trait SafeMath {
  type Output;

  fn safe_add(&self, rhs: Self::Output) -> StdResult<Self::Output, ProgramError>;
  fn safe_sub(&self, rhs: Self::Output) -> StdResult<Self::Output, ProgramError>;
  fn safe_div(&self, rhs: Self::Output) -> StdResult<Self::Output, ProgramError>;
  fn safe_mul(&self, rhs: Self::Output) -> StdResult<Self::Output, ProgramError>;
  fn safe_pow(&self, exp: u32) -> StdResult<Self::Output, ProgramError>;
}

macro_rules! safe_math {
  ($type: ident) => {
    /// $type implementation of the SafeMath trait
    impl SafeMath for $type {
      type Output = $type;

      fn safe_add(&self, rhs: Self::Output) -> StdResult<Self::Output, ProgramError> {
        match self.checked_add(rhs) {
          Some(result) => Ok(result),
          None => return Err(error!(ErrorCode::Overflow).into())
        }
      }
    
      fn safe_sub(&self, rhs: Self::Output) -> StdResult<Self::Output, ProgramError> {
        match self.checked_sub(rhs) {
          Some(result) => Ok(result),
          None => return Err(error!(ErrorCode::Underflow).into())
        }
      }

      fn safe_mul(&self, rhs: Self::Output) -> StdResult<Self::Output, ProgramError> {
        match self.checked_mul(rhs) {
          Some(result) => Ok(result),
          None => return Err(error!(ErrorCode::Underflow).into())
        }
      }

      fn safe_div(&self, rhs: Self::Output) -> StdResult<Self::Output, ProgramError> {
        match self.checked_div(rhs) {
          Some(result) => Ok(result),
          None => return Err(error!(ErrorCode::DivisionByZero).into())
        }
      }

      fn safe_pow(&self, exp: u32) -> StdResult<Self::Output, ProgramError> {
        match self.checked_pow(exp) {
          Some(result) => Ok(result),
          None => return Err(error!(ErrorCode::Overflow).into())
        }
      }
    }
  }
}

safe_math!(u128);
safe_math!(u64);
safe_math!(u32);
safe_math!(u16);
safe_math!(u8);