universe: add MultiverseRoot method to archive

Author: guggero

universe/base.go

@@ -12,6 +12,7 @@ import (
 	"github.com/davecgh/go-spew/spew"
 	"github.com/lightninglabs/taproot-assets/asset"
 	"github.com/lightninglabs/taproot-assets/fn"
+	"github.com/lightninglabs/taproot-assets/mssmt"
 	"github.com/lightninglabs/taproot-assets/proof"
 )
 
@@ -134,6 +135,69 @@ func (a *Archive) RootNodes(ctx context.Context,
 	return a.cfg.Multiverse.RootNodes(ctx, q)
 }
 
+// MultiverseRoot returns the root node of the multiverse for the specified
+// proof type. If the given list of universe IDs is non-empty, then the root
+// will be calculated just for those universes.
+func (a *Archive) MultiverseRoot(ctx context.Context, proofType ProofType,
+	filterByIDs []Identifier) (mssmt.Node, error) {
+
+	log.Debugf("Fetching multiverse root for proof type: %v", proofType)
+
+	leaveIDs, err := a.cfg.Multiverse.FetchLeaves(ctx, nil, nil, proofType)
+	if err != nil {
+		return nil, fmt.Errorf("unable to fetch multiverse leaves: %w",
+			err)
+	}
+
+	// If a filter list is provided, then we'll only include the leaves
+	// that are in the filter list.
+	includeUniverse := func(id Identifier) bool {
+		if len(filterByIDs) == 0 {
+			return true
+		}
+
+		for _, filterID := range filterByIDs {
+			if id.IsEqual(filterID) {
+				return true
+			}
+		}
+
+		return false
+	}
+
+	memStore := mssmt.NewDefaultStore()
+	tree := mssmt.NewCompactedTree(memStore)
+
+	for _, id := range leaveIDs {
+		// Only include the universe if it's in the filter list (given
+		// the filter list is non-empty).
+		if !includeUniverse(id) {
+			continue
+		}
+
+		uniRoot, err := a.cfg.Multiverse.UniverseRootNode(ctx, id)
+		if err != nil {
+			return nil, fmt.Errorf("unable to fetch universe "+
+				"root: %w", err)
+		}
+
+		rootHash := uniRoot.NodeHash()
+		rootSum := uniRoot.NodeSum()
+
+		if id.ProofType == ProofTypeIssuance {
+			rootSum = 1
+		}
+
+		uniLeaf := mssmt.NewLeafNode(rootHash[:], rootSum)
+		_, err = tree.Insert(ctx, id.Bytes(), uniLeaf)
+		if err != nil {
+			return nil, fmt.Errorf("unable to insert leaf: %w", err)
+		}
+	}
+
+	return tree.Root(ctx)
+}
+
 // UpsertProofLeaf attempts to upsert a proof for an asset issuance or transfer
 // event. This method will return an error if the passed proof is invalid. If
 // the leaf is already known, then no action is taken and the existing

universe/interface.go

@@ -84,6 +84,24 @@ func (i *Identifier) StringForLog() string {
 		i.String(), i.AssetID[:], groupKey, i.ProofType)
 }
 
+// IsEqual returns true if the two identifiers are equal.
+func (i *Identifier) IsEqual(other Identifier) bool {
+	if i == nil {
+		return false
+	}
+
+	groupKeysEqual := false
+	if i.GroupKey == nil || other.GroupKey == nil {
+		groupKeysEqual = i.GroupKey == other.GroupKey
+	} else {
+		groupKeysEqual = i.GroupKey.IsEqual(other.GroupKey)
+	}
+
+	return i.AssetID == other.AssetID &&
+		groupKeysEqual &&
+		i.ProofType == other.ProofType
+}
+
 // NewUniIDFromAsset creates a new universe ID from an asset.
 func NewUniIDFromAsset(a asset.Asset) Identifier {
 	proofType := ProofTypeTransfer
@@ -393,6 +411,13 @@ type MultiverseArchive interface {
 	// universe.
 	UniverseLeafKeys(ctx context.Context,
 		q UniverseLeafKeysQuery) ([]LeafKey, error)
+
+	// FetchLeaves returns the set of multiverse leaves for the given proof
+	// type, asset ID, and group key. If both asset ID and group key is nil,
+	// all leaves for the given proof type will be returned.
+	FetchLeaves(ctx context.Context, assetID *asset.ID,
+		groupKey *btcec.PublicKey,
+		proofType ProofType) ([]Identifier, error)
 }
 
 // Registrar is an interface that allows a caller to upsert a proof leaf in a