go4.org@v0.0.0-20230225012048-214862532bf5/readerutil/multireaderat.go (about)

     1  /*
     2  Copyright 2016 The go4 Authors
     3  
     4  Licensed under the Apache License, Version 2.0 (the "License");
     5  you may not use this file except in compliance with the License.
     6  You may obtain a copy of the License at
     7  
     8       http://www.apache.org/licenses/LICENSE-2.0
     9  
    10  Unless required by applicable law or agreed to in writing, software
    11  distributed under the License is distributed on an "AS IS" BASIS,
    12  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
    13  See the License for the specific language governing permissions and
    14  limitations under the License.
    15  */
    16  
    17  package readerutil
    18  
    19  import (
    20  	"io"
    21  	"sort"
    22  )
    23  
    24  // NewMultiReaderAt is like io.MultiReader but produces a ReaderAt
    25  // (and Size), instead of just a reader.
    26  func NewMultiReaderAt(parts ...SizeReaderAt) SizeReaderAt {
    27  	m := &multiRA{
    28  		parts: make([]offsetAndSource, 0, len(parts)),
    29  	}
    30  	var off int64
    31  	for _, p := range parts {
    32  		m.parts = append(m.parts, offsetAndSource{off, p})
    33  		off += p.Size()
    34  	}
    35  	m.size = off
    36  	return m
    37  }
    38  
    39  type offsetAndSource struct {
    40  	off int64
    41  	SizeReaderAt
    42  }
    43  
    44  type multiRA struct {
    45  	parts []offsetAndSource
    46  	size  int64
    47  }
    48  
    49  func (m *multiRA) Size() int64 { return m.size }
    50  
    51  func (m *multiRA) ReadAt(p []byte, off int64) (n int, err error) {
    52  	wantN := len(p)
    53  
    54  	// Skip past the requested offset.
    55  	skipParts := sort.Search(len(m.parts), func(i int) bool {
    56  		// This function returns whether parts[i] will
    57  		// contribute any bytes to our output.
    58  		part := m.parts[i]
    59  		return part.off+part.Size() > off
    60  	})
    61  	parts := m.parts[skipParts:]
    62  
    63  	// How far to skip in the first part.
    64  	needSkip := off
    65  	if len(parts) > 0 {
    66  		needSkip -= parts[0].off
    67  	}
    68  
    69  	for len(parts) > 0 && len(p) > 0 {
    70  		readP := p
    71  		partSize := parts[0].Size()
    72  		if int64(len(readP)) > partSize-needSkip {
    73  			readP = readP[:partSize-needSkip]
    74  		}
    75  		pn, err0 := parts[0].ReadAt(readP, needSkip)
    76  		if err0 != nil {
    77  			return n, err0
    78  		}
    79  		n += pn
    80  		p = p[pn:]
    81  		if int64(pn)+needSkip == partSize {
    82  			parts = parts[1:]
    83  		}
    84  		needSkip = 0
    85  	}
    86  
    87  	if n != wantN {
    88  		err = io.ErrUnexpectedEOF
    89  	}
    90  	return
    91  }