...
Run Format

Source file src/compress/zlib/writer.go

     1	// Copyright 2009 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 zlib
     6	
     7	import (
     8		"compress/flate"
     9		"fmt"
    10		"hash"
    11		"hash/adler32"
    12		"io"
    13	)
    14	
    15	// These constants are copied from the flate package, so that code that imports
    16	// "compress/zlib" does not also have to import "compress/flate".
    17	const (
    18		NoCompression      = flate.NoCompression
    19		BestSpeed          = flate.BestSpeed
    20		BestCompression    = flate.BestCompression
    21		DefaultCompression = flate.DefaultCompression
    22		HuffmanOnly        = flate.HuffmanOnly
    23	)
    24	
    25	// A Writer takes data written to it and writes the compressed
    26	// form of that data to an underlying writer (see NewWriter).
    27	type Writer struct {
    28		w           io.Writer
    29		level       int
    30		dict        []byte
    31		compressor  *flate.Writer
    32		digest      hash.Hash32
    33		err         error
    34		scratch     [4]byte
    35		wroteHeader bool
    36	}
    37	
    38	// NewWriter creates a new Writer.
    39	// Writes to the returned Writer are compressed and written to w.
    40	//
    41	// It is the caller's responsibility to call Close on the WriteCloser when done.
    42	// Writes may be buffered and not flushed until Close.
    43	func NewWriter(w io.Writer) *Writer {
    44		z, _ := NewWriterLevelDict(w, DefaultCompression, nil)
    45		return z
    46	}
    47	
    48	// NewWriterLevel is like NewWriter but specifies the compression level instead
    49	// of assuming DefaultCompression.
    50	//
    51	// The compression level can be DefaultCompression, NoCompression, HuffmanOnly
    52	// or any integer value between BestSpeed and BestCompression inclusive.
    53	// The error returned will be nil if the level is valid.
    54	func NewWriterLevel(w io.Writer, level int) (*Writer, error) {
    55		return NewWriterLevelDict(w, level, nil)
    56	}
    57	
    58	// NewWriterLevelDict is like NewWriterLevel but specifies a dictionary to
    59	// compress with.
    60	//
    61	// The dictionary may be nil. If not, its contents should not be modified until
    62	// the Writer is closed.
    63	func NewWriterLevelDict(w io.Writer, level int, dict []byte) (*Writer, error) {
    64		if level < HuffmanOnly || level > BestCompression {
    65			return nil, fmt.Errorf("zlib: invalid compression level: %d", level)
    66		}
    67		return &Writer{
    68			w:     w,
    69			level: level,
    70			dict:  dict,
    71		}, nil
    72	}
    73	
    74	// Reset clears the state of the Writer z such that it is equivalent to its
    75	// initial state from NewWriterLevel or NewWriterLevelDict, but instead writing
    76	// to w.
    77	func (z *Writer) Reset(w io.Writer) {
    78		z.w = w
    79		// z.level and z.dict left unchanged.
    80		if z.compressor != nil {
    81			z.compressor.Reset(w)
    82		}
    83		if z.digest != nil {
    84			z.digest.Reset()
    85		}
    86		z.err = nil
    87		z.scratch = [4]byte{}
    88		z.wroteHeader = false
    89	}
    90	
    91	// writeHeader writes the ZLIB header.
    92	func (z *Writer) writeHeader() (err error) {
    93		z.wroteHeader = true
    94		// ZLIB has a two-byte header (as documented in RFC 1950).
    95		// The first four bits is the CINFO (compression info), which is 7 for the default deflate window size.
    96		// The next four bits is the CM (compression method), which is 8 for deflate.
    97		z.scratch[0] = 0x78
    98		// The next two bits is the FLEVEL (compression level). The four values are:
    99		// 0=fastest, 1=fast, 2=default, 3=best.
   100		// The next bit, FDICT, is set if a dictionary is given.
   101		// The final five FCHECK bits form a mod-31 checksum.
   102		switch z.level {
   103		case -2, 0, 1:
   104			z.scratch[1] = 0 << 6
   105		case 2, 3, 4, 5:
   106			z.scratch[1] = 1 << 6
   107		case 6, -1:
   108			z.scratch[1] = 2 << 6
   109		case 7, 8, 9:
   110			z.scratch[1] = 3 << 6
   111		default:
   112			panic("unreachable")
   113		}
   114		if z.dict != nil {
   115			z.scratch[1] |= 1 << 5
   116		}
   117		z.scratch[1] += uint8(31 - (uint16(z.scratch[0])<<8+uint16(z.scratch[1]))%31)
   118		if _, err = z.w.Write(z.scratch[0:2]); err != nil {
   119			return err
   120		}
   121		if z.dict != nil {
   122			// The next four bytes are the Adler-32 checksum of the dictionary.
   123			checksum := adler32.Checksum(z.dict)
   124			z.scratch[0] = uint8(checksum >> 24)
   125			z.scratch[1] = uint8(checksum >> 16)
   126			z.scratch[2] = uint8(checksum >> 8)
   127			z.scratch[3] = uint8(checksum >> 0)
   128			if _, err = z.w.Write(z.scratch[0:4]); err != nil {
   129				return err
   130			}
   131		}
   132		if z.compressor == nil {
   133			// Initialize deflater unless the Writer is being reused
   134			// after a Reset call.
   135			z.compressor, err = flate.NewWriterDict(z.w, z.level, z.dict)
   136			if err != nil {
   137				return err
   138			}
   139			z.digest = adler32.New()
   140		}
   141		return nil
   142	}
   143	
   144	// Write writes a compressed form of p to the underlying io.Writer. The
   145	// compressed bytes are not necessarily flushed until the Writer is closed or
   146	// explicitly flushed.
   147	func (z *Writer) Write(p []byte) (n int, err error) {
   148		if !z.wroteHeader {
   149			z.err = z.writeHeader()
   150		}
   151		if z.err != nil {
   152			return 0, z.err
   153		}
   154		if len(p) == 0 {
   155			return 0, nil
   156		}
   157		n, err = z.compressor.Write(p)
   158		if err != nil {
   159			z.err = err
   160			return
   161		}
   162		z.digest.Write(p)
   163		return
   164	}
   165	
   166	// Flush flushes the Writer to its underlying io.Writer.
   167	func (z *Writer) Flush() error {
   168		if !z.wroteHeader {
   169			z.err = z.writeHeader()
   170		}
   171		if z.err != nil {
   172			return z.err
   173		}
   174		z.err = z.compressor.Flush()
   175		return z.err
   176	}
   177	
   178	// Close closes the Writer, flushing any unwritten data to the underlying
   179	// io.Writer, but does not close the underlying io.Writer.
   180	func (z *Writer) Close() error {
   181		if !z.wroteHeader {
   182			z.err = z.writeHeader()
   183		}
   184		if z.err != nil {
   185			return z.err
   186		}
   187		z.err = z.compressor.Close()
   188		if z.err != nil {
   189			return z.err
   190		}
   191		checksum := z.digest.Sum32()
   192		// ZLIB (RFC 1950) is big-endian, unlike GZIP (RFC 1952).
   193		z.scratch[0] = uint8(checksum >> 24)
   194		z.scratch[1] = uint8(checksum >> 16)
   195		z.scratch[2] = uint8(checksum >> 8)
   196		z.scratch[3] = uint8(checksum >> 0)
   197		_, z.err = z.w.Write(z.scratch[0:4])
   198		return z.err
   199	}
   200	

View as plain text