huffman-archiver/internal/io/buffered_io.go

package io

import "io"

const BufferSize = 1024 // bytes

type BufferedReadSeeker struct {
	Reader io.ReadSeeker
	Buf    []byte
	Size   int
	Cnt    int
	Skip   int
}

func NewBufferedReader(reader io.ReadSeeker) (res *BufferedReadSeeker) {
	res = &BufferedReadSeeker{
		Reader: reader,
		Buf:    make([]byte, BufferSize),
		Size:   0,
		Cnt:    0,
		Skip:   0,
	}

	res.UpdateBuffer()
	return
}

func (r *BufferedReadSeeker) SetLastBitsSkip(n int) {
	r.Skip = n
	r.Cnt -= r.Skip
}

var readBuf = make([]byte, BufferSize-1)

func (r *BufferedReadSeeker) UpdateBuffer() {
	if r.Cnt == 0 {
		n, _ := r.Reader.Read(readBuf)

		if n == 0 {
			r.Size = 0
			return
		}

		r.Buf[0] = r.Buf[BufferSize-1]
		for i := 1; i < BufferSize; i++ {
			r.Buf[i] = readBuf[i-1]
		}

		r.Cnt = n * 8

		r.Size = n + 1
	}
}

func (r *BufferedReadSeeker) GetOffset() int {
	return (r.Size*8 - r.Skip - r.Cnt) / 8
}

func (r *BufferedReadSeeker) IsEOF() bool {
	r.UpdateBuffer()
	return r.Size == 0
}

func (r *BufferedReadSeeker) ReadBit() (v byte) {
	v = (r.Buf[r.GetOffset()] & 0x80) >> 7
	r.Buf[r.GetOffset()] <<= 1
	r.Cnt--

	if r.Cnt == 0 {
		r.UpdateBuffer()
	}

	return
}

func (r *BufferedReadSeeker) ReadBits(n int) (v uint64) {
	if n > 64 {
		panic("cannot read more than 64 bits")
	}

	for i := 0; i < n; i++ {
		v = v<<1 | uint64(r.ReadBit())
	}

	return
}

func (r *BufferedReadSeeker) ReadByte() byte {
	return byte(r.ReadBits(8))
}

func (r *BufferedReadSeeker) ReadShort() int16 {
	return int16(r.ReadBits(16))
}

func (r *BufferedReadSeeker) SeekFromStart(offset int64) {
	if _, err := r.Reader.Seek(offset, io.SeekStart); err != nil {
		panic(err)
	}
}

type BufferedWriter struct {
	Writer io.Writer
	Buf    []byte
	Offset int
	Cnt    int
}

func NewBufferedWriter(writer io.Writer) *BufferedWriter {
	return &BufferedWriter{
		Writer: writer,
		Buf:    make([]byte, BufferSize),
		Offset: 0,
		Cnt:    0,
	}
}

func (w *BufferedWriter) WriteBit(x byte) {
	w.Buf[w.Offset] <<= 1

	if x != 0x0 {
		w.Buf[w.Offset] |= 0x1
	}

	w.Cnt++

	if w.Cnt == 8 {
		w.Offset++
		w.Cnt = 0
	}

	if w.Offset == BufferSize {
		w.Flush()
	}
}

func (w *BufferedWriter) WriteBits(x []byte) {
	for _, v := range x {
		w.WriteBit(v)
	}
}

func (w *BufferedWriter) WriteBitsFromNumber(x uint64, n int) {
	if n > 64 {
		panic("cannot read more than 64 bits")
	}

	for i := 0; i < 64; i++ {
		if 64-n <= i {
			w.WriteBit(byte((x & 0x8000000000000000) >> 63))
		}
		x <<= 1
	}
}

func (w *BufferedWriter) WriteByte(x byte) {
	w.WriteBitsFromNumber(uint64(x), 8)
}

func (w *BufferedWriter) WriteShort(x int16) {
	w.WriteBitsFromNumber(uint64(x), 16)
}

func (w *BufferedWriter) Flush() {
	for w.Cnt > 0 {
		w.WriteBit(0x0)
	}

	if _, err := w.Writer.Write(w.Buf[:w.Offset]); err != nil {
		panic(err)
	}

	w.Offset = 0
	for k := range w.Buf {
		w.Buf[k] = 0x0
	}
}