iterator for index matrix

iter_index_matrix.R

iter_index_matrix <- function(dims, chunk_size = 1L) {
  f_split <- cumprod(dims) <= chunk_size
  has_next <- TRUE
  if (all(f_split)) {
    next_element <- function() {
      if (!has_next) stop("StopIteration")
      has_next <<- FALSE
      expand.grid(lapply(dims, seq_len))
    }
  } else {
    split_dims <- split(dims, f_split)
    in_block_dims <- split_dims[["TRUE"]]
    iter_dim <- split_dims[["FALSE"]][1L]
    max_out_block_dims <- split_dims[["FALSE"]][-1L]
    out_block_dims <- rep(1L, length(max_out_block_dims))
    increment_out_block_dims <- function() {
      increment <- function(dims, max_dims) {
        if (length(dims) == 0L) {
          has_next <<- FALSE
          dims <- vector("numeric")
        } else if (dims[1L] < max_dims[1L]) {
          dims[1] <- dims[1] + 1L
        } else {
          dims <- c(1L, increment(dims[-1L], max_dims[-1L]))
        }
        dims
      }
      out_block_dims <<- increment(out_block_dims, max_out_block_dims)
      invisible(NULL)
    }
    step_size <- floor(chunk_size / prod(in_block_dims))
    inds <- c(seq.int(1L, iter_dim, by = step_size), iter_dim + 1L)
    in_block_dims <- lapply(in_block_dims, seq_len)
    i <- 2L
    next_element <- function() {
      if (!has_next) stop("StopIteration")
      start <- inds[i - 1L]
      end <- inds[i] - 1L
      args <- c(in_block_dims, list(start:end), out_block_dims)
      i <<- i + 1L
      if (i > length(inds)) {
        i <<- 2L
        increment_out_block_dims()
      }
      expand.grid(args)
    }
  }
  obj <- list(nextElem = next_element, hasNext = function() has_next)
  class(obj) <- c("ihasNext", "abstractiter", "iter")
  obj
}