diff --git a/mmap/mmap_other.go b/mmap/mmap_other.go index 799cd7c8a..ffe6bdefd 100644 --- a/mmap/mmap_other.go +++ b/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) diff --git a/mmap/mmap_test.go b/mmap/mmap_test.go index 0716cc4e1..e5a6ed94e 100644 --- a/mmap/mmap_test.go +++ b/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) + } +} diff --git a/mmap/mmap_unix.go b/mmap/mmap_unix.go index 0b92cdb68..d8746633a 100644 --- a/mmap/mmap_unix.go +++ b/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 + pos int64 } // Close closes the reader. @@ -79,6 +80,35 @@ 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) { + n, err := w.Write(r.data[r.pos:]) + r.pos += int64(n) + return int64(n), err +} + +// Read implements the io.ReadSeeker interface. +func (r *ReaderAt) Read(p []byte) (n int, err error) { + n, err = r.ReadAt(p, r.pos) + r.pos += int64(n) + return n, err +} + +// Seek implements the io.ReadSeeker interface. +func (r *ReaderAt) Seek(offset int64, whence int) (int64, error) { + switch whence { + case io.SeekStart: + r.pos = offset + case io.SeekCurrent: + r.pos += offset + case io.SeekEnd: + r.pos = int64(len(r.data)) + offset + default: + return 0, fmt.Errorf("mmap: invalid seek whence") + } + return r.pos, nil +} + // Open memory-maps the named file for reading. func Open(filename string) (*ReaderAt, error) { f, err := os.Open(filename) @@ -113,7 +143,9 @@ func Open(filename string) (*ReaderAt, error) { if err != nil { return nil, err } - r := &ReaderAt{data} + r := &ReaderAt{ + data: data, + } if debug { var p *byte if len(data) != 0 { diff --git a/mmap/mmap_windows.go b/mmap/mmap_windows.go index ea1d1cbcb..a81e8e384 100644 --- a/mmap/mmap_windows.go +++ b/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 + pos int64 } // Close closes the reader. @@ -78,6 +79,35 @@ 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) { + n, err := w.Write(r.data[r.pos:]) + r.pos += int64(n) + return int64(n), err +} + +// Read implements the io.ReadSeeker interface. +func (r *ReaderAt) Read(p []byte) (n int, err error) { + n, err = r.ReadAt(p, r.pos) + r.pos += int64(n) + return n, err +} + +// Seek implements the io.ReadSeeker interface. +func (r *ReaderAt) Seek(offset int64, whence int) (int64, error) { + switch whence { + case io.SeekStart: + r.pos = offset + case io.SeekCurrent: + r.pos += offset + case io.SeekEnd: + r.pos = int64(len(r.data)) + offset + default: + return 0, fmt.Errorf("mmap: invalid seek whence") + } + return r.pos, nil +} + // Open memory-maps the named file for reading. func Open(filename string) (*ReaderAt, error) { f, err := os.Open(filename)