1// Copyright 2015, Joe Tsai. All rights reserved.
2// Use of this source code is governed by a BSD-style
3// license that can be found in the LICENSE.md file.
4
5package prefix
6
7import (
8	"bytes"
9	"io"
10	"strings"
11)
12
13// For some of the common Readers, we wrap and extend them to satisfy the
14// compress.BufferedReader interface to improve performance.
15
16type buffer struct {
17	*bytes.Buffer
18}
19
20type bytesReader struct {
21	*bytes.Reader
22	pos int64
23	buf []byte
24	arr [512]byte
25}
26
27type stringReader struct {
28	*strings.Reader
29	pos int64
30	buf []byte
31	arr [512]byte
32}
33
34func (r *buffer) Buffered() int {
35	return r.Len()
36}
37
38func (r *buffer) Peek(n int) ([]byte, error) {
39	b := r.Bytes()
40	if len(b) < n {
41		return b, io.EOF
42	}
43	return b[:n], nil
44}
45
46func (r *buffer) Discard(n int) (int, error) {
47	b := r.Next(n)
48	if len(b) < n {
49		return len(b), io.EOF
50	}
51	return n, nil
52}
53
54func (r *bytesReader) Buffered() int {
55	r.update()
56	if r.Len() > len(r.buf) {
57		return len(r.buf)
58	}
59	return r.Len()
60}
61
62func (r *bytesReader) Peek(n int) ([]byte, error) {
63	if n > len(r.arr) {
64		return nil, io.ErrShortBuffer
65	}
66
67	// Return sub-slice of local buffer if possible.
68	r.update()
69	if len(r.buf) >= n {
70		return r.buf[:n], nil
71	}
72
73	// Fill entire local buffer, and return appropriate sub-slice.
74	cnt, err := r.ReadAt(r.arr[:], r.pos)
75	r.buf = r.arr[:cnt]
76	if cnt < n {
77		return r.arr[:cnt], err
78	}
79	return r.arr[:n], nil
80}
81
82func (r *bytesReader) Discard(n int) (int, error) {
83	var err error
84	if n > r.Len() {
85		n, err = r.Len(), io.EOF
86	}
87	r.Seek(int64(n), io.SeekCurrent)
88	return n, err
89}
90
91// update reslices the internal buffer to be consistent with the read offset.
92func (r *bytesReader) update() {
93	pos, _ := r.Seek(0, io.SeekCurrent)
94	if off := pos - r.pos; off >= 0 && off < int64(len(r.buf)) {
95		r.buf, r.pos = r.buf[off:], pos
96	} else {
97		r.buf, r.pos = nil, pos
98	}
99}
100
101func (r *stringReader) Buffered() int {
102	r.update()
103	if r.Len() > len(r.buf) {
104		return len(r.buf)
105	}
106	return r.Len()
107}
108
109func (r *stringReader) Peek(n int) ([]byte, error) {
110	if n > len(r.arr) {
111		return nil, io.ErrShortBuffer
112	}
113
114	// Return sub-slice of local buffer if possible.
115	r.update()
116	if len(r.buf) >= n {
117		return r.buf[:n], nil
118	}
119
120	// Fill entire local buffer, and return appropriate sub-slice.
121	cnt, err := r.ReadAt(r.arr[:], r.pos)
122	r.buf = r.arr[:cnt]
123	if cnt < n {
124		return r.arr[:cnt], err
125	}
126	return r.arr[:n], nil
127}
128
129func (r *stringReader) Discard(n int) (int, error) {
130	var err error
131	if n > r.Len() {
132		n, err = r.Len(), io.EOF
133	}
134	r.Seek(int64(n), io.SeekCurrent)
135	return n, err
136}
137
138// update reslices the internal buffer to be consistent with the read offset.
139func (r *stringReader) update() {
140	pos, _ := r.Seek(0, io.SeekCurrent)
141	if off := pos - r.pos; off >= 0 && off < int64(len(r.buf)) {
142		r.buf, r.pos = r.buf[off:], pos
143	} else {
144		r.buf, r.pos = nil, pos
145	}
146}
147