github.com/insolar/vanilla@v0.0.0-20201023172447-248fdf805322/iokit/prepend.go (about)

     1  // Copyright 2020 Insolar Network Ltd.
     2  // All rights reserved.
     3  // This material is licensed under the Insolar License version 1.0,
     4  // available at https://github.com/insolar/assured-ledger/blob/master/LICENSE.md.
     5  
     6  package iokit
     7  
     8  import "io"
     9  
    10  func PrependReader(b []byte, r io.Reader) *PrependedReader {
    11  	return &PrependedReader{b, r}
    12  }
    13  
    14  func PrependWriter(b []byte, w io.Writer) *PrependedWriter {
    15  	return &PrependedWriter{b, w}
    16  }
    17  
    18  var _ io.Reader = &PrependedReader{}
    19  
    20  type PrependedReader struct {
    21  	prepend []byte
    22  	R       io.Reader
    23  }
    24  
    25  func (p *PrependedReader) Read(b []byte) (int, error) {
    26  	pn := len(p.prepend)
    27  	if pn == 0 {
    28  		return p.R.Read(b)
    29  	}
    30  
    31  	cn := copy(b, p.prepend)
    32  	if cn == pn {
    33  		if pr, ok := p.R.(*PrependedReader); ok {
    34  			*p = *pr
    35  			return cn, nil
    36  		}
    37  		p.prepend = nil
    38  	} else {
    39  		p.prepend = p.prepend[cn:]
    40  	}
    41  
    42  	if cn == len(b) {
    43  		return cn, nil
    44  	}
    45  	n, err := p.R.Read(b[cn:])
    46  	return cn + n, err
    47  }
    48  
    49  var _ io.WriteCloser = &PrependedWriter{}
    50  
    51  type PrependedWriter struct {
    52  	prepend []byte
    53  	W       io.Writer
    54  }
    55  
    56  func (p *PrependedWriter) flush() error {
    57  	if len(p.prepend) == 0 {
    58  		return nil
    59  	}
    60  	fb := p.prepend
    61  	p.prepend = nil
    62  	switch n, err := p.W.Write(fb); {
    63  	case err != nil:
    64  		return err
    65  	case n != len(fb):
    66  		return io.ErrShortWrite
    67  	default:
    68  		return nil
    69  	}
    70  }
    71  
    72  func (p *PrependedWriter) Write(b []byte) (int, error) {
    73  	pn := len(p.prepend)
    74  	if pn == 0 {
    75  		return p.W.Write(b)
    76  	}
    77  	cp := cap(p.prepend) - pn
    78  	bn := len(b)
    79  
    80  	if bn < cp {
    81  		p.prepend = append(p.prepend, b...)
    82  		return bn, nil
    83  	}
    84  
    85  	p.prepend = append(p.prepend, b[:cp]...)
    86  	fb := p.prepend
    87  	p.prepend = nil
    88  	switch n, err := p.W.Write(fb); {
    89  	case err != nil:
    90  		return n, err
    91  	case n != len(fb):
    92  		return n, io.ErrShortWrite
    93  	}
    94  	// TODO flattening
    95  
    96  	if cp == len(b) {
    97  		return cp, nil
    98  	}
    99  	n, err := p.W.Write(b[cp:])
   100  	return cp + n, err
   101  }
   102  
   103  func (p *PrependedWriter) Flush() error {
   104  	if err := p.flush(); err != nil {
   105  		return err
   106  	}
   107  	if f, ok := p.W.(interface{ Flush() error }); ok {
   108  		return f.Flush()
   109  	}
   110  	return nil
   111  }
   112  
   113  func (p *PrependedWriter) Close() error {
   114  	if err := p.flush(); err != nil {
   115  		return err
   116  	}
   117  	if f, ok := p.W.(io.Closer); ok {
   118  		return f.Close()
   119  	}
   120  	return nil
   121  }