codekingpro/portable-devtools

1// Copyright 2010 The Go Authors. All rights reserved.2// Use of this source code is governed by a BSD-style3// license that can be found in the LICENSE file.4 5package io6 7type eofReader struct{}8 9func (eofReader) Read([]byte) (int, error) {10	return 0, EOF11}12 13type multiReader struct {14	readers []Reader15}16 17func (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.readers23				continue24			}25		}26		n, err = mr.readers[0].Read(p)27		if err == EOF {28			// Use eofReader instead of nil to avoid nil panic29			// after performing flatten (Issue 18232).30			mr.readers[0] = eofReader{} // permit earlier GC31			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 = nil37			}38			return39		}40	}41	return 0, EOF42}43 44func (mr *multiReader) WriteTo(w Writer) (sum int64, err error) {45	return mr.writeToWithBuffer(w, make([]byte, 1024*32))46}47 48func (mr *multiReader) writeToWithBuffer(w Writer, buf []byte) (sum int64, err error) {49	for i, r := range mr.readers {50		var n int6451		if subMr, ok := r.(*multiReader); ok { // reuse buffer with nested multiReaders52			n, err = subMr.writeToWithBuffer(w, buf)53		} else {54			n, err = copyBuffer(w, r, buf)55		}56		sum += n57		if err != nil {58			mr.readers = mr.readers[i:] // permit resume / retry after error59			return sum, err60		}61		mr.readers[i] = nil // permit early GC62	}63	mr.readers = nil64	return sum, nil65}66 67var _ WriterTo = (*multiReader)(nil)68 69// MultiReader returns a Reader that's the logical concatenation of70// the provided input readers. They're read sequentially. Once all71// inputs have returned EOF, Read will return EOF.  If any of the readers72// return a non-nil, non-EOF error, Read will return that error.73func MultiReader(readers ...Reader) Reader {74	r := make([]Reader, len(readers))75	copy(r, readers)76	return &multiReader{r}77}78 79type multiWriter struct {80	writers []Writer81}82 83func (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			return88		}89		if n != len(p) {90			err = ErrShortWrite91			return92		}93	}94	return len(p), nil95}96 97var _ StringWriter = (*multiWriter)(nil)98 99func (t *multiWriter) WriteString(s string) (n int, err error) {100	var p []byte // lazily initialized if/when needed101	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			return112		}113		if n != len(s) {114			err = ErrShortWrite115			return116		}117	}118	return len(s), nil119}120 121// MultiWriter creates a writer that duplicates its writes to all the122// 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 operation126// stops and returns the error; it does not continue down the list.127func MultiWriter(writers ...Writer) Writer {128	allWriters := make([]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