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
5package io
6
7type eofReader struct{}
8
9func (eofReader) Read([]byte) (int, error) {
10	return 0, EOF
11}
12
13type multiReader struct {
14	readers []Reader
15}
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.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
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 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
67var _ 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.
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 []Writer
81}
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			return
88		}
89		if n != len(p) {
90			err = ErrShortWrite
91			return
92		}
93	}
94	return len(p), nil
95}
96
97var _ StringWriter = (*multiWriter)(nil)
98
99func (t *multiWriter) WriteString(s string) (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.
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