peco.peco/buffer.go

package peco

import (
	"time"

	"context"

	"github.com/lestrrat-go/pdebug"
	runewidth "github.com/mattn/go-runewidth"
	"github.com/peco/peco/line"
	"github.com/peco/peco/pipeline"
	"github.com/pkg/errors"
)

func NewFilteredBuffer(src Buffer, page, perPage int) *FilteredBuffer {
	fb := FilteredBuffer{
		src: src,
	}

	start := perPage * (page - 1)

	// if for whatever reason we wanted a page that goes over the
	// capacity of the original buffer, we don't need to do any more
	// calculations. bail out
	if start > src.Size() {
		return &fb
	}

	// Copy over the selections that are applicable to this filtered buffer.
	end := start + perPage
	if end >= src.Size() {
		end = src.Size()
	}
	selection := make([]int, 0, end-start)

	lines := src.linesInRange(start, end)
	var maxcols int
	for i := start; i < end; i++ {
		selection = append(selection, i)
		cols := runewidth.StringWidth(lines[i-start].DisplayString())
		if cols > maxcols {
			maxcols = cols
		}
	}
	fb.selection = selection
	fb.maxcols = maxcols

	return &fb
}

// MaxColumn returns the max column size, which controls the amount we
// can scroll to the right
func (flb *FilteredBuffer) MaxColumn() int {
	return flb.maxcols
}

// LineAt returns the line at index `i`. Note that the i-th element
// in this filtered buffer may actually correspond to a totally
// different line number in the source buffer.
func (flb FilteredBuffer) LineAt(i int) (line.Line, error) {
	if i >= len(flb.selection) {
		return nil, errors.Errorf("specified index %d is out of range", len(flb.selection))
	}
	return flb.src.LineAt(flb.selection[i])
}

// Size returns the number of lines in the buffer
func (flb FilteredBuffer) Size() int {
	return len(flb.selection)
}

const defaultMemoryBufferCap = 1024

// NewMemoryBuffer creates a new MemoryBuffer. If cap > 0, the lines
// slice is pre-allocated with that capacity; otherwise it defaults to
// defaultMemoryBufferCap.
func NewMemoryBuffer(cap int) *MemoryBuffer {
	if cap <= 0 {
		cap = defaultMemoryBufferCap
	}
	mb := &MemoryBuffer{}
	mb.done = make(chan struct{})
	mb.lines = make([]line.Line, 0, cap)
	return mb
}

func (mb *MemoryBuffer) Size() int {
	mb.mutex.RLock()
	defer mb.mutex.RUnlock()
	return bufferSize(mb.lines)
}

func bufferSize(lines []line.Line) int {
	return len(lines)
}

func (mb *MemoryBuffer) Reset() {
	mb.mutex.Lock()
	defer mb.mutex.Unlock()
	if pdebug.Enabled {
		g := pdebug.Marker("MemoryBuffer.Reset")
		defer g.End()
	}
	mb.done = make(chan struct{})
	mb.lines = []line.Line(nil)
}

func (mb *MemoryBuffer) Done() <-chan struct{} {
	mb.mutex.RLock()
	defer mb.mutex.RUnlock()
	return mb.done
}

func (mb *MemoryBuffer) Accept(ctx context.Context, in chan interface{}, _ pipeline.ChanOutput) {
	if pdebug.Enabled {
		g := pdebug.Marker("MemoryBuffer.Accept")
		defer g.End()
	}
	defer func() {
		mb.mutex.Lock()
		close(mb.done)
		mb.mutex.Unlock()
	}()

	// batch collects lines from the channel so we can append them
	// under a single lock acquisition instead of locking per line.
	batch := make([]line.Line, 0, 256)

	start := time.Now()
	for {
		select {
		case <-ctx.Done():
			if pdebug.Enabled {
				pdebug.Printf("MemoryBuffer received context done")
			}
			return
		case v := <-in:
			switch v := v.(type) {
			case error:
				if pipeline.IsEndMark(v) {
					if pdebug.Enabled {
						pdebug.Printf("MemoryBuffer received end mark (read %d lines, %s since starting accept loop)", len(mb.lines), time.Since(start).String())
					}
					// Flush remaining batch
					if len(batch) > 0 {
						mb.mutex.Lock()
						mb.lines = append(mb.lines, batch...)
						mb.mutex.Unlock()
					}
					return
				}
			case []line.Line:
				batch = append(batch, v...)
				mb.mutex.Lock()
				mb.lines = append(mb.lines, batch...)
				mb.mutex.Unlock()
				batch = batch[:0]
			case line.Line:
				batch = append(batch, v)

				// Drain any additional ready values without blocking
			drain:
				for {
					select {
					case v2 := <-in:
						switch v2 := v2.(type) {
						case error:
							if pipeline.IsEndMark(v2) {
								if pdebug.Enabled {
									pdebug.Printf("MemoryBuffer received end mark (read %d lines, %s since starting accept loop)", len(mb.lines)+len(batch), time.Since(start).String())
								}
								mb.mutex.Lock()
								mb.lines = append(mb.lines, batch...)
								mb.mutex.Unlock()
								return
							}
						case []line.Line:
							batch = append(batch, v2...)
						case line.Line:
							batch = append(batch, v2)
						}
					default:
						break drain
					}
				}

				// Flush the batch
				mb.mutex.Lock()
				mb.lines = append(mb.lines, batch...)
				mb.mutex.Unlock()
				batch = batch[:0]
			}
		}
	}
}

// AppendLine adds a line to the buffer. This is used by the benchmark tool
// to populate a MemoryBuffer that will be used as a pipeline source.
func (mb *MemoryBuffer) AppendLine(l line.Line) {
	mb.mutex.Lock()
	mb.lines = append(mb.lines, l)
	mb.mutex.Unlock()
}

func (mb *MemoryBuffer) LineAt(n int) (line.Line, error) {
	mb.mutex.RLock()
	defer mb.mutex.RUnlock()
	return bufferLineAt(mb.lines, n)
}

func (mb *MemoryBuffer) linesInRange(start, end int) []line.Line {
	mb.mutex.RLock()
	defer mb.mutex.RUnlock()
	return mb.lines[start:end]
}

func bufferLineAt(lines []line.Line, n int) (line.Line, error) {
	if s := len(lines); s <= 0 || n >= s {
		return nil, errors.New("empty buffer")
	}

	return lines[n], nil
}

// MemoryBufferSource wraps a completed MemoryBuffer as a pipeline.Source,
// allowing previous filter results to be reused as the input for
// incremental filtering.
type MemoryBufferSource struct {
	buf *MemoryBuffer
}

// NewMemoryBufferSource creates a new MemoryBufferSource from an existing
// MemoryBuffer. The buffer should be fully populated (pipeline completed).
func NewMemoryBufferSource(buf *MemoryBuffer) *MemoryBufferSource {
	return &MemoryBufferSource{buf: buf}
}

// sourceBatchSize is the number of lines sent per batch from source to
// the filter stage. Larger batches reduce channel operations but increase
// latency to first result. 1024 is a good balance.
const sourceBatchSize = 1024

// Start iterates through the MemoryBuffer's lines and sends them in
// batches to the output channel, implementing pipeline.Source.
func (s *MemoryBufferSource) Start(ctx context.Context, out pipeline.ChanOutput) {
	defer out.SendEndMark(ctx, "end of memory buffer source")

	s.buf.mutex.RLock()
	lines := s.buf.lines
	s.buf.mutex.RUnlock()

	for i := 0; i < len(lines); i += sourceBatchSize {
		select {
		case <-ctx.Done():
			return
		default:
		}
		end := i + sourceBatchSize
		if end > len(lines) {
			end = len(lines)
		}
		out.Send(ctx, lines[i:end])
	}
}

// Reset is a no-op for MemoryBufferSource since the underlying buffer
// is immutable (from a completed pipeline run).
func (s *MemoryBufferSource) Reset() {}