github.com/anacrolix/torrent@v1.61.0/segments/index.go (about)

     1  package segments
     2  
     3  import (
     4  	"cmp"
     5  	"iter"
     6  	"slices"
     7  
     8  	g "github.com/anacrolix/generics"
     9  	"github.com/anacrolix/missinggo/v2/panicif"
    10  )
    11  
    12  func NewIndex(segments LengthIter) (ret Index) {
    13  	var start Length
    14  	for l := range segments {
    15  		ret.segments = append(ret.segments, Extent{start, l})
    16  		start += l
    17  	}
    18  	return
    19  }
    20  
    21  type Index struct {
    22  	segments []Extent
    23  }
    24  
    25  func NewIndexFromSegments(segments []Extent) Index {
    26  	return Index{segments}
    27  }
    28  
    29  // Yields segments as extents with Start relative to the previous segment's end.
    30  func (me Index) iterSegments(startIndex int) iter.Seq[Extent] {
    31  	return func(yield func(Extent) bool) {
    32  		var lastEnd g.Option[Int]
    33  		for _, cur := range me.segments[startIndex:] {
    34  			ret := Extent{
    35  				// Why ignore initial start on the first segment?
    36  				Start:  cur.Start - lastEnd.UnwrapOr(cur.Start),
    37  				Length: cur.Length,
    38  			}
    39  			lastEnd.Set(cur.End())
    40  			if !yield(ret) {
    41  				return
    42  			}
    43  		}
    44  	}
    45  }
    46  
    47  func (me Index) LocateIter(e Extent) iter.Seq2[int, Extent] {
    48  	return func(yield func(int, Extent) bool) {
    49  		// We find the first segment that ends after the start of the target extent.
    50  		first, eq := slices.BinarySearchFunc(me.segments, e.Start, func(elem Extent, target Int) int {
    51  			return cmp.Compare(elem.End(), target+1)
    52  		})
    53  		//fmt.Printf("binary search for %v in %v returned %v\n", e.Start, me.segments, first)
    54  		if first == len(me.segments) {
    55  			return
    56  		}
    57  		_ = eq
    58  		e.Start -= me.segments[first].Start
    59  		// The extent is before the first segment.
    60  		if e.Start < 0 {
    61  			e.Length += e.Start
    62  			e.Start = 0
    63  		}
    64  		i := first
    65  		for cons := range scanConsecutive(me.iterSegments(first), e) {
    66  			if !yield(i, cons) {
    67  				return
    68  			}
    69  			i++
    70  		}
    71  	}
    72  }
    73  
    74  type IndexAndOffset struct {
    75  	Index  int
    76  	Offset int64
    77  }
    78  
    79  // Returns the Extent that contains the given extent, if it exists. Panics if Extents overlap on the
    80  // offset.
    81  func (me Index) LocateOffset(off int64) (ret g.Option[IndexAndOffset]) {
    82  	// I think an Extent needs to have a non-zero to match against it? That's what this method is
    83  	// defining.
    84  	for i, e := range me.LocateIter(Extent{off, 1}) {
    85  		panicif.True(ret.Ok)
    86  		panicif.NotEq(e.Length, 1)
    87  		ret.Set(IndexAndOffset{
    88  			Index:  i,
    89  			Offset: e.Start,
    90  		})
    91  	}
    92  	return
    93  }
    94  
    95  func (me Index) Index(i int) Extent {
    96  	return me.segments[i]
    97  }