fix

io.go

@@ -40,14 +40,34 @@ func (w *writer) writeInt64(v int64) error {
 	return nil
 }
 
+type wrapByteReader struct {
+	io.Reader
+}
+
+func (br wrapByteReader) ReadByte() (byte, error) {
+	var b [1]byte
+	n, err := br.Read(b[:])
+	if n != 1 {
+		return 0, err
+	}
+	return b[0], nil
+}
+
+func toByteReader(r io.Reader) io.ByteReader {
+	if br, ok := r.(io.ByteReader); ok {
+		return br
+	}
+	return wrapByteReader{r}
+}
+
 type reader struct {
-	r   *bufio.Reader
+	r   io.ByteReader
 	err error
 }
 
 func newReader(r io.Reader) *reader {
 	return &reader{
-		r: bufio.NewReader(r),
+		r: toByteReader(r),
 	}
 }
 

trie2/trie2.go

@@ -40,16 +40,16 @@ type STrie[T any] struct {
 	values []T
 }
 
-func Freeze[T any](src *DTrie[T], copyValues bool) *STrie[T] {
-	tree := *trietree.Freeze(&src.tree)
+func (dt *DTrie[T]) Freeze(copyValues bool) *STrie[T] {
+	tree := trietree.Freeze(&dt.tree)
 	var values []T
 	if copyValues {
-		values = make([]T, len(src.values))
-		copy(values, src.values)
+		values = make([]T, len(values))
+		copy(values, dt.values)
 	} else {
-		values = src.values
+		values = dt.values
 	}
-	return &STrie[T]{tree: tree, values: values}
+	return &STrie[T]{tree: *tree, values: values}
 }
 
 func (st *STrie[T]) Marshal(w io.Writer) error {
@@ -62,17 +62,20 @@ func (st *STrie[T]) Marshal(w io.Writer) error {
 	if err := gob.NewEncoder(w).Encode(st.values); err != nil {
 		return err
 	}
-	return err
+	return nil
 }
 
 func Unmarshal[T any](r io.Reader) (*STrie[T], error) {
 	tree, err := trietree.Read(r)
 	if err != nil {
 		return nil, err
 	}
+	if len(tree.Levels) == 0 {
+		return &STrie[T]{tree: *tree}, nil
+	}
 	// read v from r then append it to values.
 	values := make([]T, 0, len(tree.Levels))
-	if err := gob.NewDecoder(r).Decode(values); err != nil {
+	if err := gob.NewDecoder(r).Decode(&values); err != nil {
 		return nil, err
 	}
 	return &STrie[T]{tree: *tree, values: values}, nil