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 bzip2
6
7import "github.com/dsnet/compress/internal/errors"
8
9// rleDone is a special "error" to indicate that the RLE stage is done.
10var rleDone = errorf(errors.Unknown, "RLE1 stage is completed")
11
12// runLengthEncoding implements the first RLE stage of bzip2. Every sequence
13// of 4..255 duplicated bytes is replaced by only the first 4 bytes, and a
14// single byte representing the repeat length. Similar to the C bzip2
15// implementation, the encoder will always terminate repeat sequences with a
16// count (even if it is the end of the buffer), and it will also never produce
17// run lengths of 256..259. The decoder can handle the latter case.
18//
19// For example, if the input was:
20//	input:  "AAAAAAABBBBCCCD"
21//
22// Then the output will be:
23//	output: "AAAA\x03BBBB\x00CCCD"
24type runLengthEncoding struct {
25	buf     []byte
26	idx     int
27	lastVal byte
28	lastCnt int
29}
30
31func (rle *runLengthEncoding) Init(buf []byte) {
32	*rle = runLengthEncoding{buf: buf}
33}
34
35func (rle *runLengthEncoding) Write(buf []byte) (int, error) {
36	for i, b := range buf {
37		if rle.lastVal != b {
38			rle.lastCnt = 0
39		}
40		rle.lastCnt++
41		switch {
42		case rle.lastCnt < 4:
43			if rle.idx >= len(rle.buf) {
44				return i, rleDone
45			}
46			rle.buf[rle.idx] = b
47			rle.idx++
48		case rle.lastCnt == 4:
49			if rle.idx+1 >= len(rle.buf) {
50				return i, rleDone
51			}
52			rle.buf[rle.idx] = b
53			rle.idx++
54			rle.buf[rle.idx] = 0
55			rle.idx++
56		case rle.lastCnt < 256:
57			rle.buf[rle.idx-1]++
58		default:
59			if rle.idx >= len(rle.buf) {
60				return i, rleDone
61			}
62			rle.lastCnt = 1
63			rle.buf[rle.idx] = b
64			rle.idx++
65		}
66		rle.lastVal = b
67	}
68	return len(buf), nil
69}
70
71func (rle *runLengthEncoding) Read(buf []byte) (int, error) {
72	for i := range buf {
73		switch {
74		case rle.lastCnt == -4:
75			if rle.idx >= len(rle.buf) {
76				return i, errorf(errors.Corrupted, "missing terminating run-length repeater")
77			}
78			rle.lastCnt = int(rle.buf[rle.idx])
79			rle.idx++
80			if rle.lastCnt > 0 {
81				break // Break the switch
82			}
83			fallthrough // Count was zero, continue the work
84		case rle.lastCnt <= 0:
85			if rle.idx >= len(rle.buf) {
86				return i, rleDone
87			}
88			b := rle.buf[rle.idx]
89			rle.idx++
90			if b != rle.lastVal {
91				rle.lastCnt = 0
92				rle.lastVal = b
93			}
94		}
95		buf[i] = rle.lastVal
96		rle.lastCnt--
97	}
98	return len(buf), nil
99}
100
101func (rle *runLengthEncoding) Bytes() []byte { return rle.buf[:rle.idx] }
102