Skip to content

Commit

Permalink
Handle #[cfg(...)] attributes on turbo tasks (vercel/turborepo#8542)
Browse files Browse the repository at this point in the history 
### Description

If a task is conditionally compiled, we must conditionally register it.

Noticed this when trying to write a test with `#[cfg(test)]`.

**Related:** vercel/turborepo#8530 adds support for registering tasks inside of `mod foo { ... }` blocks.

### Testing Instructions

Tested as part of vercel/turborepo#8529
  • Loading branch information
@bgw
bgw authored Jun 24, 2024
1 parent daed66a commit 223a0d1
Show file tree
Hide file tree
Showing 2 changed files with 129 additions and 41 deletions.
1 change: 1 addition & 0 deletions crates/turbo-tasks-build/Cargo.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -15,5 +15,6 @@ workspace = true
anyhow = { workspace = true }
cargo-lock = "8.0.2"
glob = "0.3.0"
quote = { workspace = true }
syn = { workspace = true, features = ["full"] }
turbo-tasks-macros-shared = { workspace = true }
169 changes: 128 additions & 41 deletions crates/turbo-tasks-build/src/lib.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -4,10 +4,12 @@ use std::{
fmt::{Display, Write},
fs::read_dir,
path::{PathBuf, MAIN_SEPARATOR as PATH_SEP},
sync::Arc,
};

use anyhow::{Context, Result};
use glob::glob;
use quote::ToTokens;
use syn::{
parse_quote, Attribute, Ident, Item, ItemEnum, ItemFn, ItemImpl, ItemMacro, ItemMod,
ItemStruct, ItemTrait, TraitItem, TraitItemMethod,
Expand Down Expand Up @@ -98,9 +100,18 @@ pub fn generate_register() {

let out_file = out_dir.join(filename);

let mut queue = vec![("".to_string(), entry)];

while let Some((mod_path, file_path)) = queue.pop() {
let mut queue = vec![QueueEntry {
file_path: entry,
mod_path: "".to_string(),
attributes: Vec::new(),
}];

while let Some(QueueEntry {
file_path,
mod_path,
attributes,
}) = queue.pop()
{
println!("cargo:rerun-if-changed={}", file_path.to_string_lossy());
let src = std::fs::read_to_string(&file_path).unwrap();

Expand All @@ -110,6 +121,7 @@ pub fn generate_register() {
file_path: &file_path,
prefix: &prefix,
mod_path,
attributes,

register: &mut register_code,
values: &mut values,
Expand All @@ -120,24 +132,27 @@ pub fn generate_register() {
{
Ok(file) => {
for item in file.items {
ctx.process_item(item).unwrap();
ctx.process_item(&item).unwrap();
}
}
Err(err) => println!("{}", err),
}
}

let mut values_code = String::new();
for ((mod_path, ident), (global_name, trait_idents)) in values {
for ((mod_path, ident), entry) in values {
for attribute in &entry.attributes {
values_code.push_str(attribute);
}
writeln!(
values_code,
"crate{}::{}({}, #[allow(unused_variables)] |value| {{",
mod_path,
get_register_value_type_ident(&ident),
global_name
entry.global_name,
)
.unwrap();
for trait_ident in trait_idents {
for trait_ident in entry.trait_idents {
writeln!(
values_code,
" crate{}::{}(value);",
Expand Down Expand Up @@ -183,21 +198,36 @@ pub fn rerun_if_glob(globs: &str, root: &str) {
/// (mod_path, type_ident)
type ValueKey = (String, Ident);
/// (global_name, trait_register_fns)
type ValueEntry = (String, Vec<Ident>);
struct ValueEntry {
attributes: Vec<Arc<String>>,
global_name: String,
trait_idents: Vec<Ident>,
}

struct QueueEntry {
/// The on-disk path to the file representing this module.
file_path: PathBuf,
/// The `syn::Path`-style representation of the module. Each component is
/// separated by `::`.
mod_path: String,
/// Attributes (`#[cfg(...)]`) applied to the `ItemMod`.
attributes: Vec<Arc<String>>,
}

struct RegisterContext<'a> {
queue: &'a mut Vec<(String, PathBuf)>,
queue: &'a mut Vec<QueueEntry>,

file_path: &'a PathBuf,
mod_path: String,
attributes: Vec<Arc<String>>,
prefix: &'a str,

register: &'a mut String,
values: &'a mut HashMap<ValueKey, ValueEntry>,
}

impl<'a> RegisterContext<'a> {
fn process_item(&mut self, item: Item) -> Result<()> {
fn process_item(&mut self, item: &Item) -> Result<()> {
match item {
Item::Enum(enum_item) => self.process_enum(enum_item),
Item::Fn(fn_item) => self.process_fn(fn_item),
Expand All @@ -210,16 +240,24 @@ impl<'a> RegisterContext<'a> {
}
}

fn process_enum(&mut self, enum_item: ItemEnum) -> Result<()> {
if has_attribute(&enum_item.attrs, "value") {
fn process_enum(&mut self, item: &ItemEnum) -> Result<()> {
self.with_cfg_attrs(&item.attrs, move |this| this.process_enum_inner(item))
}

fn process_enum_inner(&mut self, enum_item: &ItemEnum) -> Result<()> {
if has_turbo_attribute(&enum_item.attrs, "value") {
self.add_value(&enum_item.ident);
self.add_value_debug_impl(&enum_item.ident);
}
Ok(())
}

fn process_fn(&mut self, fn_item: ItemFn) -> Result<()> {
if has_attribute(&fn_item.attrs, "function") {
fn process_fn(&mut self, item: &ItemFn) -> Result<()> {
self.with_cfg_attrs(&item.attrs, move |this| this.process_fn_inner(item))
}

fn process_fn_inner(&mut self, fn_item: &ItemFn) -> Result<()> {
if has_turbo_attribute(&fn_item.attrs, "function") {
let ident = &fn_item.sig.ident;
let type_ident = get_native_function_ident(ident);

Expand All @@ -228,8 +266,12 @@ impl<'a> RegisterContext<'a> {
Ok(())
}

fn process_impl(&mut self, impl_item: ItemImpl) -> Result<()> {
if has_attribute(&impl_item.attrs, "value_impl") {
fn process_impl(&mut self, item: &ItemImpl) -> Result<()> {
self.with_cfg_attrs(&item.attrs, move |this| this.process_impl_inner(item))
}

fn process_impl_inner(&mut self, impl_item: &ItemImpl) -> Result<()> {
if has_turbo_attribute(&impl_item.attrs, "value_impl") {
let struct_ident = get_type_ident(&impl_item.self_ty).unwrap();

let trait_ident = impl_item
Expand All @@ -241,7 +283,7 @@ impl<'a> RegisterContext<'a> {
self.add_value_trait(&struct_ident, trait_ident);
}

for item in impl_item.items {
for item in &impl_item.items {
if let syn::ImplItem::Method(method_item) = item {
// TODO: if method_item.attrs.iter().any(|a|
// is_attribute(a,
Expand All @@ -266,46 +308,63 @@ impl<'a> RegisterContext<'a> {
Ok(())
}

fn process_mod(&mut self, mod_item: ItemMod) -> Result<()> {
if let Some((_, items)) = mod_item.content {
let mod_name = mod_item.ident.to_string();
let child_mod_path = format!("{}::{}", self.mod_path, mod_name);
fn process_mod(&mut self, item: &ItemMod) -> Result<()> {
self.with_cfg_attrs(&item.attrs, move |this| this.process_mod_inner(item))
}

fn process_mod_inner(&mut self, mod_item: &ItemMod) -> Result<()> {
let child_mod_name = mod_item.ident.to_string();
let child_mod_path = format!("{}::{}", self.mod_path, child_mod_name);
if let Some((_, items)) = &mod_item.content {
let parent_mod_path = std::mem::replace(&mut self.mod_path, child_mod_path);
for item in items {
self.process_item(item)?;
}
self.mod_path = parent_mod_path;
} else {
let name = mod_item.ident.to_string();
let parent_path = self.file_path.parent().unwrap();
let direct = parent_path.join(format!("{name}.rs"));
let parent_file_path = self.file_path.parent().unwrap();
let direct = parent_file_path.join(format!("{child_mod_name}.rs"));
if direct.exists() {
self.queue
.push((format!("{}::{name}", self.mod_path), direct));
self.queue.push(QueueEntry {
file_path: direct,
mod_path: child_mod_path,
attributes: self.attributes.clone(),
});
} else {
let nested = parent_path.join(&name).join("mod.rs");
let nested = parent_file_path.join(&child_mod_name).join("mod.rs");
if nested.exists() {
self.queue
.push((format!("{}::{name}", self.mod_path), nested));
self.queue.push(QueueEntry {
file_path: nested,
mod_path: child_mod_path,
attributes: self.attributes.clone(),
});
}
}
}
Ok(())
}

fn process_struct(&mut self, struct_item: ItemStruct) -> Result<()> {
if has_attribute(&struct_item.attrs, "value") {
fn process_struct(&mut self, item: &ItemStruct) -> Result<()> {
self.with_cfg_attrs(&item.attrs, move |this| this.process_struct_inner(item))
}

fn process_struct_inner(&mut self, struct_item: &ItemStruct) -> Result<()> {
if has_turbo_attribute(&struct_item.attrs, "value") {
self.add_value(&struct_item.ident);
self.add_value_debug_impl(&struct_item.ident);
}
Ok(())
}

fn process_trait(&mut self, trait_item: ItemTrait) -> Result<()> {
fn process_trait(&mut self, item: &ItemTrait) -> Result<()> {
self.with_cfg_attrs(&item.attrs, move |this| this.process_trait_inner(item))
}

fn process_trait_inner(&mut self, trait_item: &ItemTrait) -> Result<()> {
if let Some(attr) = trait_item
.attrs
.iter()
.find(|a| is_attribute(a, "value_trait"))
.find(|a| is_turbo_attribute(a, "value_trait"))
{
let trait_ident = &trait_item.ident;

Expand Down Expand Up @@ -343,13 +402,17 @@ impl<'a> RegisterContext<'a> {
Ok(())
}

fn process_macro(&mut self, macro_item: ItemMacro) -> Result<()> {
fn process_macro(&mut self, item: &ItemMacro) -> Result<()> {
self.with_cfg_attrs(&item.attrs, move |this| this.process_macro_inner(item))
}

fn process_macro_inner(&mut self, macro_item: &ItemMacro) -> Result<()> {
if macro_item
.mac
.path
.is_ident("__turbo_tasks_internal_primitive")
{
let input = macro_item.mac.tokens;
let input = macro_item.mac.tokens.clone();
let input = syn::parse2::<PrimitiveInput>(input).unwrap();

let ty = input.ty;
Expand All @@ -365,7 +428,7 @@ impl<'a> RegisterContext<'a> {
.path
.is_ident("__turbo_tasks_internal_generic_type")
{
let input = macro_item.mac.tokens;
let input = macro_item.mac.tokens.clone();
let input = syn::parse2::<GenericTypeInput>(input).unwrap();
let ty = input.ty;
Expand Down Expand Up @@ -398,7 +461,11 @@ impl<'a> RegisterContext<'a> {

fn add_value(&mut self, ident: &Ident) {
let key: ValueKey = (self.mod_path.clone(), ident.clone());
let value: ValueEntry = (self.get_global_name(&[ident]), Vec::new());
let value = ValueEntry {
attributes: self.attributes.clone(),
global_name: self.get_global_name(&[ident]),
trait_idents: Vec::new(),
};

assert!(
self.values.insert(key, value).is_none(),
Expand Down Expand Up @@ -444,14 +511,17 @@ impl<'a> RegisterContext<'a> {
self.file_path.display()
);
}
entry.unwrap().1.push(trait_ident.clone());
entry.unwrap().trait_idents.push(trait_ident.clone());
}

fn register(
&mut self,
type_ident: impl Display,
global_name: impl Display,
) -> std::fmt::Result {
for attribute in &self.attributes {
self.register.push_str(attribute);
}
writeln!(
self.register,
"crate{}::{}.register({});",
Expand Down Expand Up @@ -503,13 +573,24 @@ impl<'a> RegisterContext<'a> {
&["value_default"],
)
}

fn with_cfg_attrs<T>(&mut self, attrs: &[Attribute], func: impl FnOnce(&mut Self) -> T) -> T {
let orig_len = self.attributes.len();
for attr in attrs.iter().filter(|a| is_cfg_attribute(a)) {
self.attributes
.push(Arc::new(attr.to_token_stream().to_string()));
}
let ret = func(self);
self.attributes.truncate(orig_len);
ret
}
}

fn has_attribute(attrs: &[Attribute], name: &str) -> bool {
attrs.iter().any(|a| is_attribute(a, name))
fn has_turbo_attribute(attrs: &[Attribute], name: &str) -> bool {
attrs.iter().any(|a| is_turbo_attribute(a, name))
}

fn is_attribute(attr: &Attribute, name: &str) -> bool {
fn is_turbo_attribute(attr: &Attribute, name: &str) -> bool {
let path = &attr.path;
if path.leading_colon.is_some() {
return false;
Expand All @@ -524,6 +605,12 @@ fn is_attribute(attr: &Attribute, name: &str) -> bool {
}
}

fn is_cfg_attribute(attr: &Attribute) -> bool {
attr.path
.get_ident()
.is_some_and(|ident| ident == "cfg" || ident == "cfg_attr")
}

fn parse_attr_args<T>(attr: &Attribute) -> syn::Result<Option<T>>
where
T: syn::parse::Parse,
Expand Down

0 comments on commit 223a0d1

Please sign in to comment.