tx-generator: more declarative interface

bench/tx-generator/src/Cardano/Benchmarking/Compiler.hs

@@ -13,6 +13,8 @@ import           Control.Monad.Trans.RWS.CPS
 import           Data.Dependent.Sum ( (==>) )
 import           Data.DList (DList)
 import qualified Data.DList as DL
+import           Data.Text (Text)
+import qualified Data.Text as Text
 
 import           Cardano.Api
 import           Cardano.Benchmarking.Types
@@ -26,6 +28,9 @@ data CompileError where
   deriving (Show)
 type Compiler a = RWST NixServiceOptions (DList Action) Int (Except CompileError) a
 
+maxOutputsPerTx :: Int
+maxOutputsPerTx = 30
+
 type SrcWallet = WalletName
 type DstWallet = WalletName
 
@@ -46,11 +51,10 @@ compileToScript :: Compiler ()
 compileToScript = do
   initConstants
   emit . StartProtocol =<< askNixOption getNodeConfigFile
-  genesisWallet <- newWallet "genesis_wallet"
-  importGenesisFunds genesisWallet
+  genesisWallet <- importGenesisFunds
   collateralWallet <- addCollaterals genesisWallet
   splitWallet <- splittingPhase genesisWallet
-  benchmarkingPhase splitWallet collateralWallet
+  benchmarkingPhaseNew splitWallet collateralWallet
 
 initConstants :: Compiler ()
 initConstants = do
@@ -65,26 +69,34 @@ initConstants = do
     setN :: Tag v -> (NixServiceOptions -> v) -> Compiler ()
     setN key s = askNixOption s >>= setConst key
 
-importGenesisFunds :: DstWallet -> Compiler ()
-importGenesisFunds wallet = do
+importGenesisFunds :: Compiler WalletName
+importGenesisFunds = do
+  logMsg "Importing Genesis Fund."
+  wallet <- newWallet "genesis_wallet"
   era <- askNixOption _nix_era
+  fee <- askNixOption _nix_tx_fee  
   cmd1 (ReadSigningKey $ KeyName "pass-partout") _nix_sigKey
-  emit $ ImportGenesisFund era wallet LocalSocket (KeyName "pass-partout") (KeyName "pass-partout")
+  emit $ Submit era LocalSocket $ SecureGenesis fee wallet (KeyName "pass-partout") (KeyName "pass-partout")
   delay
+  logMsg "Importing Genesis Fund. Done."
+  return wallet
 
 addCollaterals :: SrcWallet -> Compiler (Maybe WalletName)
 addCollaterals src = do
   era <- askNixOption _nix_era
   isAnyPlutusMode >>= \case
     False -> return Nothing
-    True -> do      
+    True -> do
+      logMsg "Create collaterals."
       safeCollateral <- _safeCollateral <$> evilFeeMagic
       collateralWallet <- newWallet "collateral_wallet"
-      emit $ CreateChange era src LocalSocket
-               (PayToAddr (KeyName "pass-partout") collateralWallet)
-               (PayToAddr (KeyName "pass-partout") src)
-               safeCollateral
-               1
+      fee <- askNixOption _nix_tx_fee
+      let generator = Split fee src
+                        (PayToAddr (KeyName "pass-partout") collateralWallet)
+                        (PayToAddr (KeyName "pass-partout") src)
+                        [ safeCollateral ]
+      emit $ Submit era LocalSocket generator
+      logMsg "Create collaterals. Done."
       return $ Just collateralWallet
 
 splittingPhase :: SrcWallet -> Compiler DstWallet
@@ -94,70 +106,84 @@ splittingPhase srcWallet = do
   tx_fee <- askNixOption _nix_tx_fee
   era <- askNixOption _nix_era  
   minValuePerInput <- _minValuePerInput <$> evilFeeMagic
-  splitSteps <- splitSequenceWalletNames srcWallet srcWallet $ unfoldSplitSequence tx_fee minValuePerInput (tx_count * inputs_per_tx)
-  forM_ (init splitSteps) $ createChange era
-  plutus <- isAnyPlutusMode
-  (if plutus then createChangePlutus era else createChange era) $ last splitSteps
+  finalDest <- newWallet "final_split_wallet"
+  splitSteps <- splitSequenceWalletNames srcWallet finalDest $ unfoldSplitSequence tx_fee minValuePerInput (tx_count * inputs_per_tx)
+  isPlutus <- isAnyPlutusMode
+  forM_ (init splitSteps) $ createChange False era
+  createChange isPlutus era $ last splitSteps
+  return finalDest
  where
-  createChange :: AnyCardanoEra -> SplitStep -> Compiler DstWallet
-  createChange era (src, dst, value, count) = do
-     emit $ CreateChange era src LocalSocket (PayToAddr (KeyName "pass-partout") dst ) (PayToAddr  (KeyName "pass-partout") src) value count
-     delay
-     return dst
-
-  createChangePlutus :: AnyCardanoEra -> SplitStep -> Compiler DstWallet
-  createChangePlutus era (src, dst, value, count) = do
-     autoMode <- isPlutusAutoMode
-     scriptSpec <- if autoMode
-       then ScriptSpec <$> askNixOption _nix_plutusLoopScript <*> pure AutoScript
-       else do
-         executionUnits <- ExecutionUnits <$> askNixOption _nix_executionMemory <*> askNixOption _nix_executionSteps
-         debugMode <- askNixOption _nix_debugMode
-         budget <- (if debugMode then CheckScriptBudget else StaticScriptBudget)
-                     <$> (ScriptDataNumber <$> askNixOption _nix_plutusData)
-                     <*> (ScriptDataNumber <$> askNixOption _nix_plutusRedeemer)
-                     <*> pure executionUnits
-         ScriptSpec <$> askNixOption _nix_plutusScript <*> pure budget
-     emit $ CreateChange era src LocalSocket (PayToScript scriptSpec dst) (PayToScript scriptSpec src) value count
-     delay
-     return dst
+  createChange :: Bool -> AnyCardanoEra -> (SrcWallet, DstWallet, Split) -> Compiler ()
+  createChange isPlutus era (src, dst, split) = do
+    logMsg $ Text.pack $ "Splitting step: " ++ show split
+    tx_fee <- askNixOption _nix_tx_fee
+    payMode <- if isPlutus then plutusPayMode dst else return $ PayToAddr (KeyName "pass-partout") dst
+    let generator = case split of
+          SplitWithChange lovelace count -> Split tx_fee src payMode (PayToAddr  (KeyName "pass-partout") src) $ replicate count lovelace
+          FullSplits txCount -> Take txCount $ Cycle $ SplitN tx_fee src payMode maxOutputsPerTx
+    emit $ Submit era LocalSocket generator
+    delay
+    logMsg "Splitting step: Done"    
+
+  plutusPayMode :: DstWallet -> Compiler PayMode
+  plutusPayMode dst = do
+    autoMode <- isPlutusAutoMode
+    scriptSpec <- if autoMode
+      then ScriptSpec <$> askNixOption _nix_plutusLoopScript <*> pure AutoScript
+      else do
+        executionUnits <- ExecutionUnits <$> askNixOption _nix_executionMemory <*> askNixOption _nix_executionSteps
+        debugMode <- askNixOption _nix_debugMode
+        budget <- (if debugMode then CheckScriptBudget else StaticScriptBudget)
+                    <$> (ScriptDataNumber <$> askNixOption _nix_plutusData)
+                    <*> (ScriptDataNumber <$> askNixOption _nix_plutusRedeemer)
+                    <*> pure executionUnits
+        ScriptSpec <$> askNixOption _nix_plutusScript <*> pure budget
+    return $ PayToScript scriptSpec dst
 
 -- Generate src and dst wallet names for a splitSequence.
 -- testCompiler (error "opts") $ splitSequenceWalletNames (WalletName "w1") (WalletName "w2") (unfoldSplitSequence 1 1000 10000)
-type SplitStep = (SrcWallet, DstWallet, Lovelace, Int)
 
-splitSequenceWalletNames :: SrcWallet -> DstWallet -> [(Lovelace, Int)] -> Compiler [ SplitStep ]
+data Split
+  = SplitWithChange Lovelace Int
+  | FullSplits Int
+  deriving Show
+
+splitSequenceWalletNames :: SrcWallet -> DstWallet -> [Split] -> Compiler [ (SrcWallet, DstWallet, Split) ]
 splitSequenceWalletNames _src _dst [] = return []
-splitSequenceWalletNames src dst [ (val,count) ] = return [( src, dst, val, count)]
-splitSequenceWalletNames src dst ((val, count):rest) = do
-  nextDst <- newWallet "change_wallet"
-  l <- splitSequenceWalletNames dst nextDst rest
-  return $ ( src, dst, val, count) : l
-
--- Return a list of splitSteps.
-unfoldSplitSequence :: Lovelace -> Lovelace -> Int -> [(Lovelace, Int)]
-unfoldSplitSequence fee value count
-  = if count < maxOutputs
-    then [
-           -- Add an extra transaction that just contains the desired output and possible fees.
-           (value * fromIntegral count + fee, 1)
-         , (value, count )
-         ]
-    else unfoldSplitSequence fee (value * fromIntegral maxOutputs + fee) (count `div` maxOutputs + 1) ++ [ (value, count) ]
+splitSequenceWalletNames src dst [ split ] = return [( src, dst, split )]
+splitSequenceWalletNames src dst (split: rest) = do
+  tempWallet <- newWallet "change_wallet"
+  l <- splitSequenceWalletNames tempWallet dst rest
+  return $ ( src, tempWallet, split) : l
+
+unfoldSplitSequence :: Lovelace -> Lovelace -> Int -> [ Split ]
+unfoldSplitSequence fee value outputs
+  = if outputs < maxOutputsPerTx
+      then [ SplitWithChange value outputs ]
+      else
+        let txs = outputs `divCeiling` maxOutputsPerTx
+        in unfoldSplitSequence fee (value * fromIntegral maxOutputsPerTx + fee) txs ++ [ FullSplits txs ]
   where
     -- maximal number of outputs in a TX.
     -- todo: this must be in sync with Scipt/Core.hs
-    maxOutputs = 30
-
-benchmarkingPhase :: WalletName -> Maybe WalletName -> Compiler ()
-benchmarkingPhase wallet collateralWallet = do
+    divCeiling a b = case divMod a b of
+     (x, 0) -> x
+     (x, _rest) -> x+1
+
+benchmarkingPhaseNew :: WalletName -> Maybe WalletName -> Compiler ()
+benchmarkingPhaseNew wallet collateralWallet = do
   debugMode <- askNixOption _nix_debugMode
   targetNodes <- askNixOption _nix_targetNodes
   extraArgs <- evilValueMagic
   tps <- askNixOption _nix_tps
   era <- askNixOption _nix_era
-  let target = if debugMode then LocalSocket else NodeToNode targetNodes
-  emit $ RunBenchmark era wallet target (ThreadName "tx-submit-benchmark") extraArgs collateralWallet tps
+  (NumberOfTxs txCount) <- askNixOption _nix_tx_count
+  let
+    submitMode = if debugMode
+        then LocalSocket
+        else Benchmark targetNodes (ThreadName "tx-submit-benchmark") tps  extraArgs
+    generator = Take txCount $ Cycle $ BechmarkTx wallet extraArgs collateralWallet
+  emit $ Submit era submitMode generator
   unless debugMode $ do
     emit $ WaitBenchmark $ ThreadName "tx-submit-benchmark"
 
@@ -191,6 +217,9 @@ evilFeeMagic = do
 emit :: Action -> Compiler ()
 emit = tell . DL.singleton
 
+logMsg :: Text -> Compiler ()
+logMsg = emit . LogMsg
+
 cmd1 :: (v -> Action) -> (NixServiceOptions -> v) -> Compiler ()
 cmd1 cmd arg = emit . cmd =<< askNixOption arg