universe: allow FetchMultiverseLeaves to multi-query

In this commit, we modify the FetchMultiverseLeaves method to allow it
to return either all the leaves, or a subset of leaves filtered by the
`MultiverseLeafDesc` type. In the future, we can further extend this to
use postgres batch queries so we can get the resp in a single round
trip.

Author: Roasbeef

tapdb/multiverse.go

@@ -1031,54 +1031,81 @@ func (b *MultiverseStore) DeleteUniverse(ctx context.Context,
 // asset ID, and group key. If both asset ID and group key is nil, all leaves
 // for the given proof type will be returned.
 func (b *MultiverseStore) FetchLeaves(ctx context.Context,
-	assetID *asset.ID, groupKey *btcec.PublicKey,
-	proofType universe.ProofType) ([]universe.Identifier, error) {
+	universeTargets []universe.MultiverseLeafDesc,
+	proofType universe.ProofType) ([]universe.MultiverseLeaf, error) {
+
+	queries := make([]QueryMultiverseLeaves, 0, len(universeTargets))
+	switch len(universeTargets) {
+	// If we don't have any targets, then we'll have just a single query to
+	// return all universe leaves for the proof type.
+	case 0:
+		queries = append(queries, QueryMultiverseLeaves{
+			ProofType: proofType.String(),
+		})
+
+	// Otherwise, we'll do a query for each universe target specified.
+	default:
+		for _, uniTarget := range universeTargets {
+			var assetIDBytes, groupKeyBytes []byte
 
-	var assetIDBytes, groupKeyBytes []byte
-	if assetID != nil {
-		assetIDBytes = assetID[:]
-	}
-	if groupKey != nil {
-		groupKeyBytes = schnorr.SerializePubKey(groupKey)
+			uniTarget.WhenLeft(func(a asset.ID) {
+				assetIDBytes = a[:]
+			})
+			uniTarget.WhenRight(func(g btcec.PublicKey) {
+				groupKeyBytes = schnorr.SerializePubKey(&g)
+			})
+
+			queries = append(queries, QueryMultiverseLeaves{
+				ProofType: proofType.String(),
+				AssetID:   assetIDBytes,
+				GroupKey:  groupKeyBytes,
+			})
+		}
 	}
 
 	var (
 		readTx = NewBaseUniverseReadTx()
-		ids    []universe.Identifier
+		leaves []universe.MultiverseLeaf
 	)
 	dbErr := b.db.ExecTx(ctx, &readTx, func(q BaseMultiverseStore) error {
-		leaves, err := q.QueryMultiverseLeaves(
-			ctx, QueryMultiverseLeaves{
-				ProofType: proofType.String(),
-				AssetID:   assetIDBytes,
-				GroupKey:  groupKeyBytes,
-			},
-		)
-		if err != nil {
-			return err
-		}
+		leaves = nil
 
-		ids = make([]universe.Identifier, len(leaves))
-		for i, leaf := range leaves {
-			ids[i].ProofType = proofType
-			if len(leaf.AssetID) > 0 {
-				copy(ids[i].AssetID[:], leaf.AssetID)
+		for _, query := range queries {
+			dbLeaves, err := q.QueryMultiverseLeaves(ctx, query)
+			if err != nil {
+				return err
 			}
-			if len(leaf.GroupKey) > 0 {
-				ids[i].GroupKey, err = schnorr.ParsePubKey(
-					leaf.GroupKey,
-				)
-				if err != nil {
-					return err
+
+			for _, leaf := range dbLeaves {
+				var id universe.Identifier
+
+				id.ProofType = proofType
+				if len(leaf.AssetID) > 0 {
+					copy(id.AssetID[:], leaf.AssetID)
 				}
+				if len(leaf.GroupKey) > 0 {
+					id.GroupKey, err = schnorr.ParsePubKey(
+						leaf.GroupKey,
+					)
+					if err != nil {
+						return err
+					}
+				}
+
+				leaves = append(leaves, universe.MultiverseLeaf{
+					ID: id,
+					LeafNode: mssmt.NewLeafNode(
+						leaf.UniverseRootHash,
+						uint64(leaf.UniverseRootSum),
+					),
+				})
 			}
 		}
-
 		return nil
 	})
 	if dbErr != nil {
 		return nil, dbErr
 	}
 
-	return ids, nil
+	return leaves, nil
 }

tapdb/universe_test.go

@@ -116,20 +116,22 @@ func newTestUniverseWithDb(db *BaseDB,
 	return NewBaseUniverseTree(dbTxer, id), db
 }
 
-func assertIDInList(t *testing.T, leaves []universe.Identifier,
+func assertIDInList(t *testing.T, leaves []universe.MultiverseLeaf,
 	id universe.Identifier) {
 
-	require.True(t, fn.Any(leaves, func(l universe.Identifier) bool {
+	require.True(t, fn.Any(leaves, func(l universe.MultiverseLeaf) bool {
 		switch {
-		case l.AssetID != asset.ID{}:
-			return l.AssetID == id.AssetID
+		case l.ID.AssetID != asset.ID{}:
+			return l.ID.AssetID == id.AssetID
 
-		case l.GroupKey != nil:
+		case l.ID.GroupKey != nil:
 			if id.GroupKey == nil {
 				return false
 			}
 
-			return test.SchnorrKeysEqual(t, l.GroupKey, id.GroupKey)
+			return test.SchnorrKeysEqual(
+				t, l.ID.GroupKey, id.GroupKey,
+			)
 
 		default:
 			require.Fail(t, "invalid leaf")
@@ -250,12 +252,12 @@ func TestUniverseIssuanceProofs(t *testing.T) {
 
 	// The multiverse tree should be empty at this point.
 	issuanceLeaves, err := multiverse.FetchLeaves(
-		ctx, nil, nil, universe.ProofTypeIssuance,
+		ctx, nil, universe.ProofTypeIssuance,
 	)
 	require.NoError(t, err)
 	require.Len(t, issuanceLeaves, 0)
 	transferLeaves, err := multiverse.FetchLeaves(
-		ctx, nil, nil, universe.ProofTypeTransfer,
+		ctx, nil, universe.ProofTypeTransfer,
 	)
 	require.NoError(t, err)
 	require.Len(t, transferLeaves, 0)
@@ -337,7 +339,7 @@ func TestUniverseIssuanceProofs(t *testing.T) {
 	// The multiverse tree should just have a single leaf, since we inserted
 	// proofs into the same universe.
 	multiverseLeaves, err := multiverse.FetchLeaves(
-		ctx, nil, nil, id.ProofType,
+		ctx, nil, id.ProofType,
 	)
 	require.NoError(t, err)
 	require.Len(t, multiverseLeaves, 1)
@@ -887,10 +889,10 @@ func TestMultiverseRootSum(t *testing.T) {
 	}
 
 	testCases := []testCase{
-		// If we insert two transfers into a transfer tree, then the
-		// sum should be the sum of the leaf values. The leaf value
-		// here is itself the root sum of an transfer tree, or the
-		// number of transfers in a transfer tree.
+		// If we insert two transfers into a transfer tree, then the sum
+		// should be the sum of the leaf values. The leaf value here is
+		// itself the root sum of a transfer tree, or the number of
+		// transfers in a transfer tree.
 		{
 			name:      "transfer sum",
 			finalSum:  4,
@@ -922,6 +924,28 @@ func TestMultiverseRootSum(t *testing.T) {
 				},
 			},
 		},
+
+		// We also want to make sure we can insert both transfer and
+		// issuance proofs at the same time without any conflicts.
+		{
+			name:     "transfer and issuance sum",
+			finalSum: 3,
+			// By specifying this as "unspecified" we signal we want
+			// both transfer and issuance proofs. The final sum will
+			// therefore be the same for both trees.
+			proofType: universe.ProofTypeUnspecified,
+			leaves: []leaf{
+				{
+					sumAmt: 14,
+				},
+				{
+					sumAmt: 20,
+				},
+				{
+					sumAmt: 20,
+				},
+			},
+		},
 	}
 
 	runTestCase := func(t *testing.T, tc testCase) {
@@ -931,88 +955,116 @@ func TestMultiverseRootSum(t *testing.T) {
 
 		// The multiverse tree should be empty at this point.
 		issuanceLeaves, err := multiverse.FetchLeaves(
-			ctx, nil, nil, universe.ProofTypeIssuance,
+			ctx, nil, universe.ProofTypeIssuance,
 		)
 		require.NoError(t, err)
 		require.Len(t, issuanceLeaves, 0)
 		transferLeaves, err := multiverse.FetchLeaves(
-			ctx, nil, nil, universe.ProofTypeTransfer,
+			ctx, nil, universe.ProofTypeTransfer,
 		)
 		require.NoError(t, err)
 		require.Len(t, transferLeaves, 0)
 
-		leaves := make([]universe.Leaf, len(tc.leaves))
-		ids := make([]universe.Identifier, len(tc.leaves))
-		for i, testLeaf := range tc.leaves {
-			id := randUniverseID(
-				t, false, withProofType(tc.proofType),
-			)
-
-			ids[i] = id
+		ids := make([]universe.Identifier, 0, len(tc.leaves))
+		for range tc.leaves {
+			ids = append(ids, randUniverseID(t, false))
+		}
 
-			assetGen := asset.RandGenesis(t, asset.Normal)
-			leaf := randMintingLeaf(t, assetGen, id.GroupKey)
-			leaf.Amt = testLeaf.sumAmt
+		insertLeaves := func(proofType universe.ProofType) {
+			for i, testLeaf := range tc.leaves {
+				id := ids[i]
+				id.ProofType = proofType
 
-			leaves[i] = leaf
+				assetGen := asset.RandGenesis(t, asset.Normal)
+				leaf := randMintingLeaf(
+					t, assetGen, id.GroupKey,
+				)
+				leaf.Amt = testLeaf.sumAmt
 
-			targetKey := randLeafKey(t)
+				targetKey := randLeafKey(t)
 
-			// For transfer proofs, we'll modify the witness asset
-			// proof to look more like a transfer.
-			if tc.proofType == universe.ProofTypeTransfer {
-				prevWitnesses := leaf.Asset.PrevWitnesses
-				prevWitnesses[0].TxWitness = [][]byte{
-					{1}, {1}, {1},
+				// For transfer proofs, we'll modify the witness
+				// asset proof to look more like a transfer.
+				if proofType == universe.ProofTypeTransfer {
+					prevWitnesses := leaf.Asset.PrevWitnesses
+					prevWitnesses[0].TxWitness = [][]byte{
+						{1}, {1}, {1},
+					}
+					prevID := prevWitnesses[0].PrevID
+					prevID.OutPoint.Hash = [32]byte{1}
 				}
-				prevWitnesses[0].PrevID.OutPoint.Hash = [32]byte{1}
-			}
-
-			_, err := multiverse.UpsertProofLeaf(
-				ctx, id, targetKey, &leaf, nil,
-			)
-			require.NoError(t, err)
-
-			// If we should add more than one under this ID, then
-			// we'll generate another instance.
-			if tc.doubleUp {
-				targetKey = randLeafKey(t)
 
 				_, err := multiverse.UpsertProofLeaf(
 					ctx, id, targetKey, &leaf, nil,
 				)
 				require.NoError(t, err)
+
+				// If we should add more than one under this ID,
+				// then we'll generate another instance.
+				if tc.doubleUp {
+					targetKey = randLeafKey(t)
+
+					_, err := multiverse.UpsertProofLeaf(
+						ctx, id, targetKey, &leaf, nil,
+					)
+					require.NoError(t, err)
+				}
+
+				// The multiverse tree should now have one more
+				// leaf.
+				multiverseLeaves, err := multiverse.FetchLeaves(
+					ctx, nil, proofType,
+				)
+				require.NoError(t, err)
+				require.Len(t, multiverseLeaves, i+1)
+
+				// And we should actually find the leaf we just
+				// inserted.
+				assertIDInList(t, multiverseLeaves, id)
 			}
+		}
 
-			// The multiverse tree should now have one more leaf.
-			multiverseLeaves, err := multiverse.FetchLeaves(
-				ctx, nil, nil, tc.proofType,
+		checkSum := func(proofType universe.ProofType) {
+			rootNode, err := multiverse.MultiverseRootNode(
+				ctx, proofType,
 			)
 			require.NoError(t, err)
-			require.Len(t, multiverseLeaves, i+1)
 
-			// And we should actually find the leaf we just
-			// inserted.
-			assertIDInList(t, multiverseLeaves, id)
-		}
-
-		// If we fetch the root value of the tree, it should be
-		// the same as the finalSum.
-		rootNode, err := multiverse.RootNode(ctx, tc.proofType)
-		require.NoError(t, err)
+			rootNode.WhenSome(
+				func(rootNode universe.MultiverseRoot) {
+					require.EqualValues(
+						t, tc.finalSum,
+						rootNode.NodeSum(),
+					)
+				},
+			)
 
-		require.EqualValues(t, tc.finalSum, rootNode.NodeSum())
+			// We now delete the whole universe and expect the
+			// multiverse leave to also disappear.
+			id := ids[0]
+			id.ProofType = proofType
+			_, err = multiverse.DeleteUniverse(ctx, id)
+			require.NoError(t, err)
 
-		// We now delete the whole universe and expect the multiverse
-		// leave to also disappear.
-		_, err = multiverse.DeleteUniverse(ctx, ids[0])
-		require.NoError(t, err)
+			multiverseLeaves, err := multiverse.FetchLeaves(
+				ctx, nil, proofType,
+			)
+			require.NoError(t, err)
+			require.Len(t, multiverseLeaves, len(ids)-1)
+		}
 
-		multiverseLeaves, err := multiverse.FetchLeaves(
-			ctx, nil, nil, tc.proofType,
-		)
-		require.NoError(t, err)
-		require.Len(t, multiverseLeaves, len(ids)-1)
+		// If we fetch the root value of the tree, it should be the same
+		// as the finalSum.
+		if tc.proofType == universe.ProofTypeUnspecified {
+			insertLeaves(universe.ProofTypeIssuance)
+			insertLeaves(universe.ProofTypeTransfer)
+
+			checkSum(universe.ProofTypeIssuance)
+			checkSum(universe.ProofTypeTransfer)
+		} else {
+			insertLeaves(tc.proofType)
+			checkSum(tc.proofType)
+		}
 	}
 
 	for _, testCase := range testCases {