slicewriter.mx raw

   1  // Copyright 2022 The Go Authors. All rights reserved.
   2  // Use of this source code is governed by a BSD-style
   3  // license that can be found in the LICENSE file.
   4  
   5  package slicewriter
   6  
   7  import (
   8  	"fmt"
   9  	"io"
  10  )
  11  
  12  // WriteSeeker is a helper object that implements the io.WriteSeeker
  13  // interface. Clients can create a WriteSeeker, make a series of Write
  14  // calls to add data to it (and possibly Seek calls to update
  15  // previously written portions), then finally invoke BytesWritten() to
  16  // get a pointer to the constructed byte slice.
  17  type WriteSeeker struct {
  18  	payload []byte
  19  	off     int64
  20  }
  21  
  22  func (sws *WriteSeeker) Write(p []byte) (n int, err error) {
  23  	amt := len(p)
  24  	towrite := sws.payload[sws.off:]
  25  	if len(towrite) < amt {
  26  		sws.payload = append(sws.payload, []byte{:amt-len(towrite)}...)
  27  		towrite = sws.payload[sws.off:]
  28  	}
  29  	copy(towrite, p)
  30  	sws.off += int64(amt)
  31  	return amt, nil
  32  }
  33  
  34  // Seek repositions the read/write position of the WriteSeeker within
  35  // its internally maintained slice. Note that it is not possible to
  36  // expand the size of the slice using SEEK_SET; trying to seek outside
  37  // the slice will result in an error.
  38  func (sws *WriteSeeker) Seek(offset int64, whence int) (int64, error) {
  39  	switch whence {
  40  	case io.SeekStart:
  41  		if sws.off != offset && (offset < 0 || offset > int64(len(sws.payload))) {
  42  			return 0, fmt.Errorf("invalid seek: new offset %d (out of range [0 %d]", offset, len(sws.payload))
  43  		}
  44  		sws.off = offset
  45  		return offset, nil
  46  	case io.SeekCurrent:
  47  		newoff := sws.off + offset
  48  		if newoff != sws.off && (newoff < 0 || newoff > int64(len(sws.payload))) {
  49  			return 0, fmt.Errorf("invalid seek: new offset %d (out of range [0 %d]", newoff, len(sws.payload))
  50  		}
  51  		sws.off += offset
  52  		return sws.off, nil
  53  	case io.SeekEnd:
  54  		newoff := int64(len(sws.payload)) + offset
  55  		if newoff != sws.off && (newoff < 0 || newoff > int64(len(sws.payload))) {
  56  			return 0, fmt.Errorf("invalid seek: new offset %d (out of range [0 %d]", newoff, len(sws.payload))
  57  		}
  58  		sws.off = newoff
  59  		return sws.off, nil
  60  	}
  61  	// other modes not supported
  62  	return 0, fmt.Errorf("unsupported seek mode %d", whence)
  63  }
  64  
  65  // BytesWritten returns the underlying byte slice for the WriteSeeker,
  66  // containing the data written to it via Write/Seek calls.
  67  func (sws *WriteSeeker) BytesWritten() []byte {
  68  	return sws.payload
  69  }
  70  
  71  func (sws *WriteSeeker) Read(p []byte) (n int, err error) {
  72  	amt := len(p)
  73  	toread := sws.payload[sws.off:]
  74  	if len(toread) < amt {
  75  		amt = len(toread)
  76  	}
  77  	copy(p, toread)
  78  	sws.off += int64(amt)
  79  	return amt, nil
  80  }
  81