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

     1  package premul
     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  	Premul    []int
    14  	NextIndex int
    15  	Done      bool
    16  }
    17  
    18  func NewIterator(ap *shape.AP) *Iterator {
    19  	premul := make([]int, len(ap.Shape))
    20  	for i := range ap.Shape {
    21  		premul[i] = (ap.Shape[i] - 1) * ap.Stride[i]
    22  	}
    23  
    24  	return &Iterator{
    25  		AP:     ap,
    26  		Premul: premul,
    27  		Track:  make([]int, len(ap.Shape)),
    28  	}
    29  }
    30  
    31  func (it *Iterator) IsDone() bool {
    32  	return it.Done
    33  }
    34  
    35  func (it *Iterator) Next() (int, error) {
    36  	if it.Done {
    37  		return 0, io.EOF
    38  	}
    39  
    40  	last := len(it.Shape) - 1
    41  	next := it.NextIndex
    42  	result := next
    43  
    44  	// the following 3 lines causes the compiler to perform bounds check here,
    45  	// instead of being done in the loop
    46  	coord := it.Shape[:last+1]
    47  	track := it.Track[:last+1]
    48  	stride := it.Stride[:last+1]
    49  	premul := it.Premul[:last+1]
    50  	for i := last; i >= 0; i-- {
    51  		track[i]++
    52  		shapeI := coord[i]
    53  		strideI := stride[i]
    54  
    55  		if track[i] == shapeI {
    56  			if i == 0 {
    57  				it.Done = true
    58  			}
    59  			track[i] = 0
    60  			next -= premul[i]
    61  			continue
    62  		}
    63  		next += strideI
    64  		break
    65  	}
    66  	it.NextIndex = next
    67  	return result, nil
    68  }