ignoreIndexRanges argument for FORLOOP (#42)

* Add ignore argument to FORLOOP

* Make sure simplifyIndices ignores any index ranges

* Handle ignore on LHS and add better error handling

* Add tests

* Set environment when evaluating ignore argument

* Update docs

* More informative argument name

Author: kenkellner

DESCRIPTION

@@ -1,6 +1,6 @@
 Package: nimbleMacros
-Version: 0.1.1.9003
-Date: 2026-06-09
+Version: 0.1.1.9004
+Date: 2026-06-10
 Title: Macros Generating 'nimble' Code
 Authors@R: c(person("Ken", "Kellner", email="contact@kenkellner.com", 
                     role=c("cre","aut")),
@@ -23,4 +23,4 @@ URL: https://r-nimble.org
 BugReports: https://github.com/nimble-dev/nimbleMacros/issues
 Encoding: UTF-8
 VignetteBuilder: knitr
-RoxygenNote: 7.3.3
+Config/roxygen2/version: 8.0.0

R/FORLOOP.R

@@ -11,6 +11,7 @@
 #'
 #' @param code The right-hand side of a parameter declaration
 #' @param modelInfo Used internally by nimbleMacros; a list of model information such as constants and dimensions
+#' @param ignoreIndexRanges Index ranges to ignore when creating for loops; must be provided as a list of quoted code
 #' @param .env Used internally by nimbleMacros; the environment where the model was created
 #'
 #' @return NIMBLE code for a for loop or series of nested for loops.
@@ -27,7 +28,21 @@ NULL
 
 #' @export
 FORLOOP <- nimble::buildMacro(
-function(code, modelInfo, .env){
+function(code, modelInfo, ignoreIndexRanges = list(), .env){
+  # Super clunky way of extracting the ignore argument
+  # seems to be necessary to do it this way given buildMacro options
+  ignore <- eval(code[[length(code)]]$ignoreIndexRanges, envir = .env)
+  # necessary for when ignore is not explicitly included in arguments
+  if(is.null(ignore)) ignore <- list()
+  # get rid of any ignore stuff in code
+  code[[length(code)]]$ignoreIndexRanges <- NULL
+  # Check ignore structure
+  if(!is.list(ignore)) stop("ignoreIndexRanges argument must be a list", call.=FALSE)
+  if(length(ignore) > 0){
+    is_call <- sapply(ignore, is.call)
+    if(!all(is_call)) stop("Some elements of ignoreIndexRanges are not quoted code/calls", call.=FALSE)
+  }
+
   # Remove FORLOOP from the code line
   code <- removeMacroCall(code)
   LHS <- getLHS(code)
@@ -45,11 +60,13 @@ function(code, modelInfo, .env){
 
   # Subset only indexes with ranges
   idx_sub <- idx[has_range]
+  # Remove indices in ignore
+  idx_sub <- idx_sub[! idx_sub %in% ignore]
   # Create new index letters to represent those ranges in the for loops
   idx_letters <- lapply(1:length(idx_sub), function(i) modelInfo$indexCreator())
   idx_letters <- lapply(idx_letters, as.name)
   # Replace the ranges with the new corresponding letters
-  code <- replaceDeclarationIndexRanges(code, idx_letters)
+  code <- replaceDeclarationIndexRanges(code, idx_letters, ignore)
   # Then insert the ranges into the for loop structure
   idx_sub <- replaceRanges(idx_sub, idx_letters)
   # So e.g.
@@ -168,13 +185,16 @@ removeIndexAdjustments <- function(idx){
 
 # Check if index range on both sides of declaration are the same
 # Specifically, all indices on RHS have to be present in LHS
-hasMatchingIndexRanges <- function (LHS, RHS){
+# ignoreIndexRanges must be a list of quoted code e.g. list(quote(1:n))
+hasMatchingIndexRanges <- function (LHS, RHS, ignoreIndexRanges = list()){
     idx_LHS <- extractIndices(LHS)
     idx_LHS <- idx_LHS[isIndexRange(idx_LHS)]
     idx_LHS <- lapply(idx_LHS, removeIndexAdjustments)
     idx_RHS <- extractAllIndices(RHS)
     idx_RHS <- idx_RHS[isIndexRange(idx_RHS)]
     idx_RHS <- lapply(idx_RHS, removeIndexAdjustments)
+    idx_RHS <- idx_RHS[!idx_RHS %in% ignoreIndexRanges]
+    if(length(idx_RHS) == 0) return(TRUE)
     all(idx_RHS %in% idx_LHS)
 }
 
@@ -223,13 +243,16 @@ recursiveReplaceIndex <- function(code, old_idx, new_idx){
 }
 
 # Replace all indices on both sides LHS/RHS
-replaceDeclarationIndexRanges <- function(code, new_idx_list){
+# ignoreIndexRanges, if provided, must be a list of quoted code
+# e.g. list(quote(1:n), quote(1:3))
+replaceDeclarationIndexRanges <- function(code, new_idx_list, ignoreIndexRanges = list()){
   LHS <- getLHS(code)
   RHS <- getRHS(code)
   op <- code[[1]]
-  if(!hasMatchingIndexRanges(LHS, RHS)) stop("Index ranges must match")
+  if(!hasMatchingIndexRanges(LHS, RHS, ignoreIndexRanges)) stop("Index ranges must match")
   idx <- extractAllIndices(LHS)
   idx <- idx[isIndexRange(idx)]
+  idx <- idx[! idx %in% ignoreIndexRanges]
   if(length(idx) == 0) return(code)
   nidx <- length(idx)
   for (i in 1:nidx){

R/utilities.R

@@ -195,6 +195,8 @@ simplifyIndices <- function(code, new_indices){
     if(is.name(x) | is.symbol(x)) return(x)
     if(x[[1]] == "for"){    
       unique_idx <- unique(extractAllIndices(x))
+      # Ignore any remaining index *ranges*
+      unique_idx <- unique_idx[sapply(unique_idx, is.name)]
       # Ignore blank indices as in x[i,]
       unique_idx <- unique_idx[sapply(unique_idx, function(x) x!="")]
       if(length(unique_idx) > length(new_indices)){

man/FORLOOP.Rd

@@ -322,3 +322,87 @@ test_that("FORLOOP", {
 
   nimbleOptions(enableMacroComments=comments_on)
 })
+
+test_that("ignoreIndexRanges argument works", {
+  comments_on <- nimbleOptions()$enableMacroComments
+  nimbleOptions(enableMacroComments=FALSE)
+ 
+  # Check error when index on RHS not on LHS, and not ignored
+  const <- list(y=rnorm(2), x=matrix(rnorm(6), 2, 3), n=2)
+  modInfo <- list(constants = const, indexCreator = nimble:::labelFunctionCreator("i"))
+  code <- quote(y[1:n] <- FORLOOP(sum(x[1:n, 1:3])))
+  expect_error(FORLOOP$process(code, modelInfo=modInfo, .env=environment()),
+               "Index ranges must match")
+
+  # Check malformed ignore arguments
+  code <- quote(y[1:n] <- FORLOOP(sum(x[1:n, 1:3]), ignoreIndexRanges=1:3))
+  expect_error(FORLOOP$process(code, modelInfo=modInfo, .env=environment())$code,
+               "ignoreIndexRanges argument must be a list")
+
+  code <- quote(y[1:n] <- FORLOOP(sum(x[1:n, 1:3]), ignoreIndexRanges=list(1:3, 2)))
+  expect_error(FORLOOP$process(code, modelInfo=modInfo, .env=environment())$code,
+               "Some elements of ignoreIndexRanges are not quoted code/calls")
+
+  # Simple loop with one ignored range
+  code <- nimbleCode({
+    y[1:n] <- FORLOOP(sum(x[1:n, 1:3]), ignoreIndexRanges = list(quote(1:3)))
+  })
+  mod <- nimbleModel(code, constants=const)
+  expect_equal(
+    mod$getCode(),
+    quote({
+    for (i_1 in 1:n) {
+      y[i_1] <- sum(x[i_1, 1:3])
+    }
+    })
+  )
+
+  # Nested loops
+  const <- list(y=matrix(rnorm(4), 2, 2), x=array(rnorm(12), c(2,2,3)), n=2)
+  code <- nimbleCode({
+    y[1:n, 1:2] <- FORLOOP(sum(x[1:n, 1:2, 1:3]), ignoreIndexRanges = list(quote(1:3)))
+  })
+  mod <- nimbleModel(code, constants=const)
+  expect_equal(
+    mod$getCode(),
+    quote({
+    for (i_1 in 1:n) {
+        for (i_2 in 1:2) {
+            y[i_1, i_2] <- sum(x[i_1, i_2, 1:3])
+        }
+    }
+    })
+  )
+
+  # Multiple ignores
+  const <- list(y=rnorm(2), x=array(rnorm(12), c(2,2,3)), n=2)
+  code <- nimbleCode({
+    y[1:2] <- FORLOOP(sum(x[1:n, 1:2, 1:3]), ignoreIndexRanges = list(quote(1:3), quote(1:n)))
+  })
+  mod <- nimbleModel(code, constants=const)
+  expect_equal(
+    mod$getCode(),
+    quote({
+    for (i_1 in 1:2) {
+        y[i_1] <- sum(x[1:n, i_1, 1:3])
+    }
+    })
+  )
+
+  # Ignore on LHS
+  const <- list(y=matrix(rnorm(4), 2, 2), x=matrix(rnorm(4), 2, 2), n=2)
+  code <- nimbleCode({
+    y[1:n, 1:2] <- FORLOOP(x[1:n, 1:2], ignoreIndexRanges = list(quote(1:n)))
+  })
+  mod <- nimbleModel(code, constants=const)
+  expect_equal(
+    mod$getCode(),
+    quote({
+      for (i_1 in 1:2){
+        y[1:n, i_1] <- x[1:n, i_1]
+      }    
+    })
+  )
+
+  nimbleOptions(enableMacroComments=comments_on)
+})

tests/testthat/test_FORLOOP.R

@@ -152,3 +152,31 @@ test_that("Blank index slot is ignored", {
   )
 })
 
+test_that("Remaining index ranges are ignored", {
+  # Index ranges could remain in loop code based on use of ignore argument
+  # make sure simplifyForLoops doesn't replace these with a new index
+  test <- nimbleCode({
+    for (i_1 in 1:n){
+      for (i_2 in 1:2){
+        y[i_1,i_2] <- sum(x[i_1,i_2,1:m])
+      }
+    }
+
+    for(i_1 in 1:n){
+      z[i_1] <- sum(a[1:3])
+    }
+  })
+
+  out <- simplifyForLoops(test)
+
+  expect_equal(out,
+    quote({
+    for (i in 1:n) {
+        for (j in 1:2) {
+            y[i, j] <- sum(x[i, j, 1:m])
+        }
+        z[i] <- sum(a[1:3])
+    }
+    })
+  )
+})