Transforming functions with foolbox

Example: collecting symbols

To get a feeling for how the framework is used, we start with a simple example: collecting the symbols (variables) used in a function. Say we have the function f defined as below:

f <- function(x, y) {
    a <- x + y
    b <- x - y
    2 * a - b^2
}

This is a dummy-function so we don’t care what it does, but for some reason we are interested in collecting all the symbols inside it. That is, we want to collect the variables used in the implementation of the function, which are x, y, a, and b.

We can get the symbols in an expression structure by traversing it and picking all the elements where is.symbol(expr) or clang::is_symbol(expr) are TRUE. We can do this in a recursive function we write for this purpose, or we can use foolbox with a callback for symbols. To do this, we first define a callback. This is a function that must take at least one argument, expr, and allow additional arguments, .... We will use it in a traversal where it will only be called on symbols, so we can translate expr into a string whenever it is called, and we want to collect all such strings.

We build the callbacks for the traversal from the default callbacks used in an analysis, we get it using the function analysis_callbacks(), and we update the symbols callback with new function using with_symbols_callback. We can then run an analysis by setting it up—using the analyse() function and then run the callbacks using analyse_with(callbacks).

collect_symbols <- function(expr, ...) {
    list(symbols = as.character(expr))
}
callbacks <- analysis_callbacks() %>% with_symbol_callback(collect_symbols)
f %>% analyse() %>% analyse_with(callbacks)
#> $symbols
#>  [1] "{"  "<-" "a"  "+"  "x"  "y"  "<-" "b"  "-"  "x"  "y"  "-"  "*"  "a" 
#> [15] "^"  "b"

The symbols are collected in the order we see them in the function and since we do not remove duplications, we see each symbol as many times as it appears. We can simplify this a bit by post-processing the traversal results.

f %>% analyse() %>% analyse_with(callbacks) %>% 
    unlist(use.names = FALSE) %>% unique()
#>  [1] "{"  "<-" "a"  "+"  "x"  "y"  "b"  "-"  "*"  "^"

Normally, I will not bother with naming callbacks and constructing an analysis pipeline in steps, but simply construct it using a %>% pipeline like this:

f %>% analyse() %>% analyse_with(
    analysis_callbacks() %>% with_symbol_callback(
        function(expr, ...) list(symbols = as.character(expr))
    )
) %>% unlist(use.names = FALSE) %>% unique()
#>  [1] "{"  "<-" "a"  "+"  "x"  "y"  "b"  "-"  "*"  "^"

It is a matter of taste what you prefer.

I will often name a given pipeline, though, so I can reuse it later. This is something we can do by specifying the pipeline but with dot, ., as the initial input:

collect_symbols <- . %>% analyse() %>% analyse_with(
    analysis_callbacks() %>% with_symbol_callback(
        function(expr, ...) list(symbols = as.character(expr))
    )
) %>% unlist(use.names = FALSE) %>% unique()

Then, when we have a function we want to analyse, we can pipe it into the analysis:

f %>% collect_symbols()
#>  [1] "{"  "<-" "a"  "+"  "x"  "y"  "b"  "-"  "*"  "^"

The example uses an analyse() traversal. Such a traversal traverses the expression structure depth-first and invoke the callbacks on each sub-expression. What these compute are passed, bottom-up, to callbacks in the parameter bottomup. For symbols, we are at a leaf in the expression tree, so we do not get any bottomup information from the recursion, but the list we created is propagated up in the recursion. We just didn’t see it in the example because the default callbacks propagate the results up and we could use those. Using callbacks for call expressions lets us manipulate this information to construct more complicated analyses.

Example: substitution

As another example, we can consider modifying a function by replacing a given variable with a value. Here, we will use a rewrite() traversal instead of the analyse() traversal we used in the previous example. With a rewrite() traversal we have the same callbacks as in an analysis, and they are called with the same arguments (except for bottomup which is only used in analysis()), but they have to return an expression. The expression they return will be substituted into the result at the position where their expr argument sits in the input function.

We want to replace symbols with values, so we will use the symbols callback again. This time around, we write a function for the transformation because we want to specify the variable we want substituted and the value we want it substituted with. We will simply collect these as expressions and then run the function through a rewrite() pipeline. Like this:

subst_symbol <- function(fn, var, val) {
    var <- rlang::enexpr(var)
    val <- rlang::enexpr(val)
    fn %>% rewrite() %>% rewrite_with(
        rewrite_callbacks() %>% with_symbol_callback(
            function(expr, ...) if (expr == var) val else expr
        )
    )
}

Then, with a function like this:

f <- function(x, y) 2 * x + y ** 2

we can replace x with 3:

f %>% subst_symbol(x, 3)
#> function (x, y) 
#> 2 * 3 + y^2

We have replaced x with 3 inside f’s body, but not removed the formal parameter. If you want to do this, you can do it in the subst_symbol function:

subst_symbol <- function(fn, var, val) {
    var <- rlang::enexpr(var)
    val <- rlang::enexpr(val)
    fn %>% rewrite() %>% rewrite_with(
        rewrite_callbacks() %>% with_symbol_callback(
            function(expr, ...) if (expr == var) val else expr
        )
    ) %>% remove_formal_(var)
}
f %>% subst_symbol(x, 3)
#> function (y) 
#> 2 * 3 + y^2

In foolbox there is no special functionality for manipulating the function we rewrite. That can easily be done using base R functions. We focus on rewriting expressions.

Setting up callbacks

When configuring callbacks for a traversal, you should start with the default set of functions for the traversal you have in mind. You get the default callbacks for an analysis traversal using the function analysis_callbacks() and the default callbacks for a rewrite traversal using the function rewrite_callbacks(). You can then replace a callback, of each of the six kinds, using the with_type_callback() functions, as part of a pipeline. For example, the pipeline

rewrite_callbacks() %>% 
    with_call_callback(my_call_callback) %>%
    with_symbol_callback(my_symbol_callback)

will create callbacks that are defaults for rewriting for atomics, primates, pair-lists and top-down and have replaced the call and symbol callbacks with my_call_callback and my_symbol_callback, respectively.

When you use the with_type_callback functions you replace the previous callback of the same type with the new function. The existing one is not lost, however. It will be provided to the function you install via the next_cb parameter, and you can call it if you want to chain together several callbacks of the same type.

To make this concrete, consider the variable substitution from the previous section. If we try to use two different symbols callbacks to replace two different variables, we will only see the effect of the last callback we add, so in the example below, we replace y but not x.

f <- function(x, y) x + y
f %>% rewrite() %>% rewrite_with(
    rewrite_callbacks() %>% with_symbol_callback(
        function(expr, ...) if (expr == quote(x)) 1 else expr
    ) %>% with_symbol_callback(
        function(expr, ...) if (expr == quote(y)) 1 else expr
    )
)
#> function (x, y) 
#> x + 1

We can remedy this using next_cb.

f %>% rewrite() %>% rewrite_with(
    rewrite_callbacks() %>% with_symbol_callback(
        function(expr, ...) if (expr == quote(x)) 1 else expr
    ) %>% with_symbol_callback(
        function(expr, next_cb, ...) if (expr == quote(y)) 1 else next_cb(expr)
    )
)
#> function (x, y) 
#> 1 + 1

Notice here that the chaining goes in the opposite direction of how the callbacks are inserted. Whenever you add a callback with one of the with_type_callback functions you replace the previous one, so the callbacks configuration prefers the most recent changes over the older changes. This makes it easier to work with callbacks, because you can think of it as modifying a configuration. It does mean that the most recent specialisation is called first, so if you chain several callbacks, they will be invoked in the opposite order than the one they are added in. In many ways, the next_cb works for callbacks the way that NextMethod works for generic functions.

For the substitution example, you can replace multiple variables in two ways: you can make one traversal with several, chained, symbol callbacks, or you can perform several traversals, in series, each performing one substitution. The former will work as parallel substitution since while the former will allow later transformations to modify values you have inserted in previous transformations. Both approaches have their uses, and you can implement both using foolbox.

There are two other ways to add callbacks: add_topdown_callback and add_call_callback. These behave the same as with_topdown_callback and add_call_callback but allows you to specify a function they should be invoked on. They will only be called when we see call expressions to that specific function.

We can see this in use in a variation of the function we wrote earlier for extracting the symbols used in a function. Instead of extracting all symbols, we might be interested in only those that are assigned to within a function. We can get those by processing the assignments in the function; these are calls to <- and = (calls to -> are translated into calls to <- by the parser so you never see them in the expression structures).

We could capture these by installing a callback that will be called for all calls and then check if the call is to the specific function, but using add_call_callback that check is done for us. We can write the function for collecting symbols assigned to like this:

get_symbol <- function(expr, bottomup, ...) {
    bottomup <- merge_bottomup(bottomup)
    if (rlang::is_symbol(expr[[2]]))
        bottomup$variables <- c(as.character(expr[[2]]), bottomup$variables)
    bottomup
}
collect_symbols <- . %>% analyse() %>% analyse_with(
    analysis_callbacks() %>% 
        add_call_callback(`<-`, get_symbol) %>%
        add_call_callback(`=`, get_symbol)
    ) %>% unlist(use.names = FALSE) %>% unique()

This function will only extract those variables we assign to and not formal arguments or variables that are simply used in expressions.

f <- function(x, y) {
    z <- x + y
    x + y + z
}
f %>% collect_symbols()
#> [1] "z"

Initial annotation transformations

The analyse() and rewrite() functions are actually the same function and they both perform a rewrite of their input to prepare it for user-specified traversals. They perform an analysis of their input that computes which variables are assigned to in each scope in the expression-structure and which variables are bound in each scope (bound means they are either variables that are assigned to or parameters of functions in enclosing scopes). This information is added to each expression in the structure as attributes "assigned_symbols" and "bound", respectively.

Consider this function:

f <- function(x, y) {
    a <- 2 * x + y
    g <- function(z) {
        a - z
    }
    g(-x)
}

If we pipe it through rewrite() (or analyse(), the result is the same) we get an annotated function.

annotated_f <- f %>% rewrite()

At the top-level, the body of the annotated f, we see that a and g are assigned to in that scope, but that we also have x and y as bound variables, since these are parameters of f.

attr(body(annotated_f), "assigned_symbols")
#> [1] "a" "g"
attr(body(annotated_f), "bound")
#> [1] "x" "y" "a" "g"

We can extract the body of f inside f. Inside g we do not assign to any new local variables, but since g is nested inside f, all the symbols that are bound in f are also bound in g:

g_body <- body(annotated_f)[[3]][[3]][[3]]
g_body
#> {
#>     a - z
#> }
#> attr(,"assigned_symbols")
#> character(0)
#> attr(,"bound")
#> [1] "z" "x" "y" "a" "g"

It is impossible to know before runtime whether any variable that is part of an assignment call in the body of a function will actually be assigned to. So this annotation is conservative and assumes that if it is possible that a variable is assigned to, then it will also be bound. Even so, R is such a dynamic language that it is possible, in many different ways, to modify the environment of a function from outside the static expressions in its body, that the best we can do is heuristics. The heuristics used are described in the documentation for the rewrite() function that implements them:

Since R does not require that we declare local variables, and since the variables that are assigned to a local scope depend on the runtime execution of functions, we cannot determine with any certainty which variables will be assigned to in any given scope at any given program point. So the best we can do is figure out which variables are potentially assigned to. Which is what this function does.

The rules for when we are assigning to a local variable are a bit complicated. For control structures, we can assume that assignments will be to the local scope. People can change the implementation of these so it isn’t, but then they are only hurting themselves and deserve the extra pain we can give them. For other call arguments, it gets a little more complicated. With standard-evaluation, if we have an arrow assignment in a function argument, then the assignment happens in the calling scope. So we will assume this happens unless we are handling cases we know have NSE, such as with. If an assignment is inside a block, however, we will assume that NSE is in play, by default, and not consider it a local assignment.

It is quite likely that these rules will be updated in the future as I experiment more with the package, but this is what I have come up with so far.

In any case, the annotation before the user-defined traversals is used to guide the dispatching on specific functions in the add_topdown_callback and add_call_callback callbacks. It prevents these callbacks to transform local functions that just happen to share the name with one that is used in one of those two functions.

Setting up a traversal

Specifying an actual traversal, once we understand how callbacks work, is fairly simple. After calling one of the rewrite() or analyse() functions you follow up with rewrite_with() or analyse_with(). These take at least one argument, which should be a callback configuration, and might take more. If you provide more arguments to them, make them named. Then they will be passed along to all the callbacks in the traversal.

We have already seen several examples of this, so I will not provide more, but simply stress that whenever you build a transformation pipeline you have the choice between adding more callbacks to one transformation or making more transformations with the output of one becoming the input of another.

If you specify a sequence of transformations, you can pass information along with expressions using attributes. Consider, for example, this dummy-example where we inline a function by replacing calls to it with its body (ignoring the remapping of variables a real inline function should consider, but see this example for a more detailed example of inlining functions). We save, together with the transformed expression, the old function call. In a second traversal we can use this information, for example to reverse the inlining.

inline <- function(fn, f) {
    fn %>% rewrite_with(
        rewrite_callbacks() %>% add_call_callback(
            f, function(expr, ...) {
                new_expr <- body(f)
                attr(new_expr, "old_expr") <- expr
                new_expr
            }
        )
    )
}
reverse <- function(fn) {
    fn %>% rewrite_with(
        rewrite_callbacks() %>% with_call_callback(
            function(expr, ...) {
                old_expr <- attr(expr, "old_expr")
                if (!is.null(old_expr)) old_expr
                else expr
            }
        )
    )
}

We can thus run the first transformation to inline a function:

f <- function(x, y) x + y
g <- function(x) f(x, 2*x)
h <- g %>% rewrite() %>% inline(f)
h
#> function (x) 
#> x + y

And then turn around and reverse that operation because we have the old expressions stored in the tree:

h %>% rewrite() %>% reverse()
#> function (x) 
#> f(x, 2 * x)

Inlining and then reversing is, obviously, a nonsensical use of transformations, but storing information about symbols in scope is how we handle add_topdown_callback and add_call_callback callbacks correctly, and I suspect it can have other uses in static analysis.

Information-flow in traversals

The actual expression-tree traversals are handled by functions in the foolbox package, and you only have access to it via callbacks. You therefore need a way of passing information between callbacks. You have four options here:

• You can pass any fixed information along the traversal as extra parameters to the rewrite_with or analyse_with functions. Any named parameter, not on the list of parameters automatically passed to callbacks as described above, will get passed to the callbacks.
• You can provide information top-down using the topdown parameter. You can provide this parameter to rewrite_with and analyse_with to specify an initial value, and any value returned by a top-down callback will be passed to callbacks deeper in the traversal via this parameter.
• In analysis traversals, every callback must return a list. When going bottom-up in the traversal, the recursive types—pair-lists and call expressions—will have such a list for each of the constituent expressions, and this list will be passed along via the bottomup parameter. You can use the function merge_bottomup to merge this list-of-lists into a single list.
• In rewrite traversals, you can set attributes on expressions. Since the traversal is depth-first, these attributes will be available in the constituents of for composite expressions. You can use the function collect_from_args to merge information from sub-expressions.

Controlling warnings

The traversals will warn you if it encounters a call in its traversal that calls a function that is not in scope. It will also warn if you try to transform a function that might be referring to a local variable. You can turn these warnings on and off—in the annotation done in the rewrite() and analyse() functions, we have to turn off the warning for unknown functions because we do not yet know which variables are local and which are not, so we would potentially warn on functions that are known in the function’s scope, just not until runtime.

You can use the function warning_flags() to get default warnings and you can set and unset warning flags with functions named set_warn_warning-name and unset_warn_warning-name. See the documentation for warning_flags() for details.

Traversing expressions

The documentation so far has focused on rewriting and analysing whole functions, and I would imagine that would be 99% of the uses you will have for foolbox. You can, however, also rewrite and analyse expressions in themselves (although you might have to provide them with params and env arguments of some sort since these cannot be taken from a function in that case).

The expression version of rewrite and rewrite_with are rewrite_expr and rewrite_expr_with. The expression version of analyse and analyse_with are analyse_expr and analyse_expr_with. See the documentation for these functions for more details.