github.com/pyroscope-io/pyroscope@v0.37.3-0.20230725203016-5f6947968bd0/pkg/storage/dict/serialize.go (about)

     1  package dict
     2  
     3  import (
     4  	"bufio"
     5  	"bytes"
     6  	"io"
     7  
     8  	"github.com/pyroscope-io/pyroscope/pkg/util/varint"
     9  )
    10  
    11  // serialization format version. it's not very useful right now, but it will be in the future
    12  const currentVersion = 1
    13  
    14  func (t *Dict) Serialize(w io.Writer) error {
    15  	t.m.RLock()
    16  	defer t.m.RUnlock()
    17  
    18  	_, err := varint.Write(w, currentVersion)
    19  	if err != nil {
    20  		return err
    21  	}
    22  
    23  	nodes := []*trieNode{t.root}
    24  	for len(nodes) > 0 {
    25  		tn := nodes[0]
    26  		nodes = nodes[1:]
    27  
    28  		label := tn.label
    29  		_, err := varint.Write(w, uint64(len(label)))
    30  		if err != nil {
    31  			return err
    32  		}
    33  		_, err = w.Write(label)
    34  		if err != nil {
    35  			return err
    36  		}
    37  
    38  		_, err = varint.Write(w, uint64(len(tn.children)))
    39  		if err != nil {
    40  			return err
    41  		}
    42  
    43  		nodes = append(tn.children, nodes...)
    44  	}
    45  	return nil
    46  }
    47  
    48  func Deserialize(r io.Reader) (*Dict, error) {
    49  	t := New()
    50  	br := bufio.NewReader(r) // TODO if it's already a bytereader skip
    51  
    52  	// reads serialization format version, see comment at the top
    53  	_, err := varint.Read(br)
    54  	if err != nil {
    55  		return nil, err
    56  	}
    57  
    58  	parents := []*trieNode{t.root}
    59  	for len(parents) > 0 {
    60  		parent := parents[0]
    61  		parents = parents[1:]
    62  
    63  		nameLen, err := varint.Read(br)
    64  		nameBuf := make([]byte, nameLen) // TODO: maybe there are better ways to do this?
    65  		_, err = io.ReadAtLeast(br, nameBuf, int(nameLen))
    66  		if err != nil {
    67  			return nil, err
    68  		}
    69  		tn := newTrieNode(nameBuf)
    70  		parent.insert(tn)
    71  
    72  		childrenLen, err := varint.Read(br)
    73  		if err != nil {
    74  			return nil, err
    75  		}
    76  
    77  		for i := uint64(0); i < childrenLen; i++ {
    78  			parents = append([]*trieNode{tn}, parents...)
    79  		}
    80  	}
    81  
    82  	t.root = t.root.children[0]
    83  
    84  	return t, nil
    85  }
    86  
    87  func (t *Dict) Bytes() ([]byte, error) {
    88  	b := bytes.Buffer{}
    89  	if err := t.Serialize(&b); err != nil {
    90  		return nil, err
    91  	}
    92  	return b.Bytes(), nil
    93  }
    94  
    95  func FromBytes(p []byte) (*Dict, error) {
    96  	return Deserialize(bytes.NewReader(p))
    97  }