GState.hs

-- SPDX-FileCopyrightText: 2021 Oxhead Alpha
-- SPDX-License-Identifier: LicenseRef-MIT-OA

-- | Global blockchain state (emulated).

module Morley.Michelson.Runtime.GState
  (
    -- * Auxiliary types
    ContractState (..)
  , VotingPowers (..)
  , vpPick
  , vpTotal
  , mkVotingPowers
  , mkVotingPowersFromMap
  , dummyVotingPowers
  , BigMapCounter(..)
  , bigMapCounter

  -- * GState
  , GState (..)
  , gsChainIdL
  , gsImplicitAddressesL
  , gsContractAddressesL
  , gsImplicitAddressAliasesL
  , gsContractAddressAliasesL
  , gsVotingPowersL
  , gsCounterL
  , gsBigMapCounterL
  , addressesL
  , genesisAddresses
  , genesisKeyHashes
  , genesisAddress
  -- * More genesisAddresses which can be used in tests
  , genesisAddress1
  , genesisAddress2
  , genesisAddress3
  , genesisAddressN
  , genesisKeyHash
  -- * Genesis secret keys
  , genesisSecretKey
  , genesisSecrets

  , initGState
  , readGState
  , writeGState

  -- * Operations on GState
  , GStateUpdate (..)
  , GStateUpdateError (..)
  , applyUpdate
  , applyUpdates
  , extractAllContracts
  , lookupBalance
  , AddressStateFam
  ) where

import Control.Lens (at, makeLenses, (?~))
import Data.Aeson (FromJSON(..), ToJSON(..), object, withObject, (.:), (.:?), (.=))
import Data.Aeson qualified as Aeson
import Data.Aeson.Encode.Pretty qualified as Aeson
import Data.Aeson.TH (deriveJSON)
import Data.ByteString.Lazy qualified as LBS
import Data.Default (def)
import Data.Map.Strict qualified as Map
import Data.Type.Equality ((:~:)(..))
import Fmt (Buildable(build), pretty, (+|), (|+))
import System.IO.Error (IOError, isDoesNotExistError)

import Morley.Michelson.TypeCheck
  (SomeParamType(..), TcOriginatedContracts, typeCheckContractAndStorage, typeCheckingWith)
import Morley.Michelson.Typed qualified as T
import Morley.Michelson.Typed.Existential (SomeContractAndStorage(..))
import Morley.Michelson.Typed.Scope
import Morley.Michelson.Untyped (Contract, Value)
import Morley.Tezos.Address
import Morley.Tezos.Address.Alias
import Morley.Tezos.Address.Kinds
import Morley.Tezos.Core (ChainId, Mutez, divModMutezInt, dummyChainId)
import Morley.Tezos.Crypto
import Morley.Util.Aeson
import Morley.Util.Bimap (Bimap)
import Morley.Util.Bimap qualified as Bimap
import Morley.Util.Lens
import Morley.Util.Peano (ToPeano, type (>))
import Morley.Util.Sing (eqI, eqParamSing, eqParamSing2)
import Morley.Util.SizedList qualified as SL
import Morley.Util.SizedList.Types

-- | State of a contract with code.
data ContractState =
  forall cp st. (ParameterScope cp, StorageScope st) => ContractState
  { csBalance :: Mutez
  -- ^ Amount of mutez owned by this contract.
  , csContract :: T.Contract cp st
  -- ^ Contract itself.
  , csStorage :: T.Value st
  -- ^ Storage value associated with this contract.
  , csDelegate :: Maybe KeyHash
  -- ^ Delegate associated with this contract.
  }

makeLensesWith postfixLFields ''ContractState

deriving stock instance Show ContractState

instance Eq ContractState where
  ContractState b1 c1 s1 d1 == ContractState b2 c2 s2 d2 =
    b1 == b2
    && eqParamSing2 c1 c2
    && eqParamSing s1 s2
    && d1 == d2

instance ToJSON ContractState where
  toJSON ContractState{..} =
    object .
    maybe id ((:) . ("delegate" .=)) csDelegate $
    [ "balance" .= csBalance
    , "storage" .= T.untypeValue csStorage
    , "contract" .= T.convertContract csContract
    ]

-- These instance is a bit hacky because it is quite painful to
-- write proper JSON instances for typed `Instr` and `Value` so
-- we typecheck untyped representation instead of parsing.
instance FromJSON ContractState where
  parseJSON =
    withObject "contractstate" $ \o -> do
      (balance :: Mutez) <- o .: "balance"
      (uStorage :: Value) <- o .: "storage"
      (uContract :: Contract) <- o .: "contract"
      (delegate :: Maybe KeyHash) <- o .:? "delegate"
      case typeCheckingWith def $ typeCheckContractAndStorage uContract uStorage of
        Right (SomeContractAndStorage contract storage) ->
          pure $ ContractState balance contract storage delegate
        Left err -> fail $ "Unable to parse `ContractState`: " <> pretty err

instance Buildable ContractState where
  build ContractState{..} =
    "Contractstate:\n Balance: " +| csBalance |+
    "\n  Storage: " +| T.untypeValue csStorage |+
    "\n  Contract: " +| T.convertContract csContract |+
    "\n  Delegate: " +| csDelegate |+ ""

-- | Distribution of voting power among the contracts.
--
-- Voting power reflects the ability of bakers to accept, deny or pass new
-- proposals for protocol updates. I.e. each baker has its vote weight.
--
-- This datatype definition may change in future, so its internals are not
-- exported.
newtype VotingPowers
  = VotingPowers (Map KeyHash Natural)
  deriving stock (Show, Eq)

deriveJSON morleyAesonOptions ''VotingPowers

-- | Get voting power of the given address.
vpPick :: KeyHash -> VotingPowers -> Natural
vpPick key (VotingPowers distr) = Map.lookup key distr ?: 0

-- | Get total voting power.
vpTotal :: VotingPowers -> Natural
vpTotal (VotingPowers distr) = sum distr

-- | Create voting power distribution from map.
mkVotingPowersFromMap :: Map KeyHash Natural -> VotingPowers
mkVotingPowersFromMap = VotingPowers

-- | Create voting power distribution.
--
-- If some key is encountered multiple times, voting power will be summed up.
mkVotingPowers :: [(KeyHash, Natural)] -> VotingPowers
mkVotingPowers = mkVotingPowersFromMap . Map.fromListWith (+)

-- | All big_maps stored in a chain have a globally unique ID.
--
-- We use this counter to keep track of how many big_maps have been created so far,
-- and to generate new IDs whenever a new big_map is created.
newtype BigMapCounter = BigMapCounter { _bigMapCounter :: Natural }
  deriving stock (Show, Eq, Generic)
  deriving anyclass (NFData)
  deriving newtype (ToJSON, FromJSON, Num, Buildable)

makeLenses ''BigMapCounter

-- | Persistent data passed to Morley contracts which can be updated
-- as result of contract execution.
data GState = GState
  { gsChainId :: ChainId
  -- ^ Identifier of chain.
  , gsImplicitAddresses :: Map ImplicitAddress Mutez
  -- ^ All known implicit addresses and their state (i.e. balance)
  , gsContractAddresses :: Map ContractAddress ContractState
  -- ^ All known contract addresses and their state.
  , gsTxRollupAddresses :: Map TxRollupAddress ()
  -- ^ All known transaction rollup addresses and their state.
  , gsVotingPowers :: VotingPowers
  -- ^ Voting power distribution.
  , gsCounter :: GlobalCounter
  -- ^ Ever increasing operation counter.
  , gsBigMapCounter :: BigMapCounter
  , gsImplicitAddressAliases :: Bimap ImplicitAlias ImplicitAddress
  -- ^ Implicit addresses with the associated aliases/names.
  , gsContractAddressAliases :: Bimap ContractAlias ContractAddress
  -- ^ Contract addresses with the associated aliases/names.
  } deriving stock (Show, Eq)

makeLensesWith postfixLFields ''GState

deriveJSON morleyAesonOptions ''GState

-- | Number of genesis addresses, type-level
type GenesisAddressesNum = 10

-- | Convenience synonym
type GenesisList a = SizedList GenesisAddressesNum a

-- | Number of genesis addresses, term-level
genesisAddressesNum :: Natural
genesisAddressesNum = natVal @GenesisAddressesNum Proxy

-- | Secrets from which genesis addresses are derived from.
genesisSecrets :: GenesisList SecretKey
genesisSecrets = SL.generate @GenesisAddressesNum $ \i ->
    let seed = encodeUtf8 (show (i + 1) :: Text)
    in detSecretKey seed

-- | KeyHash of genesis address.
genesisKeyHashes :: GenesisList KeyHash
genesisKeyHashes = hashKey . toPublic <$> genesisSecrets

-- | Initially these addresses have a lot of money.
genesisAddresses :: GenesisList ImplicitAddress
genesisAddresses = ImplicitAddress <$> genesisKeyHashes

-- | One of genesis key hashes.
genesisKeyHash :: KeyHash
genesisKeyHash = SL.head genesisKeyHashes

-- | One of genesis addresses.
genesisAddress :: ImplicitAddress
genesisAddress = SL.head genesisAddresses

-- | Secret key assotiated with 'genesisAddress'.
genesisSecretKey :: SecretKey
genesisSecretKey = SL.head genesisSecrets

-- | More genesis addresses
genesisAddress1, genesisAddress2, genesisAddress3 :: ImplicitAddress
_ :< genesisAddress1 :< genesisAddress2 :< genesisAddress3
  :< _ = genesisAddresses

-- | More genesis addresses, via a type-level natural
--
-- > genesisAddressN @7
--
-- Note that @'genesisAddress' == genesisAddressN \@0@, @'genesisAddress1' == genesisAddressN \@1@,
-- etc.
genesisAddressN
  :: forall n. (SingIPeano n, ToPeano GenesisAddressesNum > ToPeano n ~ 'True)
  => ImplicitAddress
genesisAddressN = SL.index @n genesisAddresses

-- | Dummy 'VotingPowers'. We give all the voting power to two genesis addreses
-- as the addresses holding lot of money. Only two addresses are involved for
-- simplicity.
dummyVotingPowers :: VotingPowers
dummyVotingPowers = case genesisKeyHashes of
  k1 :< k2 :< _ -> mkVotingPowers [(k1, 50), (k2, 50)]

-- | Initial 'GState'. It's supposed to be used if no 'GState' is
-- provided. It puts plenty of money on each genesis address.
initGState :: GState
initGState =
  GState
  { gsChainId = dummyChainId
  , gsImplicitAddresses = Map.fromList
    [ (genesis, money)
    | let (money, _) = maxBound @Mutez `divModMutezInt` genesisAddressesNum
                    ?: error "Number of genesis addresses is 0"
    , genesis <- toList genesisAddresses
    ]
  , gsContractAddresses = Map.empty
  , gsTxRollupAddresses = Map.empty
  , gsVotingPowers = dummyVotingPowers
  , gsCounter = GlobalCounter 0
  , gsBigMapCounter = BigMapCounter 0
  , gsImplicitAddressAliases = Bimap.empty
  , gsContractAddressAliases = Bimap.empty
  }

data GStateParseError =
  GStateParseError String
  deriving stock (Show)

instance Exception GStateParseError where
  displayException (GStateParseError str) = "Failed to parse GState: " <> str

-- | Read 'GState' from a file.
readGState :: FilePath -> IO GState
readGState fp = (LBS.readFile fp >>= parseFile) `catch` onExc
  where
    parseFile :: LByteString -> IO GState
    parseFile lByteString =
      if null lByteString
      then pure initGState
      else (either (throwM . GStateParseError) pure . Aeson.eitherDecode') lByteString
    onExc :: IOError -> IO GState
    onExc exc
      | isDoesNotExistError exc = pure initGState
      | otherwise = throwM exc

-- | Write 'GState' to a file.
writeGState :: FilePath -> GState -> IO ()
writeGState fp gs = LBS.writeFile fp (Aeson.encodePretty' config gs)
  where
    config =
      Aeson.defConfig
      { Aeson.confTrailingNewline = True
      }

-- | Updates that can be applied to 'GState'.
data GStateUpdate where
  GSAddImplicitAddress :: ImplicitAddress -> Mutez -> GStateUpdate
  GSAddContractAddress :: ContractAddress -> ContractState -> GStateUpdate
  GSAddContractAddressAlias :: ContractAlias -> ContractAddress -> GStateUpdate
  GSSetStorageValue :: StorageScope st => ContractAddress -> T.Value st -> GStateUpdate
  GSSetBalance :: L1AddressKind kind => KindedAddress kind -> Mutez -> GStateUpdate
  GSIncrementCounter :: GStateUpdate
  GSUpdateCounter :: GlobalCounter -> GStateUpdate
  GSSetBigMapCounter :: BigMapCounter -> GStateUpdate
  GSSetDelegate :: ContractAddress -> Maybe KeyHash -> GStateUpdate

deriving stock instance Show GStateUpdate

instance Buildable GStateUpdate where
  build =
    \case
      GSAddImplicitAddress addr st ->
        "Add implicit address " +| addr |+ " with balance " +| st |+ ""
      GSAddContractAddress addr st ->
        "Add contract address " +| addr |+ " with state " +| st |+ ""
      GSAddContractAddressAlias alias addr ->
        "Add an alias " +| alias |+ " for address " +| addr |+ ""
      GSSetStorageValue addr tVal ->
        "Set storage value of address " +| addr |+ " to " +| T.untypeValue tVal |+ ""
      GSSetBalance addr balance ->
        "Set balance of address " +| addr |+ " to " +| balance |+ ""
      GSIncrementCounter ->
        "Increment internal counter after operation"
      GSUpdateCounter v ->
        "Set internal counter to " +| v |+ " after interpreting " <>
        "several 'CREATE_CONTRACT' instructions"
      GSSetBigMapCounter inc ->
        "Increment internal big_map counter by: " +| build inc
      GSSetDelegate addr key ->
        "Set delegate for contract " +| addr |+ " to " +| maybe "<nobody>" build key

data GStateUpdateError
  = GStateAddressExists Address
  | GStateUnknownAddress Address
  | GStateStorageNotMatch ContractAddress
  | GStateTxRollup TxRollupAddress GStateUpdate
  deriving stock (Show)

instance Buildable GStateUpdateError where
  build =
    \case
      GStateAddressExists addr -> "Address already exists: " <> build addr
      GStateUnknownAddress addr -> "Unknown address: " <> build addr
      GStateStorageNotMatch addr ->
        "Storage type does not match the contract in run-time state\
        \ when updating new storage value to address: " <> build addr
      GStateTxRollup addr op ->
        "Transaction rollup address " +| addr |+ " doesn't support this operation: " +| op |+ ""

-- | Apply 'GStateUpdate' to 'GState'.
applyUpdate :: GStateUpdate -> GState -> Either GStateUpdateError GState
applyUpdate =
  \case
    GSAddImplicitAddress addr st -> addImplicitAddress addr st
    GSAddContractAddress addr st -> addContractAddress addr st
    GSAddContractAddressAlias alias addr -> Right . addContractAddressAlias alias addr
    GSSetStorageValue addr newValue ->
      setStorageValue addr newValue
    GSSetBalance addr newBalance -> addr `setBalance` newBalance
    GSIncrementCounter -> Right . over gsCounterL (+1)
    GSUpdateCounter newCounter -> Right . set gsCounterL newCounter
    GSSetBigMapCounter bmCounter -> Right . set gsBigMapCounterL bmCounter
    GSSetDelegate addr key -> setDelegate addr key

-- | Apply a list of 'GStateUpdate's to 'GState'.
applyUpdates :: [GStateUpdate] -> GState -> Either GStateUpdateError GState
applyUpdates = flip (foldM (flip applyUpdate))

-- | Add an address if it hasn't been added before.
addImplicitAddress :: ImplicitAddress -> Mutez -> GState -> Either GStateUpdateError GState
addImplicitAddress addr st gs
  | addr `Map.member` (gsImplicitAddresses gs) = Left $ GStateAddressExists (MkAddress addr)
  | otherwise = Right $ gs & gsImplicitAddressesL . at addr ?~ st

-- | Add an address if it hasn't been added before.
addContractAddress :: ContractAddress -> ContractState -> GState -> Either GStateUpdateError GState
addContractAddress addr st gs
  | addr `Map.member` (gsContractAddresses gs) = Left $ GStateAddressExists (MkAddress addr)
  | otherwise = Right $ gs & gsContractAddressesL . at addr ?~ st

-- | Add an alias for given address, overwriting any existing address for the same alias.
addContractAddressAlias :: ContractAlias -> ContractAddress -> GState -> GState
addContractAddressAlias alias addr = gsContractAddressAliasesL . at alias ?~ addr

-- | Update storage value associated with given address.
setStorageValue :: forall st. (StorageScope st) =>
     ContractAddress -> T.Value st -> GState -> Either GStateUpdateError GState
setStorageValue addr newValue = updateAddressState addr modifier
  where
    modifier :: ContractState -> Either GStateUpdateError ContractState
    modifier ContractState{csStorage = _ :: T.Value st', ..} = do
      case eqI @st @st' of
        Just Refl -> Right $ ContractState{csStorage = newValue, ..}
        _ -> Left $ GStateStorageNotMatch addr

-- | Update balance value associated with given address.
setBalance
  :: forall kind. L1AddressKind kind
  => KindedAddress kind
  -> Mutez
  -> GState
  -> Either GStateUpdateError GState
setBalance addr newBalance = updateAddressState addr $
  Right . set (balanceLens addr) newBalance

lookupBalance :: forall kind. L1AddressKind kind => KindedAddress kind -> GState -> Maybe Mutez
lookupBalance addr = preview $ addressesL addr . at addr . traverse . balanceLens addr

balanceLens
  :: forall kind. L1AddressKind kind
  => KindedAddress kind
  -> Lens' (AddressStateFam kind) Mutez
balanceLens = \case
  ImplicitAddress{} -> id
  ContractAddress{} -> csBalanceL
  where _ = usingImplicitOrContractKind @kind ()

-- | Set delegate for a given address
setDelegate :: ContractAddress -> Maybe KeyHash -> GState -> Either GStateUpdateError GState
setDelegate addr key = updateAddressState addr $ pure . (csDelegateL .~ key)

type family AddressStateFam kind where
  AddressStateFam 'AddressKindImplicit = Mutez
  AddressStateFam 'AddressKindContract = ContractState
  -- TODO [#838]: support transaction rollups on the emulator
  AddressStateFam 'AddressKindTxRollup = ()

updateAddressState
  :: forall kind. KindedAddress kind
  -> (AddressStateFam kind -> Either GStateUpdateError (AddressStateFam kind))
  -> GState
  -> Either GStateUpdateError GState
updateAddressState addr f gs =
  let addrL :: Lens' GState (Maybe (AddressStateFam kind))
      addrL = addressesL addr . at addr
  in case gs ^. addrL of
    Nothing -> Left $ GStateUnknownAddress $ MkAddress addr
    Just as -> do
      newState <- f as
      return $ gs & addrL .~ Just newState

addressesL
  :: KindedAddress kind
  -> Lens' GState (Map (KindedAddress kind) (AddressStateFam kind))
addressesL = \case
  ImplicitAddress{} -> gsImplicitAddressesL
  ContractAddress{} -> gsContractAddressesL
  TxRollupAddress{} -> gsTxRollupAddressesL

-- | Retrieve all contracts stored in GState
extractAllContracts :: GState -> TcOriginatedContracts
extractAllContracts = Map.fromList . mapMaybe extractContract . toPairs . gsContractAddresses
 where
    extractContract
      :: (ContractAddress, ContractState) -> Maybe (ContractHash, SomeParamType)
    extractContract =
      \case (ContractAddress ca, ContractState{..}) ->
              Just (ca, SomeParamType $ T.cParamNotes $ csContract)