content: {-# LANGUAGE DataKinds #-}
  {-# LANGUAGE QuasiQuotes #-}
  
  module Main where
  
  import Control.Monad.State.Strict (liftIO)
  import Data.ByteString.UTF8 as BSU 
  import Data.Either
  import Data.Function
  import Data.Maybe
  import Data.String.Here
  import EVM
  import EVM.ABI
  import EVM.Effects (App, runApp)
  import EVM.Fetch qualified as Fetch
  import EVM.Solidity
  import EVM.Stepper qualified as Stepper
  import EVM.Transaction qualified
  import EVM.Types
  import EVM.UnitTest
  import Test.Tasty.Bench
  import Control.Monad.ST
  import EVM.FeeSchedule (feeSchedule)
  
  -- benchmark hevm using tasty-bench
  
  vmFromRawByteString :: App m => ByteString -> m (VM Concrete RealWorld)
  vmFromRawByteString bs = liftIO $
    bs
      & ConcreteRuntimeCode
      & RuntimeCode
      & initialContract
      & vm0
      & stToIO
      & fmap EVM.Transaction.initTx
  
  vm0 :: Contract -> ST s (VM Concrete s)
  vm0 c = makeVm $ vm0Opts c
  
  vm0Opts :: Contract -> VMOpts Concrete
  vm0Opts c =
    VMOpts
      { contract = c,
        calldata = (ConcreteBuf "", []),
        value = Lit 0xfffffffffffff, -- balance
        baseState = EmptyBase,
        address = LitAddr 0xacab,
        caller = LitAddr 0,
        origin = LitAddr 0,
        gas = 0xffffffffffffffff,
        baseFee = 0,
        priorityFee = 0,
        gaslimit = 0xffffffffffffffff,
        coinbase = LitAddr 0,
        number = Lit 0,
        timestamp = Lit 0,
        blockGaslimit = 0xffffffffffffffff,
        gasprice = 0,
        maxCodeSize = 0xffffffff,
        prevRandao = 0,
        schedule = feeSchedule,
        chainId = 1,
        create = False,
        txAccessList = mempty, -- TODO: support me soon
        allowFFI = False,
        otherContracts = [],
        freshAddresses = 0,
        beaconRoot = 0
      }
  
  vmOptsToTestVMParams :: VMOpts Concrete -> TestVMParams
  vmOptsToTestVMParams v =
    TestVMParams
      { address = v.address,
        caller = v.caller,
        origin = v.origin,
        gasCreate = v.gas,
        gasCall = v.gas,
        baseFee = v.baseFee,
        priorityFee = v.priorityFee,
        balanceCreate = case v.value of
          Lit x -> x
          _ -> 0,
        coinbase = v.coinbase,
        number = case v.number of
          Lit x -> x
          _ -> 0,
        timestamp = case v.timestamp of
          Lit x -> x
          _ -> 0,
        gaslimit = v.gaslimit,
        gasprice = v.gasprice,
        maxCodeSize = v.maxCodeSize,
        prevrandao = v.prevRandao,
        chainId = v.chainId
      }
  
  callMainForBytecode :: App m => ByteString -> m (Either EvmError (Expr 'Buf))
  callMainForBytecode bs = do
    vm <- vmFromRawByteString bs
    Stepper.interpret (Fetch.zero 0 Nothing) vm (Stepper.evm (abiCall (vmOptsToTestVMParams (vm0Opts (initialContract (RuntimeCode (ConcreteRuntimeCode bs))))) (Left ("main()", emptyAbi))) >> Stepper.execFully)
  
  benchMain :: (String, ByteString) -> Benchmark
  benchMain (name, bs) = bench name $ nfIO $ runApp $ (\x -> if isRight x then () else internalError "failed") <$> callMainForBytecode bs
  
  mkBench :: (Int -> IO a) -> [Int] -> IO [(String, a)]
  mkBench f l = mapM (\n -> (show n,) <$> f n) l
  
  main :: IO ()
  main = do
    let f (name, prog, limit) = do
          ll <- mkBench prog [2 ^ n | n :: Int <- [1 .. fromMaybe 14 limit]]
          pure $ bgroup name (benchMain <$> ll)
    let benchmarks = [ 
                       ("loop", simpleLoop, Nothing)
                     , ("primes", primes, Nothing)
                     , ("hashes", hashes, Nothing)
                     , ("hashmem", hashmem, Nothing)
                     , ("balanceTransfer", balanceTransfer, Nothing)
                     , ("funcCall", funcCall, Nothing)
                     , ("contractCreation", contractCreation, Nothing)
                     , ("contractCreationMem", contractCreationMem, Nothing)
                     , ("arrayCreationMem", arrayCreationMem, Just 9)
                     , ("mapStorage", mapStorage, Nothing)
                     , ("swapOperations", swapOperations, Nothing)
                     ]
    defaultMain =<< mapM f benchmarks
    
  
  -- Loop that adds up n numbers
  simpleLoop :: Int -> IO ByteString
  simpleLoop n = do
    let src =
          [i|
            contract A {
              function main() public {
                uint256 acc = 0;
                for (uint i = 0; i < ${n}; i++) {
                  acc += i;
                }
              }
            }
          |]
    fmap fromJust (runApp $ solcRuntime "A" src)
  
  -- Computes prime numbers and stores them up to n
  primes :: Int -> IO ByteString
  primes n = do
    let src =
          [i|
            contract A {
              mapping (uint => uint) public primes;
              function isPrime (uint n) public returns (bool) {
                if (n == 2) {
                  return true;
                }
                if (n % 2 == 0) {
                  return false;
                }
                for (uint i = 3; i * i < n; i += 2) {
                  if (n % i == 0) {
                    return false;
                  }
                }
                return true;
              }
              
              function main() public {
                uint n = 0;
                for (uint i = 0; i < ${n}; i++) {
                  if (isPrime(i)) {
                    primes[n++] = i;
                  }
                }
              }
            }
          |]
    fmap fromJust (runApp $ solcRuntime "A" src)
  
  -- Program that repeatedly hashes a value
  hashes :: Int -> IO ByteString
  hashes n = do
    let src =
          [i|
            contract A {
              function main() public {
                bytes32 h = 0;
                for (uint i = 0; i < ${n}; i++) {
                  h = keccak256(abi.encode(h));
                }
              }
            }
          |]
    fmap fromJust (runApp $ solcRuntime "A" src)
  
  -- Program that repeatedly hashes a value and stores it in a map
  hashmem :: Int -> IO ByteString
  hashmem n = do
    let src =
          [i|
            contract A {
              mapping (uint256 => uint256) public map;
              function main() public {
                uint256 h = 0;
                for (uint i = 0; i < ${n}; i++) {
                  uint256 x = h;
                  h = uint256(keccak256(abi.encode(h)));
                  map[x] = h;
                }
              }
            }
          |]
    fmap fromJust (runApp $ solcRuntime "A" src)
  
  -- Transfer ETH to an address n times
  balanceTransfer :: Int -> IO ByteString
  balanceTransfer n = do
    let src =
          [i|
            contract A {
              function main() public {
                address payable to = payable(address(0x8A8eAFb1cf62BfBeb1741769DAE1a9dd47996192)); 
                for (uint i = 0; i < ${n}; i++) {
                  to.transfer(1);
                }
              }
            }
          |]
    fmap fromJust (runApp $ solcRuntime "A" src)
  
  -- Call a public function n times
  funcCall :: Int -> IO ByteString
  funcCall n = do
    let src =
          [i|
            contract A {
              uint256 public acc;
              function f(uint256 x) public {
                acc += x;
              }
              function main() public {
                for (uint i = 0; i < ${n}; i++) {
                  f(i);
                }
              }
            }
          |]
    fmap fromJust (runApp $ solcRuntime "A" src)
  
  -- Creates n contracts
  contractCreation :: Int -> IO ByteString
  contractCreation n = do
    let src =
          [i|
            contract B { }
            contract A {
              B public b;
              function main() public {
                for (uint i = 0; i < ${n}; i++) {
                  b = new B();
                }
              }
            }
          |]
    fmap fromJust (runApp $ solcRuntime "A" src)
  
  -- Create n contracts with a string in them
  contractCreationMem :: Int -> IO ByteString
  contractCreationMem n = do
    let src =
          [i|
            contract B { string m = "the quick brown fox jumps over the lazy dog"; }
            contract A {
              B public b;
              function main() public {
                for (uint i = 0; i < ${n}; i++) {
                  b = new B();
                }
              }
            }
          |]
    fmap fromJust (runApp $ solcRuntime "A" src)
  
  -- Create large array in memory
  arrayCreationMem :: Int -> IO ByteString
  arrayCreationMem n = do
    let src =
          [i|
            struct C { int24 a; int24 b; uint256 c; uint128 d; }
            contract A {
              function work() public returns (C[] memory ret) {
                  ret = new C[](${n});
                  for(uint a = 0; a < ${n}; a++) {
                      ret[a] = C({a:1, b:1, c:0xffffffffffffff, d:5});
                  }
                  return ret;
              }
              function main() public {
                for (uint i = 0; i < ${n}; i++) {
                  C[] memory ret = work();
                }
              }
            }
          |]
    fmap fromJust (runApp $ solcRuntime "A" src)
  
  -- Create large map in storage
  mapStorage :: Int -> IO ByteString
  mapStorage n = do
    let src =
          [i|
            struct C { uint256 a; uint256 b; }
            contract A {
              mapping(uint256 => mapping(uint256 => C)) m;
              function main() public {
                for (uint i = 0; i < ${n}; i++) {
                  m[i][${n}-i] = C(i, ${n});
                }
              }
            }
          |]
    fmap fromJust (runApp $ solcRuntime "A" src)
  
  -- Do many swaps
  swapOperations :: Int -> IO ByteString
  swapOperations n = do
    let src =
          [i|
            contract A {
              function main() public pure {
                  (uint a0, uint a1, uint a2, uint a3, uint a4, uint a5, uint a6, uint a7, uint a8, uint a9, uint a10, uint a11, uint a12, uint a13, uint a14, uint a15) = (1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16);
                  for (uint i = 0; i < ${n}; i++) {
                      (a0, a15) = (a15, a0);
                      (a1, a14) = (a14, a1);
                      (a2, a13) = (a13, a2);
                      (a3, a12) = (a12, a3);
                      (a4, a11) = (a11, a4);
                      (a5, a10) = (a10, a5);
                      (a6, a9) = (a9, a6);
                      (a7, a8) = (a8, a7);
                  }
              }
          }
          |]
    fmap fromJust (runApp $ solcRuntime' "A" src True)
  