gtsocial-umbx

Unnamed repository; edit this file 'description' to name the repository.
Log | Files | Refs | README | LICENSE

writer.go (5596B)


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