tapdb: add FetchSupplyLeavesByHeight to SupplyTreeStore

In this commit, we add a new method to the SupplyTreeStore that is able
to read out the leaves of a supply tree based on a start and end height.
This will be useful for writing the new syncing state machine and the
sub-system that serves the supply tree syncer.

Author: Roasbeef

tapdb/supply_tree.go

@@ -1,14 +1,18 @@
 package tapdb
 
 import (
+	"bytes"
 	"context"
+	"database/sql"
 	"encoding/hex"
+	"errors"
 	"fmt"
 
 	"github.com/btcsuite/btcd/btcec/v2"
 	"github.com/lightninglabs/taproot-assets/asset"
 	"github.com/lightninglabs/taproot-assets/mssmt"
 	"github.com/lightninglabs/taproot-assets/proof"
+	"github.com/lightninglabs/taproot-assets/tapdb/sqlc"
 	"github.com/lightninglabs/taproot-assets/universe"
 	"github.com/lightninglabs/taproot-assets/universe/supplycommit"
 
@@ -24,6 +28,14 @@ const (
 	supplySubTreeNS = "supply-sub"
 )
 
+// allSupplyTreeTypes contains all supply sub-tree types that are commonly
+// iterated over.
+var allSupplyTreeTypes = []supplycommit.SupplySubTree{
+	supplycommit.MintTreeType,
+	supplycommit.BurnTreeType,
+	supplycommit.IgnoreTreeType,
+}
+
 // SupplyTreeStore implements the persistent storage for supply trees. It
 // manages a root supply tree per asset group, where each leaf in the root tree
 // corresponds to the root of a sub-tree (mint, burn, ignore).
@@ -202,10 +214,7 @@ func (s *SupplyTreeStore) FetchSubTrees(ctx context.Context,
 	dbErr := s.db.ExecTx(ctx, &readTx, func(db BaseUniverseStore) error {
 		// For each supply tree type, we'll fetch the corresponding
 		// sub-tree within this single transaction.
-		for _, treeType := range []supplycommit.SupplySubTree{
-			supplycommit.MintTreeType, supplycommit.BurnTreeType,
-			supplycommit.IgnoreTreeType,
-		} {
+		for _, treeType := range allSupplyTreeTypes {
 			subTree, fetchErr := fetchSubTreeInternal(
 				ctx, db, groupKey, treeType,
 			)
@@ -599,3 +608,106 @@ func (s *SupplyTreeStore) ApplySupplyUpdates(ctx context.Context,
 
 	return finalRoot, nil
 }
+
+// SupplyUpdate is a struct that holds a supply update event.
+type SupplyUpdate struct {
+	supplycommit.SupplyUpdateEvent
+}
+
+// FetchSupplyLeavesByHeight fetches all supply leaves for a given asset
+// specifier within a given block height range.
+func (s *SupplyTreeStore) FetchSupplyLeavesByHeight(ctx context.Context,
+	spec asset.Specifier, startHeight,
+	endHeight uint32) ([]SupplyUpdate, error) {
+
+	groupKey, err := spec.UnwrapGroupKeyOrErr()
+	if err != nil {
+		return nil, fmt.Errorf("group key must be "+
+			"specified for supply tree: %w", err)
+	}
+
+	var updates []SupplyUpdate
+
+	readTx := NewBaseUniverseReadTx()
+	dbErr := s.db.ExecTx(ctx, &readTx, func(db BaseUniverseStore) error {
+		for _, treeType := range allSupplyTreeTypes {
+			namespace := subTreeNamespace(groupKey, treeType)
+
+			leaves, err := db.QuerySupplyLeavesByHeight(
+				ctx, sqlc.QuerySupplyLeavesByHeightParams{
+					Namespace:   namespace,
+					StartHeight: sqlInt32(startHeight),
+					EndHeight:   sqlInt32(endHeight),
+				},
+			)
+			if err != nil {
+				if errors.Is(err, sql.ErrNoRows) {
+					continue
+				}
+
+				return fmt.Errorf("failed to query "+
+					"supply leaves: %w", err)
+			}
+
+			for _, leaf := range leaves {
+				var event supplycommit.SupplyUpdateEvent
+				switch treeType {
+				case supplycommit.MintTreeType:
+					var mintEvent supplycommit.NewMintEvent
+					err = mintEvent.Decode(
+						bytes.NewReader(
+							leaf.SupplyLeafBytes,
+						),
+					)
+					if err != nil {
+						return fmt.Errorf("failed "+
+							"to decode mint "+
+							"event: %w", err)
+					}
+
+					event = &mintEvent
+
+				case supplycommit.BurnTreeType:
+					var burnEvent supplycommit.NewBurnEvent
+					err = burnEvent.Decode(
+						bytes.NewReader(
+							leaf.SupplyLeafBytes,
+						),
+					)
+					if err != nil {
+						return fmt.Errorf("failed "+
+							"to decode burn "+
+							"event: %w", err)
+					}
+
+					event = &burnEvent
+
+				case supplycommit.IgnoreTreeType:
+					//nolint:lll
+					var ignoreEvent supplycommit.NewIgnoreEvent
+					err = ignoreEvent.Decode(
+						bytes.NewReader(
+							leaf.SupplyLeafBytes,
+						),
+					)
+					if err != nil {
+						return fmt.Errorf("failed "+
+							"to decode ignore "+
+							"event: %w", err)
+					}
+					event = &ignoreEvent
+				}
+
+				updates = append(updates, SupplyUpdate{
+					SupplyUpdateEvent: event,
+				})
+			}
+		}
+		return nil
+	})
+	if dbErr != nil {
+		return nil, dbErr
+	}
+
+	return updates, nil
+}

tapdb/supply_tree_test.go

@@ -248,7 +248,8 @@ func randIgnoreTupleGen(t *rapid.T,
 			ScriptKey: asset.ToSerialized(scriptKey.PubKey),
 			OutPoint:  op,
 		},
-		Amount: 100,
+		Amount:      100,
+		BlockHeight: rapid.Uint32Range(1, 1000).Draw(t, "block_height"),
 	}
 
 	// Create a signature for the ignore tuple.
@@ -367,6 +368,101 @@ func setupSupplyTreeTestForProps(t *testing.T) (*SupplyTreeStore,
 	return supplyStore, spec, eventGen
 }
 
+// createMintEventWithHeight creates a mint event with a specific block height.
+func createMintEventWithHeight(t *testing.T, groupKey *btcec.PublicKey,
+	height uint32) *supplycommit.NewMintEvent {
+
+	mintAsset := asset.RandAsset(t, asset.Normal)
+	mintAsset.GroupKey = &asset.GroupKey{GroupPubKey: *groupKey}
+	mintAsset.GroupKey.Witness = mintAsset.PrevWitnesses[0].TxWitness
+
+	mintProof := randProof(t, mintAsset)
+	mintProof.BlockHeight = height
+	mintProof.GroupKeyReveal = asset.NewGroupKeyRevealV0(
+		asset.ToSerialized(groupKey), nil,
+	)
+
+	var proofBuf bytes.Buffer
+	require.NoError(t, mintProof.Encode(&proofBuf))
+
+	mintLeaf := universe.Leaf{
+		GenesisWithGroup: universe.GenesisWithGroup{
+			Genesis:  mintAsset.Genesis,
+			GroupKey: mintAsset.GroupKey,
+		},
+		Asset:    &mintProof.Asset,
+		Amt:      mintProof.Asset.Amount,
+		RawProof: proofBuf.Bytes(),
+	}
+
+	mintKey := universe.AssetLeafKey{
+		BaseLeafKey: universe.BaseLeafKey{
+			OutPoint:  mintProof.OutPoint(),
+			ScriptKey: &mintProof.Asset.ScriptKey,
+		},
+		AssetID: mintProof.Asset.ID(),
+	}
+
+	return &supplycommit.NewMintEvent{
+		LeafKey:       mintKey,
+		IssuanceProof: mintLeaf,
+	}
+}
+
+// createBurnEventWithHeight creates a burn event with a specific block height.
+func createBurnEventWithHeight(t *testing.T, baseGenesis asset.Genesis,
+	groupKey *asset.GroupKey, db BatchedUniverseTree,
+	height uint32) *supplycommit.NewBurnEvent {
+
+	burnAsset := createBurnAsset(t)
+	burnAsset.Genesis = baseGenesis
+	burnAsset.GroupKey = groupKey
+
+	burnProof := randProof(t, burnAsset)
+	burnProof.BlockHeight = height
+	burnProof.GenesisReveal = &baseGenesis
+
+	// Ensure genesis exists for this burn leaf in the DB.
+	ctx := context.Background()
+	genesisPointID, err := upsertGenesisPoint(
+		ctx, db, burnAsset.Genesis.FirstPrevOut,
+	)
+	require.NoError(t, err)
+	_, err = upsertGenesis(
+		ctx, db, genesisPointID, burnAsset.Genesis,
+	)
+	require.NoError(t, err)
+
+	burnLeaf := &universe.BurnLeaf{
+		UniverseKey: universe.AssetLeafKey{
+			BaseLeafKey: universe.BaseLeafKey{
+				OutPoint:  burnProof.OutPoint(),
+				ScriptKey: &burnProof.Asset.ScriptKey,
+			},
+			AssetID: burnProof.Asset.ID(),
+		},
+		BurnProof: burnProof,
+	}
+
+	return &supplycommit.NewBurnEvent{
+		BurnLeaf: *burnLeaf,
+	}
+}
+
+// createIgnoreEventWithHeight creates an ignore event with a specific block
+// height.
+func createIgnoreEventWithHeight(t *testing.T, baseAssetID asset.ID,
+	db BatchedUniverseTree, height uint32) *supplycommit.NewIgnoreEvent {
+
+	signedTuple := randIgnoreTuple(t, db)
+	signedTuple.IgnoreTuple.Val.ID = baseAssetID
+	signedTuple.IgnoreTuple.Val.BlockHeight = height
+
+	return &supplycommit.NewIgnoreEvent{
+		SignedIgnoreTuple: signedTuple,
+	}
+}
+
 // TestSupplyTreeStoreApplySupplyUpdates tests that the ApplySupplyUpdates meets
 // a series of key invariant via property based testing.
 func TestSupplyTreeStoreApplySupplyUpdates(t *testing.T) {
@@ -538,3 +634,114 @@ func TestSupplyTreeStoreApplySupplyUpdates(t *testing.T) {
 	)
 	require.NoError(t, err)
 }
+
+// TestSupplyTreeStoreFetchSupplyLeavesByHeight tests the
+// FetchSupplyLeavesByHeight method.
+func TestSupplyTreeStoreFetchSupplyLeavesByHeight(t *testing.T) {
+	t.Parallel()
+
+	supplyStore, spec, _ := setupSupplyTreeTestForProps(t)
+	ctxb := context.Background()
+	dbTxer := supplyStore.db
+
+	groupKey, err := spec.UnwrapGroupKeyOrErr()
+	require.NoError(t, err)
+	assetID := spec.UnwrapIdToPtr()
+
+	fullGroupKey := &asset.GroupKey{
+		GroupPubKey: *groupKey,
+	}
+
+	// Create events with specific block heights, we'll use these heights
+	// below to ensure that the new leaf height is properly set/read all the
+	// way down the call stack.
+	mintEvent100 := createMintEventWithHeight(t, groupKey, 100)
+	burnEvent200 := createBurnEventWithHeight(
+		t, asset.RandGenesis(t, asset.Normal), fullGroupKey, dbTxer,
+		200,
+	)
+	ignoreEvent300 := createIgnoreEventWithHeight(t, *assetID, dbTxer, 300)
+	mintEvent400 := createMintEventWithHeight(t, groupKey, 400)
+
+	updates := []supplycommit.SupplyUpdateEvent{
+		mintEvent100, burnEvent200, ignoreEvent300, mintEvent400,
+	}
+
+	// Apply updates.
+	_, err = supplyStore.ApplySupplyUpdates(ctxb, spec, updates)
+	require.NoError(t, err)
+
+	testCases := []struct {
+		name            string
+		startHeight     uint32
+		endHeight       uint32
+		expectedCount   int
+		expectedHeights []uint32
+	}{
+		{
+			name:            "range including first",
+			startHeight:     0,
+			endHeight:       150,
+			expectedCount:   1,
+			expectedHeights: []uint32{100},
+		},
+		{
+			name:            "range including second",
+			startHeight:     150,
+			endHeight:       250,
+			expectedCount:   1,
+			expectedHeights: []uint32{200},
+		},
+		{
+			name:            "range including all",
+			startHeight:     0,
+			endHeight:       500,
+			expectedCount:   4,
+			expectedHeights: []uint32{100, 200, 300, 400},
+		},
+		{
+			name:            "exact range",
+			startHeight:     100,
+			endHeight:       400,
+			expectedCount:   4,
+			expectedHeights: []uint32{100, 200, 300, 400},
+		},
+		{
+			name:            "inner range",
+			startHeight:     101,
+			endHeight:       399,
+			expectedCount:   2,
+			expectedHeights: []uint32{200, 300},
+		},
+		{
+			name:            "range after all",
+			startHeight:     501,
+			endHeight:       1000,
+			expectedCount:   0,
+			expectedHeights: nil,
+		},
+		{
+			name:            "range before all",
+			startHeight:     0,
+			endHeight:       99,
+			expectedCount:   0,
+			expectedHeights: nil,
+		},
+	}
+
+	for _, tc := range testCases {
+		t.Run(tc.name, func(t *testing.T) {
+			leaves, err := supplyStore.FetchSupplyLeavesByHeight(
+				ctxb, spec, tc.startHeight, tc.endHeight,
+			)
+			require.NoError(t, err)
+			require.Len(t, leaves, tc.expectedCount)
+
+			var heights []uint32
+			for _, leaf := range leaves {
+				heights = append(heights, leaf.BlockHeight())
+			}
+			require.ElementsMatch(t, tc.expectedHeights, heights)
+		})
+	}
+}