add rest properties for isShared

lib/core/src/Cardano/Wallet/Primitive/Scripts.hs

@@ -48,6 +48,7 @@ import Cardano.Wallet.Primitive.AddressDiscovery.Sequential
 import Crypto.Hash.Utils
     ( blake2b224 )
 
+import qualified Data.List as L
 import qualified Data.Map.Strict as Map
 import qualified Data.Text.Encoding as T
 
@@ -63,7 +64,7 @@ isShared
 isShared script s@(SeqState !s1 !s2 !ixs !rpk !prefix !scripts) =
     let verKeysInScript = retrieveAllVerKeyHashes script
         accXPub = accountPubKey s2
-        toVerKey ix = deriveAddressPublicKey accXPub MultisigScript ix
+        toVerKey = deriveAddressPublicKey accXPub MultisigScript
         minIndex = getIndex @'Soft minBound
         scriptAddressGap = 10
         ourVerKeys =
@@ -73,13 +74,12 @@ isShared script s@(SeqState !s1 !s2 !ixs !rpk !prefix !scripts) =
             filter (\keyH -> toKeyHash keyH `elem` verKeysInScript)
             ourVerKeys
         toScriptXPub (ShelleyKey k) = ShelleyKey k
-        scriptXPubs = map toScriptXPub ourVerKeyHashesInScript
+        scriptXPubs = L.nub $ map toScriptXPub ourVerKeyHashesInScript
     in if null ourVerKeyHashesInScript then
         ([], s)
        else
         ( scriptXPubs
         , SeqState s1 s2 ixs rpk prefix (Map.insert (toScriptHash script) scriptXPubs scripts))
-  where
 
 retrieveAllVerKeyHashes :: Script -> [KeyHash]
 retrieveAllVerKeyHashes = extractVerKey []

lib/core/test/unit/Cardano/Wallet/Primitive/ScriptsSpec.hs

@@ -1,16 +1,8 @@
-{-# LANGUAGE AllowAmbiguousTypes #-}
 {-# LANGUAGE DataKinds #-}
-{-# LANGUAGE ExistentialQuantification #-}
 {-# LANGUAGE FlexibleContexts #-}
 {-# LANGUAGE FlexibleInstances #-}
-{-# LANGUAGE GADTs #-}
-{-# LANGUAGE GeneralizedNewtypeDeriving #-}
-{-# LANGUAGE KindSignatures #-}
-{-# LANGUAGE LambdaCase #-}
 {-# LANGUAGE ScopedTypeVariables #-}
-{-# LANGUAGE StandaloneDeriving #-}
 {-# LANGUAGE TypeApplications #-}
-{-# LANGUAGE UndecidableInstances #-}
 
 {-# OPTIONS_GHC -fno-warn-orphans #-}
 
@@ -23,7 +15,7 @@ import Prelude
 import Cardano.Address.Derivation
     ( XPub )
 import Cardano.Address.Script
-    ( KeyHash (..), Script (..) )
+    ( KeyHash (..), Script (..), toScriptHash )
 import Cardano.Mnemonic
     ( SomeMnemonic (..) )
 import Cardano.Wallet.Gen
@@ -85,6 +77,7 @@ import qualified Data.ByteArray as BA
 import qualified Data.ByteString.Char8 as B8
 import qualified Data.List as L
 import qualified Data.Map.Strict as Map
+import qualified Data.Set as Set
 import qualified Data.Text as T
 import qualified Data.Text.Encoding as T
 
@@ -95,34 +88,82 @@ spec = do
             (property prop_scriptFromOurVerKeys)
         it "script composed with not our verification keys should not be discovered"
             (property prop_scriptFromNotOurVerKeys)
+        it "the same script discovered twice should have the same knownScripts imprint"
+            (property prop_scriptDiscoveredTwice)
+        it "knownScripts of the sequential state is populated properly"
+            (property prop_scriptUpdatesStateProperly)
+        it "scripts with our verification keys are discovered properly"
+            (property prop_scriptsDiscovered)
 
 prop_scriptFromOurVerKeys
-    :: AccountXPubWithScript
+    :: AccountXPubWithScripts
     -> Property
-prop_scriptFromOurVerKeys (AccountXPubWithScript accXPub' script') = do
-    let sciptHashes = retrieveAllVerKeyHashes script'
+prop_scriptFromOurVerKeys (AccountXPubWithScripts accXPub' scripts') = do
+    let (script:_) = scripts'
+    let sciptKeyHashes = retrieveAllVerKeyHashes script
     let intPool = mkAddressPool accXPub' defaultAddressPoolGap []
     let extPool = mkAddressPool accXPub' defaultAddressPoolGap []
     let seqState =
             SeqState intPool extPool emptyPendingIxs dummyRewardAccount defaultPrefix Map.empty
-    let (ourSharedKeys, _) = isShared script' seqState
-    L.sort (L.nub $ map toKeyHash ourSharedKeys) === L.sort (L.nub sciptHashes)
+    let (ourSharedKeys, _) = isShared script seqState
+    L.sort (L.nub $ map toKeyHash ourSharedKeys) === L.sort (L.nub sciptKeyHashes)
 
 prop_scriptFromNotOurVerKeys
     :: ShelleyKey 'AccountK XPub
-    -> AccountXPubWithScript
+    -> AccountXPubWithScripts
     -> Property
-prop_scriptFromNotOurVerKeys accXPub' (AccountXPubWithScript _accXPub script') = do
+prop_scriptFromNotOurVerKeys accXPub' (AccountXPubWithScripts _accXPub scripts') = do
+    let (script:_) = scripts'
     let intPool = mkAddressPool accXPub' defaultAddressPoolGap []
     let extPool = mkAddressPool accXPub' defaultAddressPoolGap []
     let seqState =
             SeqState intPool extPool emptyPendingIxs dummyRewardAccount defaultPrefix Map.empty
-    let (ourSharedKeys, _) = isShared script' seqState
+    let (ourSharedKeys, _) = isShared script seqState
     ourSharedKeys === []
 
-data AccountXPubWithScript = AccountXPubWithScript
+prop_scriptUpdatesStateProperly
+    :: AccountXPubWithScripts
+    -> Property
+prop_scriptUpdatesStateProperly (AccountXPubWithScripts accXPub' scripts') = do
+    let (script:_) = scripts'
+    let sciptKeyHashes = retrieveAllVerKeyHashes script
+    let intPool = mkAddressPool accXPub' defaultAddressPoolGap []
+    let extPool = mkAddressPool accXPub' defaultAddressPoolGap []
+    let seqState =
+            SeqState intPool extPool emptyPendingIxs dummyRewardAccount defaultPrefix Map.empty
+    let (_, seqState') = isShared script seqState
+    let scriptKeyHashesInMap =
+            Set.fromList . map toKeyHash <$> Map.lookup (toScriptHash script) (knownScripts seqState')
+    scriptKeyHashesInMap === (Just $ Set.fromList (L.nub sciptKeyHashes))
+
+prop_scriptDiscoveredTwice
+    :: AccountXPubWithScripts
+    -> Property
+prop_scriptDiscoveredTwice (AccountXPubWithScripts accXPub' scripts') = do
+    let (script:_) = scripts'
+    let intPool = mkAddressPool accXPub' defaultAddressPoolGap []
+    let extPool = mkAddressPool accXPub' defaultAddressPoolGap []
+    let seqState =
+            SeqState intPool extPool emptyPendingIxs dummyRewardAccount defaultPrefix Map.empty
+    let (_, seqState') = isShared script seqState
+    let (_, seqState'') = isShared script seqState'
+    knownScripts seqState' === knownScripts seqState''
+
+prop_scriptsDiscovered
+    :: AccountXPubWithScripts
+    -> Property
+prop_scriptsDiscovered (AccountXPubWithScripts accXPub' scripts') = do
+    let intPool = mkAddressPool accXPub' defaultAddressPoolGap []
+    let extPool = mkAddressPool accXPub' defaultAddressPoolGap []
+    let seqState0 =
+            SeqState intPool extPool emptyPendingIxs dummyRewardAccount defaultPrefix Map.empty
+    let seqState = foldr (\script s -> snd $ isShared script s) seqState0 scripts'
+    let scriptHashes = Set.fromList $ Map.keys $ knownScripts seqState
+    scriptHashes === Set.fromList (map toScriptHash scripts')
+
+data AccountXPubWithScripts = AccountXPubWithScripts
     { accXPub :: ShelleyKey 'AccountK XPub
-    , script :: Script
+    , scripts :: [Script]
     } deriving (Eq, Show)
 
 defaultPrefix :: DerivationPrefix
@@ -153,12 +194,12 @@ genScript keyHashes =
         scriptTree n = do
             Positive m <- arbitrary
             let n' = n `div` (m + 1)
-            scripts <- vectorOf m (scriptTree n')
+            scripts' <- vectorOf m (scriptTree n')
             atLeast <- choose (1, fromIntegral (m + 1))
             elements
-                [ RequireAllOf scripts
-                , RequireAnyOf scripts
-                , RequireSomeOf atLeast scripts
+                [ RequireAllOf scripts'
+                , RequireAnyOf scripts'
+                , RequireSomeOf atLeast scripts'
                 ]
 
 instance Arbitrary (ShelleyKey 'AccountK XPub) where
@@ -168,17 +209,18 @@ instance Arbitrary (ShelleyKey 'AccountK XPub) where
         let rootXPrv = generateKeyFromSeed (mnemonics, Nothing) encPwd
         pure $ publicKey $ deriveAccountPrivateKey encPwd rootXPrv minBound
 
-instance Arbitrary AccountXPubWithScript where
+instance Arbitrary AccountXPubWithScripts where
     arbitrary = do
         accXPub' <- arbitrary
         keyNum <- choose (2,8)
-        let toVerKey ix = deriveAddressPublicKey accXPub' MultisigScript ix
+        let toVerKey = deriveAddressPublicKey accXPub' MultisigScript
         let minIndex = getIndex @'Soft minBound
         let verKeys =
                 map (\ix -> toVerKey (toEnum (fromInteger $ toInteger $ minIndex + ix)))
                 [0 .. keyNum]
         let verKeyHashes = map toKeyHash verKeys
-        AccountXPubWithScript accXPub' <$> genScript verKeyHashes
+        scriptsNum <- choose (1,10)
+        AccountXPubWithScripts accXPub' <$> vectorOf scriptsNum (genScript verKeyHashes)
 
 instance Arbitrary (Passphrase "raw") where
     arbitrary = do