Framework for user-defined transformations

The framework is up and running but it still needs some teststing. I also need to avoid a double-transformation. Right now, I transform when checking if a function can be transformed and again when transforming it. With refactoring, I should be able to avoid that.

A step towards resolving #5.

NAMESPACE

@@ -3,5 +3,5 @@
 export(can_loop_transform)
 export(can_loop_transform_)
 export(loop_transform)
-import(glue)
-import(rlang)
+export(transform_call)
+export(user_transform)

R/loop-transformation.R

@@ -87,26 +87,7 @@ can_transform_rec <- function(expr, fun_name, fun_call_allowed, cc) {
 }
 
 
-#' Tests if a function, provided by its name, can be transformed.
-#'
-#' This function analyses a recursive function to check if we can transform it into
-#' a loop or trampoline version with \code{\link{loop_transform}}. This version expects the
-#' function to be provided as a quosure, but see also \code{\link{can_loop_transform}}.
-#'
-#' Since this function needs to handle recursive functions, it needs to know the name of
-#' its input function, so this must be provided as a bare symbol.
-#'
-#' @param fun The function to check. Must be provided by its (bare symbol) name.
-#'
-#' @examples
-#' factorial <- function(n)
-#'     if (n <= 1) 1 else n * factorial(n - 1)
-#' factorial_acc <- function(n, acc = 1)
-#'     if (n <= 1) acc else factorial_acc(n - 1, n * acc)
-#'
-#' can_loop_transform_(rlang::quo(factorial))     # FALSE -- and prints a warning
-#' can_loop_transform_(rlang::quo(factorial_acc)) # TRUE
-#'
+#' @describeIn can_loop_transform This version expects \code{fun} to be quosure.
 #' @export
 can_loop_transform_ <- function(fun) {
     fun_name <- rlang::get_expr(fun)
@@ -134,7 +115,10 @@ can_loop_transform_ <- function(fun) {
         stop(error)
     }
 
-    callCC(function(cc) can_transform_rec(body(fun), fun_name, TRUE, cc))
+    ## FIXME: don't transform twice; we do this now both here and in the
+    ## loop transformation
+    fun_body <- user_transform(body(fun), rlang::get_env(fun))
+    callCC(function(cc) can_transform_rec(fun_body, fun_name, TRUE, cc))
 }
 
 
@@ -156,6 +140,10 @@ can_loop_transform_ <- function(fun) {
 #' can_loop_transform(factorial)     # FALSE -- and prints a warning
 #' can_loop_transform(factorial_acc) # TRUE
 #'
+#' can_loop_transform_(rlang::quo(factorial))     # FALSE -- and prints a warning
+#' can_loop_transform_(rlang::quo(factorial_acc)) # TRUE
+#'
+#' @describeIn can_loop_transform This version quotes \code{fun} itself.
 #' @export
 can_loop_transform <- function(fun) {
     fun <- rlang::enquo(fun)
@@ -355,6 +343,7 @@ loop_transform <- function(fun) {
     }
 
     fun_name <- rlang::quo_name(fun_q)
+    user_transformed_body <- user_transform(body(fun), rlang::get_env(fun_q))
     new_fun_body <- build_transformed_function(body(fun), fun_name, fun)
     rlang::new_function(
         args = formals(fun),

R/user-defined-transformations.R

@@ -0,0 +1,89 @@
+
+#' Transform a call before the tail-recursion transformation.
+#'
+#' This generic function is a hook by which you can modify how the
+#' tail-recursion transformations should handle special functions.
+#' It gives you a way to rewrite function calls to make them tail-recursive
+#' before we do any other manipulation.
+#'
+#' @param fun   The actual function. Used for dynamic dispatching.
+#' @param expr  The expression to rewrite.
+#' @return The rewritten expression.
+#'
+#' @examples
+#' my_if_else <- function(test, if_true, if_false) {
+#'     if (test) if_true else if_false
+#' }
+#' f <- function(x, y) my_if_else(x == y, x, f(y, y))
+#' f(1, 2)
+#' f(3, 3)
+#'
+#' can_loop_transform(f) # No, we can't, and we get a warning
+#'
+#' class(my_if_else) <- c("my_if_else", class(my_if_else))
+#' class(my_if_else)
+#'
+#' transform_call.my_if_else <- function(fun, expr) {
+#'     test <- expr[[2]]; if_true <- expr[[3]]; if_false <- expr[[4]]
+#'     rlang::expr(if(rlang::UQ(test)) rlang::UQ(if_true) else rlang::UQ(if_false))
+#' }
+#' transform_call(my_if_else, quote(my_if_else(x == y, x, f(y, y))))
+#'
+#' can_loop_transform(f) # Now we can, because my_if_else gets transformed
+#'
+#' @export
+transform_call <- function(fun, expr) {
+    UseMethod("transform_call")
+}
+
+#' @describeIn transform_call The default is to just return the unchanged expression.
+transform_call.default <- function(fun, expr) {
+    expr
+}
+
+#' Apply user transformations depths-first.
+#'
+#' @param expr The expression to transform -- typically a function body.
+#' @param env  The environment where functions can be found.
+#'
+#' @return Rewritten expression
+#'
+#' @examples
+#' my_if_else <- function(test, if_true, if_false) {
+#'     if (test) if_true else if_false
+#' }
+#' class(my_if_else) <- c("my_if_else", class(my_if_else))
+#' transform_call.my_if_else <- function(fun, expr) {
+#'     test <- expr[[2]]; if_true <- expr[[3]]; if_false <- expr[[4]]
+#'     rlang::expr(if(rlang::UQ(test)) rlang::UQ(if_true) else rlang::UQ(if_false))
+#' }
+#'
+#' f <- function(x, y) my_if_else(x == y, x, f(y, y))
+#' user_transform(body(f))
+#'
+#' @export
+user_transform <- function(expr, env = rlang::caller_env()) {
+    if (rlang::is_atomic(expr) || rlang::is_pairlist(expr) ||
+        rlang::is_symbol(expr) || rlang::is_primitive(expr)) {
+        expr
+
+    } else {
+        stopifnot(rlang::is_lang(expr))
+
+        fun_name <- rlang::call_name(expr)
+        if (!exists(fun_name, where = env)) {
+            error_msg <- glue::glue(
+                "The function {fun_name} was not found in the provided scope."
+            )
+            stop(simpleError(error_msg, call = expr))
+        }
+
+        fun <- get(fun_name, envir = env)
+        args <- rlang::call_args(expr)
+        for (i in seq_along(args)) {
+            expr[[i + 1]] <- user_transform(args[[i]], env)
+        }
+        transform_call(fun, expr)
+    }
+}
+

tests/testthat/test-loop-transformation.R

@@ -74,6 +74,8 @@ test_that("we simplify code-blocks", {
             factorial_acc_2(n - 1, n * acc)
         }
 
+    can_loop_transform(factorial_acc_1)
+
     transformed_1 <- loop_transform(factorial_acc_1)
     transformed_2 <- loop_transform(factorial_acc_2)
 

tests/testthat/test-user-defined-transformations.R

@@ -0,0 +1,5 @@
+context("test-user-defined-transformations.R")
+
+test_that("we call the default transformation function", {
+  skip("TODO")
+})