{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE DeriveAnyClass #-}

module SymbolicExpression where

import Data.Word
import Data.Map.Strict (Map)
import qualified Data.Map.Strict as Map
import GHC.Generics (Generic)
import Data.Bits hiding (And, Xor)

data SymExpr
  = Concrete Integer
  | Symbolic String Int
  | Add SymExpr SymExpr
  | Sub SymExpr SymExpr
  | Mul SymExpr SymExpr
  | Div SymExpr SymExpr
  | Mod SymExpr SymExpr
  | Exp SymExpr SymExpr
  | SDiv SymExpr SymExpr
  | SMod SymExpr SymExpr
  | Lt SymExpr SymExpr
  | Gt SymExpr SymExpr
  | SLt SymExpr SymExpr
  | SGt SymExpr SymExpr
  | Eq SymExpr SymExpr
  | IsZero SymExpr
  | And SymExpr SymExpr
  | Or SymExpr SymExpr
  | Xor SymExpr SymExpr
  | Not SymExpr
  | Byte SymExpr SymExpr
  | Shl SymExpr SymExpr
  | Shr SymExpr SymExpr
  | Sar SymExpr SymExpr
  | Sha3 SymExpr SymExpr
  | Address
  | Balance SymExpr
  | Origin
  | Caller
  | CallValue
  | CallDataLoad SymExpr
  | CallDataSize
  | CodeSize
  | GasPrice
  | BlockHash SymExpr
  | Coinbase
  | Timestamp
  | Number
  | Difficulty
  | GasLimit
  | ChainId
  | SelfBalance
  | BaseFee
  deriving (Eq, Ord, Generic, Show)

data Constraint
  = CTrue SymExpr
  | CFalse SymExpr
  | CEq SymExpr SymExpr
  | CNeq SymExpr SymExpr
  | CLt SymExpr SymExpr
  | CGt SymExpr SymExpr
  deriving (Eq, Ord, Show)

mkAdd :: SymExpr -> SymExpr -> SymExpr
mkAdd (Concrete 0) e = e
mkAdd e (Concrete 0) = e
mkAdd (Concrete a) (Concrete b) = Concrete ((a + b) `mod` (2^256))
mkAdd a b = Add a b

mkSub :: SymExpr -> SymExpr -> SymExpr
mkSub e (Concrete 0) = e
mkSub (Concrete a) (Concrete b) = Concrete ((a - b) `mod` (2^256))
mkSub a b = Sub a b

mkMul :: SymExpr -> SymExpr -> SymExpr
mkMul (Concrete 0) _ = Concrete 0
mkMul _ (Concrete 0) = Concrete 0
mkMul (Concrete 1) e = e
mkMul e (Concrete 1) = e
mkMul (Concrete a) (Concrete b) = Concrete ((a * b) `mod` (2^256))
mkMul a b = Mul a b

mkDiv :: SymExpr -> SymExpr -> SymExpr
mkDiv _ (Concrete 0) = Concrete 0
mkDiv e (Concrete 1) = e
mkDiv (Concrete a) (Concrete b) = if b == 0 then Concrete 0 else Concrete (a `div` b)
mkDiv a b = Div a b

mkMod :: SymExpr -> SymExpr -> SymExpr
mkMod _ (Concrete 0) = Concrete 0
mkMod (Concrete a) (Concrete b) = if b == 0 then Concrete 0 else Concrete (a `mod` b)
mkMod a b = Mod a b

mkLt :: SymExpr -> SymExpr -> SymExpr
mkLt (Concrete a) (Concrete b) = Concrete (if a < b then 1 else 0)
mkLt a b = Lt a b

mkGt :: SymExpr -> SymExpr -> SymExpr
mkGt (Concrete a) (Concrete b) = Concrete (if a > b then 1 else 0)
mkGt a b = Gt a b

mkEq :: SymExpr -> SymExpr -> SymExpr
mkEq a b | a == b = Concrete 1
mkEq (Concrete a) (Concrete b) = Concrete (if a == b then 1 else 0)
mkEq a b = Eq a b

mkIsZero :: SymExpr -> SymExpr
mkIsZero (Concrete 0) = Concrete 1
mkIsZero (Concrete _) = Concrete 0
mkIsZero e = IsZero e

mkAnd :: SymExpr -> SymExpr -> SymExpr
mkAnd (Concrete 0) _ = Concrete 0
mkAnd _ (Concrete 0) = Concrete 0
mkAnd (Concrete a) (Concrete b) = Concrete (a .&. b)
mkAnd a b = And a b

mkOr :: SymExpr -> SymExpr -> SymExpr
mkOr (Concrete a) (Concrete b) = Concrete (a .|. b)
mkOr a b = Or a b

mkXor :: SymExpr -> SymExpr -> SymExpr
mkXor (Concrete a) (Concrete b) = Concrete (xor a b)
mkXor a b = Xor a b

mkNot :: SymExpr -> SymExpr
mkNot (Concrete a) = Concrete ((2^256 - 1) - a)
mkNot e = Not e

prettyExpr :: SymExpr -> String
prettyExpr (Concrete n) = show n
prettyExpr (Symbolic name id) = name ++ "_" ++ show id
prettyExpr (Add a b) = "(" ++ prettyExpr a ++ " + " ++ prettyExpr b ++ ")"
prettyExpr (Sub a b) = "(" ++ prettyExpr a ++ " - " ++ prettyExpr b ++ ")"
prettyExpr (Mul a b) = "(" ++ prettyExpr a ++ " * " ++ prettyExpr b ++ ")"
prettyExpr (Div a b) = "(" ++ prettyExpr a ++ " / " ++ prettyExpr b ++ ")"
prettyExpr (Mod a b) = "(" ++ prettyExpr a ++ " % " ++ prettyExpr b ++ ")"
prettyExpr (Lt a b) = "(" ++ prettyExpr a ++ " < " ++ prettyExpr b ++ ")"
prettyExpr (Gt a b) = "(" ++ prettyExpr a ++ " > " ++ prettyExpr b ++ ")"
prettyExpr (Eq a b) = "(" ++ prettyExpr a ++ " == " ++ prettyExpr b ++ ")"
prettyExpr (IsZero a) = "IsZero(" ++ prettyExpr a ++ ")"
prettyExpr (And a b) = "(" ++ prettyExpr a ++ " & " ++ prettyExpr b ++ ")"
prettyExpr (Or a b) = "(" ++ prettyExpr a ++ " | " ++ prettyExpr b ++ ")"
prettyExpr (Not a) = "~" ++ prettyExpr a
prettyExpr Caller = "caller"
prettyExpr CallValue = "callvalue"
prettyExpr (CallDataLoad offset) = "calldataload(" ++ prettyExpr offset ++ ")"
prettyExpr CallDataSize = "calldatasize"
prettyExpr _ = "<expr>"

prettyConstraint :: Constraint -> String
prettyConstraint (CTrue e) = prettyExpr e ++ " != 0"
prettyConstraint (CFalse e) = prettyExpr e ++ " == 0"
prettyConstraint (CEq a b) = prettyExpr a ++ " == " ++ prettyExpr b
prettyConstraint (CNeq a b) = prettyExpr a ++ " != " ++ prettyExpr b
prettyConstraint (CLt a b) = prettyExpr a ++ " < " ++ prettyExpr b
prettyConstraint (CGt a b) = prettyExpr a ++ " > " ++ prettyExpr b