Skip to content
This repository was archived by the owner on Mar 7, 2025. It is now read-only.

Commit 3ff355f

Browse files
nigeltaonigeltao
authored
Merge pull request #51 from topos-ai/bytereader
Add ReadByte method, satisfies the io.ByteReader interface
2 parents b9440b4 + 0a27eb7 commit 3ff355f

File tree

2 files changed

+118
-30
lines changed

2 files changed

+118
-30
lines changed

‎decode.go

+53-30
Original file line numberDiff line numberDiff line change
@@ -118,32 +118,23 @@ func (r *Reader) readFull(p []byte, allowEOF bool) (ok bool) {
118118
return true
119119
}
120120

121-
// Read satisfies the io.Reader interface.
122-
func (r *Reader) Read(p []byte) (int, error) {
123-
if r.err != nil {
124-
return 0, r.err
125-
}
126-
for {
127-
if r.i < r.j {
128-
n := copy(p, r.decoded[r.i:r.j])
129-
r.i += n
130-
return n, nil
131-
}
121+
func (r *Reader) fill() error {
122+
for r.i >= r.j {
132123
if !r.readFull(r.buf[:4], true) {
133-
return 0, r.err
124+
return r.err
134125
}
135126
chunkType := r.buf[0]
136127
if !r.readHeader {
137128
if chunkType != chunkTypeStreamIdentifier {
138129
r.err = ErrCorrupt
139-
return 0, r.err
130+
return r.err
140131
}
141132
r.readHeader = true
142133
}
143134
chunkLen := int(r.buf[1]) | int(r.buf[2])<<8 | int(r.buf[3])<<16
144135
if chunkLen > len(r.buf) {
145136
r.err = ErrUnsupported
146-
return 0, r.err
137+
return r.err
147138
}
148139

149140
// The chunk types are specified at
@@ -153,31 +144,31 @@ func (r *Reader) Read(p []byte) (int, error) {
153144
// Section 4.2. Compressed data (chunk type 0x00).
154145
if chunkLen < checksumSize {
155146
r.err = ErrCorrupt
156-
return 0, r.err
147+
return r.err
157148
}
158149
buf := r.buf[:chunkLen]
159150
if !r.readFull(buf, false) {
160-
return 0, r.err
151+
return r.err
161152
}
162153
checksum := uint32(buf[0]) | uint32(buf[1])<<8 | uint32(buf[2])<<16 | uint32(buf[3])<<24
163154
buf = buf[checksumSize:]
164155

165156
n, err := DecodedLen(buf)
166157
if err != nil {
167158
r.err = err
168-
return 0, r.err
159+
return r.err
169160
}
170161
if n > len(r.decoded) {
171162
r.err = ErrCorrupt
172-
return 0, r.err
163+
return r.err
173164
}
174165
if _, err := Decode(r.decoded, buf); err != nil {
175166
r.err = err
176-
return 0, r.err
167+
return r.err
177168
}
178169
if crc(r.decoded[:n]) != checksum {
179170
r.err = ErrCorrupt
180-
return 0, r.err
171+
return r.err
181172
}
182173
r.i, r.j = 0, n
183174
continue
@@ -186,25 +177,25 @@ func (r *Reader) Read(p []byte) (int, error) {
186177
// Section 4.3. Uncompressed data (chunk type 0x01).
187178
if chunkLen < checksumSize {
188179
r.err = ErrCorrupt
189-
return 0, r.err
180+
return r.err
190181
}
191182
buf := r.buf[:checksumSize]
192183
if !r.readFull(buf, false) {
193-
return 0, r.err
184+
return r.err
194185
}
195186
checksum := uint32(buf[0]) | uint32(buf[1])<<8 | uint32(buf[2])<<16 | uint32(buf[3])<<24
196187
// Read directly into r.decoded instead of via r.buf.
197188
n := chunkLen - checksumSize
198189
if n > len(r.decoded) {
199190
r.err = ErrCorrupt
200-
return 0, r.err
191+
return r.err
201192
}
202193
if !r.readFull(r.decoded[:n], false) {
203-
return 0, r.err
194+
return r.err
204195
}
205196
if crc(r.decoded[:n]) != checksum {
206197
r.err = ErrCorrupt
207-
return 0, r.err
198+
return r.err
208199
}
209200
r.i, r.j = 0, n
210201
continue
@@ -213,15 +204,15 @@ func (r *Reader) Read(p []byte) (int, error) {
213204
// Section 4.1. Stream identifier (chunk type 0xff).
214205
if chunkLen != len(magicBody) {
215206
r.err = ErrCorrupt
216-
return 0, r.err
207+
return r.err
217208
}
218209
if !r.readFull(r.buf[:len(magicBody)], false) {
219-
return 0, r.err
210+
return r.err
220211
}
221212
for i := 0; i < len(magicBody); i++ {
222213
if r.buf[i] != magicBody[i] {
223214
r.err = ErrCorrupt
224-
return 0, r.err
215+
return r.err
225216
}
226217
}
227218
continue
@@ -230,12 +221,44 @@ func (r *Reader) Read(p []byte) (int, error) {
230221
if chunkType <= 0x7f {
231222
// Section 4.5. Reserved unskippable chunks (chunk types 0x02-0x7f).
232223
r.err = ErrUnsupported
233-
return 0, r.err
224+
return r.err
234225
}
235226
// Section 4.4 Padding (chunk type 0xfe).
236227
// Section 4.6. Reserved skippable chunks (chunk types 0x80-0xfd).
237228
if !r.readFull(r.buf[:chunkLen], false) {
238-
return 0, r.err
229+
return r.err
239230
}
240231
}
232+
233+
return nil
234+
}
235+
236+
// Read satisfies the io.Reader interface.
237+
func (r *Reader) Read(p []byte) (int, error) {
238+
if r.err != nil {
239+
return 0, r.err
240+
}
241+
242+
if err := r.fill(); err != nil {
243+
return 0, err
244+
}
245+
246+
n := copy(p, r.decoded[r.i:r.j])
247+
r.i += n
248+
return n, nil
249+
}
250+
251+
// ReadByte satisfies the io.ByteReader interface.
252+
func (r *Reader) ReadByte() (byte, error) {
253+
if r.err != nil {
254+
return 0, r.err
255+
}
256+
257+
if err := r.fill(); err != nil {
258+
return 0, err
259+
}
260+
261+
c := r.decoded[r.i]
262+
r.i++
263+
return c, nil
241264
}

‎snappy_test.go

+65
Original file line numberDiff line numberDiff line change
@@ -1037,6 +1037,71 @@ func TestReaderReset(t *testing.T) {
10371037
}
10381038
}
10391039

1040+
func TestReaderReadByte(t *testing.T) {
1041+
// Test all 32 possible sub-sequences of these 5 input slices prefixed by
1042+
// their size encoded as a uvarint.
1043+
//
1044+
// Their lengths sum to 400,000, which is over 6 times the Writer ibuf
1045+
// capacity: 6 * maxBlockSize is 393,216.
1046+
inputs := [][]byte{
1047+
bytes.Repeat([]byte{'a'}, 40000),
1048+
bytes.Repeat([]byte{'b'}, 150000),
1049+
bytes.Repeat([]byte{'c'}, 60000),
1050+
bytes.Repeat([]byte{'d'}, 120000),
1051+
bytes.Repeat([]byte{'e'}, 30000),
1052+
}
1053+
loop:
1054+
for i := 0; i < 1<<uint(len(inputs)); i++ {
1055+
var want []int
1056+
buf := new(bytes.Buffer)
1057+
w := NewBufferedWriter(buf)
1058+
p := make([]byte, binary.MaxVarintLen64)
1059+
for j, input := range inputs {
1060+
if i&(1<<uint(j)) == 0 {
1061+
continue
1062+
}
1063+
n := binary.PutUvarint(p, uint64(len(input)))
1064+
if _, err := w.Write(p[:n]); err != nil {
1065+
t.Errorf("i=%#02x: j=%d: Write Uvarint: %v", i, j, err)
1066+
continue loop
1067+
}
1068+
if _, err := w.Write(input); err != nil {
1069+
t.Errorf("i=%#02x: j=%d: Write: %v", i, j, err)
1070+
continue loop
1071+
}
1072+
want = append(want, j)
1073+
}
1074+
if err := w.Close(); err != nil {
1075+
t.Errorf("i=%#02x: Close: %v", i, err)
1076+
continue
1077+
}
1078+
r := NewReader(buf)
1079+
for _, j := range want {
1080+
size, err := binary.ReadUvarint(r)
1081+
if err != nil {
1082+
t.Errorf("i=%#02x: ReadUvarint: %v", i, err)
1083+
continue loop
1084+
}
1085+
if wantedSize := uint64(len(inputs[j])); size != wantedSize {
1086+
t.Errorf("i=%#02x: expected size %d, got %d", i, wantedSize, size)
1087+
continue loop
1088+
}
1089+
got := make([]byte, size)
1090+
if _, err := io.ReadFull(r, got); err != nil {
1091+
t.Errorf("i=%#02x: ReadFull: %v", i, err)
1092+
continue loop
1093+
}
1094+
if err := cmp(got, inputs[j]); err != nil {
1095+
t.Errorf("i=%#02x: %v", i, err)
1096+
continue
1097+
}
1098+
}
1099+
if _, err := r.ReadByte(); err != io.EOF {
1100+
t.Errorf("i=%#02x: expected size EOF, got %v", i, err)
1101+
}
1102+
}
1103+
}
1104+
10401105
func TestWriterReset(t *testing.T) {
10411106
gold := bytes.Repeat([]byte("Not all those who wander are lost;\n"), 10000)
10421107
const n = 20

0 commit comments

Comments
 (0)