writer.go raw

   1  // Copyright 2019 The Go Authors. All rights reserved.
   2  // Use of this source code is governed by a BSD-style
   3  // license that can be found in the LICENSE file.
   4  
   5  package ccitt
   6  
   7  import (
   8  	"encoding/binary"
   9  	"io"
  10  )
  11  
  12  type bitWriter struct {
  13  	w io.Writer
  14  
  15  	// order is whether to process w's bytes LSB first or MSB first.
  16  	order Order
  17  
  18  	// The high nBits bits of the bits field hold encoded bits to be written to w.
  19  	bits  uint64
  20  	nBits uint32
  21  
  22  	// bytes[:bw] holds encoded bytes not yet written to w.
  23  	// Overflow protection is ensured by using a multiple of 8 as bytes length.
  24  	bw    uint32
  25  	bytes [1024]uint8
  26  }
  27  
  28  // flushBits copies 64 bits from b.bits to b.bytes. If b.bytes is then full, it
  29  // is written to b.w.
  30  func (b *bitWriter) flushBits() error {
  31  	binary.BigEndian.PutUint64(b.bytes[b.bw:], b.bits)
  32  	b.bits = 0
  33  	b.nBits = 0
  34  	b.bw += 8
  35  	if b.bw < uint32(len(b.bytes)) {
  36  		return nil
  37  	}
  38  	b.bw = 0
  39  	if b.order != MSB {
  40  		reverseBitsWithinBytes(b.bytes[:])
  41  	}
  42  	_, err := b.w.Write(b.bytes[:])
  43  	return err
  44  }
  45  
  46  // close finalizes a bitcode stream by writing any
  47  // pending bits to bitWriter's underlying io.Writer.
  48  func (b *bitWriter) close() error {
  49  	// Write any encoded bits to bytes.
  50  	if b.nBits > 0 {
  51  		binary.BigEndian.PutUint64(b.bytes[b.bw:], b.bits)
  52  		b.bw += (b.nBits + 7) >> 3
  53  	}
  54  
  55  	if b.order != MSB {
  56  		reverseBitsWithinBytes(b.bytes[:b.bw])
  57  	}
  58  
  59  	// Write b.bw bytes to b.w.
  60  	_, err := b.w.Write(b.bytes[:b.bw])
  61  	return err
  62  }
  63  
  64  // alignToByteBoundary rounds b.nBits up to a multiple of 8.
  65  // If all 64 bits are used, flush them to bitWriter's bytes.
  66  func (b *bitWriter) alignToByteBoundary() error {
  67  	if b.nBits = (b.nBits + 7) &^ 7; b.nBits == 64 {
  68  		return b.flushBits()
  69  	}
  70  	return nil
  71  }
  72  
  73  // writeCode writes a variable length bitcode to b's underlying io.Writer.
  74  func (b *bitWriter) writeCode(bs bitString) error {
  75  	bits := bs.bits
  76  	nBits := bs.nBits
  77  	if 64-b.nBits >= nBits {
  78  		// b.bits has sufficient room for storing nBits bits.
  79  		b.bits |= uint64(bits) << (64 - nBits - b.nBits)
  80  		b.nBits += nBits
  81  		if b.nBits == 64 {
  82  			return b.flushBits()
  83  		}
  84  		return nil
  85  	}
  86  
  87  	// Number of leading bits that fill b.bits.
  88  	i := 64 - b.nBits
  89  
  90  	// Fill b.bits then flush and write remaining bits.
  91  	b.bits |= uint64(bits) >> (nBits - i)
  92  	b.nBits = 64
  93  
  94  	if err := b.flushBits(); err != nil {
  95  		return err
  96  	}
  97  
  98  	nBits -= i
  99  	b.bits = uint64(bits) << (64 - nBits)
 100  	b.nBits = nBits
 101  	return nil
 102  }
 103