Refactoring to ensure --gen-stake-delegs uses a minimal amount of memory and generates fewer files.

newhoggy · newhoggy · commit 2f68910fb930 · 2022-06-13T21:38:11.000+10:00
diff --git a/cardano-cli/src/Cardano/CLI/Shelley/Run/Address.hs b/cardano-cli/src/Cardano/CLI/Shelley/Run/Address.hs
@@ -68,27 +68,44 @@ runAddressKeyGenToFile :: AddressKeyType
                  -> VerificationKeyFile
                  -> SigningKeyFile
                  -> ExceptT ShelleyAddressCmdError IO ()
-runAddressKeyGenToFile kt (VerificationKeyFile vkeyPath) (SigningKeyFile skeyPath) =
-    case kt of
-      AddressKeyShelley         -> generateAndWriteKeyFiles AsPaymentKey
-      AddressKeyShelleyExtended -> generateAndWriteKeyFiles AsPaymentExtendedKey
-      AddressKeyByron           -> generateAndWriteKeyFiles AsByronKey
+runAddressKeyGenToFile kt vkf skf = case kt of
+  AddressKeyShelley         -> generateAndWriteKeyFiles AsPaymentKey          vkf skf
+  AddressKeyShelleyExtended -> generateAndWriteKeyFiles AsPaymentExtendedKey  vkf skf
+  AddressKeyByron           -> generateAndWriteKeyFiles AsByronKey            vkf skf
+
+generateAndWriteKeyFiles :: ()
+  => Key keyrole
+  => AsType keyrole
+  -> VerificationKeyFile
+  -> SigningKeyFile
+  -> ExceptT ShelleyAddressCmdError IO ()
+generateAndWriteKeyFiles asType vkf skf = do
+  uncurry (writePaymentKeyFiles vkf skf) =<< generatePaymentKeys asType
+
+generatePaymentKeys :: ()
+  => Key keyrole
+  => AsType keyrole
+  -> ExceptT ShelleyAddressCmdError IO (VerificationKey keyrole, SigningKey keyrole)
+generatePaymentKeys asType = do
+  skey <- liftIO $ generateSigningKey asType
+  return (getVerificationKey skey, skey)
+
+writePaymentKeyFiles :: ()
+  => Key keyrole
+  => VerificationKeyFile
+  -> SigningKeyFile
+  -> VerificationKey keyrole
+  -> SigningKey keyrole
+  -> ExceptT ShelleyAddressCmdError IO ()
+writePaymentKeyFiles (VerificationKeyFile vkeyPath) (SigningKeyFile skeyPath) vkey skey = do
+  firstExceptT ShelleyAddressCmdWriteFileError $ do
+    newExceptT $ writeFileTextEnvelope skeyPath (Just skeyDesc) skey
+    newExceptT $ writeFileTextEnvelope vkeyPath (Just vkeyDesc) vkey
   where
-    generateAndWriteKeyFiles asType = do
-      skey <- liftIO $ generateSigningKey asType
-      let vkey = getVerificationKey skey
-      firstExceptT ShelleyAddressCmdWriteFileError
-        . newExceptT
-        $ writeFileTextEnvelope skeyPath (Just skeyDesc) skey
-      firstExceptT ShelleyAddressCmdWriteFileError
-        . newExceptT
-        $ writeFileTextEnvelope vkeyPath (Just vkeyDesc) vkey
-
     skeyDesc, vkeyDesc :: TextEnvelopeDescr
     skeyDesc = "Payment Signing Key"
     vkeyDesc = "Payment Verification Key"
 
-
 runAddressKeyHash :: VerificationKeyTextOrFile
                   -> Maybe OutputFile
                   -> ExceptT ShelleyAddressCmdError IO ()
diff --git a/cardano-cli/src/Cardano/CLI/Shelley/Run/Genesis.hs b/cardano-cli/src/Cardano/CLI/Shelley/Run/Genesis.hs
@@ -80,7 +80,7 @@ import           Cardano.CLI.Shelley.Run.Node (ShelleyNodeCmdError (..), renderS
                    runNodeIssueOpCert, runNodeKeyGenCold, runNodeKeyGenKES, runNodeKeyGenVRF)
 import           Cardano.CLI.Shelley.Run.Pool (ShelleyPoolCmdError (..), renderShelleyPoolCmdError)
 import           Cardano.CLI.Shelley.Run.StakeAddress (ShelleyStakeAddressCmdError (..),
-                   renderShelleyStakeAddressCmdError, runStakeAddressKeyGenToFile, keyGenStakeAddress)
+                   renderShelleyStakeAddressCmdError, runStakeAddressKeyGenToFile)
 import           Cardano.CLI.Types
 
 import           Cardano.CLI.Byron.Delegation
@@ -703,23 +703,46 @@ runGenesisCreateStaked (GenesisDir rootdir)
         ]
 
   -- Distribute M delegates across N pools:
-  delegations :: [Delegation] <- forM distribution $ \(poolParams, index) -> do
+  delegations <- liftIO $ Lazy.forM distribution $ \(poolParams, index) -> do
     computeDelegation network stdeldir poolParams index
 
+  liftIO $ LBS.writeFile (stdeldir </> "delegations.jsonl") $ B.toLazyByteString $
+    mconcat (List.intersperse "\n" (B.lazyByteString . Aeson.encode <$> delegations))
+
+  -- NOTE The following code which reads from the same file 'delegations.jsonl' multiple times
+  -- looks like duplication, but it is not.  The file is read lazily, and it is important that
+  -- they be read multiple times because the code is streaming and reading the file multiple
+  -- times ensures that any data structures that are created as a result of the read is not
+  -- retained in memory.
+
+  !numDelegations <- fmap length $ liftIO $ LBS.lines <$> LBS.readFile (stdeldir </> "delegations.jsonl")
+
+  delegations2 <- do
+    delegationLines <- liftIO $ LBS.lines <$> LBS.readFile (stdeldir </> "delegations.jsonl")
+    return $ catMaybes $ Aeson.decode @Delegation <$> delegationLines
+
+  delegations3 <- do
+    delegationLines <- liftIO $ LBS.lines <$> LBS.readFile (stdeldir </> "delegations.jsonl")
+    return $ catMaybes $ Aeson.decode @Delegation <$> delegationLines
+
+  delegations4 <- do
+    delegationLines <- liftIO $ LBS.lines <$> LBS.readFile (stdeldir </> "delegations.jsonl")
+    return $ catMaybes $ Aeson.decode @Delegation <$> delegationLines
+
   genDlgs <- readGenDelegsMap gendir deldir
   nonDelegAddrs <- readInitialFundAddresses utxodir network
   start <- maybe (SystemStart <$> getCurrentTimePlus30) pure mStart
 
   stuffedUtxoAddrs <- liftIO $ Lazy.replicateM (fromIntegral numStuffedUtxo) genStuffedAddress
 
-  let poolMap :: Map (Ledger.KeyHash Ledger.Staking StandardCrypto) (Ledger.PoolParams StandardCrypto)
-      poolMap = Map.fromList $ mkDelegationMapEntry <$> delegations
-      delegAddrs = dInitialUtxoAddr <$> delegations
+  let stake = second Ledger._poolId . mkDelegationMapEntry <$> delegations2
+      stakePools = [ (Ledger._poolId poolParams, poolParams) | poolParams <- snd . mkDelegationMapEntry <$> delegations3 ]
+      delegAddrs = dInitialUtxoAddr <$> delegations4
       !shelleyGenesis =
         updateOutputTemplate
           -- Shelley genesis parameters
-          start genDlgs mNonDlgAmount nonDelegAddrs poolMap
-          stDlgAmount delegAddrs stuffedUtxoAddrs (toOutputTemplate template)
+          start genDlgs mNonDlgAmount (length nonDelegAddrs) nonDelegAddrs stakePools stake
+          stDlgAmount numDelegations delegAddrs stuffedUtxoAddrs (toOutputTemplate template)
 
   -- shelleyGenesis contains lazy loaded data, so using lazyToJson to serialise to avoid
   -- retaining large datastructures in memory.
@@ -735,9 +758,7 @@ runGenesisCreateStaked (GenesisDir rootdir)
     , textShow genNumUTxOKeys, " non-delegating UTxO keys, "
     , textShow genNumPools, " stake pools, "
     , textShow genNumStDelegs, " delegating UTxO keys, "
-    , textShow (length delegations), " delegation relationships, "
-    , textShow (Map.size poolMap), " delegation map entries, "
-    , textShow (length delegAddrs), " delegating addresses"
+    , textShow numDelegations, " delegation map entries, "
     ] ++
     [ mconcat
       [ ", "
@@ -855,12 +876,24 @@ createPoolCredentials dir index = do
    coldSK = SigningKeyFile $ dir </> "cold" ++ strIndex ++ ".skey"
    opCertCtr = OpCertCounterFile $ dir </> "opcert" ++ strIndex ++ ".counter"
 
-data Delegation
-  = Delegation
-    { dInitialUtxoAddr  :: AddressInEra ShelleyEra
-    , dDelegStaking     :: Ledger.KeyHash Ledger.Staking StandardCrypto
-    , dPoolParams       :: Ledger.PoolParams StandardCrypto
-    }
+data Delegation = Delegation
+  { dInitialUtxoAddr  :: AddressInEra ShelleyEra
+  , dDelegStaking     :: Ledger.KeyHash Ledger.Staking StandardCrypto
+  , dPoolParams       :: Ledger.PoolParams StandardCrypto
+  }
+
+instance ToJSON Delegation where
+  toJSON delegation = Aeson.object
+    [ "initialUtxoAddr" .= dInitialUtxoAddr delegation
+    , "delegStaking"    .= dDelegStaking delegation
+    , "poolParams"      .= dPoolParams delegation
+    ]
+
+instance FromJSON Delegation where
+  parseJSON = Aeson.withObject "Delegation" $ \v -> Delegation
+    <$> (v .: "initialUtxoAddr")
+    <*> (v .: "delegStaking")
+    <*> (v .: "poolParams")
 
 buildPool :: NetworkId -> FilePath -> Word -> ExceptT ShelleyGenesisCmdError IO (Ledger.PoolParams StandardCrypto)
 buildPool nw dir index = do
@@ -924,34 +957,13 @@ computeDelegation :: ()
   -> FilePath
   -> Ledger.PoolParams StandardCrypto
   -> Word
-  -> ExceptT ShelleyGenesisCmdError IO Delegation
-computeDelegation nw delegDir pool delegIx = do
-    let strIndex = show delegIx
-
-    let paymentVKF = VerificationKeyFile $ delegDir </> "payment" ++ strIndex ++ ".vkey"
-
-    firstExceptT ShelleyGenesisCmdAddressCmdError $ do
-      let paymentSKF = SigningKeyFile $ delegDir </> "payment" ++ strIndex ++ ".skey"
-      runAddressKeyGenToFile AddressKeyShelley paymentVKF paymentSKF
-
-    let stakingVKF = VerificationKeyFile $ delegDir </> "staking" ++ strIndex ++ ".vkey"
-
-    (_, stakeVK) <- firstExceptT ShelleyGenesisCmdStakeAddressCmdError $ do
-      -- let stakingSK = SigningKeyFile $ delegDir </> "staking" ++ strIndex ++ ".skey"
-      -- runStakeAddressKeyGenToFile stakingVKF stakingSK
-      keyGenStakeAddress
+  -> IO Delegation
+computeDelegation nw _delegDir pool _delegIx = do
+    paymentVK <- fmap getVerificationKey $ generateSigningKey AsPaymentKey
+    stakeVK   <- fmap getVerificationKey $ generateSigningKey AsStakeKey
 
-    paySVK <- firstExceptT (ShelleyGenesisCmdAddressCmdError
-                           . ShelleyAddressCmdVerificationKeyTextOrFileError) $
-                 readAddressVerificationKeyTextOrFile
-                   (VktofVerificationKeyFile paymentVKF)
-
-    initialUtxoAddr <- case paySVK of
-      APaymentVerificationKey payVK -> do
-        firstExceptT ShelleyGenesisCmdAddressCmdError $ do
-          let stakeVerifier = StakeVerifierKey . VerificationKeyFilePath $ stakingVKF
-          makeShelleyAddress nw (PaymentCredentialByKey (verificationKeyHash payVK)) <$> makeStakeAddressRef stakeVerifier
-      _ -> left $ ShelleyGenesisCmdUnexpectedAddressVerificationKey paymentVKF "APaymentVerificationKey" paySVK
+    let stakeAddressReference = StakeAddressByValue . StakeCredentialByKey . verificationKeyHash $ stakeVK
+    let initialUtxoAddr = makeShelleyAddress nw (PaymentCredentialByKey (verificationKeyHash paymentVK)) stakeAddressReference
 
     pure Delegation
       { dInitialUtxoAddr = shelleyAddressInEra initialUtxoAddr
@@ -1084,17 +1096,20 @@ updateOutputTemplate
     -> Map (Hash GenesisKey) (Hash GenesisDelegateKey, Hash VrfKey)
     -- Non-delegated initial UTxO spec:
     -> Maybe Lovelace
+    -> Int
     -> [AddressInEra ShelleyEra]
     -- Genesis staking: pools/delegation map & delegated initial UTxO spec:
-    -> Map (Ledger.KeyHash 'Ledger.Staking StandardCrypto) (Ledger.PoolParams StandardCrypto)
+    -> [(Ledger.KeyHash 'Ledger.StakePool StandardCrypto, Ledger.PoolParams StandardCrypto)]
+    -> [(Ledger.KeyHash 'Ledger.Staking StandardCrypto, Ledger.KeyHash 'Ledger.StakePool StandardCrypto)]
     -> Lovelace
-    -> [AddressInEra ShelleyEra]
+    -> Int
+    -> [AddressInEra ShelleyEra] --
     -> [AddressInEra ShelleyEra]
     -> OT.OutputShelleyGenesis StandardShelley
     -> OT.OutputShelleyGenesis StandardShelley
 updateOutputTemplate (SystemStart start)
-               genDelegMap mAmountNonDeleg utxoAddrsNonDeleg
-               poolSpecs (Lovelace amountDeleg) utxoAddrsDeleg stuffedUtxoAddrs
+               genDelegMap mAmountNonDeleg nUtxoAddrsNonDeleg utxoAddrsNonDeleg
+               pools stake (Lovelace amountDeleg) nUtxoAddrsDeleg utxoAddrsDeleg stuffedUtxoAddrs
                template = do
 
     let pparamsFromTemplate = OT.sgProtocolParams template
@@ -1105,15 +1120,16 @@ updateOutputTemplate (SystemStart start)
           , OT.sgInitialFunds = ListMap
                               [ (toShelleyAddr addr, toShelleyLovelace v)
                               | (addr, v) <-
-                                distribute (nonDelegCoin - subtractForTreasury) utxoAddrsNonDeleg ++
-                                distribute (delegCoin - subtractForTreasury)    utxoAddrsDeleg ++
-                                mkStuffedUtxo stuffedUtxoAddrs ]
+                                distribute (nonDelegCoin - subtractForTreasury) nUtxoAddrsNonDeleg  utxoAddrsNonDeleg
+                                ++
+                                distribute (delegCoin - subtractForTreasury)    nUtxoAddrsDeleg     utxoAddrsDeleg
+                                ++
+                                mkStuffedUtxo stuffedUtxoAddrs
+                                ]
           , OT.sgStaking =
-            ShelleyGenesisStaking
-              { sgsPools = Map.fromList
-                            [ (Ledger._poolId poolParams, poolParams)
-                            | poolParams <- Map.elems poolSpecs ]
-              , sgsStake = Ledger._poolId <$> poolSpecs
+            OT.OutputShelleyGenesisStaking
+              { OT.osgsPools = ListMap pools
+              , OT.osgsStake = ListMap stake
               }
           , OT.sgProtocolParams = pparamsFromTemplate
           }
@@ -1128,22 +1144,11 @@ updateOutputTemplate (SystemStart start)
     nonDelegCoin = fromIntegral (fromMaybe maximumLovelaceSupply (unLovelace <$> mAmountNonDeleg))
     delegCoin = fromIntegral amountDeleg
 
-    distribute :: Integer -> [AddressInEra ShelleyEra] -> [(AddressInEra ShelleyEra, Lovelace)]
-    distribute funds addrs =
-      fst $ List.foldl' folder ([], fromIntegral funds) addrs
-     where
-       nAddrs, coinPerAddr, splitThreshold :: Integer
-       nAddrs = fromIntegral $ length addrs
-       coinPerAddr = funds `div` nAddrs
-       splitThreshold = coinPerAddr + nAddrs
-
-       folder :: ([(AddressInEra ShelleyEra, Lovelace)], Integer)
-              -> AddressInEra ShelleyEra
-              -> ([(AddressInEra ShelleyEra, Lovelace)], Integer)
-       folder (acc, rest) addr
-         | rest > splitThreshold =
-             ((addr, Lovelace coinPerAddr) : acc, rest - coinPerAddr)
-         | otherwise = ((addr, Lovelace rest) : acc, 0)
+    distribute :: Integer -> Int -> [AddressInEra ShelleyEra] -> [(AddressInEra ShelleyEra, Lovelace)]
+    distribute funds nAddrs addrs = zipWith (,) addrs (fmap Lovelace (coinPerAddr + rest:repeat coinPerAddr))
+      where coinPerAddr :: Integer
+            coinPerAddr = funds `div` fromIntegral nAddrs
+            rest = coinPerAddr * fromIntegral nAddrs
 
     mkStuffedUtxo :: [AddressInEra ShelleyEra] -> [(AddressInEra ShelleyEra, Lovelace)]
     mkStuffedUtxo xs = (, Lovelace minUtxoVal) <$> xs
diff --git a/cardano-cli/src/Cardano/CLI/Shelley/Run/Genesis/LazyToJson.hs b/cardano-cli/src/Cardano/CLI/Shelley/Run/Genesis/LazyToJson.hs
@@ -1,13 +1,14 @@
-{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE LambdaCase #-}
 
 module Cardano.CLI.Shelley.Run.Genesis.LazyToJson
-  ( LazyToJson(..)
+  ( Aeson(..)
+  , LazyToJson(..)
   ) where
 
 import Cardano.Ledger.Address (Addr)
 import Cardano.Ledger.Crypto (StandardCrypto)
-import Data.Aeson (Value)
+import Data.Aeson (Value, ToJSON)
 import Data.Functor ((<$>))
 import Data.Monoid (Monoid(..))
 import Data.Semigroup (Semigroup(..))
@@ -38,3 +39,8 @@ instance LazyToJson a => LazyToJson [a] where
 
 instance LazyToJson (Addr StandardCrypto) where
   lazyToJson = B.lazyByteString . J.encode
+
+newtype Aeson a = Aeson a
+
+instance ToJSON a => LazyToJson (Aeson a) where
+  lazyToJson (Aeson a) = B.lazyByteString (J.encode a)
diff --git a/cardano-cli/src/Cardano/CLI/Shelley/Run/Genesis/ListMap.hs b/cardano-cli/src/Cardano/CLI/Shelley/Run/Genesis/ListMap.hs
@@ -29,15 +29,18 @@ newtype ListMap k v = ListMap
   } deriving (Eq, Show)
 
 instance ToJSONKey k => ToJSON1 (ListMap k) where
-    liftToJSON g _ = case toJSONKey of
-        ToJSONKeyText f _ -> Object . KM.fromList . unListMap . mapKeyValO f g
-        ToJSONKeyValue  f _ -> Array . V.fromList . L.map (toJSONPair f g) . unListMap
-
-    liftToEncoding g _ = case toJSONKey of
-        ToJSONKeyText _ f -> dict f g (foldrWithKey . uncurry)
-        ToJSONKeyValue _ f -> E.list (pairEncoding f) . unListMap
-      where
-        pairEncoding f (a, b) = E.list id [f a, g b]
+  liftToJSON g _ = case toJSONKey of
+    ToJSONKeyText  f _ -> Object . KM.fromList . unListMap . mapKeyValO f g
+    ToJSONKeyValue f _ -> Array . V.fromList . L.map (toJSONPair f g) . unListMap
+    where mapKeyValO :: (k1 -> k2) -> (v1 -> v2) -> ListMap k1 v1 -> ListMap k2 v2
+          mapKeyValO fk kv = ListMap . foldrWithKey (\(k, v) -> ((fk k, kv v):)) []
+          toJSONPair :: (a -> Value) -> (b -> Value) -> (a, b) -> Value
+          toJSONPair a b = liftToJSON2 a (listValue a) b (listValue b)
+
+  liftToEncoding g _ = case toJSONKey of
+    ToJSONKeyText  _ f -> dict f g (foldrWithKey . uncurry)
+    ToJSONKeyValue _ f -> E.list (pairEncoding f) . unListMap
+    where pairEncoding f (a, b) = E.list id [f a, g b]
 
 instance (ToJSON v, ToJSONKey k) => ToJSON (ListMap k v) where
     toJSON = J.toJSON1
@@ -47,14 +50,6 @@ instance (ToJSON v, ToJSONKey k) => ToJSON (ListMap k v) where
 foldrWithKey :: ((k, a) -> b -> b) -> b -> ListMap k a -> b
 foldrWithKey f z = L.foldr f z . unListMap
 
--- | Transform the keys and values of a 'M.Map'.
-mapKeyValO :: (k1 -> k2) -> (v1 -> v2) -> ListMap k1 v1 -> ListMap k2 v2
-mapKeyValO fk kv = ListMap . foldrWithKey (\(k, v) -> ((fk k, kv v):)) []
-{-# INLINE mapKeyValO #-}
-
-toJSONPair :: (a -> Value) -> (b -> Value) -> (a, b) -> Value
-toJSONPair a b = liftToJSON2 a (listValue a) b (listValue b)
-
 instance forall k v. (ToJSON k, ToJSON v) => LazyToJson (ListMap k v) where
   lazyToJson (ListMap kvs) = "{" <> mconcat (L.intersperse "," (elementLazyToJson <$> kvs)) <> "}"
     where elementLazyToJson :: (k, v) -> B.Builder
diff --git a/cardano-cli/src/Cardano/CLI/Shelley/Run/Genesis/Types.hs b/cardano-cli/src/Cardano/CLI/Shelley/Run/Genesis/Types.hs
diff --git a/cardano-cli/src/Cardano/CLI/Shelley/Run/StakeAddress.hs b/cardano-cli/src/Cardano/CLI/Shelley/Run/StakeAddress.hs
