multi.mx raw

   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 io
   6  
   7  type eofReader struct{}
   8  
   9  func (eofReader) Read([]byte) (int, error) {
  10  	return 0, EOF
  11  }
  12  
  13  type multiReader struct {
  14  	readers []Reader
  15  }
  16  
  17  func (mr *multiReader) Read(p []byte) (n int, err error) {
  18  	for len(mr.readers) > 0 {
  19  		// Optimization to flatten nested multiReaders (Issue 13558).
  20  		if len(mr.readers) == 1 {
  21  			if r, ok := mr.readers[0].(*multiReader); ok {
  22  				mr.readers = r.readers
  23  				continue
  24  			}
  25  		}
  26  		n, err = mr.readers[0].Read(p)
  27  		if err == EOF {
  28  			// Use eofReader instead of nil to avoid nil panic
  29  			// after performing flatten (Issue 18232).
  30  			mr.readers[0] = eofReader{} // permit earlier GC
  31  			mr.readers = mr.readers[1:]
  32  		}
  33  		if n > 0 || err != EOF {
  34  			if err == EOF && len(mr.readers) > 0 {
  35  				// Don't return EOF yet. More readers remain.
  36  				err = nil
  37  			}
  38  			return
  39  		}
  40  	}
  41  	return 0, EOF
  42  }
  43  
  44  func (mr *multiReader) WriteTo(w Writer) (sum int64, err error) {
  45  	return mr.writeToWithBuffer(w, []byte{:1024*32})
  46  }
  47  
  48  func (mr *multiReader) writeToWithBuffer(w Writer, buf []byte) (sum int64, err error) {
  49  	for i, r := range mr.readers {
  50  		var n int64
  51  		if subMr, ok := r.(*multiReader); ok { // reuse buffer with nested multiReaders
  52  			n, err = subMr.writeToWithBuffer(w, buf)
  53  		} else {
  54  			n, err = copyBuffer(w, r, buf)
  55  		}
  56  		sum += n
  57  		if err != nil {
  58  			mr.readers = mr.readers[i:] // permit resume / retry after error
  59  			return sum, err
  60  		}
  61  		mr.readers[i] = nil // permit early GC
  62  	}
  63  	mr.readers = nil
  64  	return sum, nil
  65  }
  66  
  67  var _ WriterTo = (*multiReader)(nil)
  68  
  69  // MultiReader returns a Reader that's the logical concatenation of
  70  // the provided input readers. They're read sequentially. Once all
  71  // inputs have returned EOF, Read will return EOF.  If any of the readers
  72  // return a non-nil, non-EOF error, Read will return that error.
  73  func MultiReader(readers ...Reader) Reader {
  74  	r := []Reader{:len(readers)}
  75  	copy(r, readers)
  76  	return &multiReader{r}
  77  }
  78  
  79  type multiWriter struct {
  80  	writers []Writer
  81  }
  82  
  83  func (t *multiWriter) Write(p []byte) (n int, err error) {
  84  	for _, w := range t.writers {
  85  		n, err = w.Write(p)
  86  		if err != nil {
  87  			return
  88  		}
  89  		if n != len(p) {
  90  			err = ErrShortWrite
  91  			return
  92  		}
  93  	}
  94  	return len(p), nil
  95  }
  96  
  97  var _ StringWriter = (*multiWriter)(nil)
  98  
  99  func (t *multiWriter) WriteString(s []byte) (n int, err error) {
 100  	var p []byte // lazily initialized if/when needed
 101  	for _, w := range t.writers {
 102  		if sw, ok := w.(StringWriter); ok {
 103  			n, err = sw.WriteString(s)
 104  		} else {
 105  			if p == nil {
 106  				p = []byte(s)
 107  			}
 108  			n, err = w.Write(p)
 109  		}
 110  		if err != nil {
 111  			return
 112  		}
 113  		if n != len(s) {
 114  			err = ErrShortWrite
 115  			return
 116  		}
 117  	}
 118  	return len(s), nil
 119  }
 120  
 121  // MultiWriter creates a writer that duplicates its writes to all the
 122  // provided writers, similar to the Unix tee(1) command.
 123  //
 124  // Each write is written to each listed writer, one at a time.
 125  // If a listed writer returns an error, that overall write operation
 126  // stops and returns the error; it does not continue down the list.
 127  func MultiWriter(writers ...Writer) Writer {
 128  	allWriters := []Writer{:0:len(writers)}
 129  	for _, w := range writers {
 130  		if mw, ok := w.(*multiWriter); ok {
 131  			allWriters = append(allWriters, mw.writers...)
 132  		} else {
 133  			allWriters = append(allWriters, w)
 134  		}
 135  	}
 136  	return &multiWriter{allWriters}
 137  }
 138