Add ability to load ledger state into sqlite for benchmarking

.gitignore

@@ -51,3 +51,4 @@ specs/**/.ghc.environment.x86_64-linux-*
 
 *.pdf
 .envrc
+ledger-state.sqlite

libs/ledger-state/app/Main.hs

@@ -1,7 +1,9 @@
+{-# LANGUAGE OverloadedStrings #-}
 module Main where
 
 import Cardano.Ledger.State.Massiv
 import Cardano.Ledger.State.UTxO
+import Cardano.Ledger.State.Query
 import Control.Monad
 import Data.IORef
 import Options.Applicative
@@ -53,7 +55,9 @@ main = do
         (header "ledger-state - Tool for analyzing ledger state")
   forM_ (optsLedgerStateBinaryFile opts) $ \fp -> do
     ls <- loadLedgerState fp
-    printNewEpochStateStats $ countNewEpochStateStats ls
+    storeLedgerState "ledger-state.sqlite" ls
+    -- nes <- loadNewEpochState fp
+    -- printNewEpochStateStats $ countNewEpochStateStats ls
   forM_ (optsUtxoJsonFile opts) $ \fp -> do
     _ <- observeMemoryOriginalMap fp
     pure ()
@@ -73,14 +77,13 @@ observeMemoryOriginalMap fp = do
   writeIORef ref $ Just utxo -- ensure utxo doesn't get GCed
   pure ref
 
-observeMemory :: FilePath -> IO (IORef (Maybe UTxOs))
+--observeMemory :: FilePath -> IO (IORef (Maybe UTxOs))
 observeMemory fp = do
   ref <- newIORef Nothing
   utxo <- loadMassivUTxO fp
   utxo `seq` putStrLn "Loaded"
   performGC
   _ <- getChar
-  printStats utxo
   writeIORef ref $ Just utxo -- ensure utxo doesn't get GCed
   pure ref
 

libs/ledger-state/bench/Memory.hs

@@ -1,21 +1,28 @@
 {-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE OverloadedStrings #-}
 
 module Main where
 
 import Weigh
 import Cardano.Ledger.State.UTxO
-import Cardano.Ledger.State.Massiv
+import Cardano.Ledger.State.Query
+import Data.Map.Strict as Map
+import qualified Data.Text as T
 
 main :: IO ()
 main = do
   let fp = "/home/lehins/Downloads/mainnet-utxo-2021-09-15.json"
+      db = "/home/lehins/iohk/cardano-ledger-specs/ledger-state.sqlite" :: T.Text
   let cols = [Case, Max, MaxOS, Live, Allocated, GCs, Check]
   -- utxo <- loadBinUTxO fp
   mainWith $ do
     setColumns cols
     -- action "loadLedgerState" $ do
     --   !_ <- loadLedgerState "/home/lehins/iohk/chain/mainnet/ledger-state.bin"
     --   pure ()
-    io "UTxO (Map TxIn ())" loadUTxO' fp
-    io "UTxO (IntMap (Map TxId ())" loadUTxOni' fp
-    io "UTxO (IntMap (HashMap TxId ())" loadUTxOhm' fp
+    io "UTxO (Map TxIn ())" (foldUTxO (\ !m !(!k, _) -> Map.insert k () m) mempty) db
+    io "UTxO (IntMap (Map TxId ())" (foldUTxO nestedInsertTxId' mempty) db
+    -- io "UTxO (IntMap (HashMap TxId ())" (foldUTxO nestedInsertHM' mempty) db
+    -- io "UTxO (Map TxIn ())" (loadUTxO') fp
+    -- io "UTxO (IntMap (Map TxId ())" loadUTxOni' fp
+    -- io "UTxO (IntMap (HashMap TxId ())" loadUTxOhm' fp

libs/ledger-state/ledger-state.cabal

@@ -44,15 +44,18 @@ common project-config
                     , deepseq
                     , persistent
                     , persistent-sqlite
+                    , persistent-template
                     , prettyprinter
                     , strict-containers
                     , small-steps
                     , text
+                    , transformers
                     , optparse-applicative
                     , heapsize
                     , massiv >= 1.0
                     , unordered-containers
                     , hashable
+                    , monad-logger
 
   ghc-options:        -Wall
                       -Wcompat
@@ -65,8 +68,10 @@ library
   import:             base, project-config
   exposed-modules:    Cardano.Ledger.State
                     , Cardano.Ledger.State.UTxO
+                    , Cardano.Ledger.State.Orphans
                     , Cardano.Ledger.State.Massiv
                     , Cardano.Ledger.State.Schema
+                    , Cardano.Ledger.State.Query
   hs-source-dirs:     src
 
 executable ledger-state
@@ -90,6 +95,8 @@ benchmark weigh
   build-depends:  base
                 , weigh
                 , ledger-state
+                , containers
+                , text
   ghc-options:   -Wall
                  -O2
                  -threaded

libs/ledger-state/src/Cardano/Ledger/State/Orphans.hs

@@ -0,0 +1,83 @@
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE StandaloneDeriving #-}
+
+module Cardano.Ledger.State.Orphans where
+
+import Cardano.Binary
+import Cardano.Crypto.Hash.Class
+import Cardano.Ledger.Alonzo.TxBody
+import Cardano.Ledger.Coin
+import Cardano.Ledger.Mary.Value
+import Cardano.Ledger.SafeHash
+import Cardano.Ledger.Shelley.CompactAddr
+import Cardano.Ledger.Shelley.LedgerState
+import Cardano.Ledger.State.UTxO
+import Cardano.Ledger.TxIn
+import Data.ByteString.Short
+import qualified Data.Text as T
+import Database.Persist
+import Database.Persist.Sqlite
+
+instance PersistField ShortByteString where
+  toPersistValue = PersistByteString . fromShort
+  fromPersistValue (PersistByteString bs) = Right $ toShort bs
+  fromPersistValue _ = Left "Unexpected type"
+
+instance PersistFieldSql ShortByteString where
+  sqlType _ = SqlBlob
+
+instance PersistField (TxId C) where
+  toPersistValue = PersistByteString . hashToBytes . extractHash . _unTxId
+  fromPersistValue (PersistByteString bs) =
+    case hashFromBytes bs of
+      Nothing -> Left "Invalid number of bytes for the hash"
+      Just h -> Right $ TxId $ unsafeMakeSafeHash h
+  fromPersistValue _ = Left "Unexpected type"
+
+instance PersistFieldSql (TxId C) where
+  sqlType _ = SqlBlob
+
+instance PersistField Coin where
+  toPersistValue = PersistInt64 . fromIntegral . unCoin
+  fromPersistValue (PersistInt64 i64) = Right $ Coin $ fromIntegral i64
+  fromPersistValue _ = Left "Unexpected type"
+
+instance PersistFieldSql Coin where
+  sqlType _ = SqlInt64
+
+instance PersistField (TxOut CurrentEra) where
+  toPersistValue = PersistByteString . serialize'
+  fromPersistValue = decodePersistValue
+
+instance PersistFieldSql (TxOut CurrentEra) where
+  sqlType _ = SqlBlob
+
+instance PersistField (PPUPState CurrentEra) where
+  toPersistValue = PersistByteString . serialize'
+  fromPersistValue = decodePersistValue
+
+instance PersistFieldSql (PPUPState CurrentEra) where
+  sqlType _ = SqlBlob
+
+instance PersistField (PState C) where
+  toPersistValue = PersistByteString . serialize'
+  fromPersistValue = decodePersistValue
+
+instance PersistFieldSql (PState C) where
+  sqlType _ = SqlBlob
+
+instance PersistField (DState C) where
+  toPersistValue = PersistByteString . serialize'
+  fromPersistValue = decodePersistValue
+
+instance PersistFieldSql (DState C) where
+  sqlType _ = SqlBlob
+
+decodePersistValue :: FromCBOR b => PersistValue -> Either T.Text b
+decodePersistValue (PersistByteString bs) =
+  case decodeFull' bs of
+    Left err -> Left $ "Could not decode: " <> T.pack (show err)
+    Right v -> Right v
+decodePersistValue _ = Left "Unexpected type"

libs/ledger-state/src/Cardano/Ledger/State/Query.hs

@@ -0,0 +1,87 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE TypeFamilies #-}
+
+module Cardano.Ledger.State.Query where
+
+import qualified Cardano.Ledger.Alonzo.TxBody as Alonzo
+import qualified Cardano.Ledger.Shelley.LedgerState as Shelley
+import qualified Cardano.Ledger.Shelley.UTxO as Shelley
+import Cardano.Ledger.State.Schema
+import Cardano.Ledger.State.UTxO
+import qualified Cardano.Ledger.TxIn as TxIn
+import Conduit
+import Control.Monad.IO.Class
+import Control.Monad.Trans.Reader
+import qualified Data.Map.Strict as Map
+import qualified Data.IntMap.Strict as IntMap
+import Database.Persist.Sqlite
+import Control.Monad.Logger (NoLoggingT(..))
+import Data.Text as T
+
+insertUTxOState ::
+  MonadIO m =>
+  Shelley.UTxOState CurrentEra ->
+  ReaderT SqlBackend m (Key UtxoState)
+insertUTxOState Shelley.UTxOState {..} = do
+  insert $
+    UtxoState
+      { utxoStateDeposited = _deposited,
+        utxoStateFees = _fees,
+        utxoStatePpups = _ppups
+      }
+
+insertUTxO ::
+  MonadIO m =>
+  Shelley.UTxO CurrentEra ->
+  Key UtxoState ->
+  ReaderT SqlBackend m ()
+insertUTxO utxo stateKey = do
+  mapM_ insertTxOut $ Map.toList (Shelley.unUTxO utxo)
+  where
+    insertTxOut (TxIn.TxIn txId txIx, out) = do
+      txKey <-
+        insert $ Tx {txInIx = fromIntegral txIx, txInId = txId, txOut = out}
+      insert_ $ UtxoEntry {utxoEntryTx = txKey, utxoEntryState = stateKey}
+
+insertLedgerState ::
+  MonadIO m => Shelley.LedgerState CurrentEra -> ReaderT SqlBackend m ()
+insertLedgerState Shelley.LedgerState {..} = do
+  stateKey <- insertUTxOState _utxoState
+  insertUTxO (Shelley._utxo _utxoState) stateKey
+  insert_ $
+    LedgerState
+      { ledgerStateUtxo = stateKey,
+        ledgerStateDstate = Shelley._dstate _delegationState,
+        ledgerStatePstate = Shelley._pstate _delegationState
+      }
+
+sourceUTxO ::
+     MonadResource m
+  => ConduitM () (TxIn.TxIn C, Alonzo.TxOut CurrentEra) (ReaderT SqlBackend m) ()
+sourceUTxO =
+  selectSource [] []
+    .| mapC (\(Entity _ Tx {..}) -> (TxIn.TxIn txInId (fromIntegral txInIx), txOut))
+
+
+foldUTxO ::
+     MonadUnliftIO m
+  => (a -> (TxIn.TxIn C, Alonzo.TxOut CurrentEra) -> a) -- ^ Folding function
+  -> a -- ^ Empty acc
+  -> Text -- ^ Path to Sqlite db
+  -> m a
+foldUTxO f m fp = runSqlite fp (runConduit (sourceUTxO .| foldlC f m))
+
+storeLedgerState ::
+     MonadUnliftIO m => Text -> Shelley.LedgerState CurrentEra -> m ()
+storeLedgerState fp ls = runSqlite fp $ do
+  runMigration migrateAll
+  insertLedgerState ls
+
+
+-- runSqlite ::
+--      MonadUnliftIO m
+--   => T.Text
+--   -> ReaderT SqlBackend (NoLoggingT m) a
+--   -> m a
+-- runSqlite dbfile = runNoLoggingT . withSqliteConn dbfile . runReaderT

libs/ledger-state/src/Cardano/Ledger/State/Schema.hs

@@ -1 +1,44 @@
+{-# LANGUAGE DataKinds #-}
+{-# LANGUAGE DerivingStrategies #-}
+{-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE GADTs #-}
+{-# LANGUAGE GeneralizedNewtypeDeriving #-}
+{-# LANGUAGE MultiParamTypeClasses #-}
+{-# LANGUAGE QuasiQuotes #-}
+{-# LANGUAGE StandaloneDeriving #-}
+{-# LANGUAGE TemplateHaskell #-}
+{-# LANGUAGE TypeFamilies #-}
+{-# LANGUAGE UndecidableInstances #-}
+
 module Cardano.Ledger.State.Schema where
+
+import qualified Cardano.Ledger.Alonzo.TxBody as Alonzo
+import Cardano.Ledger.Coin
+import qualified Cardano.Ledger.Shelley.LedgerState as Shelley
+import Cardano.Ledger.State.Orphans ()
+import Cardano.Ledger.State.UTxO
+import qualified Cardano.Ledger.TxIn as TxIn
+import Data.Word
+import Database.Persist.Sqlite
+import Database.Persist.TH
+
+share
+  [mkPersist sqlSettings, mkMigrate "migrateAll"]
+  [persistLowerCase|
+Tx
+  inIx Word64
+  inId (TxIn.TxId C)
+  out (Alonzo.TxOut CurrentEra)
+  UniqueTx inIx inId
+UtxoEntry
+  tx TxId
+  state UtxoStateId
+UtxoState
+  deposited Coin
+  fees Coin
+  ppups (Shelley.PPUPState CurrentEra)
+LedgerState
+  utxo UtxoStateId
+  dstate (Shelley.DState C)
+  pstate (Shelley.PState C)
+|]