Skip to content

Commit

Permalink
feat(turbo-tasks): Allow #[turbo_tasks::function]s to accept Resolv…
Browse files Browse the repository at this point in the history
…edVc types as arguments (#70269)

## Why?

[We want to codemod structs to use `ResolvedVc<...>` for all of their field types instead of `Vc<...>`.](https://www.notion.so/vercel/Resolved-Vcs-Vc-Lifetimes-Local-Vcs-and-Vc-Refcounts-49d666d3f9594017b5b312b87ddc5bff?pvs=4)

There are many constructor-like functions (see #70133) where we must accept an argument of type `Vc<...>`, and then explicitly call `.to_resolved().await?`.

However, internally `#[turbo_tasks::function]` arguments are guaranteed to be resolved by the time the function runs. So we can do a cheap conversion here.

## What

Instead of needing to do:

```diff
 #[turbo_tasks::value_impl]
 impl CustomProcessEnv {
     #[turbo_tasks::function]
-    pub fn new(prior: Vc<Box<dyn ProcessEnv>>, custom: Vc<EnvMap>) -> Vc<Self> {
-        CustomProcessEnv { prior, custom }.cell()
+    pub async fn new(prior: Vc<Box<dyn ProcessEnv>>, custom: Vc<EnvMap>) -> Result<Vc<Self>> {
+        let prior = prior.to_resolved().await?;
+        let custom = custom.to_resolved().await?;
+        Ok(CustomProcessEnv { prior, custom }.cell())
    }
}
```

It should now just be possible to accept `ResolvedVc` instead. The exposed function signature will be unchanged, still accepting `Vc` arguments, and a conversion will happen internally.

```diff
 #[turbo_tasks::value_impl]
 impl CustomProcessEnv {
     #[turbo_tasks::function]
-    pub fn new(prior: Vc<Box<dyn ProcessEnv>>, custom: Vc<EnvMap>) -> Vc<Self> {
+    pub fn new(prior: ResolvedVc<Box<dyn ProcessEnv>>, custom: ResolvedVc<EnvMap>) ->Vc<Self> {
         CustomProcessEnv { prior, custom }.cell()
    }
}
```

This should also work for arguments where `Vc` is inside of a `Vec` or `Option` (other collection types are not currently supported).

This PR does not support `self` arguments. That is handled by #70367.

## How

- The macro inspects the argument type and rewrites it to replace `ResolvedVc` with `Vc` to get the exposed function's signature.
- The `FromTaskInput` trait does the actual conversion.

### Why do this type expansion and conversion in the macro, and not as part of [the `TaskFn` trait](https://github.com/vercel/next.js/blob/8f9c6a86177513026ab4bc4fdc3575ca1efe025c/turbopack/crates/turbo-tasks/src/task/function.rs)?

Without [specialization](https://github.com/rust-lang/rfcs/blob/master/text/1210-impl-specialization.md) it's not possible to implement the `FromTaskInput` trait for all `TaskInput` types, as we'd end up with overlapping impls for `Option<T>` and `Vec<T>`.

There are specialization hacks ([inherent method specialization](dtolnay/case-studies#14), [autoref-specialization, and autoderef-specialization](http://lukaskalbertodt.github.io/2019/12/05/generalized-autoref-based-specialization.html)) but those hacks are mostly for macros, not for generic code:

> One thing might be worth clarifying up front: the adopted version described here does not solve *the* main limitation of autoref-based specialization, namely specializing in a generic context. For example, given `fn foo<T: Clone>()`, you cannot specialize for `T: Copy` in that function with autoref-based specialization. For these kinds of parametricity-destroying cases, “real” specialization is still required. As such, the whole autoref-based specialization technique is still mainly relevant for usage with macros.

So we need the macro to determine if a type implements `FromTaskInput` or `TaskInput`. We can't do this inside of generic function.

Aside from that, even though it's not as technically correct, expanding the types inside the macro results in *much* more readable types in rustdoc, which is why we do this in `expand_vc_return_type` as well, even though we could use a trait's associated type instead: vercel/turborepo#8096

## Test Plan

```
cargo nextest r -p turbo-tasks-memory test_resolved_vc
cargo nextest r -p turbo-tasks-macros-tests function
```

Modify some code to use this, and use `rust-analyzer`'s macro expansion feature (after telling RA to rebuild proc macros).
  • Loading branch information
bgw committed Sep 24, 2024
1 parent 908b419 commit aa445d5
Show file tree
Hide file tree
Showing 9 changed files with 389 additions and 74 deletions.
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
#![feature(arbitrary_self_types)]
#![allow(dead_code)]

use anyhow::Result;
use turbo_tasks::{ResolvedVc, Vc};

#[derive(Clone)]
#[turbo_tasks::value(resolved)]
struct ExampleStruct {
items: Vec<ResolvedVc<u32>>,
}

#[turbo_tasks::value_impl]
impl ExampleStruct {
#[turbo_tasks::function]
fn constructor_one(item: ResolvedVc<u32>) -> Vc<Self> {
ExampleStruct { items: vec![item] }.cell()
}

#[turbo_tasks::function]
fn constructor_vec(items: Vec<turbo_tasks::ResolvedVc<u32>>) -> Vc<Self> {
ExampleStruct { items }.cell()
}
}

#[turbo_tasks::value(resolved, transparent)]
struct MaybeExampleStruct(Option<ExampleStruct>);

#[turbo_tasks::function]
async fn caller_uses_unresolved_vc(items: Option<Vec<Vc<u32>>>) -> Result<Vc<MaybeExampleStruct>> {
if let Some(items) = items {
// call `constructor_vec` with `Vc` (not `ResolvedVc`)
let inner = ExampleStruct::constructor_vec(items).await?;
Ok(Vc::cell(Some((*inner).clone())))
} else {
Ok(Vc::cell(None))
}
}

fn main() {}
262 changes: 242 additions & 20 deletions turbopack/crates/turbo-tasks-macros/src/func.rs
Original file line number Diff line number Diff line change
@@ -1,24 +1,29 @@
use std::collections::HashSet;
use std::{borrow::Cow, collections::HashSet};

use proc_macro2::{Ident, Span, TokenStream};
use quote::{quote, quote_spanned, ToTokens};
use syn::{
parenthesized,
parse::{Parse, ParseStream},
parse_quote,
parse_quote, parse_quote_spanned,
punctuated::{Pair, Punctuated},
spanned::Spanned,
token::Paren,
AngleBracketedGenericArguments, Block, Expr, ExprPath, FnArg, GenericArgument, Meta, Pat,
PatIdent, PatType, Path, PathArguments, PathSegment, Receiver, ReturnType, Signature, Token,
Type, TypeGroup, TypePath, TypeTuple,
AngleBracketedGenericArguments, Block, Expr, ExprBlock, ExprPath, FnArg, GenericArgument,
Local, Meta, Pat, PatIdent, PatType, Path, PathArguments, PathSegment, Receiver, ReturnType,
Signature, Stmt, Token, Type, TypeGroup, TypePath, TypeTuple,
};

#[derive(Debug)]
pub struct TurboFn {
// signature: Signature,
// block: Block,
ident: Ident,
pub struct TurboFn<'a> {
orig_signature: &'a Signature,

/// Identifier of the exposed function (same as the original function's name).
ident: &'a Ident,

/// Identifier of the inline function (a mangled version of the original function's name).
inline_ident: Ident,

output: Type,
this: Option<Input>,
inputs: Vec<Input>,
Expand All @@ -34,14 +39,14 @@ pub struct Input {
pub ty: Type,
}

impl TurboFn {
impl TurboFn<'_> {
pub fn new(
original_signature: &Signature,
orig_signature: &Signature,
definition_context: DefinitionContext,
args: FunctionArguments,
) -> Option<TurboFn> {
if !original_signature.generics.params.is_empty() {
original_signature
if !orig_signature.generics.params.is_empty() {
orig_signature
.generics
.span()
.unwrap()
Expand All @@ -53,8 +58,8 @@ impl TurboFn {
return None;
}

if original_signature.generics.where_clause.is_some() {
original_signature
if orig_signature.generics.where_clause.is_some() {
orig_signature
.generics
.where_clause
.span()
Expand All @@ -67,7 +72,7 @@ impl TurboFn {
return None;
}

let mut raw_inputs = original_signature.inputs.iter();
let mut raw_inputs = orig_signature.inputs.iter();
let mut this = None;
let mut inputs = Vec::with_capacity(raw_inputs.len());

Expand Down Expand Up @@ -251,15 +256,25 @@ impl TurboFn {
}
}

let output = return_type_to_type(&original_signature.output);
let output = return_type_to_type(&orig_signature.output);

let orig_ident = &orig_signature.ident;
let inline_ident = Ident::new(
// Hygiene: This should use `.resolved_at(Span::def_site())`, but that's unstable, so
// instead we just pick a long, unique name
&format!("{orig_ident}_turbo_tasks_function_inline"),
orig_ident.span(),
);

Some(TurboFn {
ident: original_signature.ident.clone(),
orig_signature,
ident: orig_ident,
output,
this,
inputs,
resolved: args.resolved,
local_cells: args.local_cells.is_some(),
inline_ident,
})
}

Expand All @@ -273,8 +288,7 @@ impl TurboFn {
.chain(self.inputs.iter())
.map(|input| {
FnArg::Typed(PatType {
attrs: Default::default(),
ty: Box::new(input.ty.clone()),
attrs: Vec::new(),
pat: Box::new(Pat::Ident(PatIdent {
attrs: Default::default(),
by_ref: None,
Expand All @@ -283,6 +297,7 @@ impl TurboFn {
subpat: None,
})),
colon_token: Default::default(),
ty: Box::new(expand_task_input_type(&input.ty).into_owned()),
})
})
.collect();
Expand All @@ -304,6 +319,107 @@ impl TurboFn {
}
}

/// Signature for the "inline" function. The inline function is the function with minimal
/// changes that's called by the turbo-tasks framework during scheduling.
///
/// This is in contrast to the "exposed" function, which is the public function that the user
/// should call.
///
/// This function signature should match the name given by [`Self::inline_ident`].
///
/// A minimally wrapped version of the original function block.
pub fn inline_signature_and_block<'a>(
&self,
orig_block: &'a Block,
) -> (Signature, Cow<'a, Block>) {
let (inputs, transform_stmts): (Punctuated<_, _>, Vec<Option<_>>) = self
.orig_signature
.inputs
.iter()
.enumerate()
.map(|(idx, arg)| match arg {
FnArg::Receiver(_) => (arg.clone(), None),
FnArg::Typed(pat_type) => {
// arbitrary self types aren't `FnArg::Receiver` on syn 1.x (fixed in 2.x)
if let Pat::Ident(pat_id) = &*pat_type.pat {
// TODO: Support `self: ResolvedVc<Self>`
if pat_id.ident == "self" {
return (arg.clone(), None);
}
}
let Cow::Owned(expanded_ty) = expand_task_input_type(&pat_type.ty) else {
// common-case: skip if no type conversion is needed
return (arg.clone(), None);
};
let arg_ident = Ident::new(
&format!("arg{idx}"),
pat_type.span().resolved_at(Span::mixed_site()),
);
let arg = FnArg::Typed(PatType {
pat: Box::new(Pat::Ident(PatIdent {
attrs: Vec::new(),
by_ref: None,
mutability: None,
ident: arg_ident.clone(),
subpat: None,
})),
ty: Box::new(expanded_ty),
..pat_type.clone()
});
// convert an argument of type `FromTaskInput<T>::TaskInput` into `T`.
// essentially, replace any instances of `Vc` with `ResolvedVc`.
let orig_ty = &*pat_type.ty;
let transform_stmt = Some(Stmt::Local(Local {
attrs: Vec::new(),
let_token: Default::default(),
pat: *pat_type.pat.clone(),
init: Some((
Default::default(),
// we know the argument implements `FromTaskInput` because
// `expand_task_input_type` returned `Cow::Owned`
parse_quote_spanned! {
pat_type.span() =>
<#orig_ty as turbo_tasks::task::FromTaskInput>::from_task_input(
#arg_ident
)
},
)),
semi_token: Default::default(),
}));
(arg, transform_stmt)
}
})
.unzip();
let transform_stmts: Vec<Stmt> = transform_stmts.into_iter().flatten().collect();

let inline_signature = Signature {
ident: self.inline_ident.clone(),
inputs,
..self.orig_signature.clone()
};

let inline_block = if transform_stmts.is_empty() {
Cow::Borrowed(orig_block)
} else {
let mut stmts = transform_stmts;
stmts.push(Stmt::Expr(Expr::Block(ExprBlock {
attrs: Vec::new(),
label: None,
block: orig_block.clone(),
})));
Cow::Owned(Block {
brace_token: Default::default(),
stmts,
})
};

(inline_signature, inline_block)
}

pub fn inline_ident(&self) -> &Ident {
&self.inline_ident
}

fn input_idents(&self) -> impl Iterator<Item = &Ident> {
self.inputs.iter().map(|Input { ident, .. }| ident)
}
Expand Down Expand Up @@ -560,6 +676,112 @@ fn return_type_to_type(return_type: &ReturnType) -> Type {
}
}

/// Approximates the conversion of type `T` to `<T as FromTaskInput>::TaskInput` (in combination
/// with the `AutoFromTaskInput` specialization hack).
///
/// This expansion happens manually here for a couple reasons:
/// - While it's possible to simulate specialization of methods (with inherent impls, autoref, or
/// autoderef) there's currently no way to simulate specialization of type aliases on stable rust.
/// - Replacing arguments with types like `<T as FromTaskInput>::TaskInput` would make function
/// signatures illegible in the resulting rustdocs.
///
/// Returns `Cow::Owned` when a transformation was applied, and `Cow::Borrowed` when no change was
/// made to the input type.
fn expand_task_input_type(orig_input: &Type) -> Cow<'_, Type> {
match orig_input {
Type::Group(TypeGroup { elem, .. }) => expand_task_input_type(elem),
Type::Path(TypePath {
qself: None,
path: Path {
leading_colon,
segments,
},
}) => {
enum PathMatch {
Empty,
StdMod,
VecMod,
Vec,
OptionMod,
Option,
TurboTasksMod,
ResolvedVc,
}

let mut path_match = PathMatch::Empty;
let has_leading_colon = leading_colon.is_some();
for segment in segments {
path_match = match (has_leading_colon, path_match, &segment.ident) {
(_, PathMatch::Empty, id) if id == "std" || id == "core" || id == "alloc" => {
PathMatch::StdMod
}

(_, PathMatch::StdMod, id) if id == "vec" => PathMatch::VecMod,
(false, PathMatch::Empty | PathMatch::VecMod, id) if id == "Vec" => {
PathMatch::Vec
}

(_, PathMatch::StdMod, id) if id == "option" => PathMatch::OptionMod,
(false, PathMatch::Empty | PathMatch::OptionMod, id) if id == "Option" => {
PathMatch::Option
}

(_, PathMatch::Empty, id) if id == "turbo_tasks" => PathMatch::TurboTasksMod,
(false, PathMatch::Empty | PathMatch::TurboTasksMod, id)
if id == "ResolvedVc" =>
{
PathMatch::ResolvedVc
}

// some type we don't have an expansion for
_ => return Cow::Borrowed(orig_input),
}
}

let last_segment = segments.last().expect("non-empty");
let mut segments = Cow::Borrowed(segments);
match path_match {
PathMatch::Vec | PathMatch::Option => {
if let PathArguments::AngleBracketed(
bracketed_args @ AngleBracketedGenericArguments { args, .. },
) = &last_segment.arguments
{
if let Some(GenericArgument::Type(first_arg)) = args.first() {
match expand_task_input_type(first_arg) {
Cow::Borrowed(_) => {} // was not transformed
Cow::Owned(first_arg) => {
let mut bracketed_args = bracketed_args.clone();
*bracketed_args.args.first_mut().expect("non-empty") =
GenericArgument::Type(first_arg);
segments.to_mut().last_mut().expect("non-empty").arguments =
PathArguments::AngleBracketed(bracketed_args);
}
}
}
}
}
PathMatch::ResolvedVc => {
let args = &last_segment.arguments;
segments =
Cow::Owned(parse_quote_spanned!(segments.span() => turbo_tasks::Vc #args));
}
_ => {}
}
match segments {
Cow::Borrowed(_) => Cow::Borrowed(orig_input),
Cow::Owned(segments) => Cow::Owned(Type::Path(TypePath {
qself: None,
path: Path {
leading_colon: *leading_colon,
segments,
},
})),
}
}
_ => Cow::Borrowed(orig_input),
}
}

fn expand_vc_return_type(orig_output: &Type) -> Type {
// HACK: Approximate the expansion that we'd otherwise get from
// `<T as TaskOutput>::Return`, so that the return type shown in the rustdocs
Expand Down
Loading

0 comments on commit aa445d5

Please sign in to comment.