{-# LANGUAGE LambdaCase #-}

module Opcode where

import Data.Word
import Data.ByteString (ByteString)
import qualified Data.ByteString as BS
import Prelude hiding (LT, GT, EQ)

data Opcode
  = STOP | ADD | MUL | SUB | DIV | SDIV | MOD | SMOD
  | ADDMOD | MULMOD | EXP | SIGNEXTEND
  | LT | GT | SLT | SGT | EQ | ISZERO
  | AND | OR | XOR | NOT | BYTE | SHL | SHR | SAR
  | SHA3
  | ADDRESS | BALANCE | ORIGIN | CALLER | CALLVALUE
  | CALLDATALOAD | CALLDATASIZE | CALLDATACOPY
  | CODESIZE | CODECOPY
  | GASPRICE | EXTCODESIZE | EXTCODECOPY | RETURNDATASIZE | RETURNDATACOPY
  | EXTCODEHASH
  | BLOCKHASH | COINBASE | TIMESTAMP | NUMBER | DIFFICULTY | GASLIMIT
  | CHAINID | SELFBALANCE | BASEFEE
  | POP | MLOAD | MSTORE | MSTORE8 | SLOAD | SSTORE | JUMP | JUMPI
  | PC | MSIZE | GAS | JUMPDEST
  | PUSH1 | PUSH2 | PUSH3 | PUSH4 | PUSH5 | PUSH6 | PUSH7 | PUSH8
  | PUSH9 | PUSH10 | PUSH11 | PUSH12 | PUSH13 | PUSH14 | PUSH15 | PUSH16
  | PUSH17 | PUSH18 | PUSH19 | PUSH20 | PUSH21 | PUSH22 | PUSH23 | PUSH24
  | PUSH25 | PUSH26 | PUSH27 | PUSH28 | PUSH29 | PUSH30 | PUSH31 | PUSH32
  | DUP1 | DUP2 | DUP3 | DUP4 | DUP5 | DUP6 | DUP7 | DUP8
  | DUP9 | DUP10 | DUP11 | DUP12 | DUP13 | DUP14 | DUP15 | DUP16
  | SWAP1 | SWAP2 | SWAP3 | SWAP4 | SWAP5 | SWAP6 | SWAP7 | SWAP8
  | SWAP9 | SWAP10 | SWAP11 | SWAP12 | SWAP13 | SWAP14 | SWAP15 | SWAP16
  | LOG0 | LOG1 | LOG2 | LOG3 | LOG4
  | CREATE | CALL | CALLCODE | RETURN | DELEGATECALL | CREATE2
  | STATICCALL | REVERT | INVALID | SELFDESTRUCT
  | UNKNOWN Word8
  deriving (Show, Eq)

decodeOpcode :: Word8 -> Opcode
decodeOpcode = \case
  0x00 -> STOP
  0x01 -> ADD
  0x02 -> MUL
  0x03 -> SUB
  0x04 -> DIV
  0x05 -> SDIV
  0x06 -> MOD
  0x07 -> SMOD
  0x08 -> ADDMOD
  0x09 -> MULMOD
  0x0a -> EXP
  0x0b -> SIGNEXTEND
  0x10 -> LT
  0x11 -> GT
  0x12 -> SLT
  0x13 -> SGT
  0x14 -> EQ
  0x15 -> ISZERO
  0x16 -> AND
  0x17 -> OR
  0x18 -> XOR
  0x19 -> NOT
  0x1a -> BYTE
  0x1b -> SHL
  0x1c -> SHR
  0x1d -> SAR
  0x20 -> SHA3
  0x30 -> ADDRESS
  0x31 -> BALANCE
  0x32 -> ORIGIN
  0x33 -> CALLER
  0x34 -> CALLVALUE
  0x35 -> CALLDATALOAD
  0x36 -> CALLDATASIZE
  0x37 -> CALLDATACOPY
  0x38 -> CODESIZE
  0x39 -> CODECOPY
  0x3a -> GASPRICE
  0x3b -> EXTCODESIZE
  0x3c -> EXTCODECOPY
  0x3d -> RETURNDATASIZE
  0x3e -> RETURNDATACOPY
  0x3f -> EXTCODEHASH
  0x40 -> BLOCKHASH
  0x41 -> COINBASE
  0x42 -> TIMESTAMP
  0x43 -> NUMBER
  0x44 -> DIFFICULTY
  0x45 -> GASLIMIT
  0x46 -> CHAINID
  0x47 -> SELFBALANCE
  0x48 -> BASEFEE
  0x50 -> POP
  0x51 -> MLOAD
  0x52 -> MSTORE
  0x53 -> MSTORE8
  0x54 -> SLOAD
  0x55 -> SSTORE
  0x56 -> JUMP
  0x57 -> JUMPI
  0x58 -> PC
  0x59 -> MSIZE
  0x5a -> GAS
  0x5b -> JUMPDEST
  0x60 -> PUSH1
  0x61 -> PUSH2
  0x62 -> PUSH3
  0x63 -> PUSH4
  0x64 -> PUSH5
  0x65 -> PUSH6
  0x66 -> PUSH7
  0x67 -> PUSH8
  0x68 -> PUSH9
  0x69 -> PUSH10
  0x6a -> PUSH11
  0x6b -> PUSH12
  0x6c -> PUSH13
  0x6d -> PUSH14
  0x6e -> PUSH15
  0x6f -> PUSH16
  0x70 -> PUSH17
  0x71 -> PUSH18
  0x72 -> PUSH19
  0x73 -> PUSH20
  0x74 -> PUSH21
  0x75 -> PUSH22
  0x76 -> PUSH23
  0x77 -> PUSH24
  0x78 -> PUSH25
  0x79 -> PUSH26
  0x7a -> PUSH27
  0x7b -> PUSH28
  0x7c -> PUSH29
  0x7d -> PUSH30
  0x7e -> PUSH31
  0x7f -> PUSH32
  0x80 -> DUP1
  0x81 -> DUP2
  0x82 -> DUP3
  0x83 -> DUP4
  0x84 -> DUP5
  0x85 -> DUP6
  0x86 -> DUP7
  0x87 -> DUP8
  0x88 -> DUP9
  0x89 -> DUP10
  0x8a -> DUP11
  0x8b -> DUP12
  0x8c -> DUP13
  0x8d -> DUP14
  0x8e -> DUP15
  0x8f -> DUP16
  0x90 -> SWAP1
  0x91 -> SWAP2
  0x92 -> SWAP3
  0x93 -> SWAP4
  0x94 -> SWAP5
  0x95 -> SWAP6
  0x96 -> SWAP7
  0x97 -> SWAP8
  0x98 -> SWAP9
  0x99 -> SWAP10
  0x9a -> SWAP11
  0x9b -> SWAP12
  0x9c -> SWAP13
  0x9d -> SWAP14
  0x9e -> SWAP15
  0x9f -> SWAP16
  0xa0 -> LOG0
  0xa1 -> LOG1
  0xa2 -> LOG2
  0xa3 -> LOG3
  0xa4 -> LOG4
  0xf0 -> CREATE
  0xf1 -> CALL
  0xf2 -> CALLCODE
  0xf3 -> RETURN
  0xf4 -> DELEGATECALL
  0xf5 -> CREATE2
  0xfa -> STATICCALL
  0xfd -> REVERT
  0xfe -> INVALID
  0xff -> SELFDESTRUCT
  b    -> UNKNOWN b

pushBytes :: Opcode -> Int
pushBytes PUSH1 = 1
pushBytes PUSH2 = 2
pushBytes PUSH3 = 3
pushBytes PUSH4 = 4
pushBytes PUSH5 = 5
pushBytes PUSH6 = 6
pushBytes PUSH7 = 7
pushBytes PUSH8 = 8
pushBytes PUSH9 = 9
pushBytes PUSH10 = 10
pushBytes PUSH11 = 11
pushBytes PUSH12 = 12
pushBytes PUSH13 = 13
pushBytes PUSH14 = 14
pushBytes PUSH15 = 15
pushBytes PUSH16 = 16
pushBytes PUSH17 = 17
pushBytes PUSH18 = 18
pushBytes PUSH19 = 19
pushBytes PUSH20 = 20
pushBytes PUSH21 = 21
pushBytes PUSH22 = 22
pushBytes PUSH23 = 23
pushBytes PUSH24 = 24
pushBytes PUSH25 = 25
pushBytes PUSH26 = 26
pushBytes PUSH27 = 27
pushBytes PUSH28 = 28
pushBytes PUSH29 = 29
pushBytes PUSH30 = 30
pushBytes PUSH31 = 31
pushBytes PUSH32 = 32
pushBytes _ = 0

isPush :: Opcode -> Bool
isPush op = pushBytes op > 0

dupDepth :: Opcode -> Maybe Int
dupDepth DUP1 = Just 1
dupDepth DUP2 = Just 2
dupDepth DUP3 = Just 3
dupDepth DUP4 = Just 4
dupDepth DUP5 = Just 5
dupDepth DUP6 = Just 6
dupDepth DUP7 = Just 7
dupDepth DUP8 = Just 8
dupDepth DUP9 = Just 9
dupDepth DUP10 = Just 10
dupDepth DUP11 = Just 11
dupDepth DUP12 = Just 12
dupDepth DUP13 = Just 13
dupDepth DUP14 = Just 14
dupDepth DUP15 = Just 15
dupDepth DUP16 = Just 16
dupDepth _ = Nothing

swapDepth :: Opcode -> Maybe Int
swapDepth SWAP1 = Just 1
swapDepth SWAP2 = Just 2
swapDepth SWAP3 = Just 3
swapDepth SWAP4 = Just 4
swapDepth SWAP5 = Just 5
swapDepth SWAP6 = Just 6
swapDepth SWAP7 = Just 7
swapDepth SWAP8 = Just 8
swapDepth SWAP9 = Just 9
swapDepth SWAP10 = Just 10
swapDepth SWAP11 = Just 11
swapDepth SWAP12 = Just 12
swapDepth SWAP13 = Just 13
swapDepth SWAP14 = Just 14
swapDepth SWAP15 = Just 15
swapDepth SWAP16 = Just 16
swapDepth _ = Nothing

readBytecode :: ByteString -> Int -> Maybe Word8
readBytecode code pc
  | pc >= 0 && pc < BS.length code = Just (BS.index code pc)
  | otherwise = Nothing

readBytes :: ByteString -> Int -> Int -> Integer
readBytes code start n = go 0 0
  where
    go acc i
      | i >= n = acc
      | otherwise = case readBytecode code (start + i) of
          Just b -> go (acc * 256 + fromIntegral b) (i + 1)
          Nothing -> acc