github.com/egonelbre/exp@v0.0.0-20240430123955-ed1d3aa93911/niterator/onearrrevadvance/iterator.go (about)

     1  package onearrrevadvance
     2  
     3  import (
     4  	"io"
     5  
     6  	"github.com/egonelbre/exp/niterator/shape"
     7  )
     8  
     9  type Index struct {
    10  	Track   uint32
    11  	Advance uint32
    12  	Shape   uint32
    13  }
    14  
    15  type Iterator struct {
    16  	*shape.AP
    17  
    18  	Track     []Index
    19  	NextIndex uint32
    20  	Done      bool
    21  }
    22  
    23  func NewIterator(ap *shape.AP) *Iterator {
    24  	track := make([]Index, len(ap.Shape))
    25  
    26  	last := len(ap.Shape) - 1
    27  	stride := ap.Stride[:last+1]
    28  	shape := ap.Shape[:last+1]
    29  
    30  	track[0].Track = uint32(shape[last])
    31  	track[0].Advance = uint32(stride[last])
    32  	track[0].Shape = uint32(shape[last])
    33  
    34  	for i := 1; i < last+1; i++ {
    35  		track[i].Track = uint32(shape[last-i])
    36  		track[i].Shape = uint32(shape[last-i])
    37  		track[i].Advance = uint32(stride[last-i] - stride[last-i+1]*shape[last-i+1])
    38  	}
    39  
    40  	return &Iterator{
    41  		AP:    ap,
    42  		Track: track,
    43  	}
    44  }
    45  
    46  func (it *Iterator) IsDone() bool {
    47  	return it.Done
    48  }
    49  
    50  func (it *Iterator) Next() (int, error) {
    51  	if it.Done {
    52  		return 0, io.EOF
    53  	}
    54  
    55  	next := it.NextIndex
    56  	result := next
    57  	track := it.Track
    58  	for i := range track {
    59  		x := &track[i]
    60  		x.Track--
    61  		next += x.Advance
    62  		if x.Track > 0 {
    63  			it.NextIndex = next
    64  			return int(result), nil
    65  		}
    66  		x.Track = x.Shape
    67  	}
    68  
    69  	it.Done = true
    70  	it.NextIndex = next
    71  	return int(result), nil
    72  }