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		"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 that satisfies writes by compressing data
    38	// 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	// writeHeader writes the ZLIB header.
    74	func (z *Writer) writeHeader() (err error) {
    75		z.wroteHeader = true
    76		// ZLIB has a two-byte header (as documented in RFC 1950).
    77		// The first four bits is the CINFO (compression info), which is 7 for the default deflate window size.
    78		// The next four bits is the CM (compression method), which is 8 for deflate.
    79		z.scratch[0] = 0x78
    80		// The next two bits is the FLEVEL (compression level). The four values are:
    81		// 0=fastest, 1=fast, 2=default, 3=best.
    82		// The next bit, FDICT, is set if a dictionary is given.
    83		// The final five FCHECK bits form a mod-31 checksum.
    84		switch z.level {
    85		case 0, 1:
    86			z.scratch[1] = 0 << 6
    87		case 2, 3, 4, 5:
    88			z.scratch[1] = 1 << 6
    89		case 6, -1:
    90			z.scratch[1] = 2 << 6
    91		case 7, 8, 9:
    92			z.scratch[1] = 3 << 6
    93		default:
    94			panic("unreachable")
    95		}
    96		if z.dict != nil {
    97			z.scratch[1] |= 1 << 5
    98		}
    99		z.scratch[1] += uint8(31 - (uint16(z.scratch[0])<<8+uint16(z.scratch[1]))%31)
   100		if _, err = z.w.Write(z.scratch[0:2]); err != nil {
   101			return err
   102		}
   103		if z.dict != nil {
   104			// The next four bytes are the Adler-32 checksum of the dictionary.
   105			checksum := adler32.Checksum(z.dict)
   106			z.scratch[0] = uint8(checksum >> 24)
   107			z.scratch[1] = uint8(checksum >> 16)
   108			z.scratch[2] = uint8(checksum >> 8)
   109			z.scratch[3] = uint8(checksum >> 0)
   110			if _, err = z.w.Write(z.scratch[0:4]); err != nil {
   111				return err
   112			}
   113		}
   114		z.compressor, err = flate.NewWriterDict(z.w, z.level, z.dict)
   115		if err != nil {
   116			return err
   117		}
   118		z.digest = adler32.New()
   119		return nil
   120	}
   121	
   122	// Write writes a compressed form of p to the underlying io.Writer. The
   123	// compressed bytes are not necessarily flushed until the Writer is closed or
   124	// explicitly flushed.
   125	func (z *Writer) Write(p []byte) (n int, err error) {
   126		if !z.wroteHeader {
   127			z.err = z.writeHeader()
   128		}
   129		if z.err != nil {
   130			return 0, z.err
   131		}
   132		if len(p) == 0 {
   133			return 0, nil
   134		}
   135		n, err = z.compressor.Write(p)
   136		if err != nil {
   137			z.err = err
   138			return
   139		}
   140		z.digest.Write(p)
   141		return
   142	}
   143	
   144	// Flush flushes the Writer to its underlying io.Writer.
   145	func (z *Writer) Flush() error {
   146		if !z.wroteHeader {
   147			z.err = z.writeHeader()
   148		}
   149		if z.err != nil {
   150			return z.err
   151		}
   152		z.err = z.compressor.Flush()
   153		return z.err
   154	}
   155	
   156	// Calling Close does not close the wrapped io.Writer originally passed to NewWriter.
   157	func (z *Writer) Close() error {
   158		if !z.wroteHeader {
   159			z.err = z.writeHeader()
   160		}
   161		if z.err != nil {
   162			return z.err
   163		}
   164		z.err = z.compressor.Close()
   165		if z.err != nil {
   166			return z.err
   167		}
   168		checksum := z.digest.Sum32()
   169		// ZLIB (RFC 1950) is big-endian, unlike GZIP (RFC 1952).
   170		z.scratch[0] = uint8(checksum >> 24)
   171		z.scratch[1] = uint8(checksum >> 16)
   172		z.scratch[2] = uint8(checksum >> 8)
   173		z.scratch[3] = uint8(checksum >> 0)
   174		_, z.err = z.w.Write(z.scratch[0:4])
   175		return z.err
   176	}