tree: add asset-aware materialization context

lib/tree/asset_tree_context.go

@@ -0,0 +1,162 @@
+package tree
+
+import (
+	"fmt"
+	"math"
+
+	"github.com/btcsuite/btcd/wire/v2"
+)
+
+// AssetTreeContext contains the data needed to materialize and sign an asset
+// tree. Builders populate the context before sharing the tree.
+type AssetTreeContext struct {
+	// Nodes do not have inputs until materialization, so the structure pass
+	// records amounts by node first.
+	amountsByNode map[*Node]uint64
+
+	// Input outpoints remain stable when paths are extracted or serialized.
+	amountsByInput map[wire.OutPoint]uint64
+
+	leafRootsByInput map[wire.OutPoint][]byte
+	tweaksByInput    map[wire.OutPoint][]byte
+	packagesByInput  map[wire.OutPoint][]byte
+	assetRef         string
+}
+
+// NewAssetTreeContext returns an empty asset tree context.
+func NewAssetTreeContext() *AssetTreeContext {
+	return &AssetTreeContext{
+		amountsByNode:    make(map[*Node]uint64),
+		amountsByInput:   make(map[wire.OutPoint]uint64),
+		leafRootsByInput: make(map[wire.OutPoint][]byte),
+		tweaksByInput:    make(map[wire.OutPoint][]byte),
+		packagesByInput:  make(map[wire.OutPoint][]byte),
+	}
+}
+
+// SetNodeAssetAmount records a node's subtree asset amount.
+func (c *AssetTreeContext) SetNodeAssetAmount(node *Node, amount uint64) {
+	c.amountsByNode[node] = amount
+	if node.Input != (wire.OutPoint{}) {
+		// Extracted paths contain cloned nodes, but preserve their
+		// inputs.
+		c.amountsByInput[node.Input] = amount
+	}
+}
+
+// NodeAssetAmount returns a node's subtree asset amount.
+func (c *AssetTreeContext) NodeAssetAmount(node *Node) uint64 {
+	if amount, ok := c.amountsByNode[node]; ok {
+		return amount
+	}
+
+	// Path extraction clones nodes but preserves their inputs.
+	return c.amountsByInput[node.Input]
+}
+
+// SetSigningTweak records the taproot tweak for a tree transaction.
+func (c *AssetTreeContext) SetSigningTweak(input wire.OutPoint, tweak []byte) {
+	c.tweaksByInput[input] = append([]byte(nil), tweak...)
+}
+
+// SigningTweak returns the taproot tweak for a tree transaction.
+func (c *AssetTreeContext) SigningTweak(input wire.OutPoint) []byte {
+	return append([]byte(nil), c.tweaksByInput[input]...)
+}
+
+// SetSealedPackage records a tree transaction's sealed transfer package.
+func (c *AssetTreeContext) SetSealedPackage(input wire.OutPoint, pkg []byte) {
+	c.packagesByInput[input] = append([]byte(nil), pkg...)
+}
+
+// SealedPackage returns a tree transaction's sealed transfer package.
+func (c *AssetTreeContext) SealedPackage(input wire.OutPoint) []byte {
+	return append([]byte(nil), c.packagesByInput[input]...)
+}
+
+// SetAssetRef records the tree's asset reference.
+func (c *AssetTreeContext) SetAssetRef(ref string) {
+	c.assetRef = ref
+}
+
+// AssetRef returns the tree's asset reference.
+func (c *AssetTreeContext) AssetRef() string {
+	return c.assetRef
+}
+
+// IsEmpty reports whether the context carries no asset data at all.
+func (c *AssetTreeContext) IsEmpty() bool {
+	if c == nil {
+		return true
+	}
+
+	return len(c.amountsByNode) == 0 && len(c.amountsByInput) == 0 &&
+		len(c.leafRootsByInput) == 0 && len(c.tweaksByInput) == 0 &&
+		len(c.packagesByInput) == 0 && c.assetRef == ""
+}
+
+// tweakLookup returns the signing tweak for a node's input.
+func (c *AssetTreeContext) tweakLookup() func(*Node) []byte {
+	return func(node *Node) []byte {
+		return c.SigningTweak(node.Input)
+	}
+}
+
+// aggregateAssetAmounts records the asset amount of every subtree.
+func aggregateAssetAmounts(node *Node, leafAssets map[*Node]uint64,
+	ctx *AssetTreeContext) (uint64, error) {
+
+	if len(node.Children) == 0 {
+		amount := leafAssets[node]
+		ctx.SetNodeAssetAmount(node, amount)
+
+		return amount, nil
+	}
+
+	var total uint64
+	for _, child := range node.Children {
+		childAmount, err := aggregateAssetAmounts(
+			child, leafAssets, ctx,
+		)
+		if err != nil {
+			return 0, err
+		}
+
+		if childAmount > math.MaxUint64-total {
+			return 0, fmt.Errorf("asset amount overflow " +
+				"aggregating subtree totals")
+		}
+		total += childAmount
+	}
+
+	ctx.SetNodeAssetAmount(node, total)
+
+	return total, nil
+}
+
+// SetLeafAssetRoot records an asset leaf's commitment root.
+func (c *AssetTreeContext) SetLeafAssetRoot(input wire.OutPoint, root []byte) {
+	if c == nil || len(root) == 0 {
+		return
+	}
+
+	c.leafRootsByInput[input] = append([]byte(nil), root...)
+}
+
+// LeafAssetRoot returns an asset leaf's commitment root.
+func (c *AssetTreeContext) LeafAssetRoot(input wire.OutPoint) []byte {
+	if c == nil {
+		return nil
+	}
+
+	return append([]byte(nil), c.leafRootsByInput[input]...)
+}
+
+// LeafAssetRoot returns an asset leaf's commitment root.
+func (t *Tree) LeafAssetRoot(input wire.OutPoint) []byte {
+	if t == nil {
+		return nil
+	}
+
+	return t.AssetContext.LeafAssetRoot(input)
+}

lib/tree/asset_tree_context_test.go

@@ -0,0 +1,172 @@
+package tree
+
+import (
+	"bytes"
+	"math"
+	"testing"
+
+	"github.com/btcsuite/btcd/btcec/v2"
+	"github.com/btcsuite/btcd/wire/v2"
+	"github.com/stretchr/testify/require"
+)
+
+// assetTestLeaves returns leaves with the given asset amounts.
+func assetTestLeaves(t *testing.T, amounts ...uint64) []LeafDescriptor {
+	t.Helper()
+
+	leaves := make([]LeafDescriptor, len(amounts))
+	for i, amount := range amounts {
+		_, owner := createTestKey(t)
+		leaves[i] = LeafDescriptor{
+			PkScript: []byte{
+				0x51,
+				byte(i),
+			},
+			Amount:      10_000,
+			CoSignerKey: owner,
+			AssetAmount: amount,
+		}
+	}
+
+	return leaves
+}
+
+// TestBuildStructureAssetAggregation tests subtree amount calculation.
+func TestBuildStructureAssetAggregation(t *testing.T) {
+	t.Parallel()
+
+	_, operatorKey := createTestKey(t)
+	cfg := StructureConfig{OperatorKey: operatorKey, Radix: 2}
+
+	t.Run("btc-only trees have no asset context", func(t *testing.T) {
+		t.Parallel()
+
+		structure, err := BuildStructure(
+			assetTestLeaves(t, 0, 0, 0), cfg,
+		)
+		require.NoError(t, err)
+		require.Nil(t, structure.AssetContext)
+	})
+
+	t.Run("subtree totals cover every node", func(t *testing.T) {
+		t.Parallel()
+
+		structure, err := BuildStructure(
+			assetTestLeaves(t, 800, 200, 500), cfg,
+		)
+		require.NoError(t, err)
+		require.NotNil(t, structure.AssetContext)
+
+		assetCtx := structure.AssetContext
+		require.EqualValues(
+			t, 1_500, assetCtx.NodeAssetAmount(structure.Root),
+		)
+
+		// Every node's total must equal the sum of its children,
+		// and each leaf must carry one of the input amounts.
+		var leafTotal uint64
+		for node := range structure.Root.NodesIter() {
+			if len(node.Children) == 0 {
+				leafTotal += assetCtx.NodeAssetAmount(node)
+				continue
+			}
+
+			var childSum uint64
+			for _, child := range node.Children {
+				childSum += assetCtx.NodeAssetAmount(child)
+			}
+			require.Equal(
+				t, childSum, assetCtx.NodeAssetAmount(node),
+			)
+		}
+		require.EqualValues(t, 1_500, leafTotal)
+	})
+
+	t.Run("mixed asset and btc leaves", func(t *testing.T) {
+		t.Parallel()
+
+		structure, err := BuildStructure(
+			assetTestLeaves(t, 300, 0), cfg,
+		)
+		require.NoError(t, err)
+		require.NotNil(t, structure.AssetContext)
+		require.EqualValues(
+			t, 300,
+			structure.AssetContext.NodeAssetAmount(structure.Root),
+		)
+	})
+
+	t.Run("overflow rejected", func(t *testing.T) {
+		t.Parallel()
+
+		_, err := BuildStructure(
+			assetTestLeaves(t, math.MaxUint64, 1), cfg,
+		)
+		require.ErrorContains(t, err, "overflow")
+	})
+}
+
+// TestAssetTreeContextTweakLookup tests outpoint lookup and byte ownership.
+func TestAssetTreeContextTweakLookup(t *testing.T) {
+	t.Parallel()
+
+	assetCtx := NewAssetTreeContext()
+	input := wire.OutPoint{Index: 3}
+	input.Hash[0] = 0xAB
+
+	tweak := bytes.Repeat([]byte{0x07}, 32)
+	assetCtx.SetSigningTweak(input, tweak)
+	tweak[0] = 0xFF
+
+	lookup := assetCtx.tweakLookup()
+	got := lookup(&Node{Input: input})
+	require.Equal(t, bytes.Repeat([]byte{0x07}, 32), got)
+	got[0] = 0xEE
+	require.Equal(t, byte(0x07), lookup(&Node{Input: input})[0])
+	require.Nil(t, lookup(&Node{Input: wire.OutPoint{Index: 9}}))
+
+	pkg := []byte{0x01, 0x02}
+	assetCtx.SetSealedPackage(input, pkg)
+	pkg[0] = 0xEE
+	gotPkg := assetCtx.SealedPackage(input)
+	require.Equal(t, []byte{0x01, 0x02}, gotPkg)
+	gotPkg[0] = 0xDD
+	require.Equal(t, byte(0x01), assetCtx.SealedPackage(input)[0])
+
+	root := []byte{0x03, 0x04}
+	assetCtx.SetLeafAssetRoot(input, root)
+	root[0] = 0xCC
+	gotRoot := assetCtx.LeafAssetRoot(input)
+	require.Equal(t, []byte{0x03, 0x04}, gotRoot)
+	gotRoot[0] = 0xBB
+	require.Equal(t, byte(0x03), assetCtx.LeafAssetRoot(input)[0])
+	require.False(t, assetCtx.IsEmpty())
+}
+
+// TestAssetTreeContextEmpty tests empty context detection.
+func TestAssetTreeContextEmpty(t *testing.T) {
+	t.Parallel()
+
+	var nilContext *AssetTreeContext
+	require.True(t, nilContext.IsEmpty())
+
+	assetCtx := NewAssetTreeContext()
+	require.True(t, assetCtx.IsEmpty())
+
+	assetCtx.SetLeafAssetRoot(wire.OutPoint{Index: 1}, []byte{0x01})
+	require.False(t, assetCtx.IsEmpty())
+}
+
+// TestAssetTreeContextAmountsByIdentity tests node identity keying.
+func TestAssetTreeContextAmountsByIdentity(t *testing.T) {
+	t.Parallel()
+
+	assetCtx := NewAssetTreeContext()
+	_, key := createTestKey(t)
+	nodeA := &Node{CoSigners: []*btcec.PublicKey{key}}
+	nodeB := &Node{CoSigners: []*btcec.PublicKey{key}}
+
+	assetCtx.SetNodeAssetAmount(nodeA, 42)
+	require.EqualValues(t, 42, assetCtx.NodeAssetAmount(nodeA))
+	require.Zero(t, assetCtx.NodeAssetAmount(nodeB))
+}

lib/tree/leaf.go

@@ -19,4 +19,7 @@ type LeafDescriptor struct {
 	// CoSignerKey is the public key of the leaf owner who must participate
 	// in signing this leaf's transaction along with the operator.
 	CoSignerKey *btcec.PublicKey
+
+	// AssetAmount is the number of asset units in the leaf output.
+	AssetAmount uint64
 }