From d49e885129403f38d9a62e798f6b71ebf62e0ff0 Mon Sep 17 00:00:00 2001
From: Alexander Lyon <arlyon@me.com>
Date: Fri, 8 Mar 2024 08:11:12 +0000
Subject: [PATCH] feat: add mechanism to bypass turbo engine for annotated
 methods

This PR still needs some work, since we will need to transparently
forward VCs or eagerly evaluate them. We need to ensure that the parent
task is still lazy, but that we are just cutting out task tracking
between tasks.
---
 crates/turbo-tasks-build/src/lib.rs           | 19 +++++++++++++-
 .../turbo-tasks-macros/src/function_macro.rs  | 26 +++++++++++++++++++
 2 files changed, 44 insertions(+), 1 deletion(-)

diff --git a/crates/turbo-tasks-build/src/lib.rs b/crates/turbo-tasks-build/src/lib.rs
index 55d7e5640e170d..d34fbf2dcb2038 100644
--- a/crates/turbo-tasks-build/src/lib.rs
+++ b/crates/turbo-tasks-build/src/lib.rs
@@ -21,6 +21,21 @@ use turbo_tasks_macros_shared::{
     get_type_ident, GenericTypeInput, PrimitiveInput, ValueTraitArguments,
 };
 
+/// the same module exists in turbo-tasks-macros but due to its size we
+/// just copy the relevant parts in both places
+mod ignored {
+    use std::{cell::OnceCell, collections::HashSet};
+    // a newline-separated list of task names to opt-out of turbo tracking
+    const IGNORE_TASKS_RAW: Option<&str> = std::option_env!("TURBO_IGNORE_TASKS");
+    const IGNORE_TASKS: OnceCell<HashSet<&'static str>> = OnceCell::new();
+
+    pub fn task_ignored(name: &str) -> bool {
+        IGNORE_TASKS
+            .get_or_init(|| IGNORE_TASKS_RAW.unwrap_or_default().split(',').collect())
+            .contains(name)
+    }
+}
+
 pub fn generate_register() {
     println!("cargo:rerun-if-changed=build.rs");
 
@@ -221,7 +236,9 @@ impl<'a> RegisterContext<'a> {
     }
 
     fn process_fn(&mut self, fn_item: ItemFn) -> Result<()> {
-        if has_attribute(&fn_item.attrs, "function") {
+        if has_attribute(&fn_item.attrs, "function")
+            && !ignored::task_ignored(&fn_item.sig.ident.to_string())
+        {
             let ident = &fn_item.sig.ident;
             let type_ident = get_native_function_ident(ident);
 
diff --git a/crates/turbo-tasks-macros/src/function_macro.rs b/crates/turbo-tasks-macros/src/function_macro.rs
index 4d805de9ad711e..207f9f424fd876 100644
--- a/crates/turbo-tasks-macros/src/function_macro.rs
+++ b/crates/turbo-tasks-macros/src/function_macro.rs
@@ -6,9 +6,35 @@ use turbo_tasks_macros_shared::{get_native_function_id_ident, get_native_functio
 
 use crate::func::{DefinitionContext, NativeFn, TurboFn};
 
+/// the same module exists in turbo-tasks-build but due to its size we
+/// just copy the relevant parts in both places
+mod ignored {
+    use std::{cell::OnceCell, collections::HashSet};
+    // a newline-separated list of task names to opt-out of turbo tracking
+    const IGNORE_TASKS_RAW: Option<&str> = std::option_env!("TURBO_IGNORE_TASKS");
+    const IGNORE_TASKS: OnceCell<HashSet<&'static str>> = OnceCell::new();
+
+    pub fn task_ignored(name: &str) -> bool {
+        IGNORE_TASKS
+            .get_or_init(|| IGNORE_TASKS_RAW.unwrap_or_default().split(',').collect())
+            .contains(name)
+    }
+}
+
 pub fn function(_args: TokenStream, input: TokenStream) -> TokenStream {
     let item = parse_macro_input!(input as ItemFn);
 
+    // TODO: ideally here we would be able to match on the entire module path
+    //       however macros are evaluated in a different context and we don't
+    //       have access to it. in the mean time, just use the function name.
+    //       the likelihood that two functions with the same name exist is low
+    //       and, if we do need to disable just one of them, can be mitigated
+    //       through naming
+    let name = item.sig.ident.to_string();
+    if ignored::task_ignored(&name) {
+        return quote! { #item }.into();
+    }
+
     let ItemFn {
         attrs,
         vis,