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

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