mmap: implement io.WriterTo and io.ReadSeeker

Author: costela

mmap/mmap_other.go

@@ -9,6 +9,7 @@ package mmap
 
 import (
 	"fmt"
+	"io"
 	"os"
 )
 
@@ -46,6 +47,21 @@ func (r *ReaderAt) ReadAt(p []byte, off int64) (int, error) {
 	return r.f.ReadAt(p, off)
 }
 
+// WriteTo implements the io.WriterTo interface.
+func (r *ReaderAt) WriteTo(w io.Writer) (int64, error) {
+	return r.f.WriterTo(w)
+}
+
+// Read implements the io.ReadSeeker interface.
+func (r *ReaderAt) Read(p []byte) (n int, err error) {
+	return r.f.Read(p)
+}
+
+// Seek implements the io.ReadSeeker interface.
+func (r *ReaderAt) Seek(offset int64, whence int) (int64, error) {
+	return r.f.Seek(offset, whence)
+}
+
 // Open memory-maps the named file for reading.
 func Open(filename string) (*ReaderAt, error) {
 	f, err := os.Open(filename)

mmap/mmap_test.go

@@ -32,3 +32,120 @@ func TestOpen(t *testing.T) {
 		t.Fatalf("\ngot  %q\nwant %q", string(got), string(want))
 	}
 }
+
+func TestSeekRead(t *testing.T) {
+	const filename = "mmap_test.go"
+	r, err := Open(filename)
+	if err != nil {
+		t.Fatalf("Open: %v", err)
+	}
+	buf := make([]byte, 1)
+	if _, err := r.Seek(0, io.SeekStart); err != nil {
+		t.Fatalf("Seek: %v", err)
+	}
+	n, err := r.Read(buf)
+	if err != nil {
+		t.Fatalf("Read: %v", err)
+	}
+	if n != 1 {
+		t.Fatalf("Read: got %d bytes, want 1", n)
+	}
+	if buf[0] != '/' { // first comment slash
+		t.Fatalf("Read: got %q, want '/'", buf[0])
+	}
+	if _, err := r.Seek(1, io.SeekCurrent); err != nil {
+		t.Fatalf("Seek: %v", err)
+	}
+	n, err = r.Read(buf)
+	if err != nil {
+		t.Fatalf("Read: %v", err)
+	}
+	if n != 1 {
+		t.Fatalf("Read: got %d bytes, want 1", n)
+	}
+	if buf[0] != ' ' { // space after comment
+		t.Fatalf("Read: got %q, want ' '", buf[0])
+	}
+	if _, err := r.Seek(-1, io.SeekEnd); err != nil {
+		t.Fatalf("Seek: %v", err)
+	}
+	n, err = r.Read(buf)
+	if err != nil {
+		t.Fatalf("Read: %v", err)
+	}
+	if n != 1 {
+		t.Fatalf("Read: got %d bytes, want 1", n)
+	}
+	if buf[0] != '\n' { // last newline
+		t.Fatalf("Read: got %q, want newline", buf[0])
+	}
+	if _, err := r.Seek(0, io.SeekEnd); err != nil {
+		t.Fatalf("Seek: %v", err)
+	}
+	if _, err := r.Read(buf); err != io.EOF {
+		t.Fatalf("Read: expected EOF, got %v", err)
+	}
+}
+
+func TestWriterTo_idempotency(t *testing.T) {
+	const filename = "mmap_test.go"
+	r, err := Open(filename)
+	if err != nil {
+		t.Fatalf("Open: %v", err)
+	}
+	buf := bytes.NewBuffer(make([]byte, 0, len(r.data)))
+	// first run
+	n, err := r.WriteTo(buf)
+	if err != nil {
+		t.Fatalf("WriteTo: %v", err)
+	}
+	if n != int64(len(r.data)) {
+		t.Fatalf("WriteTo: got %d bytes, want %d", n, len(r.data))
+	}
+	if !bytes.Equal(buf.Bytes(), r.data) {
+		t.Fatalf("WriteTo: got %q, want %q", buf.Bytes(), r.data)
+	}
+	// second run
+	n, err = r.WriteTo(buf)
+	if err != nil {
+		t.Fatalf("WriteTo: %v", err)
+	}
+	if n != 0 {
+		t.Fatalf("WriteTo: got %d bytes, want %d", n, 0)
+	}
+	if !bytes.Equal(buf.Bytes(), r.data) {
+		t.Fatalf("WriteTo: got %q, want %q", buf.Bytes(), r.data)
+	}
+}
+
+func BenchmarkMmapCopy(b *testing.B) {
+	var f io.ReadSeeker
+
+	// mmap some big-ish file; will only work on unix-like OSs.
+	r, err := Open("/proc/self/exe")
+	if err != nil {
+		b.Fatalf("Open: %v", err)
+	}
+
+	// Sanity check: ensure we will run into the io.Copy optimization when using the NEW code above.
+	var _ io.WriterTo = r
+
+	// f = io.NewSectionReader(r, 0, int64(len(r.data))) // old
+	f = r // new
+
+	buf := bytes.NewBuffer(make([]byte, 0, len(r.data)))
+	// "Hide" the ReaderFrom interface by wrapping the writer.
+	// Otherwise we skew the results by optimizing the wrong side.
+	writer := struct{ io.Writer }{buf}
+
+	b.ReportAllocs()
+	b.ResetTimer()
+
+	for i := 0; i < b.N; i++ {
+		_, _ = f.Seek(0, io.SeekStart)
+		buf.Reset()
+
+		n, _ := io.Copy(writer, f)
+		b.SetBytes(n)
+	}
+}

mmap/mmap_unix.go

@@ -31,6 +31,7 @@ const debug = false
 // not safe to call Close and reading methods concurrently.
 type ReaderAt struct {
 	data []byte
+	*io.SectionReader
 }
 
 // Close closes the reader.
@@ -79,6 +80,20 @@ func (r *ReaderAt) ReadAt(p []byte, off int64) (int, error) {
 	return n, nil
 }
 
+// WriteTo implements the io.WriterTo interface.
+func (r *ReaderAt) WriteTo(w io.Writer) (int64, error) {
+	pos, err := r.Seek(0, io.SeekCurrent)
+	if err != nil {
+		return 0, fmt.Errorf("mmap: WriteTo: %w", err)
+	}
+	n, err := w.Write(r.data[pos:])
+	_, seekErr := r.Seek(int64(n), io.SeekCurrent)
+	if seekErr != nil {
+		panic("cannot happen") // neither errWhence nor errOffset are possible with current usage
+	}
+	return int64(n), err
+}
+
 // Open memory-maps the named file for reading.
 func Open(filename string) (*ReaderAt, error) {
 	f, err := os.Open(filename)
@@ -113,7 +128,10 @@ func Open(filename string) (*ReaderAt, error) {
 	if err != nil {
 		return nil, err
 	}
-	r := &ReaderAt{data}
+	r := &ReaderAt{
+		data: data,
+	}
+	r.SectionReader = io.NewSectionReader(r, 0, size)
 	if debug {
 		var p *byte
 		if len(data) != 0 {

mmap/mmap_windows.go

@@ -30,6 +30,7 @@ const debug = false
 // not safe to call Close and reading methods concurrently.
 type ReaderAt struct {
 	data []byte
+	*io.SectionReader
 }
 
 // Close closes the reader.
@@ -78,6 +79,20 @@ func (r *ReaderAt) ReadAt(p []byte, off int64) (int, error) {
 	return n, nil
 }
 
+// WriteTo implements the io.WriterTo interface.
+func (r *ReaderAt) WriteTo(w io.Writer) (int64, error) {
+	pos, err := r.Seek(0, io.SeekCurrent)
+	if err != nil {
+		return 0, fmt.Errorf("mmap: WriteTo: %w", err)
+	}
+	n, err := w.Write(r.data[pos:])
+	_, seekErr := r.Seek(int64(n), io.SeekCurrent)
+	if seekErr != nil {
+		panic("cannot happen") // neither errWhence nor errOffset are possible with current usage
+	}
+	return int64(n), err
+}
+
 // Open memory-maps the named file for reading.
 func Open(filename string) (*ReaderAt, error) {
 	f, err := os.Open(filename)
@@ -121,6 +136,7 @@ func Open(filename string) (*ReaderAt, error) {
 	data := unsafe.Slice((*byte)(unsafe.Pointer(ptr)), size)
 
 	r := &ReaderAt{data: data}
+	r.SectionReader = io.NewSectionReader(r, 0, size)
 	if debug {
 		var p *byte
 		if len(data) != 0 {