add AST walker and update tests

Author: ben-schwen

R/data.table.R

@@ -3342,9 +3342,12 @@ is_constantish = function(q, check_singleton=FALSE) {
       "as.complex", "as.logical", "as.Date", "as.POSIXct", "as.factor")
 }
 
+.gforce_ops = c("+", "-", "*", "/", "^", "%%", "%/%")
+
 .gforce_ok = function(q, x, envir=parent.frame(2L)) {
   if (is.N(q)) return(TRUE) # For #334
   if (!is.call(q)) return(FALSE)  # plain columns are not gforce-able since they might not aggregate (see test 104.1)
+  if (q %iscall% "(") return(.gforce_ok(q[[2L]], x, envir))
 
   q1 = .get_gcall(q)
   if (!is.null(q1)) {
@@ -3362,7 +3365,7 @@ is_constantish = function(q, check_singleton=FALSE) {
   }
 
   # check if arithmetic operator -> recursively validate ALL branches (like in AST)
-  if (is.symbol(q[[1L]]) && q[[1L]] %chin% c("+", "-", "*", "/", "^", "%%", "%/%")) {
+  if (is.symbol(q[[1L]]) && q[[1L]] %chin% .gforce_ops) {
     for (i in 2:length(q)) {
       if (!.gforce_ok(q[[i]], x, envir)) return(FALSE)
     }
@@ -3374,6 +3377,10 @@ is_constantish = function(q, check_singleton=FALSE) {
 
 .gforce_jsub = function(q, names_x, envir=parent.frame(2L)) {
   if (!is.call(q)) return(q)
+  if (q %iscall% "(") {
+    q[[2L]] = .gforce_jsub(q[[2L]], names_x, envir)
+    return(q)
+  }
 
   q1 = .get_gcall(q)
   if (!is.null(q1)) {
@@ -3390,7 +3397,7 @@ is_constantish = function(q, check_singleton=FALSE) {
   }
 
   # if arithmetic operator, recursively substitute its operands. we know what branches are valid from .gforce_ok
-  if (is.symbol(q[[1L]]) && q[[1L]] %chin% c("+", "-", "*", "/", "^", "%%", "%/%")) {
+  if (is.symbol(q[[1L]]) && q[[1L]] %chin% .gforce_ops) {
     for (i in 2:length(q)) {
       q[[i]] = .gforce_jsub(q[[i]], names_x, envir)
     }

inst/tests/tests.Rraw

@@ -21219,20 +21219,20 @@ test(2344.04, key(DT[, .(V4 = c("b", "a"), V2, V5 = c("y", "x"), V1)]), c("V1",
 # conversions should not turn gforce off #2934
 # lapply gforce should also work without .SD #5032
 # support arithmetic in j with gforce #3815
-dt = data.table(a=1:4, b=1:2)
 out = c("GForce FALSE", "GForce FALSE", "GForce TRUE")
+dt = data.table(a=1:4, b=1:2)
 test(2345.01, levels=0:2, dt[, max(as.character(a)), by=b, verbose=TRUE], data.table(b=1:2, V1=c("3","4")), output=out)
 test(2345.02, levels=0:2, dt[, max(as.numeric(a)), by=b, verbose=TRUE], data.table(b=1:2, V1=c(3,4)), output=out)
-test(2345.03, levels=0:2, dt[, Map(sum, .SD), b, verbose=TRUE], dt[, lapply(.SD, sum), b], output=out)
+dt = data.table(a=1:4, b=1:2)
+test(2345.11, levels=0:2, dt[, Map(sum, .SD), b, verbose=TRUE], dt[, lapply(.SD, sum), b], output=out)
+test(2345.12, levels=0:2, dt[, Map(sum, .SD, .SD), by=b, verbose=TRUE], output="GForce FALSE")
 dt = data.table(a = NA_integer_, b = 1:2, c = c(TRUE, FALSE))
-test(2345.04, levels=0:2, dt[, Map(weighted.mean, .SD, na.rm=c), b, .SDcols="a", verbose=TRUE], data.table(b=1:2, a=c(NaN, NA_real_)), output="GForce FALSE")
-test(2345.05, levels=0:2, dt[,list(weighted.mean(a, na.rm=c)), b, verbose=TRUE], data.table(b=1:2, V1=c(NaN, NA_real_)), output="GForce FALSE")
+test(2345.13, levels=0:2, dt[, Map(weighted.mean, .SD, na.rm=c), b, .SDcols="a", verbose=TRUE], data.table(b=1:2, a=c(NaN, NA_real_)), output="GForce FALSE")
+test(2345.14, levels=0:2, dt[,list(weighted.mean(a, na.rm=c)), b, verbose=TRUE], data.table(b=1:2, V1=c(NaN, NA_real_)), output="GForce FALSE")
 dt = data.table(a=1:2, b=1, c=1:4)
-out = c("GForce FALSE", "lapply optimization changed j", "GForce TRUE")
-test(2345.06, levels=0:2, dt[, lapply(list(b, c), sum), by=a, verbose=TRUE], output=out)
-test(2345.07, levels=0:2, dt[, c(list(sum(b), sum(c))), by=a, verbose=TRUE], output=out)
-test(2345.08, levels=0:2, names(dt[, lapply(list(b, c), sum), by=a]))
-dt = data.table(a=1:4, b=1:2)
-out = c("GForce FALSE", "GForce FALSE", "GForce TRUE")
-test(2345.09, levels=0:2, dt[, .(max(a)-min(a)), by=b, verbose=TRUE], output=out)
+test(2345.21, levels=0:2, dt[, lapply(list(b, c), sum), by=a, verbose=TRUE], output=out)
+test(2345.22, levels=0:2, dt[, c(list(sum(b), sum(c))), by=a, verbose=TRUE], output=out)
+test(2345.23, levels=0:2, names(dt[, lapply(list(b, c), sum), by=a]))
 dt = data.table(a=1:4, b=1:2)
+test(2345.31, levels=0:2, dt[, .(max(a)-min(a)), by=b, verbose=TRUE], output=out)
+test(2345.32, levels=0:2, dt[, .((max(a) - min(a)) / (max(a) + min(a))), by=b, verbose=TRUE], data.table(b=1:2, V1=c(0.5, 1/3)), output=out)