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

     1  package unrollinreverseadvance
     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     [4]Index
    19  	NextIndex uint32
    20  	Done      bool
    21  }
    22  
    23  func NewIterator(ap *shape.AP) *Iterator {
    24  	track := [4]Index{}
    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 := int(next)
    57  
    58  	x := &it.Track[0]
    59  	x.Track--
    60  	next += x.Advance
    61  	if x.Track > 0 {
    62  		it.NextIndex = next
    63  		return result, nil
    64  	}
    65  	x.Track = x.Shape
    66  
    67  	x = &it.Track[1]
    68  	x.Track--
    69  	next += x.Advance
    70  	if x.Track > 0 {
    71  		it.NextIndex = next
    72  		return result, nil
    73  	}
    74  	x.Track = x.Shape
    75  
    76  	x = &it.Track[2]
    77  	x.Track--
    78  	next += x.Advance
    79  	if x.Track > 0 {
    80  		it.NextIndex = next
    81  		return result, nil
    82  	}
    83  	x.Track = x.Shape
    84  
    85  	x = &it.Track[3]
    86  	x.Track--
    87  	next += x.Advance
    88  	if x.Track > 0 {
    89  		it.NextIndex = next
    90  		return result, nil
    91  	}
    92  	it.Done = true
    93  	it.NextIndex = next
    94  
    95  	return result, nil
    96  }