feat: added partial sequence helper

Author: f7f376a1fcd0d0e11a10ed1b6577c9

dag/dag.go

@@ -709,9 +709,137 @@ func (d *Dag) GetPartial(start, end int) (*Dag, error) {
 	return partialDag, nil
 }
 
+// getPartialLeafSequence returns an ordered sequence of leaves for transmission from a partial DAG
+// This is an internal method used by GetLeafSequence when dealing with partial DAGs
+func (d *Dag) getPartialLeafSequence() []*TransmissionPacket {
+	var sequence []*TransmissionPacket
+
+	// Get the root leaf
+	rootLeaf := d.Leafs[d.Root]
+	if rootLeaf == nil {
+		return sequence // Return empty sequence if root leaf is missing
+	}
+
+	// First, build a map of proofs organized by parent hash and child hash
+	// This will allow us to look up the proof for a specific child when creating its packet
+	proofMap := make(map[string]map[string]*ClassicTreeBranch)
+
+	// Populate the proof map from all leaves in the partial DAG
+	for _, leaf := range d.Leafs {
+		if len(leaf.Proofs) > 0 {
+			// Create an entry for this parent if it doesn't exist
+			if _, exists := proofMap[leaf.Hash]; !exists {
+				proofMap[leaf.Hash] = make(map[string]*ClassicTreeBranch)
+			}
+
+			// Add all proofs from this leaf to the map
+			for childHash, proof := range leaf.Proofs {
+				proofMap[leaf.Hash][childHash] = proof
+			}
+		}
+	}
+
+	// Now perform BFS traversal similar to the full DAG method
+	visited := make(map[string]bool)
+
+	// Start with the root
+	rootLeafClone := rootLeaf.Clone()
+
+	// We need to preserve the original links for the root leaf
+	// because they're part of its identity and hash calculation
+	originalLinks := make(map[string]string)
+	for k, v := range rootLeaf.Links {
+		originalLinks[k] = v
+	}
+
+	// Clear links for transmission (they'll be reconstructed on the receiving end)
+	rootLeafClone.Links = make(map[string]string)
+
+	// We need to preserve these fields for verification
+	// but clear proofs for the root packet - they'll be sent with child packets
+	originalMerkleRoot := rootLeafClone.ClassicMerkleRoot
+	originalLatestLabel := rootLeafClone.LatestLabel
+	originalLeafCount := rootLeafClone.LeafCount
+
+	rootLeafClone.Proofs = nil
+
+	// Restore the critical fields
+	rootLeafClone.ClassicMerkleRoot = originalMerkleRoot
+	rootLeafClone.LatestLabel = originalLatestLabel
+	rootLeafClone.LeafCount = originalLeafCount
+
+	rootPacket := &TransmissionPacket{
+		Leaf:       rootLeafClone,
+		ParentHash: "", // Root has no parent
+		Proofs:     make(map[string]*ClassicTreeBranch),
+	}
+	sequence = append(sequence, rootPacket)
+	visited[d.Root] = true
+
+	// Restore the original links for the root leaf in the DAG
+	rootLeaf.Links = originalLinks
+
+	// BFS traversal
+	queue := []string{d.Root}
+	for len(queue) > 0 {
+		current := queue[0]
+		queue = queue[1:]
+
+		currentLeaf := d.Leafs[current]
+
+		// Sort links for deterministic order
+		var sortedLinks []string
+		for _, link := range currentLeaf.Links {
+			sortedLinks = append(sortedLinks, link)
+		}
+		sort.Strings(sortedLinks)
+
+		// Process each child
+		for _, childHash := range sortedLinks {
+			if !visited[childHash] {
+				childLeaf := d.Leafs[childHash]
+				if childLeaf == nil {
+					continue // Skip if child leaf doesn't exist in this partial DAG
+				}
+
+				// Clone the leaf and clear its links for transmission
+				leafClone := childLeaf.Clone()
+				leafClone.Links = make(map[string]string)
+				leafClone.Proofs = nil // Clear proofs from the leaf
+
+				packet := &TransmissionPacket{
+					Leaf:       leafClone,
+					ParentHash: current,
+					Proofs:     make(map[string]*ClassicTreeBranch),
+				}
+
+				// Add the proof for this specific child from the proof map
+				if parentProofs, exists := proofMap[current]; exists {
+					if proof, hasProof := parentProofs[childHash]; hasProof {
+						packet.Proofs[childHash] = proof
+					}
+				}
+
+				sequence = append(sequence, packet)
+				visited[childHash] = true
+				queue = append(queue, childHash)
+			}
+		}
+	}
+
+	return sequence
+}
+
 // GetLeafSequence returns an ordered sequence of leaves for transmission
 // Each packet contains a leaf, its parent hash, and any proofs needed for verification
 func (d *Dag) GetLeafSequence() []*TransmissionPacket {
+	// Check if this is a partial DAG
+	if d.IsPartial() {
+		// Use specialized method for partial DAGs
+		return d.getPartialLeafSequence()
+	}
+
+	// Original implementation for complete DAGs
 	var sequence []*TransmissionPacket
 	visited := make(map[string]bool)
 

dag/leaves.go

@@ -516,6 +516,7 @@ func (leaf *DagLeaf) Clone() *DagLeaf {
 		CurrentLinkCount:  leaf.CurrentLinkCount,
 		LatestLabel:       leaf.LatestLabel,
 		LeafCount:         leaf.LeafCount,
+		ParentHash:        leaf.ParentHash,
 		Links:             make(map[string]string),
 		AdditionalData:    make(map[string]string),
 		Proofs:            make(map[string]*ClassicTreeBranch),
@@ -534,38 +535,8 @@ func (leaf *DagLeaf) Clone() *DagLeaf {
 		}
 	}
 
-	// Copy root-specific fields if this is the root leaf
-	if leaf.Hash == leaf.ParentHash || leaf.Hash == GetHash(leaf.Hash) {
-		cloned.LatestLabel = leaf.LatestLabel
-		cloned.LeafCount = leaf.LeafCount
-		cloned.ParentHash = cloned.Hash // Root is its own parent
-	} else {
-		cloned.ParentHash = leaf.ParentHash
-	}
-
-	// If leaf has multiple children according to CurrentLinkCount,
-	// we need to handle its merkle tree state
-	if leaf.CurrentLinkCount > 1 {
-		if len(leaf.Links) > 1 {
-			// Build merkle tree with current links
-			builder := merkle_tree.CreateTree()
-			for l, h := range leaf.Links {
-				builder.AddLeaf(l, h)
-			}
-			merkleTree, leafMap, err := builder.Build()
-			if err == nil {
-				cloned.MerkleTree = merkleTree
-				cloned.LeafMap = leafMap
-				cloned.ClassicMerkleRoot = merkleTree.Root
-			}
-		} else {
-			// Clear merkle tree if we don't have enough links to rebuild it
-			cloned.MerkleTree = nil
-			cloned.LeafMap = nil
-			// But keep ClassicMerkleRoot as it's part of the leaf's identity
-			cloned.ClassicMerkleRoot = leaf.ClassicMerkleRoot
-		}
-	}
+	// MerkleTree and LeafMap are not deep-copied because they're regenerated when needed
+	// But we preserve the ClassicMerkleRoot which is part of the leaf's identity
 
 	return cloned
 }

dag/transmission_test.go

@@ -75,3 +75,83 @@ func TestLeafByLeafTransmission(t *testing.T) {
 
 	t.Log("Successfully verified full DAG after discarding proofs")
 }
+
+func TestPartialDagTransmission(t *testing.T) {
+	testDir, err := os.MkdirTemp("", "dag_partial_transmission_test_*")
+	if err != nil {
+		t.Fatalf("Failed to create temp directory: %v", err)
+	}
+	defer os.RemoveAll(testDir)
+
+	GenerateDummyDirectory(filepath.Join(testDir, "input"), 3, 5, 2, 4)
+
+	originalDag, err := CreateDag(filepath.Join(testDir, "input"), true)
+	if err != nil {
+		t.Fatalf("Failed to create DAG: %v", err)
+	}
+
+	err = originalDag.Verify()
+	if err != nil {
+		t.Fatalf("Original DAG verification failed: %v", err)
+	}
+
+	partialDag, err := originalDag.GetPartial(0, 3)
+	if err != nil {
+		t.Fatalf("Failed to get partial DAG: %v", err)
+	}
+
+	err = partialDag.Verify()
+	if err != nil {
+		t.Fatalf("Partial DAG verification failed: %v", err)
+	}
+
+	if !partialDag.IsPartial() {
+		t.Fatal("DAG not recognized as partial")
+	}
+
+	sequence := partialDag.GetLeafSequence()
+	if len(sequence) == 0 {
+		t.Fatal("No transmission packets generated from partial DAG")
+	}
+
+	t.Logf("Generated %d transmission packets from partial DAG with %d leaves",
+		len(sequence), len(partialDag.Leafs))
+
+	receiverDag := &Dag{
+		Root:  partialDag.Root,
+		Leafs: make(map[string]*DagLeaf),
+	}
+
+	for i, p := range sequence {
+		bytes, err := p.ToCBOR()
+		if err != nil {
+			t.Fatalf("Failed to serialize packet from partial DAG")
+		}
+
+		packet, err := TransmissionPacketFromCBOR(bytes)
+		if err != nil {
+			t.Fatalf("Failed to deserialize packet from partial DAG")
+		}
+
+		receiverDag.ApplyTransmissionPacket(packet)
+
+		err = receiverDag.Verify()
+		if err != nil {
+			t.Fatalf("Verification failed after packet %d from partial DAG: %v", i, err)
+		}
+
+		t.Logf("Successfully verified after packet %d from partial DAG, DAG now has %d leaves",
+			i, len(receiverDag.Leafs))
+	}
+
+	if len(receiverDag.Leafs) != len(partialDag.Leafs) {
+		t.Fatalf("Receiver DAG has %d leaves, expected %d (same as partial DAG)",
+			len(receiverDag.Leafs), len(partialDag.Leafs))
+	}
+
+	if !receiverDag.IsPartial() {
+		t.Fatal("Reconstructed DAG not recognized as partial")
+	}
+
+	t.Log("Successfully transmitted and verified partial DAG")
+}