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

     1  package onearrrevspecializeadvance
     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  	Stride  uint32
    14  }
    15  
    16  type Iterator struct {
    17  	*shape.AP
    18  
    19  	Track     [4]Index
    20  	NextIndex uint32
    21  	Done      bool
    22  }
    23  
    24  func NewIterator(ap *shape.AP) *Iterator {
    25  	it := &Iterator{}
    26  	it.AP = ap
    27  
    28  	last := len(ap.Shape) - 1
    29  	stride := ap.Stride[:last+1]
    30  	shape := ap.Shape[:last+1]
    31  	track := &it.Track
    32  
    33  	(*track)[0].Track = uint32(shape[last])
    34  	(*track)[0].Advance = uint32(stride[last])
    35  	(*track)[0].Shape = uint32(shape[last])
    36  
    37  	for i := 1; i < last+1; i++ {
    38  		(*track)[i].Track = uint32(shape[last-i])
    39  		(*track)[i].Shape = uint32(shape[last-i])
    40  		(*track)[i].Advance = uint32(stride[last-i] - stride[last-i+1]*shape[last-i+1])
    41  	}
    42  
    43  	return it
    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  	for i := range it.Track {
    58  		x := &it.Track[i]
    59  		x.Track--
    60  		next += x.Advance
    61  		if x.Track > 0 {
    62  			it.NextIndex = next
    63  			return int(result), nil
    64  		}
    65  		x.Track = x.Shape
    66  	}
    67  
    68  	it.Done = true
    69  	it.NextIndex = next
    70  	return int(result), nil
    71  }