Writer.hs

{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}

module Haskoin.Store.Database.Writer (WriterT, runWriter) where

import Control.Monad.Reader (ReaderT (..))
import qualified Control.Monad.Reader as R
import Data.HashMap.Strict (HashMap)
import qualified Data.HashMap.Strict as M
import Data.List (sortOn)
import Data.Ord (Down (..))
import Data.Tuple (swap)
import Database.RocksDB (BatchOp, DB)
import Database.RocksDB.Query (
    deleteOp,
    deleteOpCF,
    insertOp,
    insertOpCF,
    writeBatch,
 )
import Haskoin (
    Address,
    BlockHash,
    BlockHeight,
    Network,
    OutPoint (..),
    TxHash,
    headerHash,
    txHash,
 )
import Haskoin.Store.Common
import Haskoin.Store.Data
import Haskoin.Store.Database.Reader
import Haskoin.Store.Database.Types
import UnliftIO (
    MonadIO,
    TVar,
    atomically,
    liftIO,
    modifyTVar,
    newTVarIO,
    readTVarIO,
 )

data Writer = Writer
    { getReader :: !DatabaseReader
    , getState :: !(TVar Memory)
    }

type WriterT = ReaderT Writer

instance MonadIO m => StoreReadBase (WriterT m) where
    getNetwork = getNetworkI
    getBestBlock = getBestBlockI
    getBlocksAtHeight = getBlocksAtHeightI
    getBlock = getBlockI
    getTxData = getTxDataI
    getSpender = getSpenderI
    getUnspent = getUnspentI
    getBalance = getBalanceI
    getMempool = getMempoolI

data Memory = Memory
    { hNet ::
        !(Maybe Network)
    , hBest ::
        !(Maybe (Maybe BlockHash))
    , hBlock ::
        !(HashMap BlockHash (Maybe BlockData))
    , hHeight ::
        !(HashMap BlockHeight [BlockHash])
    , hTx ::
        !(HashMap TxHash (Maybe TxData))
    , hSpender ::
        !(HashMap OutPoint (Maybe Spender))
    , hUnspent ::
        !(HashMap OutPoint (Maybe Unspent))
    , hBalance ::
        !(HashMap Address (Maybe Balance))
    , hAddrTx ::
        !(HashMap (Address, TxRef) (Maybe ()))
    , hAddrOut ::
        !(HashMap (Address, BlockRef, OutPoint) (Maybe OutVal))
    , hMempool ::
        !(HashMap TxHash UnixTime)
    }
    deriving (Eq, Show)

instance MonadIO m => StoreWrite (WriterT m) where
    setBest h =
        ReaderT $ \Writer{getState = s} ->
            liftIO . atomically . modifyTVar s $
                setBestH h
    insertBlock b =
        ReaderT $ \Writer{getState = s} ->
            liftIO . atomically . modifyTVar s $
                insertBlockH b
    setBlocksAtHeight h g =
        ReaderT $ \Writer{getState = s} ->
            liftIO . atomically . modifyTVar s $
                setBlocksAtHeightH h g
    insertTx t =
        ReaderT $ \Writer{getState = s} ->
            liftIO . atomically . modifyTVar s $
                insertTxH t
    insertSpender p s' =
        ReaderT $ \Writer{getState = s} ->
            liftIO . atomically . modifyTVar s $
                insertSpenderH p s'
    deleteSpender p =
        ReaderT $ \Writer{getState = s} ->
            liftIO . atomically . modifyTVar s $
                deleteSpenderH p
    insertAddrTx a t =
        ReaderT $ \Writer{getState = s} ->
            liftIO . atomically . modifyTVar s $
                insertAddrTxH a t
    deleteAddrTx a t =
        ReaderT $ \Writer{getState = s} ->
            liftIO . atomically . modifyTVar s $
                deleteAddrTxH a t
    insertAddrUnspent a u =
        ReaderT $ \Writer{getState = s} ->
            liftIO . atomically . modifyTVar s $
                insertAddrUnspentH a u
    deleteAddrUnspent a u =
        ReaderT $ \Writer{getState = s} ->
            liftIO . atomically . modifyTVar s $
                deleteAddrUnspentH a u
    addToMempool x t =
        ReaderT $ \Writer{getState = s} ->
            liftIO . atomically . modifyTVar s $
                addToMempoolH x t
    deleteFromMempool x =
        ReaderT $ \Writer{getState = s} ->
            liftIO . atomically . modifyTVar s $
                deleteFromMempoolH x
    setBalance b =
        ReaderT $ \Writer{getState = s} ->
            liftIO . atomically . modifyTVar s $
                setBalanceH b
    insertUnspent h =
        ReaderT $ \Writer{getState = s} ->
            liftIO . atomically . modifyTVar s $
                insertUnspentH h
    deleteUnspent p =
        ReaderT $ \Writer{getState = s} ->
            liftIO . atomically . modifyTVar s $
                deleteUnspentH p

getLayered ::
    MonadIO m =>
    (Memory -> Maybe a) ->
    DatabaseReaderT m a ->
    WriterT m a
getLayered f g =
    ReaderT $ \Writer{getReader = db, getState = tmem} ->
        f <$> readTVarIO tmem >>= \case
            Just x -> return x
            Nothing -> runReaderT g db

runWriter ::
    MonadIO m =>
    DatabaseReader ->
    WriterT m a ->
    m a
runWriter bdb@DatabaseReader{databaseHandle = db} f = do
    mem <- runReaderT getMempool bdb
    hm <- newTVarIO (newMemory mem)
    x <- R.runReaderT f Writer{getReader = bdb, getState = hm}
    mem' <- readTVarIO hm
    let ops = hashMapOps db mem'
    writeBatch db ops
    return x

hashMapOps :: DB -> Memory -> [BatchOp]
hashMapOps db mem =
    bestBlockOp (hBest mem)
        <> blockHashOps db (hBlock mem)
        <> blockHeightOps db (hHeight mem)
        <> txOps db (hTx mem)
        <> spenderOps db (hSpender mem)
        <> balOps db (hBalance mem)
        <> addrTxOps db (hAddrTx mem)
        <> addrOutOps db (hAddrOut mem)
        <> mempoolOp (hMempool mem)
        <> unspentOps db (hUnspent mem)

bestBlockOp :: Maybe (Maybe BlockHash) -> [BatchOp]
bestBlockOp Nothing = []
bestBlockOp (Just Nothing) = [deleteOp BestKey]
bestBlockOp (Just (Just b)) = [insertOp BestKey b]

blockHashOps :: DB -> HashMap BlockHash (Maybe BlockData) -> [BatchOp]
blockHashOps db = map (uncurry f) . M.toList
  where
    f k (Just d) = insertOpCF (blockCF db) (BlockKey k) d
    f k Nothing = deleteOpCF (blockCF db) (BlockKey k)

blockHeightOps :: DB -> HashMap BlockHeight [BlockHash] -> [BatchOp]
blockHeightOps db = map (uncurry f) . M.toList
  where
    f = insertOpCF (heightCF db) . HeightKey

txOps :: DB -> HashMap TxHash (Maybe TxData) -> [BatchOp]
txOps db = map (uncurry f) . M.toList
  where
    f k (Just t) = insertOpCF (txCF db) (TxKey k) t
    f k Nothing = deleteOpCF (txCF db) (TxKey k)

spenderOps :: DB -> HashMap OutPoint (Maybe Spender) -> [BatchOp]
spenderOps db = map (uncurry f) . M.toList
  where
    f o (Just s) = insertOpCF (spenderCF db) (SpenderKey o) s
    f o Nothing = deleteOpCF (spenderCF db) (SpenderKey o)

balOps :: DB -> HashMap Address (Maybe Balance) -> [BatchOp]
balOps db = map (uncurry f) . M.toList
  where
    f a (Just b) = insertOpCF (balanceCF db) (BalKey a) (balanceToVal b)
    f a Nothing = deleteOpCF (balanceCF db) (BalKey a)

addrTxOps :: DB -> HashMap (Address, TxRef) (Maybe ()) -> [BatchOp]
addrTxOps db = map (uncurry f) . M.toList
  where
    f (a, t) (Just ()) = insertOpCF (addrTxCF db) (AddrTxKey a t) ()
    f (a, t) Nothing = deleteOpCF (addrTxCF db) (AddrTxKey a t)

addrOutOps ::
    DB ->
    HashMap (Address, BlockRef, OutPoint) (Maybe OutVal) ->
    [BatchOp]
addrOutOps db = map (uncurry f) . M.toList
  where
    f (a, b, p) (Just l) =
        insertOpCF
            (addrOutCF db)
            ( AddrOutKey
                { addrOutKeyA = a
                , addrOutKeyB = b
                , addrOutKeyP = p
                }
            )
            l
    f (a, b, p) Nothing =
        deleteOpCF
            (addrOutCF db)
            AddrOutKey
                { addrOutKeyA = a
                , addrOutKeyB = b
                , addrOutKeyP = p
                }

mempoolOp :: HashMap TxHash UnixTime -> [BatchOp]
mempoolOp = return . insertOp MemKey . sortOn Down . map swap . M.toList

unspentOps :: DB -> HashMap OutPoint (Maybe Unspent) -> [BatchOp]
unspentOps db = map (uncurry f) . M.toList
  where
    f p (Just u) =
        insertOpCF (unspentCF db) (UnspentKey p) (snd (unspentToVal u))
    f p Nothing =
        deleteOpCF (unspentCF db) (UnspentKey p)

getNetworkI :: MonadIO m => WriterT m Network
getNetworkI = getLayered hNet getNetwork

getBestBlockI :: MonadIO m => WriterT m (Maybe BlockHash)
getBestBlockI = getLayered getBestBlockH getBestBlock

getBlocksAtHeightI :: MonadIO m => BlockHeight -> WriterT m [BlockHash]
getBlocksAtHeightI bh =
    getLayered (getBlocksAtHeightH bh) (getBlocksAtHeight bh)

getBlockI :: MonadIO m => BlockHash -> WriterT m (Maybe BlockData)
getBlockI bh = getLayered (getBlockH bh) (getBlock bh)

getTxDataI :: MonadIO m => TxHash -> WriterT m (Maybe TxData)
getTxDataI th = getLayered (getTxDataH th) (getTxData th)

getSpenderI :: MonadIO m => OutPoint -> WriterT m (Maybe Spender)
getSpenderI op = getLayered (getSpenderH op) (getSpender op)

getBalanceI :: MonadIO m => Address -> WriterT m (Maybe Balance)
getBalanceI a = getLayered (getBalanceH a) (getBalance a)

getUnspentI :: MonadIO m => OutPoint -> WriterT m (Maybe Unspent)
getUnspentI op = getLayered (getUnspentH op) (getUnspent op)

getMempoolI :: MonadIO m => WriterT m [(UnixTime, TxHash)]
getMempoolI =
    ReaderT $ \Writer{getState = tmem} ->
        getMempoolH <$> readTVarIO tmem

newMemory :: [(UnixTime, TxHash)] -> Memory
newMemory mem =
    Memory
        { hNet = Nothing
        , hBest = Nothing
        , hBlock = M.empty
        , hHeight = M.empty
        , hTx = M.empty
        , hSpender = M.empty
        , hUnspent = M.empty
        , hBalance = M.empty
        , hAddrTx = M.empty
        , hAddrOut = M.empty
        , hMempool = M.fromList (map swap mem)
        }

getBestBlockH :: Memory -> Maybe (Maybe BlockHash)
getBestBlockH = hBest

getBlocksAtHeightH :: BlockHeight -> Memory -> Maybe [BlockHash]
getBlocksAtHeightH h = M.lookup h . hHeight

getBlockH :: BlockHash -> Memory -> Maybe (Maybe BlockData)
getBlockH h = M.lookup h . hBlock

getTxDataH :: TxHash -> Memory -> Maybe (Maybe TxData)
getTxDataH t = M.lookup t . hTx

getSpenderH :: OutPoint -> Memory -> Maybe (Maybe Spender)
getSpenderH op db = M.lookup op (hSpender db)

getBalanceH :: Address -> Memory -> Maybe (Maybe Balance)
getBalanceH a = M.lookup a . hBalance

getMempoolH :: Memory -> [(UnixTime, TxHash)]
getMempoolH = sortOn Down . map swap . M.toList . hMempool

setBestH :: BlockHash -> Memory -> Memory
setBestH h db = db{hBest = Just (Just h)}

insertBlockH :: BlockData -> Memory -> Memory
insertBlockH bd db =
    db
        { hBlock =
            M.insert
                (headerHash (blockDataHeader bd))
                (Just bd)
                (hBlock db)
        }

setBlocksAtHeightH :: [BlockHash] -> BlockHeight -> Memory -> Memory
setBlocksAtHeightH hs g db =
    db{hHeight = M.insert g hs (hHeight db)}

insertTxH :: TxData -> Memory -> Memory
insertTxH tx db =
    db{hTx = M.insert (txHash (txData tx)) (Just tx) (hTx db)}

insertSpenderH :: OutPoint -> Spender -> Memory -> Memory
insertSpenderH op s db =
    db{hSpender = M.insert op (Just s) (hSpender db)}

deleteSpenderH :: OutPoint -> Memory -> Memory
deleteSpenderH op db =
    db{hSpender = M.insert op Nothing (hSpender db)}

setBalanceH :: Balance -> Memory -> Memory
setBalanceH bal db =
    db{hBalance = M.insert (balanceAddress bal) (Just bal) (hBalance db)}

insertAddrTxH :: Address -> TxRef -> Memory -> Memory
insertAddrTxH a tr db =
    db{hAddrTx = M.insert (a, tr) (Just ()) (hAddrTx db)}

deleteAddrTxH :: Address -> TxRef -> Memory -> Memory
deleteAddrTxH a tr db =
    db{hAddrTx = M.insert (a, tr) Nothing (hAddrTx db)}

insertAddrUnspentH :: Address -> Unspent -> Memory -> Memory
insertAddrUnspentH a u db =
    let k = (a, unspentBlock u, unspentPoint u)
        v =
            OutVal
                { outValAmount = unspentAmount u
                , outValScript = unspentScript u
                }
     in db{hAddrOut = M.insert k (Just v) (hAddrOut db)}

deleteAddrUnspentH :: Address -> Unspent -> Memory -> Memory
deleteAddrUnspentH a u db =
    let k = (a, unspentBlock u, unspentPoint u)
     in db{hAddrOut = M.insert k Nothing (hAddrOut db)}

addToMempoolH :: TxHash -> UnixTime -> Memory -> Memory
addToMempoolH h t db =
    db{hMempool = M.insert h t (hMempool db)}

deleteFromMempoolH :: TxHash -> Memory -> Memory
deleteFromMempoolH h db =
    db{hMempool = M.delete h (hMempool db)}

getUnspentH :: OutPoint -> Memory -> Maybe (Maybe Unspent)
getUnspentH op db = M.lookup op (hUnspent db)

insertUnspentH :: Unspent -> Memory -> Memory
insertUnspentH u db =
    let k = fst (unspentToVal u)
     in db{hUnspent = M.insert k (Just u) (hUnspent db)}

deleteUnspentH :: OutPoint -> Memory -> Memory
deleteUnspentH op db =
    db{hUnspent = M.insert op Nothing (hUnspent db)}