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

View as plain text