From 27a594bf3505ec8f187c8295ca17fae78ffcbedc Mon Sep 17 00:00:00 2001 From: Roman Atachiants Date: Sun, 12 Dec 2021 20:20:24 +0400 Subject: [PATCH] Add Write/Read Self --- reader.go | 9 ++++++++- source.go | 30 +++++++++++++++--------------- source_test.go | 23 +++++++++++++++++++++++ writer.go | 26 +++++++++++++++++++++----- writer_test.go | 18 ++++++++++++++++++ 5 files changed, 85 insertions(+), 21 deletions(-) diff --git a/reader.go b/reader.go index 20c9ba4..c2ce686 100644 --- a/reader.go +++ b/reader.go @@ -27,7 +27,7 @@ func NewReader(src io.Reader) *Reader { } // Offset returns the number of bytes read through this reader. -func (r *Reader) Offset() uint { +func (r *Reader) Offset() int64 { return r.src.Offset() } @@ -182,6 +182,13 @@ func (r *Reader) ReadText(v encoding.TextUnmarshaler) error { return v.UnmarshalText(b) } +// ReadSelf uses the provider io.ReaderFrom in order to read the data from +// the source reader. +func (r *Reader) ReadSelf(v io.ReaderFrom) error { + _, err := v.ReadFrom(r) + return err +} + // --------------------------- Strings --------------------------- // ReadString a string prefixed with a variable-size integer size. diff --git a/source.go b/source.go index e4382b5..1afffb0 100644 --- a/source.go +++ b/source.go @@ -27,7 +27,7 @@ type source interface { Slice(n int) (buffer []byte, err error) ReadUvarint() (uint64, error) ReadVarint() (int64, error) - Offset() uint + Offset() int64 } // newSource figures out the most efficient source to use for the provided type @@ -53,7 +53,7 @@ func newSource(r io.Reader) source { // sliceSource implements a source that only reads from a slice type sliceSource struct { buffer []byte - offset uint64 // current reading index + offset int64 } // newSliceSource returns a new source reading from b. @@ -62,24 +62,24 @@ func newSliceSource(b []byte) *sliceSource { } // Offset returns the number of bytes read through this reader. -func (r *sliceSource) Offset() uint { - return uint(r.offset) +func (r *sliceSource) Offset() int64 { + return r.offset } // Read implements the io.Reader interface. func (r *sliceSource) Read(b []byte) (n int, err error) { - if r.offset >= uint64(len(r.buffer)) { + if r.offset >= int64(len(r.buffer)) { return 0, io.EOF } n = copy(b, r.buffer[r.offset:]) - r.offset += uint64(n) + r.offset += int64(n) return } // ReadByte implements the io.ByteReader interface. func (r *sliceSource) ReadByte() (byte, error) { - if r.offset >= uint64(len(r.buffer)) { + if r.offset >= int64(len(r.buffer)) { return 0, io.EOF } @@ -93,12 +93,12 @@ func (r *sliceSource) ReadByte() (byte, error) { // returns a sub-slice pointing to the same array. Since this requires access // to the underlying data, this is only available for our default source. func (r *sliceSource) Slice(n int) ([]byte, error) { - if r.offset+uint64(n) > uint64(len(r.buffer)) { + if r.offset+int64(n) > int64(len(r.buffer)) { return nil, io.EOF } cur := r.offset - r.offset += uint64(n) + r.offset += int64(n) return r.buffer[cur:r.offset], nil } @@ -106,7 +106,7 @@ func (r *sliceSource) Slice(n int) ([]byte, error) { func (r *sliceSource) ReadUvarint() (uint64, error) { var x uint64 for s := 0; s < maxVarintLen64; s += 7 { - if r.offset >= uint64(len(r.buffer)) { + if r.offset >= int64(len(r.buffer)) { return 0, io.EOF } @@ -140,7 +140,7 @@ type streamSource struct { io.Reader io.ByteReader scratch []byte - offset uint64 + offset int64 } // newStreamSource returns a new stream source @@ -163,14 +163,14 @@ func newStreamSource(r io.Reader) *streamSource { } // Offset returns the number of bytes read through this reader. -func (r *streamSource) Offset() uint { - return uint(r.offset) +func (r *streamSource) Offset() int64 { + return r.offset } // Read implements the io.Reader interface. func (r *streamSource) Read(b []byte) (int, error) { n, err := r.Reader.Read(b) - r.offset += uint64(n) + r.offset += int64(n) return n, err } @@ -189,7 +189,7 @@ func (r *streamSource) Slice(n int) ([]byte, error) { // Read from the stream into our scratch buffer n, err := io.ReadAtLeast(r.Reader, r.scratch[:n], n) - r.offset += uint64(n) + r.offset += int64(n) return r.scratch[:n], err } diff --git a/source_test.go b/source_test.go index 5e95222..eb6c4f3 100644 --- a/source_test.go +++ b/source_test.go @@ -112,3 +112,26 @@ func (w *limitWriter) Write(p []byte) (int, error) { return w.buffer.Write(p) } + +func (w *limitWriter) Close() error { + return nil +} + +// --------------------------- Self Reader/Writer --------------------------- + +type person struct { + Name string +} + +func (p *person) WriteTo(dst io.Writer) (int64, error) { + w := NewWriter(dst) + err := w.WriteString(p.Name) + return w.Offset(), err +} + +func (p *person) ReadFrom(src io.Reader) (int64, error) { + r := NewReader(src) + name, err := r.ReadString() + p.Name = name + return r.Offset(), err +} diff --git a/writer.go b/writer.go index 82ba287..55a5160 100644 --- a/writer.go +++ b/writer.go @@ -13,7 +13,7 @@ import ( type Writer struct { scratch [10]byte out io.Writer - offset uint64 + offset int64 } // NewWriter creates a new stream writer. @@ -34,8 +34,8 @@ func (w *Writer) Reset(out io.Writer) { } // Offset returns the number of bytes written through this writer. -func (w *Writer) Offset() uint { - return uint(w.offset) +func (w *Writer) Offset() int64 { + return w.offset } // --------------------------- io.Writer --------------------------- @@ -44,17 +44,26 @@ func (w *Writer) Offset() uint { // souurce. func (w *Writer) Write(p []byte) (int, error) { n, err := w.out.Write(p) - w.offset += uint64(n) + w.offset += int64(n) return n, err } // Write writes the contents of p into the buffer. func (w *Writer) write(p []byte) error { n, err := w.out.Write(p) - w.offset += uint64(n) + w.offset += int64(n) return err } +// Close closes the writer's underlying stream and return its error. If the +// underlying stream is not an io.Closer, it is a no-op. +func (w *Writer) Close() error { + if closer, ok := w.out.(io.Closer); ok { + return closer.Close() + } + return nil +} + // --------------------------- Unsigned Integers --------------------------- // WriteUvarint writes a variable size unsigned integer @@ -187,6 +196,13 @@ func (w *Writer) WriteText(v encoding.TextMarshaler) error { return err } +// WriteSelf uses the provider io.WriterTo in order to write the data into +// the destination writer. +func (w *Writer) WriteSelf(v io.WriterTo) error { + _, err := v.WriteTo(w) + return err +} + // --------------------------- Strings --------------------------- // WriteString writes a string prefixed with a variable-size integer. diff --git a/writer_test.go b/writer_test.go index 5db61a8..9ac977e 100644 --- a/writer_test.go +++ b/writer_test.go @@ -146,6 +146,18 @@ var Fixtures = map[string]struct { Buffer: []byte{0x14, 0x31, 0x39, 0x37, 0x30, 0x2d, 0x30, 0x31, 0x2d, 0x30, 0x31, 0x54, 0x30, 0x30, 0x3a, 0x30, 0x31, 0x3a, 0x30, 0x30, 0x5a}, Value: time.Unix(60, 0).UTC(), }, + "person": { + Encode: func(w *Writer) error { + return w.WriteSelf(&person{Name: "Roman"}) + }, + Decode: func(r *Reader) (interface{}, error) { + var out person + err := r.ReadSelf(&out) + return out, err + }, + Buffer: []byte{0x5, 0x52, 0x6f, 0x6d, 0x61, 0x6e}, + Value: person{Name: "Roman"}, + }, } func TestWrite(t *testing.T) { @@ -174,6 +186,12 @@ func TestNewWriter(t *testing.T) { w1 := NewWriter(bytes.NewBuffer(nil)) w2 := NewWriter(w1) assert.Equal(t, w1, w2) + assert.NoError(t, w1.Close()) +} + +func TestWriterClose(t *testing.T) { + w := NewWriter(new(limitWriter)) + assert.NoError(t, w.Close()) } // assertWrite asserts a single write operation