memoise in rstats with a limit on how large an object can be in the cache

memoise-with-size-limit.R

#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
#' A version of 'memoise::memoise' with limits on individual object size
#'
#' @param f     Function of which to create a memoised copy.
#' @param ... optional variables specified as formulas with no RHS to use as
#' additional restrictions on caching. See Examples for usage.
#' @param envir Environment of the returned function.
#' @param cache Cache function.
#' @param object_size_limit maximum size of objects stored in cache.
#'        Default: 1048576 bytes (1MB)
#'
#'
#' @import memoise
#' @importFrom stats setNames
#' @importFrom digest digest
#~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
memoise_with_size_limit <- function (f, ..., envir = environment(f), 
                                     cache = memoise::cache_memory(), 
                                     object_size_limit = 1048576L) {
  f_formals <- formals(args(f))
  if (memoise::is.memoised(f)) {
    stop("`f` must not be memoised.", call. = FALSE)
  }
  f_formal_names <- names(f_formals)
  f_formal_name_list <- lapply(f_formal_names, as.name)
  init_call_args <- setNames(f_formal_name_list, f_formal_names)
  init_call <- memoise:::make_call(quote(`_f`), init_call_args)
  memoise:::validate_formulas(...)
  additional <- list(...)
  memo_f <- eval(bquote(function(...) {
    called_args <- as.list(match.call())[-1]
    default_args <- Filter(function(x) !identical(x, quote(expr = )),
                           as.list(formals()))
    default_args <- default_args[setdiff(names(default_args),
                                         names(called_args))]
    args <- c(lapply(called_args, eval, parent.frame()),
              lapply(default_args, eval, envir = environment()))
    hash <- `_cache`$digest(c(body(`_f`), args, lapply(`_additional`,
                                                       function(x) eval(x[[2L]], environment(x)))))
    if (`_cache`$has_key(hash)) {
      res <- `_cache`$get(hash)
    }
    else {
      res <- withVisible(.(init_call))
      #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
      # Check size and only store if < object_size_limit
      #~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
      if (pryr::object_size(res) < .(object_size_limit)) {
        `_cache`$set(hash, res)
      }
    }
    if (res$visible) {
      res$value
    }
    else {
      invisible(res$value)
    }
  }, as.environment(list(init_call = init_call, object_size_limit = object_size_limit))))
  formals(memo_f) <- f_formals
  attr(memo_f, "memoised") <- TRUE
  if (is.null(envir)) {
    envir <- baseenv()
  }
  memo_f_env <- new.env(parent = envir)
  memo_f_env$`_cache` <- cache
  memo_f_env$`_f` <- f
  memo_f_env$`_additional` <- additional
  environment(memo_f) <- memo_f_env
  class(memo_f) <- c("memoised", "function")
  memo_f
}