module Test.Ouroboros.Network.Testnet.Node.ChainDB
  ( ChainDB (..)
  , SelectChain (..)
  , newChainDB
  , addBlock
  , getBlockPointSet
  ) where

import Control.Concurrent.Class.MonadSTM (MonadSTM (..))
import Data.Coerce (coerce)
import Data.Map (Map)
import Data.Map qualified as Map
import Data.Set (Set)
import Data.Set qualified as Set
import Ouroboros.Network.AnchoredFragment (Point)
import Ouroboros.Network.Block (ChainHash (..), HasFullHeader, HasHeader,
           blockHash, blockPoint, blockPrevHash)
import Ouroboros.Network.Mock.Chain (Chain (..), selectChain)
import Ouroboros.Network.Mock.Chain qualified as Chain

-- | ChainDB is an in memory store of all fetched (downloaded) blocks.
--
-- This type holds an index mapping previous hashes to their blocks (i.e. if a
-- block "A" has prevHash "H" then the entry "H -> [A]" exists in the map) and
-- the current version of the longest chain.
--
-- Used to simulate real world ChainDB, it offers the invariant that
-- 'cdbLongestChainVar' is always the longest known chain of downloaded blocks.
-- Whenever a node generates a new block it gets added here, and whenever it gets
-- a block via block fetch it gets added here as well. Everytime 'addBlock' is
-- called the possibly new longest chain gets computed, since the API is atomic
-- we can guarantee that in each moment ChainDB has the current longest chain.
--
-- This type is used in diffusion simulation.
--
data ChainDB block m = ChainDB { forall block (m :: * -> *).
ChainDB block m -> TVar m (Map (ChainHash block) [block])
cdbIndexVar :: TVar m (Map (ChainHash block) [block]),
                                 forall block (m :: * -> *). ChainDB block m -> TVar m (Chain block)
cdbLongestChainVar :: TVar m (Chain block)
                               }

-- | Constructs a new ChainDB, the index has only 1 value which is the
-- 'GenesisHash' but this hash does not map to any block.
--
newChainDB :: MonadSTM m => m (ChainDB block m)
newChainDB :: forall (m :: * -> *) block. MonadSTM m => m (ChainDB block m)
newChainDB = do
  indexVar <- Map (ChainHash block) [block]
-> m (TVar m (Map (ChainHash block) [block]))
forall a. a -> m (TVar m a)
forall (m :: * -> *) a. MonadSTM m => a -> m (TVar m a)
newTVarIO (ChainHash block -> [block] -> Map (ChainHash block) [block]
forall k a. k -> a -> Map k a
Map.singleton ChainHash block
forall {k} (b :: k). ChainHash b
GenesisHash [])
  longestChain <- newTVarIO Genesis
  return (ChainDB indexVar longestChain)

-- | Adds a block to ChainDB.
--
-- This function also recomputes the longest chain with the new block
-- information.
--
addBlock :: (MonadSTM m, HasFullHeader block)
         => block -> ChainDB block m -> STM m ()
addBlock :: forall (m :: * -> *) block.
(MonadSTM m, HasFullHeader block) =>
block -> ChainDB block m -> STM m ()
addBlock block
block chainDB :: ChainDB block m
chainDB@(ChainDB TVar m (Map (ChainHash block) [block])
indexVar TVar m (Chain block)
lchainVar) = do
  TVar m (Map (ChainHash block) [block])
-> (Map (ChainHash block) [block] -> Map (ChainHash block) [block])
-> STM m ()
forall a. TVar m a -> (a -> a) -> STM m ()
forall (m :: * -> *) a.
MonadSTM m =>
TVar m a -> (a -> a) -> STM m ()
modifyTVar' TVar m (Map (ChainHash block) [block])
indexVar ((Map (ChainHash block) [block] -> Map (ChainHash block) [block])
 -> STM m ())
-> (Map (ChainHash block) [block] -> Map (ChainHash block) [block])
-> STM m ()
forall a b. (a -> b) -> a -> b
$ \Map (ChainHash block) [block]
index ->
    case ChainHash block -> Map (ChainHash block) [block] -> Maybe [block]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup (block -> ChainHash block
forall b. HasFullHeader b => b -> ChainHash b
blockPrevHash block
block) Map (ChainHash block) [block]
index of
      Maybe [block]
Nothing -> ([block] -> [block] -> [block])
-> ChainHash block
-> [block]
-> Map (ChainHash block) [block]
-> Map (ChainHash block) [block]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith [block] -> [block] -> [block]
forall a. [a] -> [a] -> [a]
(++) ChainHash block
forall {k} (b :: k). ChainHash b
GenesisHash [block
block] Map (ChainHash block) [block]
index
      Just [block]
_  -> ([block] -> [block] -> [block])
-> ChainHash block
-> [block]
-> Map (ChainHash block) [block]
-> Map (ChainHash block) [block]
forall k a. Ord k => (a -> a -> a) -> k -> a -> Map k a -> Map k a
Map.insertWith [block] -> [block] -> [block]
forall a. [a] -> [a] -> [a]
(++) (block -> ChainHash block
forall b. HasFullHeader b => b -> ChainHash b
blockPrevHash block
block) [block
block] Map (ChainHash block) [block]
index
  longestChain <- ChainDB block m -> STM m (Chain block)
forall block (m :: * -> *).
(HasHeader block, MonadSTM m) =>
ChainDB block m -> STM m (Chain block)
getLongestChain ChainDB block m
chainDB
  writeTVar lchainVar longestChain

-- | Constructs the block Point set of all downloaded blocks
--
getBlockPointSet :: (MonadSTM m, HasHeader block)
                 => ChainDB block m -> STM m (Set (Point block))
getBlockPointSet :: forall (m :: * -> *) block.
(MonadSTM m, HasHeader block) =>
ChainDB block m -> STM m (Set (Point block))
getBlockPointSet (ChainDB TVar m (Map (ChainHash block) [block])
indexVar TVar m (Chain block)
_) = do
  index <- TVar m (Map (ChainHash block) [block])
-> STM m (Map (ChainHash block) [block])
forall a. TVar m a -> STM m a
forall (m :: * -> *) a. MonadSTM m => TVar m a -> STM m a
readTVar TVar m (Map (ChainHash block) [block])
indexVar
  return (foldMap (Set.fromList . map blockPoint) index)

-- | Computes the longest chain from Genesis
--
getLongestChain :: (HasHeader block, MonadSTM m)
                => ChainDB block m
                -> STM m (Chain block)
getLongestChain :: forall block (m :: * -> *).
(HasHeader block, MonadSTM m) =>
ChainDB block m -> STM m (Chain block)
getLongestChain (ChainDB TVar m (Map (ChainHash block) [block])
indexVar TVar m (Chain block)
_) = do
    index <- TVar m (Map (ChainHash block) [block])
-> STM m (Map (ChainHash block) [block])
forall a. TVar m a -> STM m a
forall (m :: * -> *) a. MonadSTM m => TVar m a -> STM m a
readTVar TVar m (Map (ChainHash block) [block])
indexVar
    return (go Nothing Genesis index)
  where
    go :: HasHeader block
       => Maybe block
       -> Chain block
       -> Map (ChainHash block) [block]
       -> Chain block
    go :: forall block.
HasHeader block =>
Maybe block
-> Chain block -> Map (ChainHash block) [block] -> Chain block
go Maybe block
mbblock Chain block
chain Map (ChainHash block) [block]
m =
      let hash :: ChainHash block
hash = ChainHash block
-> (block -> ChainHash block) -> Maybe block -> ChainHash block
forall b a. b -> (a -> b) -> Maybe a -> b
maybe ChainHash block
forall {k} (b :: k). ChainHash b
GenesisHash (HeaderHash block -> ChainHash block
forall {k} (b :: k). HeaderHash b -> ChainHash b
BlockHash (HeaderHash block -> ChainHash block)
-> (block -> HeaderHash block) -> block -> ChainHash block
forall b c a. (b -> c) -> (a -> b) -> a -> c
. block -> HeaderHash block
forall b. HasHeader b => b -> HeaderHash b
blockHash) Maybe block
mbblock
       in case ChainHash block -> Map (ChainHash block) [block] -> Maybe [block]
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup ChainHash block
hash Map (ChainHash block) [block]
m of
            Maybe [block]
Nothing -> Chain block -> (block -> Chain block) -> Maybe block -> Chain block
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Chain block
forall block. Chain block
Genesis (block -> Chain block -> Chain block
forall block.
HasHeader block =>
block -> Chain block -> Chain block
`Chain.addBlock` Chain block
chain) Maybe block
mbblock
            Just [block]
blocks ->
              let longestChain :: Chain block
longestChain = SelectChain block -> Chain block
forall block. SelectChain block -> Chain block
getSelectedChain
                               (SelectChain block -> Chain block)
-> SelectChain block -> Chain block
forall a b. (a -> b) -> a -> b
$ (block -> SelectChain block) -> [block] -> SelectChain block
forall m a. Monoid m => (a -> m) -> [a] -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap (\block
b -> Chain block -> SelectChain block
forall block. Chain block -> SelectChain block
SelectChain
                                              (Chain block -> SelectChain block)
-> Chain block -> SelectChain block
forall a b. (a -> b) -> a -> b
$ Maybe block
-> Chain block -> Map (ChainHash block) [block] -> Chain block
forall block.
HasHeader block =>
Maybe block
-> Chain block -> Map (ChainHash block) [block] -> Chain block
go (block -> Maybe block
forall a. a -> Maybe a
Just block
b) Chain block
chain Map (ChainHash block) [block]
m)
                                         [block]
blocks
               in Chain block -> (block -> Chain block) -> Maybe block -> Chain block
forall b a. b -> (a -> b) -> Maybe a -> b
maybe Chain block
longestChain (block -> Chain block -> Chain block
forall block.
HasHeader block =>
block -> Chain block -> Chain block
`Chain.addBlock` Chain block
longestChain) Maybe block
mbblock

-- | Chain selection as a 'Monoid'.
--
newtype SelectChain block = SelectChain { forall block. SelectChain block -> Chain block
getSelectedChain :: Chain block }

instance HasHeader block => Semigroup (SelectChain block) where
    <> :: SelectChain block -> SelectChain block -> SelectChain block
(<>) = ((Chain block -> Chain block -> Chain block)
-> SelectChain block -> SelectChain block -> SelectChain block
forall {block}.
(Chain block -> Chain block -> Chain block)
-> SelectChain block -> SelectChain block -> SelectChain block
forall a b. Coercible a b => a -> b
coerce :: (     Chain block ->       Chain block ->       Chain block)
                   -> SelectChain block -> SelectChain block -> SelectChain block)
           Chain block -> Chain block -> Chain block
forall block.
HasHeader block =>
Chain block -> Chain block -> Chain block
selectChain

instance HasHeader block => Monoid (SelectChain block) where
    mempty :: SelectChain block
mempty = Chain block -> SelectChain block
forall block. Chain block -> SelectChain block
SelectChain Chain block
forall block. Chain block
Genesis