Merge #3073

3073: Refactor and simplify `SeqState n k` r=HeinrichApfelmus a=HeinrichApfelmus

### Issue number

ADP-1308

### Overview

Previous work in epic ADP-1043 introduced delta encodings, DBVars, and an embedding of the wallet state and its delta encodings into a database table. It's time to integrate these tools with the wallet code. To facilitate code review, the integration proceeds in a sequence of refactorings that do not change functionality and pass all unit tests.

In this step, we continue refactoring the address discovery state. Here, we refactor and simplify the `SeqState n k` type to use the new abstract data type `Pool addr ix`, which aids with BIP-44 style address discovery.

### Details

* The testing module `PoolSpec` now also provides a shrinker `shrinkPool` for use in the old testing module `Cardano.Wallet.Primitive.AddressDiscovery.SequentialSpec`.
* The property tests pertaining to the address discovery aspects of `SeqState` are superseded by the more general unit tests in `Cardano.Wallet.Address.PoolSpec`.
 
### Comments

* Merge PR ##3068 before this one, because this pull request is based on the branch of the former.


Co-authored-by: Heinrich Apfelmus <heinrich.apfelmus@iohk.io>

lib/core/bench/db-bench.hs

@@ -87,7 +87,6 @@ import Cardano.Wallet.Primitive.AddressDerivation
     , Passphrase (..)
     , PaymentAddress (..)
     , PersistPrivateKey
-    , Role (..)
     , WalletKey (..)
     )
 import Cardano.Wallet.Primitive.AddressDerivation.Byron
@@ -97,14 +96,12 @@ import Cardano.Wallet.Primitive.AddressDerivation.Shelley
 import Cardano.Wallet.Primitive.AddressDiscovery.Random
     ( DerivationPath, RndState (..), mkRndState )
 import Cardano.Wallet.Primitive.AddressDiscovery.Sequential
-    ( AddressPool
-    , DerivationPrefix (..)
-    , ParentContext (..)
+    ( DerivationPrefix (..)
+    , SeqAddressPool (..)
     , SeqState (..)
     , coinTypeAda
     , defaultAddressPoolGap
-    , emptyPendingIxs
-    , mkAddressPool
+    , mkSeqStateFromAccountXPub
     , mkSeqStateFromRootXPrv
     , purposeCIP1852
     )
@@ -179,6 +176,8 @@ import Data.Functor
     ( ($>) )
 import Data.Functor.Identity
     ( Identity (..) )
+import Data.List
+    ( foldl' )
 import Data.Map.Strict
     ( Map )
 import Data.Maybe
@@ -216,6 +215,7 @@ import UnliftIO.Temporary
 
 import qualified Cardano.BM.Configuration.Model as CM
 import qualified Cardano.BM.Data.BackendKind as CM
+import qualified Cardano.Wallet.Address.Pool as AddressPool
 import qualified Cardano.Wallet.Primitive.AddressDerivation.Byron as Byron
 import qualified Cardano.Wallet.Primitive.Types.TokenBundle as TokenBundle
 import qualified Cardano.Wallet.Primitive.Types.TokenMap as TokenMap
@@ -359,31 +359,26 @@ bgroupWriteSeqState db = bgroup "SeqState"
             pure cps
         cps :: [WalletBench]
         cps =
-            [ snd $ initWallet (withMovingSlot i block0) $
-                SeqState
-                    (mkIntPool a i)
-                    (mkExtPool a i)
-                    emptyPendingIxs
-                    rewardAccount
-                    defaultPrefix
+            [ snd $ initWallet (withMovingSlot i block0) $ mkSeqState a i
             | i <- [1..n]
             ]
 
 benchPutSeqState :: DBLayerBench -> [WalletBench] -> IO ()
 benchPutSeqState DBLayer{..} cps = do
     unsafeRunExceptT $ mapExceptT atomically $ mapM_ (putCheckpoint testWid) cps
 
-mkExtPool :: Int -> Int -> AddressPool  'UtxoExternal ShelleyKey
-mkExtPool numAddrs i =
-    mkAddressPool @'Mainnet (ParentContextUtxo ourAccount) defaultAddressPoolGap addrs
-  where
-    addrs = [ force (mkAddress i j, Unused) | j <- [1..numAddrs] ]
-
-mkIntPool :: Int -> Int -> AddressPool 'UtxoInternal ShelleyKey
-mkIntPool numAddrs i =
-    mkAddressPool @'Mainnet (ParentContextUtxo ourAccount) defaultAddressPoolGap addrs
+mkSeqState :: Int -> Int -> SeqState 'Mainnet ShelleyKey
+mkSeqState numAddrs _ = s 
+    { internalPool = fillPool (internalPool s)
+    , externalPool = fillPool (externalPool s)
+    }
   where
-    addrs = [ force (mkAddress i j, Unused) | j <- [1..numAddrs] ]
+    s = mkSeqStateFromAccountXPub @'Mainnet
+        ourAccount purposeCIP1852 defaultAddressPoolGap
+    fillPool (SeqAddressPool pool0) = SeqAddressPool $
+        foldl' (\p ix -> AddressPool.update (gen ix) p) pool0 [0 .. numAddrs-1]
+      where
+        gen ix = AddressPool.generator pool0 $ toEnum ix
 
 ----------------------------------------------------------------------------
 -- Wallet State (Random Scheme) Benchmarks

lib/core/src/Cardano/Wallet.hs

@@ -737,6 +737,7 @@ createIcarusWallet
         , PaymentAddress n k
         , k ~ IcarusKey
         , s ~ SeqState n k
+        , Typeable n
         )
     => ctx
     -> WalletId

lib/core/src/Cardano/Wallet/Api/Server.hs

@@ -352,6 +352,7 @@ import Cardano.Wallet.Primitive.AddressDiscovery.Sequential
     ( DerivationPrefix (..)
     , SeqState (..)
     , defaultAddressPoolGap
+    , getGap
     , mkSeqStateFromAccountXPub
     , mkSeqStateFromRootXPrv
     , purposeCIP1852
@@ -748,7 +749,6 @@ postWallet
         , MkKeyFingerprint k (Proxy n, k 'AddressK XPub)
         , MkKeyFingerprint k Address
         , WalletKey k
-        , Bounded (Index (AddressIndexDerivationType k) 'AddressK)
         , HasDBFactory s k ctx
         , HasWorkerRegistry s k ctx
         , IsOurs s RewardAccount
@@ -776,7 +776,6 @@ postShelleyWallet
         , MkKeyFingerprint k (Proxy n, k 'AddressK XPub)
         , MkKeyFingerprint k Address
         , WalletKey k
-        , Bounded (Index (AddressIndexDerivationType k) 'AddressK)
         , HasDBFactory s k ctx
         , HasWorkerRegistry s k ctx
         , IsOurs s RewardAccount
@@ -816,6 +815,7 @@ postAccountWallet
         , HasWorkerRegistry s k ctx
         , IsOurs s RewardAccount
         , (k == SharedKey) ~ 'False
+        , Typeable n
         )
     => ctx
     -> MkApiWallet ctx s w
@@ -873,7 +873,7 @@ mkShelleyWallet ctx wid cp meta pending progress = do
     let available = availableBalance pending cp
     let total = totalBalance pending reward cp
     pure ApiWallet
-        { addressPoolGap = ApiT $ getState cp ^. #externalPool . #gap
+        { addressPoolGap = ApiT $ getGap $ getState cp ^. #externalPool
         , balance = ApiWalletBalance
             { available = coinToQuantity (available ^. #coin)
             , total = coinToQuantity (total ^. #coin)
@@ -1269,6 +1269,7 @@ postIcarusWallet
         , k ~ IcarusKey
         , HasWorkerRegistry s k ctx
         , PaymentAddress n IcarusKey
+        , Typeable n
         )
     => ctx
     -> ByronWalletPostData '[12,15,18,21,24]
@@ -1289,6 +1290,7 @@ postTrezorWallet
         , k ~ IcarusKey
         , HasWorkerRegistry s k ctx
         , PaymentAddress n IcarusKey
+        , Typeable n
         )
     => ctx
     -> ByronWalletPostData '[12,15,18,21,24]
@@ -1309,6 +1311,7 @@ postLedgerWallet
         , k ~ IcarusKey
         , HasWorkerRegistry s k ctx
         , PaymentAddress n IcarusKey
+        , Typeable n
         )
     => ctx
     -> ByronWalletPostData '[12,15,18,21,24]

lib/core/src/Cardano/Wallet/DB/Sqlite/AddressBook.hs

@@ -3,7 +3,6 @@
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE RankNTypes #-}
 {-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE TypeApplications #-}
 {-# LANGUAGE TypeOperators #-}
 {-# LANGUAGE UndecidableInstances #-}
 
@@ -18,30 +17,21 @@ module Cardano.Wallet.DB.Sqlite.AddressBook
     ( AddressBookIso (..)
     , Prologue (..)
     , Discoveries (..)
-    , SeqAddressList (..)
+    , SeqAddressMap (..)
     )
   where
 
 import Prelude
 
-import Cardano.Address.Derivation
-    ( XPub )
 import Cardano.Wallet.Primitive.AddressDerivation
     ( Depth (..)
     , DerivationType (..)
     , Index (..)
     , KeyFingerprint (..)
-    , MkKeyFingerprint (..)
-    , MkKeyFingerprint (..)
-    , NetworkDiscriminant (..)
-    , PaymentAddress (..)
     , Role (..)
-    , SoftDerivation (..)
     )
 import Cardano.Wallet.Primitive.AddressDerivation.SharedKey
     ( SharedKey (..) )
-import Cardano.Wallet.Primitive.AddressDiscovery
-    ( GetPurpose )
 import Cardano.Wallet.Primitive.Types.Address
     ( Address (..), AddressState (..) )
 import Data.Generics.Internal.VL
@@ -50,18 +40,13 @@ import Data.Kind
     ( Type )
 import Data.Map.Strict
     ( Map )
-import Data.Proxy
-    ( Proxy (..) )
 import Data.Type.Equality
     ( type (==) )
-import Data.Typeable
-    ( Typeable )
 
 import qualified Cardano.Wallet.Address.Pool as AddressPool
 import qualified Cardano.Wallet.Primitive.AddressDiscovery.Random as Rnd
 import qualified Cardano.Wallet.Primitive.AddressDiscovery.Sequential as Seq
 import qualified Cardano.Wallet.Primitive.AddressDiscovery.Shared as Shared
-import qualified Cardano.Wallet.Primitive.Types.Address as W
 import qualified Data.Map.Strict as Map
 
 {-------------------------------------------------------------------------------
@@ -83,9 +68,6 @@ class AddressBookIso s where
 {-------------------------------------------------------------------------------
     Sequential address book
 -------------------------------------------------------------------------------}
--- | Sequential list of addresses
-newtype SeqAddressList (c :: Role) = SeqAddressList [(W.Address, W.AddressState)]
-
 -- piggy-back on SeqState existing instance, to simulate the same behavior.
 instance AddressBookIso (Seq.SeqState n k)
     => AddressBookIso (Seq.SeqAnyState n k p)
@@ -99,56 +81,47 @@ instance AddressBookIso (Seq.SeqState n k)
         in  iso from2 to2
 
 -- | Isomorphism for sequential address book.
-instance 
-    ( MkKeyFingerprint key (Proxy n, key 'AddressK XPub)
-    , GetPurpose key
-    , SoftDerivation key
-    , PaymentAddress n key
-    , (key == SharedKey) ~ 'False
-    ) => AddressBookIso (Seq.SeqState n key)
+instance ( (key == SharedKey) ~ 'False ) => AddressBookIso (Seq.SeqState n key)
   where
     data Prologue (Seq.SeqState n key)
         = SeqPrologue (Seq.SeqState n key)
         -- Trick: We keep the type, but we empty the discovered addresses
     data Discoveries (Seq.SeqState n key)
         = SeqDiscoveries
-            (SeqAddressList 'UtxoInternal)
-            (SeqAddressList 'UtxoExternal)
+            (SeqAddressMap 'UtxoInternal key)
+            (SeqAddressMap 'UtxoExternal key)
 
     addressIso = iso from to
       where
-        from (Seq.SeqState int ext a b c) =
-            ( SeqPrologue $ Seq.SeqState (emptyPool int) (emptyPool ext) a b c
-            , SeqDiscoveries (toDiscoveries @n int) (toDiscoveries @n ext)
-            )
-        to (SeqPrologue (Seq.SeqState int ext a b c), SeqDiscoveries ints exts)
-          = Seq.SeqState (fromDiscoveries @n int ints) (fromDiscoveries @n ext exts) a b c
-
--- | Extract the discovered addresses from an address pool.
-toDiscoveries
-    :: forall (n :: NetworkDiscriminant) (c :: Role) key.
-    ( GetPurpose key
-    , PaymentAddress n key
-    , Typeable c
-    ) => Seq.AddressPool c key -> SeqAddressList c
-toDiscoveries pool = SeqAddressList
-    [ (a,st) | (a,st,_) <- Seq.addresses (liftPaymentAddress @n) pool ]
-
--- | Fill an empty address pool with addresses.
-fromDiscoveries
-    :: forall (n :: NetworkDiscriminant) (c :: Role) key.
-    ( MkKeyFingerprint key (Proxy n, key 'AddressK XPub)
-    , MkKeyFingerprint key Address
-    , SoftDerivation key
-    , Typeable c
-    ) => Seq.AddressPool c key -> SeqAddressList c -> Seq.AddressPool c key
-fromDiscoveries ctx (SeqAddressList addrs) =
-    Seq.mkAddressPool @n (Seq.context ctx) (Seq.gap ctx) addrs
-
--- | Remove all discovered addresses from an address pool,
--- but keep context.
-emptyPool :: Seq.AddressPool c key -> Seq.AddressPool c key
-emptyPool pool = pool{ Seq.indexedKeys = Map.empty }
+        from (Seq.SeqState int ext a b c d) =
+            let int0 = clear int
+                ext0 = clear ext
+            in  ( SeqPrologue $ Seq.SeqState int0 ext0 a b c d
+                , SeqDiscoveries (addresses int) (addresses ext)
+                )
+        to (SeqPrologue (Seq.SeqState int0 ext0 a b c d), SeqDiscoveries ints exts)
+          = Seq.SeqState (loadUnsafe int0 ints) (loadUnsafe ext0 exts) a b c d
+
+-- | Address data from sequential address pool.
+-- The phantom type parameter @c@ prevents mixing up
+-- the internal with the external pool.
+newtype SeqAddressMap (c :: Role) (key :: Depth -> Type -> Type) = SeqAddressMap
+        ( Map
+            (KeyFingerprint "payment" key)
+            (Index 'Soft 'AddressK, AddressState)
+        )
+
+clear :: Seq.SeqAddressPool c k -> Seq.SeqAddressPool c k
+clear = Seq.SeqAddressPool . AddressPool.clear . Seq.getPool
+
+addresses :: Seq.SeqAddressPool c k -> SeqAddressMap c k
+addresses = SeqAddressMap . AddressPool.addresses . Seq.getPool
+
+loadUnsafe
+    :: Seq.SeqAddressPool c k
+    -> SeqAddressMap c k -> Seq.SeqAddressPool c k
+loadUnsafe (Seq.SeqAddressPool pool0) (SeqAddressMap addrs) =
+    Seq.SeqAddressPool $ AddressPool.loadUnsafe pool0 addrs
 
 {-------------------------------------------------------------------------------
     Shared key address book