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

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