diff --git a/bitset.go b/bitset.go index 164544b..06b7482 100644 --- a/bitset.go +++ b/bitset.go @@ -895,8 +895,22 @@ func (b *BitSet) WriteTo(stream io.Writer) (int64, error) { } // Write set - err = binary.Write(stream, binaryOrder, b.set) - return int64(b.BinaryStorageSize()), err + // current implementation of bufio.Writer is more memory efficient than + // binary.Write for large set + writer := bufio.NewWriter(stream) + var item = make([]byte, 8) // for serializing uint64 + var n = binary.Size(uint64(0)) // number of bytes written + for _, x := range b.set { + binaryOrder.PutUint64(item, x) + nn, err := writer.Write(item) + if err != nil { + return int64(n+nn), err + } + n += nn + } + + err = writer.Flush() + return int64(n), err } // ReadFrom reads a BitSet from a stream written using WriteTo @@ -915,9 +929,20 @@ func (b *BitSet) ReadFrom(stream io.Reader) (int64, error) { } // Read remaining bytes as set - err = binary.Read(stream, binaryOrder, newset.set) - if err != nil { - return 0, err + // current implementation bufio.Reader is more memory efficient than + // binary.Read for large set + reader := bufio.NewReader(stream) + i := 0 + var item = make([]byte, 8) // one uint64 + for { + if _, err := reader.Read(item); err != nil { + if err == io.EOF { + break // done + } + return 0, err + } + newset.set[i] = binaryOrder.Uint64(item) + i++ } *b = *newset @@ -927,25 +952,18 @@ func (b *BitSet) ReadFrom(stream io.Reader) (int64, error) { // MarshalBinary encodes a BitSet into a binary form and returns the result. func (b *BitSet) MarshalBinary() ([]byte, error) { var buf bytes.Buffer - writer := bufio.NewWriter(&buf) - - _, err := b.WriteTo(writer) + _, err := b.WriteTo(&buf) if err != nil { return []byte{}, err } - err = writer.Flush() - return buf.Bytes(), err } // UnmarshalBinary decodes the binary form generated by MarshalBinary. func (b *BitSet) UnmarshalBinary(data []byte) error { buf := bytes.NewReader(data) - reader := bufio.NewReader(buf) - - _, err := b.ReadFrom(reader) - + _, err := b.ReadFrom(buf) return err }