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 (..)
  , ImplicitState (..)
  , TicketKey (..)
  , toTicketKey
  , isBalanceL
  , isDelegateL
  , isTicketsL
  , 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(..), FromJSONKey, ToJSON(..), ToJSONKey, 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.Constraint ((\\))
import Data.Default (def)
import Data.Map.Strict qualified as Map
import Data.Type.Equality ((:~:)(..))
import Fmt (Buildable(build), nameF, pretty, unlinesF, (+|), (|+))
import System.IO.Error (IOError, isDoesNotExistError)

import Morley.Michelson.TypeCheck
  (SomeParamType(..), TcOriginatedContracts, tcStrict, 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.Michelson.Untyped qualified as U
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.Named
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

data ImplicitState = ImplicitState
  { isBalance :: Mutez
    -- ^ Implicit address balance.
  , isTickets :: HashMap TicketKey Natural
    -- ^ Implicit account's ticket balances.
  , isDelegate :: Maybe KeyHash
    -- ^ Delegate, if set. Implicit address can have a delegate set to itself,
    -- then it's considered "registered delegate" and can be set as delegate for
    -- other addresses. It's impossible to "unregister" a delegate, so once
    -- delegate is set to the address itself, it can't be changed.
  } deriving stock (Show, Eq)

-- | A triple of ticketer, value and type, which uniquely defines a ticket.
newtype TicketKey = TicketKey (Address, U.Value, U.Ty)
  deriving newtype (Eq, Show, ToJSON, FromJSON, FromJSONKey, ToJSONKey)

instance Buildable TicketKey where
  build (TicketKey (addr, val, ty)) =
    "(" +| addr |+ ", " +| val |+ ", " +| ty |+ ")"

instance Hashable TicketKey where
  hashWithSalt s = hashWithSalt s . toJSON

-- | Convert a typed ticket value to 'TicketKey' and amount.
toTicketKey :: forall t. T.Value ('T.TTicket t) -> (TicketKey, Natural)
toTicketKey (T.VTicket ticketer value amount) =
  (TicketKey (ticketer, T.untypeValue value, T.mkUType $ T.starNotes @t), amount)
  \\ T.valueTypeSanity value
  \\ T.comparableImplies (Proxy @t)

deriveJSON morleyAesonOptions ''ImplicitState
makeLensesWith postfixLFields ''ImplicitState

instance Buildable ImplicitState where
  build ImplicitState{..} = nameF "Implicit state" $ unlinesF
    [ nameF "Balance" $ build isBalance
    , nameF "Delegate" $ build isDelegate
    ]

-- | 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{ tcStrict=False } $ 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 ImplicitState
  -- ^ All known implicit addresses and their state (i.e. balance)
  , gsContractAddresses :: Map ContractAddress ContractState
  -- ^ All known contract addresses and their state.
  , gsSmartRollupAddresses :: Map SmartRollupAddress ()
  -- ^ All known smart 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 addresses
-- 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, ImplicitState money mempty Nothing)
    | let (money, _) = maxBound @Mutez `divModMutezInt` genesisAddressesNum
                    ?: error "Number of genesis addresses is 0"
    , genesis <- toList genesisAddresses
    ]
  , gsContractAddresses = Map.empty
  , gsSmartRollupAddresses = 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 -> [(TicketKey, Natural)] -> GStateUpdate
  GSAddContractAddress :: ContractAddress -> ContractState -> GStateUpdate
  GSAddContractAddressAlias :: ContractAlias -> ContractAddress -> GStateUpdate
  GSSetStorageValue :: StorageScope st => ContractAddress -> T.Value st -> GStateUpdate
  GSSetBalance :: L1AddressKind kind => KindedAddress kind -> Mutez -> GStateUpdate
  GSAddTickets :: ImplicitAddress -> TicketKey -> Natural -> GStateUpdate
  GSRemoveTickets :: ImplicitAddress -> TicketKey -> Natural -> GStateUpdate
  GSIncrementCounter :: GStateUpdate
  GSUpdateCounter :: GlobalCounter -> GStateUpdate
  GSSetBigMapCounter :: BigMapCounter -> GStateUpdate
  GSSetDelegate :: L1AddressKind kind => KindedAddress kind -> 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 " +| addr |+ " to " +| maybe "<nobody>" build key
      GSAddTickets addr key amount ->
        "Add " +| amount |+ " tickets to " +| addr |+ " type " +| key |+ ""
      GSRemoveTickets addr key amount ->
        "Remove " +| amount |+ " tickets from " +| addr |+ " type " +| key |+ ""

data GStateUpdateError
  = GStateAddressExists Address
  | GStateUnknownAddress Address
  | GStateStorageNotMatch ContractAddress
  | GStateNotDelegate ImplicitAddress
  | GStateCantDeleteDelegate ImplicitAddress
  | GStateNoBLSDelegate Address KeyHash
  | GStateAlreadySetDelegate L1Address (Maybe KeyHash)
  | GStateInsufficientTickets ImplicitAddress TicketKey ("needs" :! Natural) ("has" :! Natural)
  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
    GStateNotDelegate addr -> "Address " +| addr |+ " is not registered as delegate."
    GStateNoBLSDelegate addr kh ->
      nameF ("Can not set delegate for " +| addr |+ "") $
        nameF "tz4 addresses can't be delegates" $
          build kh
    GStateCantDeleteDelegate addr ->
      "Delegate deletion is forbidden for " +| addr |+ ""
    GStateAlreadySetDelegate addr kh ->
      "Already set delegate for " +| addr |+ " to " <> maybe "<nothing>" build kh
    GStateInsufficientTickets addr key (arg #needs -> needs) (arg #has -> has) ->
      "Address " +| addr |+ "does not have enough tickets of type " +| key |+ ": has " +| has
        |+ ", needs " +| needs |+ ""

-- | Apply 'GStateUpdate' to 'GState'.
applyUpdate :: GStateUpdate -> GState -> Either GStateUpdateError GState
applyUpdate =
  \case
    GSAddImplicitAddress addr st tickets -> addImplicitAddress addr st tickets
    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
    GSAddTickets addr key amount -> addTickets addr key amount
    GSRemoveTickets addr key amount -> removeTickets addr key amount

-- | Update ticket balance value associated with given address.
addTickets
  :: ImplicitAddress
  -> TicketKey
  -> Natural
  -> GState
  -> Either GStateUpdateError GState
addTickets addr key num = updateAddressState addr $
  Right . (isTicketsL . at key %~ Just . maybe num (+ num))

-- | Update ticket balance value associated with given address.
removeTickets
  :: ImplicitAddress
  -> TicketKey
  -> Natural
  -> GState
  -> Either GStateUpdateError GState
removeTickets addr key num = updateAddressState addr \st ->
  case st ^. isTicketsL . at key of
    Just n | n >= num -> Right $ st & isTicketsL . at key .~ case n - num of
      0 -> Nothing
      n' -> Just n'
    n -> Left $ GStateInsufficientTickets addr key ! #needs num ! #has (fromMaybe 0 n)

-- | 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 -> [(TicketKey, Natural)] -> GState -> Either GStateUpdateError GState
addImplicitAddress addr st tickets gs
  | addr `Map.member` (gsImplicitAddresses gs) = Left $ GStateAddressExists (MkAddress addr)
  | otherwise = Right $
      gs & gsImplicitAddressesL . at addr ?~ ImplicitState st (fromList tickets) Nothing

-- | 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{} -> isBalanceL
  ContractAddress{} -> csBalanceL
  where _ = usingImplicitOrContractKind @kind ()

-- | Check if address delegate is set to itself, i.e. it's "registered delegate"
-- in network terms. "Registered delegates" can't change their delegate. On the
-- other hand, those can be set as delegates for other addresses.
isRegisteredDelegate :: ImplicitAddress -> GState -> Bool
isRegisteredDelegate addr@(ImplicitAddress kh) gs =
  gs ^? gsImplicitAddressesL . at addr . traverse . isDelegateL == Just (Just kh)

-- | Set delegate for a given address
setDelegate
  :: forall kind. L1AddressKind kind
  => KindedAddress kind -> Maybe KeyHash -> GState -> Either GStateUpdateError GState
setDelegate addr newDelegate gs
  | Just oldDelegate@Just{} <- gs ^? addressesL addr . at addr . traverse . delegateLens
  , oldDelegate == newDelegate
  -- network fails when new delegate == old delegate == Just kh, so we do likewise
  -- note it _doesn't_ fail when new == old == Nothing
  = Left $ GStateAlreadySetDelegate (Constrained addr) newDelegate
  | ImplicitAddress{} <- addr
  , isRegisteredDelegate addr gs
  -- implicit addresses that are registered delegates can't change delegates
  = Left $ GStateCantDeleteDelegate addr
  | Just h@Hash{..} <- newDelegate
  , HashKey KeyTypeBLS <- hTag
  = Left $ GStateNoBLSDelegate (Constrained addr) h
  | Just kh <- newDelegate
  , let keyAddr = ImplicitAddress kh
  , not $ isRegisteredDelegate keyAddr gs
  -- can't set delegate to an address not registered as delegate
  , MkAddress keyAddr /= MkAddress addr -- but implicit contract can set delegate to itself
  = Left $ GStateNotDelegate keyAddr
  | otherwise = updateAddressState addr (pure . (delegateLens .~ newDelegate)) gs
  where
    delegateLens :: Lens' (AddressStateFam kind) (Maybe KeyHash)
    delegateLens = case addr of
      ContractAddress{} -> csDelegateL
      ImplicitAddress{} -> isDelegateL

type family AddressStateFam kind where
  AddressStateFam 'AddressKindImplicit = ImplicitState
  AddressStateFam 'AddressKindContract = ContractState
  AddressStateFam 'AddressKindSmartRollup = ()

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
  SmartRollupAddress{} -> gsSmartRollupAddressesL

-- | 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)