package buffer import ( "io" ) type block struct { buf []byte next int // index in pool plus one active bool } type bufferPool struct { pool []block head int // index in pool plus one tail int // index in pool plus one pos int // byte pos in tail } func (z *bufferPool) swap(oldBuf []byte, size int) []byte { // find new buffer that can be reused swap := -1 for i := 0; i < len(z.pool); i++ { if !z.pool[i].active && size <= cap(z.pool[i].buf) { swap = i break } } if swap == -1 { // no free buffer found for reuse if z.tail == 0 && z.pos >= len(oldBuf) && size <= cap(oldBuf) { // but we can reuse the current buffer! z.pos -= len(oldBuf) return oldBuf[:0] } // allocate new z.pool = append(z.pool, block{make([]byte, 0, size), 0, true}) swap = len(z.pool) - 1 } newBuf := z.pool[swap].buf // put current buffer into pool z.pool[swap] = block{oldBuf, 0, true} if z.head != 0 { z.pool[z.head-1].next = swap + 1 } z.head = swap + 1 if z.tail == 0 { z.tail = swap + 1 } return newBuf[:0] } func (z *bufferPool) free(n int) { z.pos += n // move the tail over to next buffers for z.tail != 0 && z.pos >= len(z.pool[z.tail-1].buf) { z.pos -= len(z.pool[z.tail-1].buf) newTail := z.pool[z.tail-1].next z.pool[z.tail-1].active = false // after this, any thread may pick up the inactive buffer, so it can't be used anymore z.tail = newTail } if z.tail == 0 { z.head = 0 } } // StreamLexer is a buffered reader that allows peeking forward and shifting, taking an io.Reader. // It keeps data in-memory until Free, taking a byte length, is called to move beyond the data. type StreamLexer struct { r io.Reader err error pool bufferPool buf []byte start int // index in buf pos int // index in buf prevStart int free int } // NewStreamLexer returns a new StreamLexer for a given io.Reader with a 4kB estimated buffer size. // If the io.Reader implements Bytes, that buffer is used instead. func NewStreamLexer(r io.Reader) *StreamLexer { return NewStreamLexerSize(r, defaultBufSize) } // NewStreamLexerSize returns a new StreamLexer for a given io.Reader and estimated required buffer size. // If the io.Reader implements Bytes, that buffer is used instead. func NewStreamLexerSize(r io.Reader, size int) *StreamLexer { // if reader has the bytes in memory already, use that instead if buffer, ok := r.(interface { Bytes() []byte }); ok { return &StreamLexer{ err: io.EOF, buf: buffer.Bytes(), } } return &StreamLexer{ r: r, buf: make([]byte, 0, size), } } func (z *StreamLexer) read(pos int) byte { if z.err != nil { return 0 } // free unused bytes z.pool.free(z.free) z.free = 0 // get new buffer c := cap(z.buf) p := pos - z.start + 1 if 2*p > c { // if the token is larger than half the buffer, increase buffer size c = 2*c + p } d := len(z.buf) - z.start buf := z.pool.swap(z.buf[:z.start], c) copy(buf[:d], z.buf[z.start:]) // copy the left-overs (unfinished token) from the old buffer // read in new data for the rest of the buffer var n int for pos-z.start >= d && z.err == nil { n, z.err = z.r.Read(buf[d:cap(buf)]) d += n } pos -= z.start z.pos -= z.start z.start, z.buf = 0, buf[:d] if pos >= d { return 0 } return z.buf[pos] } // Err returns the error returned from io.Reader. It may still return valid bytes for a while though. func (z *StreamLexer) Err() error { if z.err == io.EOF && z.pos < len(z.buf) { return nil } return z.err } // Free frees up bytes of length n from previously shifted tokens. // Each call to Shift should at one point be followed by a call to Free with a length returned by ShiftLen. func (z *StreamLexer) Free(n int) { z.free += n } // Peek returns the ith byte relative to the end position and possibly does an allocation. // Peek returns zero when an error has occurred, Err returns the error. // TODO: inline function func (z *StreamLexer) Peek(pos int) byte { pos += z.pos if uint(pos) < uint(len(z.buf)) { // uint for BCE return z.buf[pos] } return z.read(pos) } // PeekRune returns the rune and rune length of the ith byte relative to the end position. func (z *StreamLexer) PeekRune(pos int) (rune, int) { // from unicode/utf8 c := z.Peek(pos) if c < 0xC0 { return rune(c), 1 } else if c < 0xE0 { return rune(c&0x1F)<<6 | rune(z.Peek(pos+1)&0x3F), 2 } else if c < 0xF0 { return rune(c&0x0F)<<12 | rune(z.Peek(pos+1)&0x3F)<<6 | rune(z.Peek(pos+2)&0x3F), 3 } return rune(c&0x07)<<18 | rune(z.Peek(pos+1)&0x3F)<<12 | rune(z.Peek(pos+2)&0x3F)<<6 | rune(z.Peek(pos+3)&0x3F), 4 } // Move advances the position. func (z *StreamLexer) Move(n int) { z.pos += n } // Pos returns a mark to which can be rewinded. func (z *StreamLexer) Pos() int { return z.pos - z.start } // Rewind rewinds the position to the given position. func (z *StreamLexer) Rewind(pos int) { z.pos = z.start + pos } // Lexeme returns the bytes of the current selection. func (z *StreamLexer) Lexeme() []byte { return z.buf[z.start:z.pos] } // Skip collapses the position to the end of the selection. func (z *StreamLexer) Skip() { z.start = z.pos } // Shift returns the bytes of the current selection and collapses the position to the end of the selection. // It also returns the number of bytes we moved since the last call to Shift. This can be used in calls to Free. func (z *StreamLexer) Shift() []byte { if z.pos > len(z.buf) { // make sure we peeked at least as much as we shift z.read(z.pos - 1) } b := z.buf[z.start:z.pos] z.start = z.pos return b } // ShiftLen returns the number of bytes moved since the last call to ShiftLen. This can be used in calls to Free because it takes into account multiple Shifts or Skips. func (z *StreamLexer) ShiftLen() int { n := z.start - z.prevStart z.prevStart = z.start return n }