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

     1  package basic
     2  
     3  import (
     4  	"io"
     5  
     6  	"github.com/egonelbre/exp/niterator/shape"
     7  )
     8  
     9  type Iterator struct {
    10  	*shape.AP
    11  
    12  	Track     []int
    13  	NextIndex int
    14  	Done      bool
    15  }
    16  
    17  func NewIterator(ap *shape.AP) *Iterator {
    18  	return &Iterator{
    19  		AP:    ap,
    20  		Track: make([]int, len(ap.Shape)),
    21  	}
    22  }
    23  
    24  func (it *Iterator) IsDone() bool {
    25  	return it.Done
    26  }
    27  
    28  func (it *Iterator) Next() (int, error) {
    29  	if it.Done {
    30  		return 0, io.EOF
    31  	}
    32  
    33  	last := len(it.Shape) - 1
    34  	next := it.NextIndex
    35  	result := next
    36  
    37  	// the following 3 lines causes the compiler to perform bounds check here,
    38  	// instead of being done in the loop
    39  	coord := it.Shape[:last+1]
    40  	track := it.Track[:last+1]
    41  	stride := it.Stride[:last+1]
    42  	for i := last; i >= 0; i-- {
    43  		track[i]++
    44  		shapeI := coord[i]
    45  		strideI := stride[i]
    46  
    47  		if track[i] == shapeI {
    48  			if i == 0 {
    49  				it.Done = true
    50  			}
    51  			track[i] = 0
    52  			next -= (shapeI - 1) * strideI
    53  			continue
    54  		}
    55  		next += strideI
    56  		break
    57  	}
    58  	it.NextIndex = next
    59  	return result, nil
    60  }