R/meta.R

In duckdblabs/duckplyr: A 'DuckDB'-Backed Version of 'dplyr'

call_stack <- collections::stack()
pre_code_cache <- collections::queue()
code_cache <- collections::queue()
ext_cache <- collections::dict()
macro_cache <- collections::dict()
df_cache <- collections::dict()
rel_cache <- collections::dict()

meta_call <- function(name) {
  meta_call_start(name)
  withr::defer_parent(meta_call_end())
}

meta_call_start <- function(name) {
  call_stack$push(name)
}

meta_call_end <- function() {
  call_stack$pop()
}

meta_call_current <- function() {
  if (call_stack$size() == 0) {
    return(NULL)
  }
  call_stack$peek()
}

meta_clear <- function() {
  pre_code_cache$clear()
  code_cache$clear()
  ext_cache$clear()
  macro_cache$clear()
  df_cache$clear()
  rel_cache$clear()
}

meta_pre_record <- function(call) {
  pre_code_cache$push(new_prom_fun({{ call }}))
  invisible()
}

meta_record <- function(call) {
  code_cache$push(new_prom_fun({{ call }}))
  invisible()
}

meta_replay <- function(add_pre_code = TRUE) {
  if (add_pre_code) {
    con_exprs <- list(
      expr(duckdb <- asNamespace("duckdb")),
      expr(drv <- duckdb::duckdb()),
      expr(con <- DBI::dbConnect(drv)),
      expr(experimental <- !!(Sys.getenv("DUCKPLYR_EXPERIMENTAL") == "TRUE"))
    )
    con_code <- map(con_exprs, constructive::deparse_call)
    pre_code <- c(
      con_code,
      map(pre_code_cache$as_list(), ~ .x())
    )
  } else {
    pre_code <- NULL
  }

  # HACK
  count <- rel_cache$size()
  res_name <- sym(paste0("rel", count))
  res_mat_expr <- expr(duckdb$rel_to_altrep(!!res_name))
  res_code <- map(list(res_name, res_mat_expr), constructive::deparse_call)

  all_code <- c(
    pre_code,
    map(code_cache$as_list(), ~ .x()),
    res_code
  )

  walk(all_code, print)
}

meta_replay_to_fun_code <- function() {
  code <- utils::capture.output(meta_replay(add_pre_code = FALSE))
  code <- c(
    paste0("function(con, experimental) {"),
    paste0("  ", code),
    "}"
  )

  # Trailing newline
  paste0(code, "\n", collapse = "")
}

meta_replay_to_fun <- function(text = meta_replay_to_fun_code()) {
  eval(parse(text = text)[[1]])
}

meta_replay_to_fun_file <- function(name) {
  code <- paste0(
    "# Generated by meta_replay_to_fun_file(), do not edit by hand\n", name, " <- ",
    meta_replay_to_fun_code()
  )

  path <- file.path("R", paste0(name, ".R"))

  brio::write_file(code, path)
}

meta_replay_to_file <- function(path, extra = character()) {
  code <- utils::capture.output(meta_replay())
  writeLines(c(extra, code), path)
}

meta_replay_to_new_doc <- function() {
  code <- utils::capture.output(meta_replay())
  rstudioapi::documentNew(code, execute = TRUE)
}

meta_replay_to_reprex <- function(...) {
  code <- utils::capture.output(meta_replay())
  reprex::reprex(input = code, ...)
}

meta_eval <- function() {
  code <- utils::capture.output(meta_replay())
  eval(parse(text = code))
}

meta_ext_register <- function(name = "rfuns") {
  if (ext_cache$has(name)) {
    return(invisible())
  }

  stopifnot(identical(name, "rfuns"))

  ext_install_expr <- expr(invisible(
    duckdb$rapi_load_rfuns(drv@database_ref)
  ))
  meta_pre_record(constructive::deparse_call(ext_install_expr))

  ext_cache$set(name, TRUE)
  invisible()
}

meta_macro_register <- function(name) {
  macro <- duckplyr_macros[name]
  if (is.na(macro)) {
    return(invisible())
  }

  if (macro_cache$has(name)) {
    return(invisible())
  }

  # Register functions from the rfuns extension
  # Can't use '^"r_' because of the way the macro is defined
  if (grepl('"r_', macro)) {
    meta_ext_register()
  }

  macro_expr <- expr(invisible(
    DBI::dbExecute(con, !!paste0('CREATE MACRO "', names(macro), '"', macro))
  ))
  meta_pre_record(constructive::deparse_call(macro_expr))

  macro_cache$set(name, TRUE)
  invisible()
}

meta_df_register <- function(df) {
  if (df_cache$has(df)) {
    return(invisible(df_cache$get(df)))
  }

  count <- df_cache$size()
  name <- sym(paste0("df", count + 1))

  df_expr <- NULL

  if (Sys.getenv("DUCKPLYR_META_GLOBAL") == "TRUE") {
    global_dfs <- mget(ls(.GlobalEnv), .GlobalEnv, mode = "list", ifnotfound = list(NULL))

    for (df_name in names(global_dfs)) {
      global_df <- global_dfs[[df_name]]
      # FIXME: Does this also work with pointer comparison?
      if (identical(df, global_df)) {
        df_expr <- sym(df_name)
        break
      }
    }
  }

  df_cache$set(df, name)

  if (is.null(df_expr)) {
    class(df) <- setdiff(class(df), "duckplyr_df")
    meta_record(constructive::construct_multi(list2(!!name := df)))
  } else {
    # Changes df in-place!
    class(df) <- setdiff(class(df), "duckplyr_df")
    meta_record(constructive::deparse_call(expr(!!name <- !!df_expr)))
  }

  invisible(name)
}

meta_rel_register_df <- function(rel, df) {
  if (Sys.getenv("DUCKPLYR_META_SKIP") == "TRUE") {
    return(invisible())
  }

  df_name <- meta_df_register(df)
  # Expect experimental argument from outside
  rel_expr <- expr(duckdb$rel_from_df(con, !!df_name, experimental = experimental))
  meta_rel_register(rel, rel_expr)
}

meta_rel_register_file <- function(rel, path, table_function, options) {
  if (Sys.getenv("DUCKPLYR_META_SKIP") == "TRUE") {
    return(invisible())
  }

  rel_expr <- expr(
    duckdb$rel_from_table_function(con, !!table_function, list(!!path), list(!!!options))
  )
  meta_rel_register(rel, rel_expr)
}

meta_rel_register <- function(rel, rel_expr) {
  if (Sys.getenv("DUCKPLYR_META_SKIP") == "TRUE") {
    return(invisible())
  }

  force(rel_expr)

  count <- rel_cache$size()
  name <- sym(paste0("rel", count + 1))

  current_call <- meta_call_current()
  if (!is.null(current_call)) {
    # FIXME: This is probably too convoluted
    meta_record(constructive::deparse_call(expr(!!current_call)))
  }

  # https://github.com/cynkra/constructive/issues/102
  meta_record(constructive::deparse_call(expr(!!name <- !!rel_expr)))

  obj <- list(rel = rel, name = name, df = df)
  hash <- deparse(rel)

  rel_cache$set(hash, obj)
  invisible()
}

meta_rel_get <- function(rel) {
  hash <- deparse(rel)

  if (!rel_cache$has(hash)) {
    rel_out <- paste(utils::capture.output(print(rel), type = "message"), collapse = "\n")
    cli::cli_abort(c(
      "duckplyr: internal: hash not found",
      i = "hash: {hash}",
      i = "relation: {rel_out}"
    ))
  }

  rel_cache$get(hash)
}

new_prom_fun <- function(code) {
  quo <- enquo(code)

  valid <- FALSE
  out <- NULL

  function() {
    if (!valid) {
      out <<- eval_tidy(quo)
      valid <<- TRUE
    }
    out
  }
}