hamlab/internal/space/codec.go

package space

import (
	"bufio"
	"encoding/binary"
	"errors"
	"io"
)

// The tree crosses from the helper process to the daemon as a preorder stream of varints:
// a few bytes per node, decoded without reflection, so a million nodes take well under a
// second either way.

const codecMagic = "HLSPACE1"

func Encode(w io.Writer, root *Node) error {
	bw := bufio.NewWriterSize(w, 1<<20)
	if _, err := bw.WriteString(codecMagic); err != nil {
		return err
	}
	var buf [binary.MaxVarintLen64]byte
	putU := func(v uint64) error { n := binary.PutUvarint(buf[:], v); _, err := bw.Write(buf[:n]); return err }
	putI := func(v int64) error { n := binary.PutVarint(buf[:], v); _, err := bw.Write(buf[:n]); return err }
	var enc func(n *Node) error
	enc = func(n *Node) error {
		if err := putU(uint64(len(n.Name))); err != nil {
			return err
		}
		if _, err := bw.WriteString(n.Name); err != nil {
			return err
		}
		flags := uint64(0)
		if n.Dir {
			flags |= 1
		}
		if n.Mount {
			flags |= 2
		}
		if n.Denied {
			flags |= 4
		}
		for _, f := range []func() error{
			func() error { return putU(flags) },
			func() error { return putI(n.Size) },
			func() error { return putI(n.Apparent) },
			func() error { return putI(n.MTime) },
			func() error { return putU(uint64(n.Files)) },
			func() error { return putU(uint64(n.Dirs)) },
			func() error { return putU(uint64(n.SmallN)) },
			func() error { return putI(n.SmallSize) },
			func() error { return putU(uint64(len(n.Children))) },
		} {
			if err := f(); err != nil {
				return err
			}
		}
		for _, c := range n.Children {
			if err := enc(c); err != nil {
				return err
			}
		}
		return nil
	}
	if err := enc(root); err != nil {
		return err
	}
	return bw.Flush()
}

func Decode(r io.Reader) (*Node, error) {
	br := bufio.NewReaderSize(r, 1<<20)
	magic := make([]byte, len(codecMagic))
	if _, err := io.ReadFull(br, magic); err != nil || string(magic) != codecMagic {
		return nil, errors.New("space: not a scan stream")
	}
	var dec func() (*Node, error)
	dec = func() (*Node, error) {
		l, err := binary.ReadUvarint(br)
		if err != nil {
			return nil, err
		}
		if l > 1<<16 {
			return nil, errors.New("space: bad name length")
		}
		name := make([]byte, l)
		if _, err := io.ReadFull(br, name); err != nil {
			return nil, err
		}
		n := &Node{Name: string(name)}
		flags, err := binary.ReadUvarint(br)
		if err != nil {
			return nil, err
		}
		n.Dir, n.Mount, n.Denied = flags&1 != 0, flags&2 != 0, flags&4 != 0
		if n.Size, err = binary.ReadVarint(br); err != nil {
			return nil, err
		}
		if n.Apparent, err = binary.ReadVarint(br); err != nil {
			return nil, err
		}
		if n.MTime, err = binary.ReadVarint(br); err != nil {
			return nil, err
		}
		u, err := binary.ReadUvarint(br)
		if err != nil {
			return nil, err
		}
		n.Files = uint32(u)
		if u, err = binary.ReadUvarint(br); err != nil {
			return nil, err
		}
		n.Dirs = uint32(u)
		if u, err = binary.ReadUvarint(br); err != nil {
			return nil, err
		}
		n.SmallN = uint32(u)
		if n.SmallSize, err = binary.ReadVarint(br); err != nil {
			return nil, err
		}
		cnt, err := binary.ReadUvarint(br)
		if err != nil {
			return nil, err
		}
		if cnt > 1<<24 {
			return nil, errors.New("space: bad child count")
		}
		if cnt > 0 {
			n.Children = make([]*Node, 0, cnt)
			for i := uint64(0); i < cnt; i++ {
				c, err := dec()
				if err != nil {
					return nil, err
				}
				n.Children = append(n.Children, c)
			}
		}
		return n, nil
	}
	return dec()
}