Skip to content

Commit

Permalink
universe: allow FetchMultiverseLeaves to multi-query
Browse files Browse the repository at this point in the history
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.
  • Loading branch information
Roasbeef authored and guggero committed Dec 12, 2023
1 parent 6c55912 commit 6744d9c
Show file tree
Hide file tree
Showing 2 changed files with 180 additions and 101 deletions.
91 changes: 59 additions & 32 deletions tapdb/multiverse.go
Original file line number Diff line number Diff line change
Expand Up @@ -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
}
190 changes: 121 additions & 69 deletions tapdb/universe_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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")
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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) {
Expand All @@ -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 {
Expand Down

0 comments on commit 6744d9c

Please sign in to comment.