Promote 'ChangeChain' from the term-level to type-level on 'AddressPool'

src/Cardano/Wallet.hs

@@ -51,8 +51,7 @@ import Cardano.Wallet.DB
 import Cardano.Wallet.Network
     ( NetworkLayer (..) )
 import Cardano.Wallet.Primitive.AddressDerivation
-    ( ChangeChain (..)
-    , Depth (RootK)
+    ( Depth (RootK)
     , Key
     , Passphrase
     , XPrv
@@ -219,9 +218,9 @@ mkWalletLayer db network = WalletLayer
         let accXPrv =
                 deriveAccountPrivateKey mempty rootXPrv minBound
         let extPool =
-                mkAddressPool (publicKey accXPrv) (gap w) ExternalChain []
+                mkAddressPool (publicKey accXPrv) (gap w) []
         let intPool =
-                mkAddressPool (publicKey accXPrv) minBound InternalChain []
+                mkAddressPool (publicKey accXPrv) minBound []
         let wid =
                 WalletId (digest $ publicKey rootXPrv)
         let checkpoint = initWallet $ SeqState

src/Cardano/Wallet/Primitive/AddressDerivation.hs

@@ -105,6 +105,8 @@ import Data.Text
     ( Text )
 import Data.Text.Class
     ( FromText (..), TextDecodingError (..), ToText (..) )
+import Data.Typeable
+    ( Typeable )
 import Data.Word
     ( Word32 )
 import Fmt
@@ -361,7 +363,7 @@ instance MonadRandom ((->) (Passphrase "salt")) where
 data ChangeChain
     = InternalChain
     | ExternalChain
-    deriving (Generic, Show, Eq)
+    deriving (Generic, Typeable, Show, Eq)
 
 instance NFData ChangeChain
 

src/Cardano/Wallet/Primitive/AddressDiscovery.hs

@@ -1,9 +1,12 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
 {-# LANGUAGE BangPatterns #-}
 {-# LANGUAGE DataKinds #-}
 {-# LANGUAGE DeriveGeneric #-}
 {-# LANGUAGE DerivingStrategies #-}
 {-# LANGUAGE FlexibleInstances #-}
+{-# LANGUAGE KindSignatures #-}
 {-# LANGUAGE LambdaCase #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeApplications #-}
 
 -- |
@@ -74,8 +77,12 @@ import Data.Map.Strict
     ( Map )
 import Data.Maybe
     ( isJust )
+import Data.Proxy
+    ( Proxy (..) )
 import Data.Text.Class
     ( FromText (..), TextDecodingError (..), ToText (..) )
+import Data.Typeable
+    ( Typeable, typeRep )
 import Data.Word
     ( Word8 )
 import GHC.Generics
@@ -163,29 +170,43 @@ defaultAddressPoolGap =
 --
 -- >>> mkAddressPool xpub gap changeChain mempty
 -- AddressPool { }
-data AddressPool = AddressPool
+data AddressPool (chain :: ChangeChain) = AddressPool
     { accountPubKey
         :: !(Key 'AccountK XPub)
         -- ^ Corresponding key for the pool (a pool is tied to only one account)
     , gap
         :: !AddressPoolGap
         -- ^ The actual gap for the pool. This can't change for a given pool.
-    , changeChain
-        :: !ChangeChain
-        -- ^ Whether this pool tracks addrs on the internal or external chain
     , indexedAddresses
         :: !(Map Address (Index 'Soft 'AddressK))
     } deriving (Generic, Show, Eq)
 
-instance NFData AddressPool
+instance NFData (AddressPool chain)
+
+-- | Bring a 'ChangeChain' type back to the term-level. This requires a type
+-- application and either a scoped type variable, or an explicit passing of a
+-- 'ChangeChain'.
+--
+-- >>> changeChain @'ExternalChain
+-- ExternalChain
+--
+-- >>> changeChain @chain
+-- ...
+changeChain :: forall (chain :: ChangeChain). Typeable chain => ChangeChain
+changeChain =
+    case typeRep (Proxy :: Proxy chain) of
+        t | t == typeRep (Proxy :: Proxy 'InternalChain) ->
+            InternalChain
+        _ ->
+            ExternalChain
 
 -- | Get all addresses in the pool, sorted from the first address discovered,
 -- up until the next one.
 --
 -- In practice, we always have:
 --
 -- > mkAddressPool key g cc (addresses pool) == pool
-addresses :: AddressPool -> [Address]
+addresses :: AddressPool chain -> [Address]
 addresses = map fst . L.sortOn snd . Map.toList . indexedAddresses
 
 -- | Create a new Address pool from a list of addresses. Note that, the list is
@@ -194,16 +215,15 @@ addresses = map fst . L.sortOn snd . Map.toList . indexedAddresses
 -- The pool will grow from the start if less than @g :: AddressPoolGap@ are
 -- given, such that, there are always @g@ undiscovered addresses in the pool.
 mkAddressPool
-    :: Key 'AccountK XPub
+    :: forall chain. Typeable chain
+    => Key 'AccountK XPub
     -> AddressPoolGap
-    -> ChangeChain
     -> [Address]
-    -> AddressPool
-mkAddressPool key g cc addrs = AddressPool
+    -> AddressPool chain
+mkAddressPool key g addrs = AddressPool
     { accountPubKey = key
     , gap = g
-    , changeChain = cc
-    , indexedAddresses = nextAddresses key g cc minBound <>
+    , indexedAddresses = nextAddresses key g (changeChain @chain) minBound <>
         Map.fromList (zip addrs [minBound..maxBound])
     }
 
@@ -213,9 +233,10 @@ mkAddressPool key g cc addrs = AddressPool
 -- possible that the pool is not amended at all - this happens in the case that
 -- an address is discovered 'far' from the edge.
 lookupAddress
-    :: Address
-    -> AddressPool
-    -> (Maybe (Index 'Soft 'AddressK), AddressPool)
+    :: Typeable chain
+    => Address
+    -> AddressPool chain
+    -> (Maybe (Index 'Soft 'AddressK), AddressPool chain)
 lookupAddress !target !pool =
     case Map.lookup target (indexedAddresses pool) of
         Just ix ->
@@ -226,9 +247,10 @@ lookupAddress !target !pool =
 -- | If an address is discovered near the edge, we extend the address sequence,
 -- otherwise we return the pool untouched.
 extendAddressPool
-    :: Index 'Soft 'AddressK
-    -> AddressPool
-    -> AddressPool
+    :: forall chain. Typeable chain
+    => Index 'Soft 'AddressK
+    -> AddressPool chain
+    -> AddressPool chain
 extendAddressPool !ix !pool
     | isOnEdge  = pool { indexedAddresses = indexedAddresses pool <> next }
     | otherwise = pool
@@ -238,7 +260,7 @@ extendAddressPool !ix !pool
     next = if ix == maxBound then mempty else nextAddresses
         (accountPubKey pool)
         (gap pool)
-        (changeChain pool)
+        (changeChain @chain)
         (succ ix)
 
 -- | Compute the pool extension from a starting index
@@ -296,7 +318,7 @@ updatePendingIxs ix (PendingIxs ixs) =
 -- exchanges who care less about privacy / not-reusing addresses than
 -- regular users.
 nextChangeIndex
-    :: AddressPool
+    :: AddressPool 'InternalChain
     -> PendingIxs
     -> (Index 'Soft 'AddressK, PendingIxs)
 nextChangeIndex pool (PendingIxs ixs) =
@@ -322,9 +344,9 @@ nextChangeIndex pool (PendingIxs ixs) =
 -------------------------------------------------------------------------------}
 
 data SeqState = SeqState
-    { internalPool :: !AddressPool
+    { internalPool :: !(AddressPool 'InternalChain)
         -- ^ Addresses living on the 'InternalChain'
-    , externalPool :: !AddressPool
+    , externalPool :: !(AddressPool 'ExternalChain)
         -- ^ Addresses living on the 'ExternalChain'
     , pendingChangeIxs :: !PendingIxs
         -- ^ Indexes from the internal pool that have been used in pending
@@ -380,23 +402,24 @@ class AddressScheme s where
 instance AddressScheme SeqState where
     keyFrom addr (rootPrv, pwd) (SeqState !s1 !s2 _) =
         let
-            xPrv1 = lookupAndDeriveXPrv s1 InternalChain
-            xPrv2 = lookupAndDeriveXPrv s2 ExternalChain
+            xPrv1 = lookupAndDeriveXPrv s1
+            xPrv2 = lookupAndDeriveXPrv s2
             xPrv = xPrv1 <|> xPrv2
         in
             xPrv
       where
         lookupAndDeriveXPrv
-            :: AddressPool
-            -> ChangeChain
+            :: forall chain. Typeable chain
+            => AddressPool chain
             -> Maybe (Key 'AddressK XPrv)
-        lookupAndDeriveXPrv pool chain =
+        lookupAndDeriveXPrv pool =
             let
                 -- We are assuming there is only one account
                 accountPrv = deriveAccountPrivateKey pwd rootPrv minBound
                 (addrIx, _) = lookupAddress addr pool
+                cc = changeChain @chain
             in
-                deriveAddressPrivateKey pwd accountPrv chain <$> addrIx
+                deriveAddressPrivateKey pwd accountPrv cc <$> addrIx
 
     -- | We pick indexes in sequence from the first known available index (i.e.
     -- @length addrs - gap@) but we do not generate _new change addresses_. As a

test/unit/Cardano/Wallet/Primitive/AddressDiscoverySpec.hs

@@ -1,4 +1,7 @@
+{-# LANGUAGE AllowAmbiguousTypes #-}
 {-# LANGUAGE DataKinds #-}
+{-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE TypeApplications #-}
 
 {-# OPTIONS_GHC -fno-warn-orphans #-}
@@ -47,6 +50,8 @@ import Data.Proxy
     ( Proxy (..) )
 import Data.Text.Class
     ( TextDecodingError (..), fromText )
+import Data.Typeable
+    ( Typeable )
 import Data.Word
     ( Word8 )
 import Test.Hspec
@@ -90,15 +95,25 @@ spec = do
         it "defaultAddressPoolGap is valid"
             (property prop_defaultValid)
 
-    describe "AddressPool" $ do
+    describe "AddressPool ExternalChain" $ do
         it "'lookupAddressPool' extends the pool by a maximum of 'gap'"
-            (checkCoverage prop_poolGrowWithinGap)
+            (checkCoverage (prop_poolGrowWithinGap @'ExternalChain))
         it "'addresses' preserves the address order"
-            (checkCoverage prop_roundtripMkAddressPool)
+            (checkCoverage (prop_roundtripMkAddressPool @'ExternalChain))
         it "An AddressPool always contains at least 'gap pool' addresses"
-            (property prop_poolAtLeastGapAddresses)
+            (property (prop_poolAtLeastGapAddresses @'ExternalChain))
         it "Our addresses are eventually discovered"
-            (property prop_poolEventuallyDiscoverOurs)
+            (property (prop_poolEventuallyDiscoverOurs @'ExternalChain))
+
+    describe "AddressPool InternalChain" $ do
+        it "'lookupAddressPool' extends the pool by a maximum of 'gap'"
+            (checkCoverage (prop_poolGrowWithinGap @'InternalChain))
+        it "'addresses' preserves the address order"
+            (checkCoverage (prop_roundtripMkAddressPool @'InternalChain))
+        it "An AddressPool always contains at least 'gap pool' addresses"
+            (property (prop_poolAtLeastGapAddresses @'InternalChain))
+        it "Our addresses are eventually discovered"
+            (property (prop_poolEventuallyDiscoverOurs @'InternalChain))
 
     describe "AddressPoolGap - Text Roundtrip" $ do
         textRoundtrip $ Proxy @AddressPoolGap
@@ -117,6 +132,7 @@ spec = do
         where
             err :: String
             err = "An address pool gap must be a natural number between 10 and 100"
+
 {-------------------------------------------------------------------------------
                         Properties for AddressPoolGap
 -------------------------------------------------------------------------------}
@@ -164,7 +180,8 @@ prop_roundtripEnumGap g =
 
 -- | After a lookup, a property should never grow more than its gap value.
 prop_poolGrowWithinGap
-    :: (AddressPool, Address)
+    :: (Typeable chain)
+    => (AddressPool chain, Address)
     -> Property
 prop_poolGrowWithinGap (pool, addr) =
     cover 10 (isJust $ fst res) "pool hit" prop
@@ -181,19 +198,20 @@ prop_poolGrowWithinGap (pool, addr) =
 
 -- | A pool gives back its addresses in correct order and can be reconstructed
 prop_roundtripMkAddressPool
-    :: AddressPool
+    :: (Typeable chain)
+    => AddressPool chain
     -> Property
 prop_roundtripMkAddressPool pool =
     ( mkAddressPool
         (accountPubKey pool)
         (gap pool)
-        (changeChain pool)
         (addresses pool)
     ) === pool
 
 -- | A pool always contains a number of addresses at least equal to its gap
 prop_poolAtLeastGapAddresses
-    :: AddressPool
+    :: (Typeable chain)
+    => AddressPool chain
     -> Property
 prop_poolAtLeastGapAddresses pool =
     property prop
@@ -202,13 +220,14 @@ prop_poolAtLeastGapAddresses pool =
 
 -- | Our addresses are eventually discovered
 prop_poolEventuallyDiscoverOurs
-    :: (AddressPoolGap, ChangeChain, Address)
+    :: forall (chain :: ChangeChain). (Typeable chain)
+    => (AddressPoolGap, Address)
     -> Property
-prop_poolEventuallyDiscoverOurs (g, cc, addr) =
+prop_poolEventuallyDiscoverOurs (g, addr) =
     addr `elem` ours ==> withMaxSuccess 10 $ property prop
   where
-    ours = take 25 (ourAddresses cc)
-    pool = flip execState (mkAddressPool ourAccount g cc mempty) $
+    ours = take 25 (ourAddresses (changeChain @chain))
+    pool = flip execState (mkAddressPool @chain ourAccount g mempty) $
         forM ours (state . lookupAddress)
     prop = (fromEnum <$> fst (lookupAddress addr pool)) === elemIndex addr ours
 
@@ -254,26 +273,24 @@ instance Arbitrary Address where
             let xprv = unsafeGenerateKeyFromSeed (bytes, mempty) mempty
             return $ keyToAddress $ publicKey xprv
 
-instance Arbitrary AddressPool where
+instance Typeable chain => Arbitrary (AddressPool chain) where
     shrink pool =
         let
             key = accountPubKey pool
             g = gap pool
-            cc = changeChain pool
             addrs = addresses pool
         in case length addrs of
             k | k == fromEnum g && g == minBound ->
                 []
             k | k == fromEnum g && g > minBound ->
-                [ mkAddressPool key minBound cc [] ]
+                [ mkAddressPool key minBound [] ]
             k ->
-                [ mkAddressPool key minBound cc []
-                , mkAddressPool key g cc []
-                , mkAddressPool key g cc (take (k - (fromEnum g `div` 5)) addrs)
+                [ mkAddressPool key minBound []
+                , mkAddressPool key g []
+                , mkAddressPool key g (take (k - (fromEnum g `div` 5)) addrs)
                 ]
     arbitrary = do
         g <- arbitrary
         n <- choose (0, 2 * fromEnum g)
-        cc <- arbitrary
-        let addrs = take n (ourAddresses cc)
-        return $ mkAddressPool ourAccount g cc addrs
+        let addrs = take n (ourAddresses (changeChain @chain))
+        return $ mkAddressPool ourAccount g addrs