libgo/go/archive/zip/register.go

   1 // Copyright 2010 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 zip
   6
   7 import (
   8         "compress/flate"
   9         "errors"
  10         "io"
  11         "sync"
  12 )
  13
  14 // A Compressor returns a new compressing writer, writing to w.
  15 // The WriteCloser's Close method must be used to flush pending data to w.
  16 // The Compressor itself must be safe to invoke from multiple goroutines
  17 // simultaneously, but each returned writer will be used only by
  18 // one goroutine at a time.
  19 type Compressor func(w io.Writer) (io.WriteCloser, error)
  20
  21 // A Decompressor returns a new decompressing reader, reading from r.
  22 // The ReadCloser's Close method must be used to release associated resources.
  23 // The Decompressor itself must be safe to invoke from multiple goroutines
  24 // simultaneously, but each returned reader will be used only by
  25 // one goroutine at a time.
  26 type Decompressor func(r io.Reader) io.ReadCloser
  27
  28 var flateWriterPool sync.Pool
  29
  30 func newFlateWriter(w io.Writer) io.WriteCloser {
  31         fw, ok := flateWriterPool.Get().(*flate.Writer)
  32         if ok {
  33                 fw.Reset(w)
  34         } else {
  35                 fw, _ = flate.NewWriter(w, 5)
  36         }
  37         return &pooledFlateWriter{fw: fw}
  38 }
  39
  40 type pooledFlateWriter struct {
  41         mu sync.Mutex // guards Close and Write
  42         fw *flate.Writer
  43 }
  44
  45 func (w *pooledFlateWriter) Write(p []byte) (n int, err error) {
  46         w.mu.Lock()
  47         defer w.mu.Unlock()
  48         if w.fw == nil {
  49                 return 0, errors.New("Write after Close")
  50         }
  51         return w.fw.Write(p)
  52 }
  53
  54 func (w *pooledFlateWriter) Close() error {
  55         w.mu.Lock()
  56         defer w.mu.Unlock()
  57         var err error
  58         if w.fw != nil {
  59                 err = w.fw.Close()
  60                 flateWriterPool.Put(w.fw)
  61                 w.fw = nil
  62         }
  63         return err
  64 }
  65
  66 var flateReaderPool sync.Pool
  67
  68 func newFlateReader(r io.Reader) io.ReadCloser {
  69         fr, ok := flateReaderPool.Get().(io.ReadCloser)
  70         if ok {
  71                 fr.(flate.Resetter).Reset(r, nil)
  72         } else {
  73                 fr = flate.NewReader(r)
  74         }
  75         return &pooledFlateReader{fr: fr}
  76 }
  77
  78 type pooledFlateReader struct {
  79         mu sync.Mutex // guards Close and Read
  80         fr io.ReadCloser
  81 }
  82
  83 func (r *pooledFlateReader) Read(p []byte) (n int, err error) {
  84         r.mu.Lock()
  85         defer r.mu.Unlock()
  86         if r.fr == nil {
  87                 return 0, errors.New("Read after Close")
  88         }
  89         return r.fr.Read(p)
  90 }
  91
  92 func (r *pooledFlateReader) Close() error {
  93         r.mu.Lock()
  94         defer r.mu.Unlock()
  95         var err error
  96         if r.fr != nil {
  97                 err = r.fr.Close()
  98                 flateReaderPool.Put(r.fr)
  99                 r.fr = nil
 100         }
 101         return err
 102 }
 103
 104 var (
 105         compressors   sync.Map // map[uint16]Compressor
 106         decompressors sync.Map // map[uint16]Decompressor
 107 )
 108
 109 func init() {
 110         compressors.Store(Store, Compressor(func(w io.Writer) (io.WriteCloser, error) { return &nopCloser{w}, nil }))
 111         compressors.Store(Deflate, Compressor(func(w io.Writer) (io.WriteCloser, error) { return newFlateWriter(w), nil }))
 112
 113         decompressors.Store(Store, Decompressor(io.NopCloser))
 114         decompressors.Store(Deflate, Decompressor(newFlateReader))
 115 }
 116
 117 // RegisterDecompressor allows custom decompressors for a specified method ID.
 118 // The common methods Store and Deflate are built in.
 119 func RegisterDecompressor(method uint16, dcomp Decompressor) {
 120         if _, dup := decompressors.LoadOrStore(method, dcomp); dup {
 121                 panic("decompressor already registered")
 122         }
 123 }
 124
 125 // RegisterCompressor registers custom compressors for a specified method ID.
 126 // The common methods Store and Deflate are built in.
 127 func RegisterCompressor(method uint16, comp Compressor) {
 128         if _, dup := compressors.LoadOrStore(method, comp); dup {
 129                 panic("compressor already registered")
 130         }
 131 }
 132
 133 func compressor(method uint16) Compressor {
 134         ci, ok := compressors.Load(method)
 135         if !ok {
 136                 return nil
 137         }
 138         return ci.(Compressor)
 139 }
 140
 141 func decompressor(method uint16) Decompressor {
 142         di, ok := decompressors.Load(method)
 143         if !ok {
 144                 return nil
 145         }
 146         return di.(Decompressor)
 147 }