R devel - Aug 2021 - JIT compiler does not compile closures with custom environments

I have encountered a behavior of R?s JIT compiler that I can?t quite figure out.
Consider the following code:


   f_global <- function(x) {
     for(i in 1:10000) x <- x + 1
     x
   }

   f_env <- local({
    function(x) {
      for(i in 1:10000) x <- x + 1
      x
    }
   })

   compiler::enableJIT(3)

  bench::mark(f_global(0), f_env(0))
  # 1 f_global(0)    103?s 107.61?s     8770.    11.4KB      0    4384     0
  # 2 f_env(0)       1.1ms   1.42ms      712.        0B     66.3   290    27
  
Inspecting the closures shows that f_global has been byte-compiled while f_env
has not been byte-compiled. Furthermore, if I assign a new environment to
f_global (e.g. via environment(f_global) <- new.env()), it won?t be
byte-compiled either.

However, if I have a function returning a closure, that closure does get
byte-compiled:

  f_closure <- (function() {
    function(x) {
      for(i in 1:10000) x <- x + 1
     x
   }
  })()

  bench::mark(f_closure(0))
  # 1 f_closure(0)    105?s    109?s     8625.        0B     2.01  4284     1   
497ms

What is going on here? Both f_closure and f_env have non-global environments.
Why is one JIT-compiled, but not the other? Is there a way to ensure that
functions defined in environments will be JIT-compiled?

Thanks, 

Taras

On 18/08/2021 9:00 a.m., Taras Zakharko wrote:
> I have encountered a behavior of R?s JIT compiler that I can?t quite figure
out. Consider the following code:
> 
> 
>     f_global <- function(x) {
>       for(i in 1:10000) x <- x + 1
>       x
>     }
> 
>     f_env <- local({
>      function(x) {
>        for(i in 1:10000) x <- x + 1
>        x
>      }
>     })
> 
>     compiler::enableJIT(3)
> 
>    bench::mark(f_global(0), f_env(0))
>    # 1 f_global(0)    103?s 107.61?s     8770.    11.4KB      0    4384    
0
>    # 2 f_env(0)       1.1ms   1.42ms      712.        0B     66.3   290   
27
>    
> Inspecting the closures shows that f_global has been byte-compiled while
f_env has not been byte-compiled. Furthermore, if I assign a new environment to
f_global (e.g. via environment(f_global) <- new.env()), it won?t be
byte-compiled either.
> 
> However, if I have a function returning a closure, that closure does get
byte-compiled:
> 
>    f_closure <- (function() {
>      function(x) {
>        for(i in 1:10000) x <- x + 1
>       x
>     }
>    })()
> 
>    bench::mark(f_closure(0))
>    # 1 f_closure(0)    105?s    109?s     8625.        0B     2.01  4284   
1      497ms
> 
> What is going on here? Both f_closure and f_env have non-global
environments. Why is one JIT-compiled, but not the other? Is there a way to
ensure that functions defined in environments will be JIT-compiled?
About what is going on in f_closure:  I think the anonymous factory

function() {
       function(x) {
         for(i in 1:10000) x <- x + 1
        x
      }
     }

got byte compiled before first use, and that compiled its result.  That 
seems to be what this code indicates:

   f_closure <- (function() {
   res <- function(x) {
   for(i in 1:10000) x <- x + 1
   x
   }; print(res); res
   })()
   #> function(x) {
   #> for(i in 1:10000) x <- x + 1
   #> x
   #> }
   #> <bytecode: 0x7fb43ec3aa70>
   #> <environment: 0x7fb441117ac0>

But even if that's true, it doesn't address the bigger question of why 
f_global and f_env are treated differently.

Duncan Murdoch