memio - memio.go

  // Package memio implements Read, Write, Seek, Close and other io methods for a byte slice.
  package memio // import "vimagination.zapto.org/memio"
  
  import (
  	"bytes"
  	"errors"
  	"io"
  )
  
  const (
  	seekSet = iota
  	seekCurr
  	seekEnd
  )
  
  // ErrClosed is an error returned when trying to perform an operation after using Close().
  var ErrClosed = errors.New("operation not permitted when closed")
  
  // ReadMem holds a byte slice that can be used for many io interfaces
  type ReadMem struct {
  	*bytes.Reader
  }
  
  // Open uses a byte slice for reading. Implements io.Reader, io.Seeker,
  // io.Closer, io.ReaderAt, io.ByteReader and io.WriterTo.
  func Open(data []byte) ReadMem {
  	return ReadMem{bytes.NewReader(data)}
  }
  
  // Close is a no-op func the lets ReadMem implement interfaces that require a
  // Close method
  func (ReadMem) Close() error {
  	return nil
  }
  
  // Peek reads the next n bytes without advancing the position
  func (r ReadMem) Peek(n int) ([]byte, error) {
  	pos, _ := r.Seek(0, seekCurr)
  	buf := make([]byte, n)
  	_, err := r.ReadAt(buf, pos)
  	return buf, err
  }
  
  // WriteMem holds a pointer to a byte slice and allows numerous io interfaces
  // to be used with it.
  type WriteMem struct {
  	data *[]byte
  	pos  int
  }
  
  // Create uses a byte slice for writing. Implements io.Writer, io.Seeker,
  // io.Closer, io.WriterAt, io.ByteWriter and io.ReaderFrom.
  func Create(data *[]byte) *WriteMem {
  	return &WriteMem{data, 0}
  }
  
  // Write is an implementation of the io.Writer interface
  func (b *WriteMem) Write(p []byte) (int, error) {
  	if b.data == nil {
  		return 0, ErrClosed
  	}
  	b.setSize(b.pos + len(p))
  	n := copy((*b.data)[b.pos:], p)
  	b.pos += n
  	return n, nil
  }
  
  // WriteAt is an implementation of the io.WriterAt interface
  func (b *WriteMem) WriteAt(p []byte, off int64) (int, error) {
  	if b.data == nil {
  		return 0, ErrClosed
  	}
  	b.setSize(int(off) + len(p))
  	return copy((*b.data)[off:], p), nil
  }
  
  // WriteByte is an implementation of the io.WriteByte interface
  func (b *WriteMem) WriteByte(c byte) error {
  	if b.data == nil {
  		return ErrClosed
  	}
  	b.setSize(b.pos + 1)
  	(*b.data)[b.pos] = c
  	b.pos++
  	return nil
  }
  
  // ReadFrom is an implementation of the io.ReaderFrom interface
  func (b *WriteMem) ReadFrom(f io.Reader) (int64, error) {
  	if b.data == nil {
  		return 0, ErrClosed
  	}
  	var (
  		c   int64
  		n   int
  		err error
  	)
  	buf := make([]byte, 1024)
  	for {
  		n, err = f.Read(buf)
  		if n > 0 {
  			c += int64(n)
  			b.setSize(b.pos + n)
  			copy((*b.data)[b.pos:], buf[:n])
  			b.pos += n
  		}
  		if err != nil {
  			if err == io.EOF {
  				err = nil
  			}
  			break
  		}
  	}
  	return c, err
  }
  
  // Seek is an implementation of the io.Seeker interface
  func (b *WriteMem) Seek(offset int64, whence int) (int64, error) {
  	if b.data == nil {
  		return 0, ErrClosed
  	}
  	switch whence {
  	case seekSet:
  		b.pos = int(offset)
  	case seekCurr:
  		b.pos += int(offset)
  	case seekEnd:
  		b.pos = len(*b.data) + int(offset)
  	}
  	if b.pos < 0 {
  		b.pos = 0
  	}
  	return int64(b.pos), nil
  }
  
  // Close is an implementation of the io.Closer interface
  func (b *WriteMem) Close() error {
  	b.data = nil
  	return nil
  }
  
  func (b *WriteMem) setSize(end int) {
  	if end > len(*b.data) {
  		if end < cap(*b.data) {
  			*b.data = (*b.data)[:end]
  		} else {
  			var newData []byte
  			if len(*b.data) < 512 {
  				newData = make([]byte, end, end<<1)
  			} else {
  				newData = make([]byte, end, end+(end>>2))
  			}
  			copy(newData, *b.data)
  			*b.data = newData
  		}
  	}
  }
  
  // Truncate changes the length of the byte slice to the given amount
  func (b *WriteMem) Truncate(s int64) error {
  	if l := int64(len(*b.data)); l > s {
  		copy((*b.data)[s:], make([]byte, l-s))
  		*b.data = (*b.data)[:s]
  	} else if l < s {
  		b.setSize(int(s))
  	}
  	return nil
  }
  
  // WriteString writes a string to the underlying memory
  func (b *WriteMem) WriteString(s string) (int, error) {
  	return b.Write([]byte(s))
  }