Skip to content

Commit

Permalink
mmap: implement io.WriterTo and io.ReadSeeker
Browse files Browse the repository at this point in the history
  • Loading branch information
costela committed Jan 15, 2024
1 parent db7319d commit ab75851
Show file tree
Hide file tree
Showing 4 changed files with 196 additions and 1 deletion.
16 changes: 16 additions & 0 deletions mmap/mmap_other.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,7 @@ package mmap

import (
"fmt"
"io"
"os"
)

Expand Down Expand Up @@ -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)
Expand Down
117 changes: 117 additions & 0 deletions mmap/mmap_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -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)
}
}
34 changes: 33 additions & 1 deletion mmap/mmap_unix.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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 {
Expand Down
30 changes: 30 additions & 0 deletions mmap/mmap_windows.go
Original file line number Diff line number Diff line change
Expand Up @@ -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.
Expand Down Expand Up @@ -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)
Expand Down

0 comments on commit ab75851

Please sign in to comment.