Type safe storage

concordium-consensus/src/Concordium/Afgjort/Finalize.hs

@@ -73,6 +73,7 @@ import Concordium.GlobalState.BlockPointer
 import Concordium.GlobalState.BlockState
 import Concordium.GlobalState.Finalization
 import Concordium.GlobalState.Parameters
+import Concordium.GlobalState.Persistent.BlobStore (MBSStore)
 import Concordium.GlobalState.Transactions
 import Concordium.GlobalState.TreeState
 import Concordium.Kontrol
@@ -1158,6 +1159,8 @@ newtype ActiveFinalizationM (pv :: ProtocolVersion) (r :: Type) (s :: Type) (m :
           SkovQueryMonad
         )
 
+type instance MBSStore (ActiveFinalizationM pv r s m) = MBSStore m
+
 deriving instance (TokenStateOperations ts m) => TokenStateOperations ts (ActiveFinalizationM pv r s m)
 deriving instance (PLTQuery bs ts m) => PLTQuery bs ts (ActiveFinalizationM pv r s m)
 deriving instance (MonadProtocolVersion m) => MonadProtocolVersion (ActiveFinalizationM pv r s m)

concordium-consensus/src/Concordium/Birk/Bake.hs

@@ -200,8 +200,13 @@ class (SkovMonad m, FinalizationMonad m) => BakerMonad m where
     tryBake :: BakerIdentity -> Slot -> m BakeResult
 
 instance
-    (FinalizationMonad (SkovT pv h c m), MonadIO m, SkovMonad (SkovT pv h c m), TreeStateMonad (SkovT pv h c m), OnSkov (SkovT pv h c m)) =>
-    BakerMonad (SkovT pv h c m)
+    ( FinalizationMonad (SkovT store pv h c m),
+      MonadIO m,
+      SkovMonad (SkovT store pv h c m),
+      TreeStateMonad (SkovT store pv h c m),
+      OnSkov (SkovT store pv h c m)
+    ) =>
+    BakerMonad (SkovT store pv h c m)
     where
     bakeForSlot = doBakeForSlot
     tryBake = doTryBake

concordium-consensus/src/Concordium/GlobalState.hs

@@ -1,5 +1,6 @@
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE DerivingVia #-}
+{-# LANGUAGE ExistentialQuantification #-}
 {-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeFamilies #-}
@@ -52,10 +53,12 @@ instance Show GlobalStateInitException where
 instance Exception GlobalStateInitException
 
 -- | The read-only context type associated with a global state configuration.
-type GSContext pv = PersistentBlockStateContext pv
+type GSContext store pv = PersistentBlockStateContext store pv
 
 -- | The (mutable) state type associated with a global state configuration.
-type GSState pv = SkovPersistentData pv
+type GSState store pv = SkovPersistentData store pv
+
+data InitialisedState pv = forall store. InitialisedState (GSContext store pv) (GSState store pv)
 
 -- | Generate context and state from the initial configuration if the state
 --  exists already. This may have 'IO' side effects to set up any necessary
@@ -68,7 +71,10 @@ type GSState pv = SkovPersistentData pv
 --  Note that even if the state is successfully loaded it is not in a usable
 --  state for an active consensus and must be activated before. Use
 --  'activateGlobalState' for that.
-initialiseExistingGlobalState :: forall pv. (IsProtocolVersion pv) => SProtocolVersion pv -> GlobalStateConfig -> LogIO (Maybe (GSContext pv, GSState pv))
+initialiseExistingGlobalState ::
+    forall pv.
+    (IsProtocolVersion pv) =>
+    SProtocolVersion pv -> GlobalStateConfig -> LogIO (Maybe (InitialisedState pv))
 initialiseExistingGlobalState _ GlobalStateConfig{..} = do
     -- check if all the necessary database files exist
     existingDB <- checkExistingDatabase dtdbTreeStateDirectory dtdbBlockStateFile
@@ -84,13 +90,14 @@ initialiseExistingGlobalState _ GlobalStateConfig{..} = do
                 skovData <-
                     runLoggerT (loadSkovPersistentData dtdbRuntimeParameters dtdbTreeStateDirectory pbsc) logm
                         `onException` closeBlobStore pbscBlobStore
-                return (Just (pbsc, skovData))
+                return (Just $ InitialisedState pbsc skovData)
         else return Nothing
 
 -- | Initialize a 'PersistentBlockStateContext' via the provided
 --  'GlobalStateConfig'.
 --  This function attempts to create a new blob store.
-initializePersistentBlockStateContext :: GlobalStateConfig -> IO (PersistentBlockStateContext pv)
+initializePersistentBlockStateContext ::
+    GlobalStateConfig -> IO (PersistentBlockStateContext store pv)
 initializePersistentBlockStateContext GlobalStateConfig{..} = do
     pbscBlobStore <- createBlobStore dtdbBlockStateFile
     pbscAccountCache <- newAccountCache (rpAccountsCacheSize dtdbRuntimeParameters)
@@ -118,15 +125,15 @@ migrateExistingState ::
     -- | The configuration.
     GlobalStateConfig ->
     -- | Global state context for the state we are migrating from.
-    GSContext oldpv ->
+    GSContext oldstore oldpv ->
     -- | The state of the chain we are migrating from. See documentation above for assumptions.
-    GSState oldpv ->
+    GSState oldstore oldpv ->
     -- | Auxiliary migration data.
     StateMigrationParameters oldpv pv ->
     -- | Regenesis data for the new chain. This is in effect the genesis block of the new chain.
     Regenesis pv ->
     -- | The return value is the context and state for the new chain.
-    LogIO (GSContext pv, GSState pv)
+    LogIO (InitialisedState pv)
 migrateExistingState gsc@GlobalStateConfig{..} oldPbsc oldState migration genData = do
     pbsc <- liftIO $ initializePersistentBlockStateContext gsc
     newInitialBlockState <- flip runBlobStoreT oldPbsc . flip runBlobStoreT pbsc $ do
@@ -148,12 +155,14 @@ migrateExistingState gsc@GlobalStateConfig{..} oldPbsc oldState migration genDat
     isd <-
         runReaderT (runPersistentBlockStateMonad initGS) pbsc
             `onException` liftIO (destroyBlobStore (pbscBlobStore pbsc))
-    return (pbsc, isd)
+    return (InitialisedState pbsc isd)
 
 -- | Initialise new global state with the given genesis. If the state already
 --  exists this will raise an exception. It is not necessary to call 'activateGlobalState'
 --  on the generated state, as this will establish the necessary invariants.
-initialiseNewGlobalState :: (IsProtocolVersion pv, IsConsensusV0 pv) => GenesisData pv -> GlobalStateConfig -> LogIO (GSContext pv, GSState pv)
+initialiseNewGlobalState ::
+    (IsProtocolVersion pv, IsConsensusV0 pv) =>
+    GenesisData pv -> GlobalStateConfig -> LogIO (InitialisedState pv)
 initialiseNewGlobalState genData gsc@GlobalStateConfig{..} = do
     pbsc@PersistentBlockStateContext{..} <- liftIO $ initializePersistentBlockStateContext gsc
     let initGS = do
@@ -169,22 +178,27 @@ initialiseNewGlobalState genData gsc@GlobalStateConfig{..} = do
     isd <-
         runReaderT (runPersistentBlockStateMonad initGS) pbsc
             `onException` liftIO (destroyBlobStore pbscBlobStore)
-    return (pbsc, isd)
+    return (InitialisedState pbsc isd)
 
 -- | Either initialise an existing state, or if it does not exist, initialise a new one with the given genesis.
-initialiseGlobalState :: forall pv. (IsProtocolVersion pv, IsConsensusV0 pv) => GenesisData pv -> GlobalStateConfig -> LogIO (GSContext pv, GSState pv)
+initialiseGlobalState ::
+    forall pv.
+    (IsProtocolVersion pv, IsConsensusV0 pv) =>
+    GenesisData pv -> GlobalStateConfig -> LogIO (InitialisedState pv)
 initialiseGlobalState gd cfg =
     initialiseExistingGlobalState (protocolVersion @pv) cfg >>= \case
         Nothing -> initialiseNewGlobalState gd cfg
         Just config -> return config
 
 -- | Establish all the necessary invariants so that the state can be used by
 --  consensus. This should only be called once per initialised state.
-activateGlobalState :: (IsProtocolVersion pv) => Proxy pv -> GSContext pv -> GSState pv -> LogIO (GSState pv)
+activateGlobalState ::
+    (IsProtocolVersion pv) =>
+    Proxy pv -> GSContext store pv -> GSState store pv -> LogIO (GSState store pv)
 activateGlobalState _ = activateSkovPersistentData
 
 -- | Shutdown the global state.
-shutdownGlobalState :: SProtocolVersion pv -> GSContext pv -> GSState pv -> IO ()
+shutdownGlobalState :: SProtocolVersion pv -> GSContext store pv -> GSState store pv -> IO ()
 shutdownGlobalState _ PersistentBlockStateContext{..} st = do
     closeBlobStore pbscBlobStore
     closeSkovPersistentData st

concordium-consensus/src/Concordium/GlobalState/AccountMap.hs

@@ -69,33 +69,33 @@ newtype AccountMap (pv :: ProtocolVersion) fix = AccountMap
     }
 
 -- | The account map to be used in the persistent block state.
-type PersistentAccountMap pv = AccountMap pv BufferedFix
+type PersistentAccountMap store pv = AccountMap pv (BufferedFix store)
 
 -- | See documentation of @migratePersistentBlockState@.
 migratePersistentAccountMap ::
     (BlobStorable m AccountIndex, BlobStorable (t m) AccountIndex, MonadTrans t) =>
-    PersistentAccountMap oldpv ->
-    t m (PersistentAccountMap pv)
+    PersistentAccountMap (MBSStore m) oldpv ->
+    t m (PersistentAccountMap (MBSStore (t m)) pv)
 migratePersistentAccountMap (AccountMap am) = AccountMap <$> Trie.migrateTrieN True return am
 
 -- | The account map that is purely in memory and used in the basic block state.
 type PureAccountMap pv = AccountMap pv Fix
 
 -- Necessary state storage instances for the persistent map. The pure one is not
 -- stored so does not need the related instances.
-instance (MonadBlobStore m) => Cacheable m (PersistentAccountMap pv) where
+instance (MonadBlobStore m, store ~ MBSStore m) => Cacheable m (PersistentAccountMap store pv) where
     cache (AccountMap am) = AccountMap <$> cache am
     {-# INLINE cache #-}
 
-instance (MonadBlobStore m) => BlobStorable m (PersistentAccountMap pv) where
+instance (MonadBlobStore m, store ~ MBSStore m) => BlobStorable m (PersistentAccountMap store pv) where
     storeUpdate (AccountMap am) = second AccountMap <$> storeUpdate am
     {-# INLINE storeUpdate #-}
 
     load = fmap AccountMap <$> load
     {-# INLINE load #-}
 
 -- | Convert a pure account map to the persistent one.
-toPersistent :: (MonadBlobStore m) => PureAccountMap pv -> m (PersistentAccountMap pv)
+toPersistent :: (MonadBlobStore m) => PureAccountMap pv -> m (PersistentAccountMap (MBSStore m) pv)
 toPersistent = fmap AccountMap . Trie.fromTrie . unAccountMap
 
 -- Aliases for reducing constraint repetition in methods below.

concordium-consensus/src/Concordium/GlobalState/BlockState.hs

@@ -337,38 +337,38 @@ class (BlockStateTypes m, Monad m) => AccountOperations m where
 --  state. At the end of contract execution the mutable state is "frozen", which
 --  converts it to the persistent version, retaining as much sharing as possible
 --  with the previous version.
-type family UpdatableContractState (v :: Wasm.WasmVersion) = ty | ty -> v where
-    UpdatableContractState GSWasm.V0 = Wasm.ContractState
-    UpdatableContractState GSWasm.V1 = StateV1.MutableState
+type family UpdatableContractState store (v :: Wasm.WasmVersion) = ty | ty -> v where
+    UpdatableContractState store GSWasm.V0 = Wasm.ContractState
+    UpdatableContractState store GSWasm.V1 = StateV1.MutableState store
 
 -- | An external representation of the persistent (i.e., frozen) contract state.
 --  This is used to pass this state through FFI for queries and should not be
 --  used during contract execution in the scheduler since it's considered an
 --  implementation detail and needs to be used together with the correct loader
 --  callback. Higher-level abstractions should be used in the scheduler.
-type family ExternalContractState (v :: Wasm.WasmVersion) = ty | ty -> v where
-    ExternalContractState GSWasm.V0 = Wasm.ContractState
-    ExternalContractState GSWasm.V1 = StateV1.PersistentState
+type family ExternalContractState store (v :: Wasm.WasmVersion) = ty | ty -> v where
+    ExternalContractState store GSWasm.V0 = Wasm.ContractState
+    ExternalContractState store GSWasm.V1 = StateV1.PersistentState store
 
 class (BlockStateTypes m, Monad m) => ContractStateOperations m where
     -- | Convert a persistent state to a mutable one that can be updated by the
     --  scheduler. This function must generate independent mutable states for
     --  each invocation, where independent means that updates to different
     --  versions are __not__ reflected in others.
-    thawContractState :: ContractState m v -> m (UpdatableContractState v)
+    thawContractState :: ContractState m v -> m (UpdatableContractState (MBSStore m) v)
 
     -- | Convert a persistent state to its external representation that can be
     --  passed through FFI. The state should be used together with the
     --  callbacks returned by 'getV1StateContext'.
-    externalContractState :: ContractState m v -> m (ExternalContractState v)
+    externalContractState :: ContractState m v -> m (ExternalContractState (MBSStore m) v)
 
     -- | Get the callback to allow loading the contract state. Contracts are
     --  executed on the other end of FFI, and state is managed by Haskell, this
     --  gives access to state across the FFI boundary.
     --
     --  V0 state is a simple byte-array which is copied over the FFI boundary, so
     --  it does not require an analogous construct.
-    getV1StateContext :: m LoadCallback
+    getV1StateContext :: m (LoadCallback (MBSStore m))
 
     -- | Size of the persistent V0 state. The way charging is done for V0
     --  contracts requires us to get this information when loading the state __at
@@ -793,15 +793,15 @@ mintTotal MintAmounts{..} = mintBakingReward + mintFinalizationReward + mintDeve
 --  to simplify function API. Thus values are immediately deconstructed.
 --  It is parameterized by the concrete instrumented module @im@ and the
 --  WasmVersion @v@.
-data NewInstanceData im v = NewInstanceData
+data NewInstanceData store im v = NewInstanceData
     { -- | Name of the init method used to initialize the contract.
       nidInitName :: Wasm.InitName,
       -- | Receive functions suitable for this instance.
       nidEntrypoints :: Set.Set Wasm.ReceiveName,
       -- | Module interface that contains the code of the contract.
       nidInterface :: GSWasm.ModuleInterfaceA im,
       -- | Initial state of the instance.
-      nidInitialState :: UpdatableContractState v,
+      nidInitialState :: UpdatableContractState store v,
       -- | Initial balance.
       nidInitialAmount :: Amount,
       -- | Owner/creator of the instance.
@@ -873,7 +873,7 @@ class (BlockStateQuery m, PLTQuery (UpdatableBlockState m) (MutableTokenState m)
     bsoCreateAccount :: UpdatableBlockState m -> GlobalContext -> AccountAddress -> AccountCredential -> m (Maybe (Account m), UpdatableBlockState m)
 
     -- | Add a new smart contract instance to the state.
-    bsoPutNewInstance :: forall v. (Wasm.IsWasmVersion v) => UpdatableBlockState m -> NewInstanceData (InstrumentedModuleRef m v) v -> m (ContractAddress, UpdatableBlockState m)
+    bsoPutNewInstance :: forall v. (Wasm.IsWasmVersion v) => UpdatableBlockState m -> NewInstanceData (MBSStore m) (InstrumentedModuleRef m v) v -> m (ContractAddress, UpdatableBlockState m)
 
     -- | Add the module to the global state. If a module with the given address
     --  already exists return @False@.
@@ -937,7 +937,7 @@ class (BlockStateQuery m, PLTQuery (UpdatableBlockState m) (MutableTokenState m)
         UpdatableBlockState m ->
         ContractAddress ->
         AmountDelta ->
-        Maybe (UpdatableContractState v) ->
+        Maybe (UpdatableContractState (MBSStore m) v) ->
         Maybe (GSWasm.ModuleInterfaceA (InstrumentedModuleRef m v), Set.Set Wasm.ReceiveName) ->
         m (UpdatableBlockState m)
 
@@ -1812,7 +1812,7 @@ class (BlockStateOperations m, FixedSizeSerialization (BlockStateRef m)) => Bloc
 
     -- | Retrieve the callback that is needed to read state that is not in
     --  memory. This is needed for using V1 contract state.
-    blockStateLoadCallback :: m LoadCallback
+    blockStateLoadCallback :: m (LoadCallback (MBSStore m))
 
     -- | Shut down any caches used by the block state. This is used to free
     --  up the memory in the case where the block state is no longer being

concordium-consensus/src/Concordium/GlobalState/ContractStateFFIHelpers.hs

@@ -17,20 +17,20 @@ data Vec
 --  vector that should be passed to the Rust runtime.
 type LoadCallbackType = Word64 -> IO (Ptr Vec)
 
-type LoadCallback = FunPtr LoadCallbackType
+type LoadCallback store = FunPtr LoadCallbackType
 
 -- | Callback for writing to the blob store from the provided buffer. The
 --  arguments are the buffer where the data is and the amount of data to write.
 --  It is assumed that the buffer has sufficient size. The return value is the
 --  location where data was written.
 type StoreCallbackType = Ptr Word8 -> CSize -> IO Word64
 
-type StoreCallback = FunPtr StoreCallbackType
+type StoreCallback store = FunPtr StoreCallbackType
 
 -- | Wrappers for making callbacks from Haskell functions or closures.
-foreign import ccall "wrapper" createLoadCallback :: LoadCallbackType -> IO LoadCallback
+foreign import ccall "wrapper" createLoadCallback :: LoadCallbackType -> IO (LoadCallback store)
 
-foreign import ccall "wrapper" createStoreCallback :: StoreCallbackType -> IO StoreCallback
+foreign import ccall "wrapper" createStoreCallback :: StoreCallbackType -> IO (StoreCallback store)
 
 -- | Allocate and return a Rust vector that contains the given data.
 foreign import ccall "copy_to_vec_ffi" copyToRustVec :: Ptr Word8 -> CSize -> IO (Ptr Vec)
@@ -39,9 +39,9 @@ foreign import ccall "copy_to_vec_ffi" copyToRustVec :: Ptr Word8 -> CSize -> IO
 --  implementation which never stores any data in the backing store. NOINLINE
 --  here ensures that only a single instance of callbacks is allocated.
 {-# NOINLINE errorLoadCallback #-}
-errorLoadCallback :: LoadCallback
+errorLoadCallback :: LoadCallback store
 errorLoadCallback = unsafePerformIO $ createLoadCallback (\_location -> error "Error load callback invoked, and it should not have been.")
 
 -- | Deallocate the callbacks. This should generally be called to not leak memory.
-freeErrorCallback :: LoadCallback -> IO ()
+freeErrorCallback :: LoadCallback store -> IO ()
 freeErrorCallback = freeHaskellFunPtr