YAPC::EU 2018 Glasgow Update!

Task/Bitwise-IO/Go/bitwise-io-1.go

@@ -0,0 +1,203 @@
+// Package bit provides bit-wise IO to an io.Writer and from an io.Reader.
+package bit
+
+import (
+	"bufio"
+	"errors"
+	"io"
+)
+
+// Order specifies the bit ordering within a byte stream.
+type Order int
+
+const (
+	// LSB is for Least Significant Bits first
+	LSB Order = iota
+	// MSB is for Most  Significant Bits first
+	MSB
+)
+
+// ==== Writing / Encoding ====
+
+type writer interface {
+	io.ByteWriter
+	Flush() error
+}
+
+// Writer implements bit-wise writing to an io.Writer.
+type Writer struct {
+	w     writer
+	order Order
+	write func(uint32, uint) error // writeLSB or writeMSB
+	bits  uint32
+	nBits uint
+	err   error
+}
+
+// writeLSB writes `width` bits of `c` in LSB order.
+func (w *Writer) writeLSB(c uint32, width uint) error {
+	w.bits |= c << w.nBits
+	w.nBits += width
+	for w.nBits >= 8 {
+		if err := w.w.WriteByte(uint8(w.bits)); err != nil {
+			return err
+		}
+		w.bits >>= 8
+		w.nBits -= 8
+	}
+	return nil
+}
+
+// writeMSB writes `width` bits of `c` in MSB order.
+func (w *Writer) writeMSB(c uint32, width uint) error {
+	w.bits |= c << (32 - width - w.nBits)
+	w.nBits += width
+	for w.nBits >= 8 {
+		if err := w.w.WriteByte(uint8(w.bits >> 24)); err != nil {
+			return err
+		}
+		w.bits <<= 8
+		w.nBits -= 8
+	}
+	return nil
+}
+
+// WriteBits writes up to 16 bits of `c` to the underlying writer.
+// Even for MSB ordering the bits are taken from the lower bits of `c`.
+// (e.g. WriteBits(0x0f,4) writes four 1 bits).
+func (w *Writer) WriteBits(c uint16, width uint) error {
+	if w.err == nil {
+		w.err = w.write(uint32(c), width)
+	}
+	return w.err
+}
+
+var errClosed = errors.New("bit reader/writer is closed")
+
+// Close closes the writer, flushing any pending output.
+// It does not close the underlying writer.
+func (w *Writer) Close() error {
+	if w.err != nil {
+		if w.err == errClosed {
+			return nil
+		}
+		return w.err
+	}
+	// Write the final bits (zero padded).
+	if w.nBits > 0 {
+		if w.order == MSB {
+			w.bits >>= 24
+		}
+		if w.err = w.w.WriteByte(uint8(w.bits)); w.err != nil {
+			return w.err
+		}
+	}
+	w.err = w.w.Flush()
+	if w.err != nil {
+		return w.err
+	}
+
+	// Make any future calls to Write return errClosed.
+	w.err = errClosed
+	return nil
+}
+
+// NewWriter returns a new bit Writer that writes completed bytes to `w`.
+func NewWriter(w io.Writer, order Order) *Writer {
+	bw := &Writer{order: order}
+	switch order {
+	case LSB:
+		bw.write = bw.writeLSB
+	case MSB:
+		bw.write = bw.writeMSB
+	default:
+		bw.err = errors.New("bit writer: unknown order")
+		return bw
+	}
+	if byteWriter, ok := w.(writer); ok {
+		bw.w = byteWriter
+	} else {
+		bw.w = bufio.NewWriter(w)
+	}
+	return bw
+}
+
+// ==== Reading / Decoding ====
+
+// Reader implements bit-wise reading from an io.Reader.
+type Reader struct {
+	r     io.ByteReader
+	bits  uint32
+	nBits uint
+	read  func(width uint) (uint16, error) // readLSB or readMSB
+	err   error
+}
+
+func (r *Reader) readLSB(width uint) (uint16, error) {
+	for r.nBits < width {
+		x, err := r.r.ReadByte()
+		if err != nil {
+			return 0, err
+		}
+		r.bits |= uint32(x) << r.nBits
+		r.nBits += 8
+	}
+	bits := uint16(r.bits & (1<<width - 1))
+	r.bits >>= width
+	r.nBits -= width
+	return bits, nil
+}
+
+func (r *Reader) readMSB(width uint) (uint16, error) {
+	for r.nBits < width {
+		x, err := r.r.ReadByte()
+		if err != nil {
+			return 0, err
+		}
+		r.bits |= uint32(x) << (24 - r.nBits)
+		r.nBits += 8
+	}
+	bits := uint16(r.bits >> (32 - width))
+	r.bits <<= width
+	r.nBits -= width
+	return bits, nil
+}
+
+// ReadBits reads up to 16 bits from the underlying reader.
+func (r *Reader) ReadBits(width uint) (uint16, error) {
+	var bits uint16
+	if r.err == nil {
+		bits, r.err = r.read(width)
+	}
+	return bits, r.err
+}
+
+// Close closes the reader.
+// It does not close the underlying reader.
+func (r *Reader) Close() error {
+	if r.err != nil && r.err != errClosed {
+		return r.err
+	}
+	r.err = errClosed
+	return nil
+}
+
+// NewReader returns a new bit Reader that reads bytes from `r`.
+func NewReader(r io.Reader, order Order) *Reader {
+	br := new(Reader)
+	switch order {
+	case LSB:
+		br.read = br.readLSB
+	case MSB:
+		br.read = br.readMSB
+	default:
+		br.err = errors.New("bit writer: unknown order")
+		return br
+	}
+	if byteReader, ok := r.(io.ByteReader); ok {
+		br.r = byteReader
+	} else {
+		br.r = bufio.NewReader(r)
+	}
+	return br
+}