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

View as plain text