untangle marshal

basics_test.go

@@ -88,6 +88,18 @@ func validate(t *testing.T, actual ipld.Node, expected *pb.PBNode) {
 func runTest(t *testing.T, bytsHex string, expected *pb.PBNode) {
 	byts, _ := hex.DecodeString(bytsHex)
 
+	roundTrip := func(t *testing.T, node ipld.Node) {
+		var buf bytes.Buffer
+		if err := Marshal(node, &buf); err != nil {
+			t.Fatal(err)
+		}
+
+		// fmt.Printf("CMP\n\tFrom: %v\n\tTo:   %v\n", hex.EncodeToString(byts), hex.EncodeToString(buf.Bytes()))
+		if bytes.Compare(buf.Bytes(), byts) != 0 {
+			t.Fatal("Round-trip resulted in different bytes")
+		}
+	}
+
 	t.Run("basicnode", func(t *testing.T) {
 		nb := basicnode.Prototype__Map{}.NewBuilder()
 		err := Unmarshal(nb, bytes.NewReader(byts))
@@ -97,15 +109,7 @@ func runTest(t *testing.T, bytsHex string, expected *pb.PBNode) {
 
 		node := nb.Build()
 		validate(t, node, expected)
-
-		var buf bytes.Buffer
-		if err := Marshal(node, &buf); err != nil {
-			t.Fatal(err)
-		}
-
-		if bytes.Compare(buf.Bytes(), byts) != 0 {
-			t.Fatal("Round-trip resulted in different bytes")
-		}
+		roundTrip(t, node)
 	})
 
 	t.Run("typed", func(t *testing.T) {
@@ -114,7 +118,9 @@ func runTest(t *testing.T, bytsHex string, expected *pb.PBNode) {
 		if err != nil {
 			t.Fatal(err)
 		}
-		validate(t, nb.Build(), expected)
+		node := nb.Build()
+		validate(t, node, expected)
+		roundTrip(t, node)
 	})
 }
 

marshal.go

@@ -3,133 +3,195 @@ package dagpb
 import (
 	"fmt"
 	"io"
+	math_bits "math/bits"
+	"sort"
 
+	"github.com/ipfs/go-cid"
 	ipld "github.com/ipld/go-ipld-prime"
 	cidlink "github.com/ipld/go-ipld-prime/linking/cid"
-	pb "github.com/rvagg/go-dagpb/pb"
 )
 
+type pbLink struct {
+	hash     cid.Cid
+	name     string
+	hasName  bool
+	tsize    uint64
+	hasTsize bool
+}
+
 func Marshal(inNode ipld.Node, out io.Writer) error {
+	// Wrap in a typed node for some basic schema form checking
 	builder := Type.PBNode.NewBuilder()
 	if err := builder.AssignNode(inNode); err != nil {
 		return err
 	}
 	node := builder.Build()
 
-	curField := pb.TypeLinks
-	var linksIter ipld.ListIterator
-	var link ipld.Node
-
-	tokenSource := func() (pb.Token, error) {
-		if curField == pb.TypeLinks {
 	links, err := node.LookupByString("Links")
 	if err != nil {
-				return pb.Token{}, err
-			}
-			if links.Length() == 0 {
-				curField = pb.TypeData
-			} else {
-				curField = pb.TypeHash
-				linksIter = links.ListIterator()
-				_, link, err = linksIter.Next()
-				if err != nil {
-					return pb.Token{}, err
-				}
-			}
+		return err
 	}
+	if links.Length() > 0 {
+		// collect links into a slice so we can properly sort for encoding
+		pbLinks := make([]pbLink, links.Length())
 
-		if curField == pb.TypeData {
-			curField = pb.TypeEnd
-
-			d, err := node.LookupByString("Data")
-			if err != nil {
-				return pb.Token{}, err
-			}
-			if !d.IsAbsent() {
-				b, err := d.AsBytes()
+		linksIter := links.ListIterator()
+		for !linksIter.Done() {
+			ii, link, err := linksIter.Next()
 			if err != nil {
-					return pb.Token{}, err
-				}
-				return pb.Token{Type: pb.TypeData, Bytes: b}, nil
-			}
-		}
-
-		if curField == pb.TypeEnd {
-			return pb.Token{Type: pb.TypeEnd}, nil
+				return err
 			}
 
-		for {
-			if curField == pb.TypeHash {
-				curField = pb.TypeName
-
+			{ // Hash
 				d, err := link.LookupByString("Hash")
 				if err != nil {
-					return pb.Token{}, err
+					return err
 				}
+				// TODO:
+				// 		return 0, fmt.Errorf("invalid DAG-PB form (link must have a Hash)")
 				l, err := d.AsLink()
 				if err != nil {
-					return pb.Token{}, err
+					return err
+				}
+				if err != nil {
+					return err
 				}
-				if cl, ok := l.(cidlink.Link); ok {
-					return pb.Token{Type: pb.TypeHash, Cid: &cl.Cid}, nil
+				cl, ok := l.(cidlink.Link)
+				if !ok {
+					return fmt.Errorf("unexpected Link type [%v]", l)
 				}
-				return pb.Token{}, fmt.Errorf("unexpected Link type [%v]", l)
+				pbLinks[ii].hash = cl.Cid
 			}
 
-			if curField == pb.TypeName {
-				curField = pb.TypeTSize
-
+			{ // Name
 				nameNode, err := link.LookupByString("Name")
 				if err != nil {
-					return pb.Token{}, err
+					return err
 				}
 				if !nameNode.IsAbsent() {
 					name, err := nameNode.AsString()
 					if err != nil {
-						return pb.Token{}, err
+						return err
 					}
-					return pb.Token{Type: pb.TypeName, Bytes: []byte(name)}, nil
+					pbLinks[ii].name = name
+					pbLinks[ii].hasName = true
 				}
 			}
 
-			if curField == pb.TypeTSize {
-				curField = pb.TypeLinkEnd
-
+			{ // Tsize
 				tsizeNode, err := link.LookupByString("Tsize")
 				if err != nil {
-					return pb.Token{}, err
+					return err
 				}
 				if !tsizeNode.IsAbsent() {
 					tsize, err := tsizeNode.AsInt()
 					if err != nil {
-						return pb.Token{}, err
+						return err
 					}
 					if tsize < 0 {
-						return pb.Token{}, fmt.Errorf("Link has negative Tsize value [%v]", tsize)
+						return fmt.Errorf("Link has negative Tsize value [%v]", tsize)
+					}
+					utsize := uint64(tsize)
+					pbLinks[ii].tsize = utsize
+					pbLinks[ii].hasTsize = true
+				}
+			}
+		} // for
+
+		sortLinks(pbLinks)
+		for _, link := range pbLinks {
+			size := link.encodedSize()
+			chunk := make([]byte, size+sizeOfVarint(uint64(size))+1)
+			chunk[0] = 0x12
+			offset := encodeVarint(chunk, 1, uint64(size))
+			wrote, err := link.marshal(chunk, offset)
+			if err != nil {
+				return err
 			}
-					return pb.Token{Type: pb.TypeTSize, Int: uint64(tsize)}, nil
+			if wrote != size {
+				return fmt.Errorf("bad PBLink marshal, wrote wrong number of bytes")
 			}
+			out.Write(chunk)
 		}
+	} // if links
 
-			if curField == pb.TypeLinkEnd {
-				if linksIter.Done() {
-					curField = pb.TypeData
-				} else {
-					curField = pb.TypeHash
-					var err error
-					_, link, err = linksIter.Next()
+	data, err := node.LookupByString("Data")
 	if err != nil {
-						return pb.Token{}, err
+		return err
 	}
+	if !data.IsAbsent() {
+		byts, err := data.AsBytes()
+		if err != nil {
+			return err
 		}
-				return pb.Token{Type: pb.TypeLinkEnd}, nil
+		size := uint64(len(byts))
+		lead := make([]byte, sizeOfVarint(size)+1)
+		lead[0] = 0xa
+		encodeVarint(lead, 1, size)
+		out.Write(lead)
+		out.Write(byts)
 	}
 
-			if curField != pb.TypeHash {
-				return pb.Token{}, fmt.Errorf("unexpected and invalid token state")
+	return nil
+}
+
+func (link pbLink) encodedSize() (n int) {
+	l := link.hash.ByteLen()
+	n += 1 + l + sizeOfVarint(uint64(l))
+	if link.hasName {
+		l = len(link.name)
+		n += 1 + l + sizeOfVarint(uint64(l))
 	}
+	if link.hasTsize {
+		n += 1 + sizeOfVarint(uint64(link.tsize))
 	}
+	return n
+}
+
+func (link pbLink) marshal(data []byte, offset int) (int, error) {
+	base := offset
+	data[offset] = 0xa
+	byts := link.hash.Bytes()
+	offset = encodeVarint(data, offset+1, uint64(len(byts)))
+	copy(data[offset:], byts)
+	offset += len(byts)
+	if link.hasName {
+		data[offset] = 0x12
+		offset = encodeVarint(data, offset+1, uint64(len(link.name)))
+		copy(data[offset:], link.name)
+		offset += len(link.name)
+	}
+	if link.hasTsize {
+		data[offset] = 0x18
+		offset = encodeVarint(data, offset+1, uint64(link.tsize))
+	}
+	return offset - base, nil
+}
+
+func encodeVarint(data []byte, offset int, v uint64) int {
+	for v >= 1<<7 {
+		data[offset] = uint8(v&0x7f | 0x80)
+		v >>= 7
+		offset++
 	}
+	data[offset] = uint8(v)
+	return offset + 1
+}
+
+func sortLinks(links []pbLink) {
+	sort.Stable(pbLinkSlice(links))
+}
+
+type pbLinkSlice []pbLink
+
+func (ls pbLinkSlice) Len() int           { return len(ls) }
+func (ls pbLinkSlice) Swap(a, b int)      { ls[a], ls[b] = ls[b], ls[a] }
+func (ls pbLinkSlice) Less(a, b int) bool { return pbLinkLess(ls[a], ls[b]) }
+
+func pbLinkLess(a pbLink, b pbLink) bool {
+	return a.name < b.name
+}
 
-	return pb.Marshal(out, tokenSource)
+func sizeOfVarint(x uint64) (n int) {
+	return (math_bits.Len64(x|1) + 6) / 7
 }

pb/dagpb.go

@@ -60,7 +60,11 @@ func NewPBLink(name string, c cid.Cid, tsize uint64) *PBLink {
 }
 
 func (node *PBNode) SortLinks() {
-	sort.Stable(pbLinkSlice(node.Links))
+	SortLinks(node.Links)
+}
+
+func SortLinks(links []*PBLink) {
+	sort.Stable(pbLinkSlice(links))
 }
 
 type pbLinkSlice []*PBLink

pb/marshal.go

@@ -11,9 +11,9 @@ import (
 // Marshal TODO
 func Marshal(out io.Writer, tokenSource func() (Token, error)) error {
 	writeLead := func(wire byte, size uint64) {
-		lead := make([]byte, sizeOfVarint(size)+1)
+		lead := make([]byte, SizeOfVarint(size)+1)
 		lead[0] = wire
-		encodeVarint(lead, len(lead), size)
+		EncodeVarint(lead, len(lead), size)
 		out.Write(lead)
 	}
 
@@ -33,10 +33,10 @@ func Marshal(out io.Writer, tokenSource func() (Token, error)) error {
 			writeLead(0xa, uint64(len(tok.Bytes)))
 			out.Write(tok.Bytes)
 		case TypeLinkEnd:
-			l := link.size()
+			l := link.EncodedSize()
 			writeLead(0x12, uint64(l))
 			chunk := make([]byte, l)
-			wrote, err := link.marshal(chunk)
+			wrote, err := link.Marshal(chunk)
 			if err != nil {
 				return err
 			}
@@ -59,50 +59,50 @@ func Marshal(out io.Writer, tokenSource func() (Token, error)) error {
 }
 
 func MarshalPBNode(node *PBNode) ([]byte, error) {
-	if err := node.validate(); err != nil {
+	if err := node.Validate(); err != nil {
 		return nil, err
 	}
-	size := node.size()
+	size := node.EncodedSize()
 	data := make([]byte, size)
 
 	i := len(data)
 	if node.Data != nil {
 		i -= len(node.Data)
 		copy(data[i:], node.Data)
-		i = encodeVarint(data, i, uint64(len(node.Data))) - 1
+		i = EncodeVarint(data, i, uint64(len(node.Data))) - 1
 		data[i] = 0xa
 	}
 	if len(node.Links) > 0 {
 		for index := len(node.Links) - 1; index >= 0; index-- {
-			size, err := node.Links[index].marshal(data[:i])
+			size, err := node.Links[index].Marshal(data[:i])
 			if err != nil {
 				return nil, err
 			}
 			i -= size
-			i = encodeVarint(data, i, uint64(size)) - 1
+			i = EncodeVarint(data, i, uint64(size)) - 1
 			data[i] = 0x12
 		}
 	}
 	return data[:size], nil
 }
 
-func (link *PBLink) marshal(data []byte) (int, error) {
+func (link *PBLink) Marshal(data []byte) (int, error) {
 	i := len(data)
 	if link.Tsize != nil {
-		i = encodeVarint(data, i, uint64(*link.Tsize)) - 1
+		i = EncodeVarint(data, i, uint64(*link.Tsize)) - 1
 		data[i] = 0x18
 	}
 	if link.Name != nil {
 		i -= len(*link.Name)
 		copy(data[i:], *link.Name)
-		i = encodeVarint(data, i, uint64(len(*link.Name))) - 1
+		i = EncodeVarint(data, i, uint64(len(*link.Name))) - 1
 		data[i] = 0x12
 	}
 	if link.Hash != nil {
 		byts := link.Hash.Bytes()
 		i -= len(byts)
 		copy(data[i:], byts)
-		i = encodeVarint(data, i, uint64(len(byts))) - 1
+		i = EncodeVarint(data, i, uint64(len(byts))) - 1
 		data[i] = 0xa
 	} else {
 		return 0, fmt.Errorf("invalid DAG-PB form (link must have a Hash)")
@@ -110,7 +110,7 @@ func (link *PBLink) marshal(data []byte) (int, error) {
 	return len(data) - i, nil
 }
 
-func (node *PBNode) validate() error {
+func (node *PBNode) Validate() error {
 	if node == nil {
 		return fmt.Errorf("PBNode not defined")
 	}
@@ -132,45 +132,45 @@ func (node *PBNode) validate() error {
 	return nil
 }
 
-func (link *PBLink) size() (n int) {
+func (link *PBLink) EncodedSize() (n int) {
 	if link == nil {
 		return 0
 	}
 	var l int
 	if link.Hash != nil {
 		l = link.Hash.ByteLen()
-		n += 1 + l + sizeOfVarint(uint64(l))
+		n += 1 + l + SizeOfVarint(uint64(l))
 	}
 	if link.Name != nil {
 		l = len(*link.Name)
-		n += 1 + l + sizeOfVarint(uint64(l))
+		n += 1 + l + SizeOfVarint(uint64(l))
 	}
 	if link.Tsize != nil {
-		n += 1 + sizeOfVarint(uint64(*link.Tsize))
+		n += 1 + SizeOfVarint(uint64(*link.Tsize))
 	}
 	return n
 }
 
-func (node *PBNode) size() (n int) {
+func (node *PBNode) EncodedSize() (n int) {
 	if node == nil {
 		return 0
 	}
 	var l int
 	if node.Data != nil {
 		l = len(node.Data)
-		n += 1 + l + sizeOfVarint(uint64(l))
+		n += 1 + l + SizeOfVarint(uint64(l))
 	}
 	if len(node.Links) > 0 {
 		for _, e := range node.Links {
-			l = e.size()
-			n += 1 + l + sizeOfVarint(uint64(l))
+			l = e.EncodedSize()
+			n += 1 + l + SizeOfVarint(uint64(l))
 		}
 	}
 	return n
 }
 
-func encodeVarint(data []byte, offset int, v uint64) int {
-	offset -= sizeOfVarint(v)
+func EncodeVarint(data []byte, offset int, v uint64) int {
+	offset -= SizeOfVarint(v)
 	base := offset
 	for v >= 1<<7 {
 		data[offset] = uint8(v&0x7f | 0x80)
@@ -181,6 +181,6 @@ func encodeVarint(data []byte, offset int, v uint64) int {
 	return base
 }
 
-func sizeOfVarint(x uint64) (n int) {
+func SizeOfVarint(x uint64) (n int) {
 	return (math_bits.Len64(x|1) + 6) / 7
 }