NDIter.cr

require "benchmark"

struct IterFast
  include Iterator(UInt64)

  @track : Pointer(UInt64)
  @shape : Pointer(UInt64)
  @strides : Pointer(UInt64)
  @ndims : Int32
  @next_index : UInt64
  @last_index : UInt64
  @done : Bool = false
  @contiguous : Bool
  @size : Int32

  def initialize(shape : Array(Int32), strides : Array(Int32), @contiguous : Bool)
    @ndims = shape.size
    @size = shape.reduce { |i, j| i * j } * strides[-1]
    @track = Pointer(UInt64).malloc(@ndims, UInt64.new(0))
    @shape = Pointer(UInt64).malloc(@ndims) { |i| UInt64.new(shape[i]) }
    @strides = Pointer(UInt64).malloc(@ndims) { |i| UInt64.new(strides[i]) }
    @next_index = 0
    @last_index = 0
  end

  def rewind
    @done = false
    @last_index = 0
    @next_index = 0
  end

  def next
    if @contiguous
      onednext
    else
      ndnext
    end
  end

  def onednext
    @last_index = @next_index
    @next_index += @strides[@ndims - 1]
    if @next_index > @size
      return stop
    end
    @last_index
  end

  def ndnext
    if @done
      return stop
    end

    u = @ndims - 1
    ni = @next_index
    @last_index = ni

    u.step(to: 0, by: -1) do |i|
      @track[i] += 1
      shape_i = @shape[i]
      stride_i = @strides[i]

      if @track[i] == shape_i
        if i == 0
          @done = true
        end
        @track[i] = 0
        ni -= (shape_i - 1) * stride_i
        next
      end
      ni += stride_i
      break
    end

    @next_index = ni
    @last_index
  end
end

Benchmark.ips do |b|
  b.report("Contiguous") do
    a = 0
    t = IterFast.new([100, 100, 100], [10000, 100, 1], true)
    t.each { |i| a = i }
  end

  b.report("nd") do
    a = 0
    t = IterFast.new([100, 100, 100], [10000, 100, 1], false)
    t.each { |i| a = i }
  end
end