Skip to content

Commit d33ee59

Browse files
feat: partial transmission helpers
1 parent a86abc8 commit d33ee59

6 files changed

+345
-14
lines changed

dag/dag.go

+101-12
Original file line numberDiff line numberDiff line change
@@ -464,14 +464,16 @@ func (d *Dag) IterateDag(processLeaf func(leaf *DagLeaf, parent *DagLeaf) error)
464464
return iterate(d.Root, nil)
465465
}
466466

467-
// IsPartial returns true if this DAG is a partial DAG (has pruned links)
467+
// IsPartial returns true if this DAG is a partial DAG (has fewer leaves than the total count)
468468
func (d *Dag) IsPartial() bool {
469-
for _, leaf := range d.Leafs {
470-
if len(leaf.Links) < leaf.CurrentLinkCount {
471-
return true
472-
}
469+
// Get the root leaf
470+
rootLeaf := d.Leafs[d.Root]
471+
if rootLeaf == nil {
472+
return true // If root leaf is missing, it's definitely partial
473473
}
474-
return false
474+
475+
// Check if the number of leaves in the DAG matches the total leaf count
476+
return len(d.Leafs) < rootLeaf.LeafCount
475477
}
476478

477479
// pruneIrrelevantLinks removes links that aren't needed for partial verification
@@ -707,11 +709,98 @@ func (d *Dag) GetPartial(start, end int) (*Dag, error) {
707709
return partialDag, nil
708710
}
709711

710-
// Helper function to get keys from a map for debugging
711-
func getMapKeys(m map[string]*ClassicTreeBranch) []string {
712-
keys := make([]string, 0, len(m))
713-
for k := range m {
714-
keys = append(keys, k)
712+
// GetLeafSequence returns an ordered sequence of leaves for transmission
713+
// Each packet contains a leaf, its parent hash, and any proofs needed for verification
714+
func (d *Dag) GetLeafSequence() []*TransmissionPacket {
715+
var sequence []*TransmissionPacket
716+
visited := make(map[string]bool)
717+
718+
rootLeaf := d.Leafs[d.Root]
719+
if rootLeaf == nil {
720+
return sequence
721+
}
722+
723+
totalLeafCount := rootLeaf.LeafCount
724+
725+
rootLeafClone := rootLeaf.Clone()
726+
rootLeafClone.Links = make(map[string]string)
727+
728+
rootPacket := &TransmissionPacket{
729+
Leaf: rootLeafClone,
730+
ParentHash: "",
731+
Proofs: make(map[string]*ClassicTreeBranch),
732+
}
733+
sequence = append(sequence, rootPacket)
734+
visited[d.Root] = true
735+
736+
queue := []string{d.Root}
737+
for len(queue) > 0 && len(sequence) <= totalLeafCount {
738+
current := queue[0]
739+
queue = queue[1:]
740+
741+
currentLeaf := d.Leafs[current]
742+
743+
var sortedLinks []string
744+
for _, link := range currentLeaf.Links {
745+
sortedLinks = append(sortedLinks, link)
746+
}
747+
sort.Strings(sortedLinks)
748+
749+
for _, childHash := range sortedLinks {
750+
if !visited[childHash] && len(sequence) <= totalLeafCount {
751+
branch, err := d.buildVerificationBranch(d.Leafs[childHash])
752+
if err != nil {
753+
continue
754+
}
755+
756+
leafClone := d.Leafs[childHash].Clone()
757+
leafClone.Links = make(map[string]string)
758+
759+
packet := &TransmissionPacket{
760+
Leaf: leafClone,
761+
ParentHash: current,
762+
Proofs: make(map[string]*ClassicTreeBranch),
763+
}
764+
765+
for _, pathNode := range branch.Path {
766+
if pathNode.Proofs != nil {
767+
for k, v := range pathNode.Proofs {
768+
packet.Proofs[k] = v
769+
}
770+
}
771+
}
772+
773+
sequence = append(sequence, packet)
774+
visited[childHash] = true
775+
queue = append(queue, childHash)
776+
}
777+
}
778+
}
779+
780+
return sequence
781+
}
782+
783+
func (d *Dag) ApplyTransmissionPacket(packet *TransmissionPacket) {
784+
d.Leafs[packet.Leaf.Hash] = packet.Leaf
785+
786+
if packet.ParentHash != "" {
787+
if parent, exists := d.Leafs[packet.ParentHash]; exists {
788+
label := GetLabel(packet.Leaf.Hash)
789+
if label != "" {
790+
parent.Links[label] = packet.Leaf.Hash
791+
}
792+
}
793+
}
794+
795+
for leafHash, proof := range packet.Proofs {
796+
for _, leaf := range d.Leafs {
797+
if leaf.HasLink(leafHash) {
798+
if leaf.Proofs == nil {
799+
leaf.Proofs = make(map[string]*ClassicTreeBranch)
800+
}
801+
leaf.Proofs[leafHash] = proof
802+
break
803+
}
804+
}
715805
}
716-
return keys
717806
}

dag/leaves.go

+2-2
Original file line numberDiff line numberDiff line change
@@ -42,8 +42,8 @@ func (b *DagLeafBuilder) AddLink(label string, hash string) {
4242
}
4343

4444
func (b *DagBuilder) GetLatestLabel() string {
45-
var result string = "1"
46-
var latestLabel int64 = 1
45+
var result string = "0"
46+
var latestLabel int64 = 0
4747
for hash := range b.Leafs {
4848
label := GetLabel(hash)
4949

dag/serialize.go

+102
Original file line numberDiff line numberDiff line change
@@ -29,6 +29,13 @@ type SerializableDagLeaf struct {
2929
StoredProofs map[string]*ClassicTreeBranch `json:"stored_proofs,omitempty" cbor:"stored_proofs,omitempty"`
3030
}
3131

32+
// SerializableTransmissionPacket is a minimal version of TransmissionPacket for efficient serialization
33+
type SerializableTransmissionPacket struct {
34+
Leaf *SerializableDagLeaf
35+
ParentHash string
36+
Proofs map[string]*ClassicTreeBranch `json:"proofs,omitempty" cbor:"proofs,omitempty"`
37+
}
38+
3239
// ToSerializable converts a Dag to its serializable form
3340
func (dag *Dag) ToSerializable() *SerializableDag {
3441
serializable := &SerializableDag{
@@ -187,3 +194,98 @@ func FromJSON(data []byte) (*Dag, error) {
187194
}
188195
return FromSerializable(&serializable), nil
189196
}
197+
198+
// ToSerializable converts a TransmissionPacket to its serializable form
199+
func (packet *TransmissionPacket) ToSerializable() *SerializableTransmissionPacket {
200+
serializable := &SerializableTransmissionPacket{
201+
Leaf: packet.Leaf.ToSerializable(),
202+
ParentHash: packet.ParentHash,
203+
Proofs: make(map[string]*ClassicTreeBranch),
204+
}
205+
206+
// Copy proofs
207+
if packet.Proofs != nil {
208+
for k, v := range packet.Proofs {
209+
serializable.Proofs[k] = v
210+
}
211+
}
212+
213+
return serializable
214+
}
215+
216+
// TransmissionPacketFromSerializable reconstructs a TransmissionPacket from its serializable form
217+
func TransmissionPacketFromSerializable(s *SerializableTransmissionPacket) *TransmissionPacket {
218+
// Create a DagLeaf from the serializable leaf
219+
leaf := &DagLeaf{
220+
Hash: s.Leaf.Hash,
221+
ItemName: s.Leaf.ItemName,
222+
Type: s.Leaf.Type,
223+
ContentHash: s.Leaf.ContentHash,
224+
Content: s.Leaf.Content,
225+
ClassicMerkleRoot: s.Leaf.ClassicMerkleRoot,
226+
CurrentLinkCount: s.Leaf.CurrentLinkCount,
227+
LatestLabel: s.Leaf.LatestLabel,
228+
LeafCount: s.Leaf.LeafCount,
229+
Links: make(map[string]string),
230+
AdditionalData: make(map[string]string),
231+
Proofs: make(map[string]*ClassicTreeBranch),
232+
}
233+
234+
// Copy and sort links
235+
leaf.Links = sortMapByKeys(s.Leaf.Links)
236+
237+
// Copy and sort additional data
238+
leaf.AdditionalData = sortMapByKeys(s.Leaf.AdditionalData)
239+
240+
// Copy stored proofs
241+
if s.Leaf.StoredProofs != nil {
242+
for k, v := range s.Leaf.StoredProofs {
243+
leaf.Proofs[k] = v
244+
}
245+
}
246+
247+
packet := &TransmissionPacket{
248+
Leaf: leaf,
249+
ParentHash: s.ParentHash,
250+
Proofs: make(map[string]*ClassicTreeBranch),
251+
}
252+
253+
// Copy proofs
254+
if s.Proofs != nil {
255+
for k, v := range s.Proofs {
256+
packet.Proofs[k] = v
257+
}
258+
}
259+
260+
return packet
261+
}
262+
263+
// ToCBOR serializes a TransmissionPacket to CBOR format
264+
func (packet *TransmissionPacket) ToCBOR() ([]byte, error) {
265+
serializable := packet.ToSerializable()
266+
return cbor.Marshal(serializable)
267+
}
268+
269+
// ToJSON serializes a TransmissionPacket to JSON format
270+
func (packet *TransmissionPacket) ToJSON() ([]byte, error) {
271+
serializable := packet.ToSerializable()
272+
return json.MarshalIndent(serializable, "", " ")
273+
}
274+
275+
// TransmissionPacketFromCBOR deserializes a TransmissionPacket from CBOR format
276+
func TransmissionPacketFromCBOR(data []byte) (*TransmissionPacket, error) {
277+
var serializable SerializableTransmissionPacket
278+
if err := cbor.Unmarshal(data, &serializable); err != nil {
279+
return nil, err
280+
}
281+
return TransmissionPacketFromSerializable(&serializable), nil
282+
}
283+
284+
// TransmissionPacketFromJSON deserializes a TransmissionPacket from JSON format
285+
func TransmissionPacketFromJSON(data []byte) (*TransmissionPacket, error) {
286+
var serializable SerializableTransmissionPacket
287+
if err := json.Unmarshal(data, &serializable); err != nil {
288+
return nil, err
289+
}
290+
return TransmissionPacketFromSerializable(&serializable), nil
291+
}

dag/serialize_test.go

+57
Original file line numberDiff line numberDiff line change
@@ -131,4 +131,61 @@ func TestSerialization(t *testing.T) {
131131
t.Errorf("Failed to recreate directory from deserialized DAG: %v", err)
132132
}
133133
})
134+
135+
t.Run("TransmissionPacket", func(t *testing.T) {
136+
// Get a sequence of transmission packets
137+
sequence := originalDag.GetLeafSequence()
138+
if len(sequence) == 0 {
139+
t.Fatal("No transmission packets generated")
140+
}
141+
142+
// Test the first packet
143+
packet := sequence[0]
144+
145+
// Serialize to JSON
146+
jsonData, err := packet.ToJSON()
147+
if err != nil {
148+
t.Fatalf("Failed to serialize TransmissionPacket to JSON: %v", err)
149+
}
150+
151+
// Deserialize from JSON
152+
deserializedPacket, err := TransmissionPacketFromJSON(jsonData)
153+
if err != nil {
154+
t.Fatalf("Failed to deserialize TransmissionPacket from JSON: %v", err)
155+
}
156+
157+
// Verify the deserialized packet
158+
if packet.Leaf.Hash != deserializedPacket.Leaf.Hash {
159+
t.Errorf("Leaf hash mismatch: expected %s, got %s", packet.Leaf.Hash, deserializedPacket.Leaf.Hash)
160+
}
161+
if packet.ParentHash != deserializedPacket.ParentHash {
162+
t.Errorf("Parent hash mismatch: expected %s, got %s", packet.ParentHash, deserializedPacket.ParentHash)
163+
}
164+
if len(packet.Proofs) != len(deserializedPacket.Proofs) {
165+
t.Errorf("Proofs count mismatch: expected %d, got %d", len(packet.Proofs), len(deserializedPacket.Proofs))
166+
}
167+
168+
// Serialize to CBOR
169+
cborData, err := packet.ToCBOR()
170+
if err != nil {
171+
t.Fatalf("Failed to serialize TransmissionPacket to CBOR: %v", err)
172+
}
173+
174+
// Deserialize from CBOR
175+
deserializedPacket, err = TransmissionPacketFromCBOR(cborData)
176+
if err != nil {
177+
t.Fatalf("Failed to deserialize TransmissionPacket from CBOR: %v", err)
178+
}
179+
180+
// Verify the deserialized packet
181+
if packet.Leaf.Hash != deserializedPacket.Leaf.Hash {
182+
t.Errorf("Leaf hash mismatch: expected %s, got %s", packet.Leaf.Hash, deserializedPacket.Leaf.Hash)
183+
}
184+
if packet.ParentHash != deserializedPacket.ParentHash {
185+
t.Errorf("Parent hash mismatch: expected %s, got %s", packet.ParentHash, deserializedPacket.ParentHash)
186+
}
187+
if len(packet.Proofs) != len(deserializedPacket.Proofs) {
188+
t.Errorf("Proofs count mismatch: expected %d, got %d", len(packet.Proofs), len(deserializedPacket.Proofs))
189+
}
190+
})
134191
}

dag/transmission_test.go

+77
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,77 @@
1+
package dag
2+
3+
import (
4+
"os"
5+
"path/filepath"
6+
"testing"
7+
)
8+
9+
func TestLeafByLeafTransmission(t *testing.T) {
10+
testDir, err := os.MkdirTemp("", "dag_transmission_test_*")
11+
if err != nil {
12+
t.Fatalf("Failed to create temp directory: %v", err)
13+
}
14+
defer os.RemoveAll(testDir)
15+
16+
GenerateDummyDirectory(filepath.Join(testDir, "input"), 3, 5, 2, 3)
17+
18+
originalDag, err := CreateDag(filepath.Join(testDir, "input"), true)
19+
if err != nil {
20+
t.Fatalf("Failed to create DAG: %v", err)
21+
}
22+
23+
err = originalDag.Verify()
24+
if err != nil {
25+
t.Fatalf("Original DAG verification failed: %v", err)
26+
}
27+
28+
sequence := originalDag.GetLeafSequence()
29+
30+
if len(sequence) == 0 {
31+
t.Fatal("No transmission packets generated")
32+
}
33+
34+
t.Logf("Generated %d transmission packets", len(sequence))
35+
36+
receiverDag := &Dag{
37+
Root: originalDag.Root,
38+
Leafs: make(map[string]*DagLeaf),
39+
}
40+
41+
for i, p := range sequence {
42+
bytes, err := p.ToCBOR()
43+
if err != nil {
44+
t.Fatalf("Failed to serialize packet")
45+
}
46+
47+
packet, err := TransmissionPacketFromCBOR(bytes)
48+
if err != nil {
49+
t.Fatalf("Failed to deserialize packet")
50+
}
51+
52+
receiverDag.ApplyTransmissionPacket(packet)
53+
54+
err = receiverDag.Verify()
55+
if err != nil {
56+
t.Fatalf("Verification failed after packet %d: %v", i, err)
57+
}
58+
59+
t.Logf("Successfully verified after packet %d, DAG now has %d leaves", i, len(receiverDag.Leafs))
60+
}
61+
62+
if len(receiverDag.Leafs) != len(originalDag.Leafs) {
63+
t.Fatalf("Receiver DAG has %d leaves, expected %d",
64+
len(receiverDag.Leafs), len(originalDag.Leafs))
65+
}
66+
67+
for _, leaf := range receiverDag.Leafs {
68+
leaf.Proofs = nil
69+
}
70+
71+
err = receiverDag.Verify()
72+
if err != nil {
73+
t.Fatalf("Full DAG verification after discarding proofs failed: %v", err)
74+
}
75+
76+
t.Log("Successfully verified full DAG after discarding proofs")
77+
}

0 commit comments

Comments
 (0)