bufat - bufat.go

  // Package bufat provides a buffered io.ReaderAt implementation.
  package bufat // import "vimagination.zapto.org/bufat"
  
  import (
  	"errors"
  	"io"
  	"io/fs"
  	"iter"
  
  	"vimagination.zapto.org/cache"
  )
  
  type reader struct {
  	r                         io.ReaderAt
  	length, offset, blockSize int64
  	cache                     cache.Cache[int64, string]
  }
  
  // Buffer wraps an io.ReaderAt with a caching layer.
  //
  // By default, the cache will store 1,024 4,096 byte blocks in an LRU cache.
  //
  // The defaults can be changed by specifying various options to this function.
  func Buffer(r io.ReaderAt, os ...Option) io.ReaderAt {
  	o := opts{length: -1}
  
  	for _, opt := range os {
  		opt(&o)
  	}
  
  	if o.chunkSize <= 0 {
  		o.chunkSize = 1 << 12
  	}
  
  	if o.chunks <= 0 {
  		o.chunks = 1 << 10
  	}
  
  	var c cache.Cache[int64, string]
  
  	if o.mru {
  		c = cache.NewMRU[int64, string](o.chunks)
  	} else {
  		c = cache.NewLRU[int64, string](o.chunks)
  	}
  
  	return &reader{
  		r:         r,
  		length:    o.length,
  		offset:    o.offset,
  		blockSize: o.chunkSize,
  		cache:     c,
  	}
  }
  
  // ReadAt implements the io.ReaderAt interface.
  func (r *reader) ReadAt(p []byte, n int64) (int, error) {
  	n += r.offset
  
  	if r.length >= 0 && n > r.length {
  		return 0, io.EOF
  	} else if n < 0 {
  		return 0, fs.ErrInvalid
  	}
  
  	if r.length >= 0 {
  		p = p[:min(int64(len(p)), r.length-n)]
  	}
  
  	blocks, err := r.getBlocks(n, int64(len(p)))
  	if err != nil {
  		return 0, err
  	}
  
  	if n%r.blockSize > int64(len(blocks[0])) {
  		return 0, io.EOF
  	}
  
  	blocks[0] = blocks[0][n%r.blockSize:]
  
  	l := len(p)
  
  	for _, block := range blocks {
  		p = p[copy(p, block):]
  	}
  
  	return l, nil
  }
  
  func (r *reader) getBlocks(start, length int64) ([]string, error) {
  	blocks, requests := r.getExistingBlocks(start, length)
  	if len(requests) == 0 {
  		return blocks, nil
  	}
  
  	if err := r.setNewBlocks(blocks, requests, start); err != nil {
  		return nil, err
  	}
  
  	return blocks, nil
  }
  
  type requests [][2]int64
  
  func (r requests) Iter() iter.Seq[int64] {
  	return func(yield func(int64) bool) {
  		for _, request := range r {
  			for b := request[0]; b <= request[1]; b++ {
  				if !yield(b) {
  					return
  				}
  			}
  		}
  	}
  }
  
  func (r *reader) getExistingBlocks(start, length int64) ([]string, requests) {
  	var (
  		blocks   []string
  		requests requests
  	)
  
  	startBlock := start / r.blockSize
  
  	for block := startBlock; block <= (start+length-1)/r.blockSize; block++ {
  		b, ok := r.cache.Get(block)
  		if ok {
  			blocks = append(blocks, b)
  
  			continue
  		}
  
  		blocks = append(blocks, "")
  
  		if len(requests) > 0 && requests[len(requests)-1][1]+1 == block {
  			requests[len(requests)-1][1] = block
  		} else {
  			requests = append(requests, [2]int64{block, block})
  		}
  	}
  
  	return blocks, requests
  }
  
  func (r *reader) setNewBlocks(blocks []string, requests requests, start int64) error {
  	startBlock := start / r.blockSize
  	buf := make([]byte, r.blockSize)
  
  	for block := range requests.Iter() {
  		n, err := io.ReadFull(io.NewSectionReader(r.r, block*r.blockSize, r.blockSize), buf)
  		if err != nil && !errors.Is(err, io.ErrUnexpectedEOF) {
  			return err
  		}
  
  		data := string(buf[:n])
  		blocks[block-startBlock] = data
  
  		r.cache.Set(block, data)
  	}
  
  	return nil
  }