...
Run Format

Source file src/compress/lzw/writer.go

Documentation: compress/lzw

     1  // Copyright 2011 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 lzw
     6  
     7  import (
     8  	"bufio"
     9  	"errors"
    10  	"fmt"
    11  	"io"
    12  )
    13  
    14  // A writer is a buffered, flushable writer.
    15  type writer interface {
    16  	io.ByteWriter
    17  	Flush() error
    18  }
    19  
    20  // An errWriteCloser is an io.WriteCloser that always returns a given error.
    21  type errWriteCloser struct {
    22  	err error
    23  }
    24  
    25  func (e *errWriteCloser) Write([]byte) (int, error) {
    26  	return 0, e.err
    27  }
    28  
    29  func (e *errWriteCloser) Close() error {
    30  	return e.err
    31  }
    32  
    33  const (
    34  	// A code is a 12 bit value, stored as a uint32 when encoding to avoid
    35  	// type conversions when shifting bits.
    36  	maxCode     = 1<<12 - 1
    37  	invalidCode = 1<<32 - 1
    38  	// There are 1<<12 possible codes, which is an upper bound on the number of
    39  	// valid hash table entries at any given point in time. tableSize is 4x that.
    40  	tableSize = 4 * 1 << 12
    41  	tableMask = tableSize - 1
    42  	// A hash table entry is a uint32. Zero is an invalid entry since the
    43  	// lower 12 bits of a valid entry must be a non-literal code.
    44  	invalidEntry = 0
    45  )
    46  
    47  // encoder is LZW compressor.
    48  type encoder struct {
    49  	// w is the writer that compressed bytes are written to.
    50  	w writer
    51  	// order, write, bits, nBits and width are the state for
    52  	// converting a code stream into a byte stream.
    53  	order Order
    54  	write func(*encoder, uint32) error
    55  	bits  uint32
    56  	nBits uint
    57  	width uint
    58  	// litWidth is the width in bits of literal codes.
    59  	litWidth uint
    60  	// hi is the code implied by the next code emission.
    61  	// overflow is the code at which hi overflows the code width.
    62  	hi, overflow uint32
    63  	// savedCode is the accumulated code at the end of the most recent Write
    64  	// call. It is equal to invalidCode if there was no such call.
    65  	savedCode uint32
    66  	// err is the first error encountered during writing. Closing the encoder
    67  	// will make any future Write calls return errClosed
    68  	err error
    69  	// table is the hash table from 20-bit keys to 12-bit values. Each table
    70  	// entry contains key<<12|val and collisions resolve by linear probing.
    71  	// The keys consist of a 12-bit code prefix and an 8-bit byte suffix.
    72  	// The values are a 12-bit code.
    73  	table [tableSize]uint32
    74  }
    75  
    76  // writeLSB writes the code c for "Least Significant Bits first" data.
    77  func (e *encoder) writeLSB(c uint32) error {
    78  	e.bits |= c << e.nBits
    79  	e.nBits += e.width
    80  	for e.nBits >= 8 {
    81  		if err := e.w.WriteByte(uint8(e.bits)); err != nil {
    82  			return err
    83  		}
    84  		e.bits >>= 8
    85  		e.nBits -= 8
    86  	}
    87  	return nil
    88  }
    89  
    90  // writeMSB writes the code c for "Most Significant Bits first" data.
    91  func (e *encoder) writeMSB(c uint32) error {
    92  	e.bits |= c << (32 - e.width - e.nBits)
    93  	e.nBits += e.width
    94  	for e.nBits >= 8 {
    95  		if err := e.w.WriteByte(uint8(e.bits >> 24)); err != nil {
    96  			return err
    97  		}
    98  		e.bits <<= 8
    99  		e.nBits -= 8
   100  	}
   101  	return nil
   102  }
   103  
   104  // errOutOfCodes is an internal error that means that the encoder has run out
   105  // of unused codes and a clear code needs to be sent next.
   106  var errOutOfCodes = errors.New("lzw: out of codes")
   107  
   108  // incHi increments e.hi and checks for both overflow and running out of
   109  // unused codes. In the latter case, incHi sends a clear code, resets the
   110  // encoder state and returns errOutOfCodes.
   111  func (e *encoder) incHi() error {
   112  	e.hi++
   113  	if e.hi == e.overflow {
   114  		e.width++
   115  		e.overflow <<= 1
   116  	}
   117  	if e.hi == maxCode {
   118  		clear := uint32(1) << e.litWidth
   119  		if err := e.write(e, clear); err != nil {
   120  			return err
   121  		}
   122  		e.width = e.litWidth + 1
   123  		e.hi = clear + 1
   124  		e.overflow = clear << 1
   125  		for i := range e.table {
   126  			e.table[i] = invalidEntry
   127  		}
   128  		return errOutOfCodes
   129  	}
   130  	return nil
   131  }
   132  
   133  // Write writes a compressed representation of p to e's underlying writer.
   134  func (e *encoder) Write(p []byte) (n int, err error) {
   135  	if e.err != nil {
   136  		return 0, e.err
   137  	}
   138  	if len(p) == 0 {
   139  		return 0, nil
   140  	}
   141  	if maxLit := uint8(1<<e.litWidth - 1); maxLit != 0xff {
   142  		for _, x := range p {
   143  			if x > maxLit {
   144  				e.err = errors.New("lzw: input byte too large for the litWidth")
   145  				return 0, e.err
   146  			}
   147  		}
   148  	}
   149  	n = len(p)
   150  	code := e.savedCode
   151  	if code == invalidCode {
   152  		// The first code sent is always a literal code.
   153  		code, p = uint32(p[0]), p[1:]
   154  	}
   155  loop:
   156  	for _, x := range p {
   157  		literal := uint32(x)
   158  		key := code<<8 | literal
   159  		// If there is a hash table hit for this key then we continue the loop
   160  		// and do not emit a code yet.
   161  		hash := (key>>12 ^ key) & tableMask
   162  		for h, t := hash, e.table[hash]; t != invalidEntry; {
   163  			if key == t>>12 {
   164  				code = t & maxCode
   165  				continue loop
   166  			}
   167  			h = (h + 1) & tableMask
   168  			t = e.table[h]
   169  		}
   170  		// Otherwise, write the current code, and literal becomes the start of
   171  		// the next emitted code.
   172  		if e.err = e.write(e, code); e.err != nil {
   173  			return 0, e.err
   174  		}
   175  		code = literal
   176  		// Increment e.hi, the next implied code. If we run out of codes, reset
   177  		// the encoder state (including clearing the hash table) and continue.
   178  		if err1 := e.incHi(); err1 != nil {
   179  			if err1 == errOutOfCodes {
   180  				continue
   181  			}
   182  			e.err = err1
   183  			return 0, e.err
   184  		}
   185  		// Otherwise, insert key -> e.hi into the map that e.table represents.
   186  		for {
   187  			if e.table[hash] == invalidEntry {
   188  				e.table[hash] = (key << 12) | e.hi
   189  				break
   190  			}
   191  			hash = (hash + 1) & tableMask
   192  		}
   193  	}
   194  	e.savedCode = code
   195  	return n, nil
   196  }
   197  
   198  // Close closes the encoder, flushing any pending output. It does not close or
   199  // flush e's underlying writer.
   200  func (e *encoder) Close() error {
   201  	if e.err != nil {
   202  		if e.err == errClosed {
   203  			return nil
   204  		}
   205  		return e.err
   206  	}
   207  	// Make any future calls to Write return errClosed.
   208  	e.err = errClosed
   209  	// Write the savedCode if valid.
   210  	if e.savedCode != invalidCode {
   211  		if err := e.write(e, e.savedCode); err != nil {
   212  			return err
   213  		}
   214  		if err := e.incHi(); err != nil && err != errOutOfCodes {
   215  			return err
   216  		}
   217  	}
   218  	// Write the eof code.
   219  	eof := uint32(1)<<e.litWidth + 1
   220  	if err := e.write(e, eof); err != nil {
   221  		return err
   222  	}
   223  	// Write the final bits.
   224  	if e.nBits > 0 {
   225  		if e.order == MSB {
   226  			e.bits >>= 24
   227  		}
   228  		if err := e.w.WriteByte(uint8(e.bits)); err != nil {
   229  			return err
   230  		}
   231  	}
   232  	return e.w.Flush()
   233  }
   234  
   235  // NewWriter creates a new io.WriteCloser.
   236  // Writes to the returned io.WriteCloser are compressed and written to w.
   237  // It is the caller's responsibility to call Close on the WriteCloser when
   238  // finished writing.
   239  // The number of bits to use for literal codes, litWidth, must be in the
   240  // range [2,8] and is typically 8. Input bytes must be less than 1<<litWidth.
   241  func NewWriter(w io.Writer, order Order, litWidth int) io.WriteCloser {
   242  	var write func(*encoder, uint32) error
   243  	switch order {
   244  	case LSB:
   245  		write = (*encoder).writeLSB
   246  	case MSB:
   247  		write = (*encoder).writeMSB
   248  	default:
   249  		return &errWriteCloser{errors.New("lzw: unknown order")}
   250  	}
   251  	if litWidth < 2 || 8 < litWidth {
   252  		return &errWriteCloser{fmt.Errorf("lzw: litWidth %d out of range", litWidth)}
   253  	}
   254  	bw, ok := w.(writer)
   255  	if !ok {
   256  		bw = bufio.NewWriter(w)
   257  	}
   258  	lw := uint(litWidth)
   259  	return &encoder{
   260  		w:         bw,
   261  		order:     order,
   262  		write:     write,
   263  		width:     1 + lw,
   264  		litWidth:  lw,
   265  		hi:        1<<lw + 1,
   266  		overflow:  1 << (lw + 1),
   267  		savedCode: invalidCode,
   268  	}
   269  }
   270  

View as plain text