Function Components & Hooks V2 (#2401)

* Make a use_hook hook with the new Hook trait.

* Implement Lifetime.

* Rewrites function signature.

* Only apply lifetime if there're other lifetimes.

* Cleanup signature rewrite logic.

* Rewrite hook body.

* Port some built-in hooks.

* Finish porting all built-in hooks.

* Port tests.

* Fix tests.

* Migrate to macro-based hooks.

* Fix HookContext, add tests on non-possible locations.

* Fix stderr for trybuild.

* Add 1 more test case.

* Adjust doc location.

* Pretty print hook signature.

* Fix Items & std::ops::Fn*.

* Add use_memo.

* Optimise Implementation of hooks.

* Use Box to capture function value only.

* Detect whether needs boxing.

* Add args if boxing not needed.

* Enforce hook number.

* Deduplicate use_effect.

* Optimise Implementation.

* Update documentation.

* Fix website test. Strip BoxedHook implementation from it.

* Allow doc string.

* Workaround doc tests.

* Optimise codebase & documentation.

* Fix website test.

* Reduce implementation complexity.

* Destructor is no more.

* Documentation and macros.

* Reduce heap allocation and hook complexity.

* Remove Queue as well.

* Prefer Generics.

* Fix typo.

* Remove more allocations.

* Add comments.

* Remove outdated comment.

* Bare Function Pointer for better code size.

examples/function_todomvc/src/hooks/use_bool_toggle.rs

@@ -1,6 +1,6 @@
 use std::ops::Deref;
 use std::rc::Rc;
-use yew::{use_state_eq, UseStateHandle};
+use yew::prelude::*;
 
 #[derive(Clone)]
 pub struct UseBoolToggleHandle {
@@ -47,6 +47,7 @@ impl Deref for UseBoolToggleHandle {
 /// <button {onclick}>{ "Click me" }</button>
 /// ...
 /// ```
+#[hook]
 pub fn use_bool_toggle(default: bool) -> UseBoolToggleHandle {
     let state = use_state_eq(|| default);
 

examples/simple_ssr/src/main.rs

@@ -57,6 +57,7 @@ impl PartialEq for UuidState {
     }
 }
 
+#[hook]
 fn use_random_uuid() -> SuspensionResult<Uuid> {
     let s = use_state(UuidState::new);
 

examples/suspense/src/use_sleep.rs

@@ -28,6 +28,7 @@ impl Reducible for SleepState {
     }
 }
 
+#[hook]
 pub fn use_sleep() -> SuspensionResult<Rc<dyn Fn()>> {
     let sleep_state = use_reducer(SleepState::new);
 

packages/yew-agent/src/hooks.rs

@@ -29,6 +29,7 @@ where
 ///
 /// Takes a callback as the only argument. The callback will be updated on every render to make
 /// sure captured values (if any) are up to date.
+#[hook]
 pub fn use_bridge<T, F>(on_output: F) -> UseBridgeHandle<T>
 where
     T: Bridged,

packages/yew-macro/Cargo.toml

@@ -21,7 +21,9 @@ lazy_static = "1"
 proc-macro-error = "1"
 proc-macro2 = "1"
 quote = "1"
-syn = { version = "1", features = ["full", "extra-traits"] }
+syn = { version = "1", features = ["full", "extra-traits", "visit-mut"] }
+once_cell = "1"
+prettyplease = "0.1.1"
 
 # testing
 [dev-dependencies]

packages/yew-macro/src/function_component.rs

@@ -3,7 +3,11 @@ use quote::{format_ident, quote, ToTokens};
 use syn::parse::{Parse, ParseStream};
 use syn::punctuated::Punctuated;
 use syn::token::{Comma, Fn};
-use syn::{Attribute, Block, FnArg, Generics, Ident, Item, ItemFn, ReturnType, Type, Visibility};
+use syn::{
+    visit_mut, Attribute, Block, FnArg, Generics, Ident, Item, ItemFn, ReturnType, Type, Visibility,
+};
+
+use crate::hook::BodyRewriter;
 
 #[derive(Clone)]
 pub struct FunctionComponent {
@@ -169,7 +173,7 @@ fn print_fn(func_comp: FunctionComponent, use_fn_name: bool) -> TokenStream {
         fn_token,
         name,
         attrs,
-        block,
+        mut block,
         return_type,
         generics,
         arg,
@@ -184,9 +188,14 @@ fn print_fn(func_comp: FunctionComponent, use_fn_name: bool) -> TokenStream {
         Ident::new("inner", Span::mixed_site())
     };
 
+    let ctx_ident = Ident::new("ctx", Span::mixed_site());
+
+    let mut body_rewriter = BodyRewriter::default();
+    visit_mut::visit_block_mut(&mut body_rewriter, &mut *block);
+
     quote! {
         #(#attrs)*
-        #fn_token #name #ty_generics (#arg) -> #return_type
+        #fn_token #name #ty_generics (#ctx_ident: &mut ::yew::functional::HookContext, #arg) -> #return_type
         #where_clause
         {
             #block
@@ -241,6 +250,8 @@ pub fn function_component_impl(
         Ident::new("inner", Span::mixed_site())
     };
 
+    let ctx_ident = Ident::new("ctx", Span::mixed_site());
+
     let quoted = quote! {
         #[doc(hidden)]
         #[allow(non_camel_case_types)]
@@ -253,10 +264,10 @@ pub fn function_component_impl(
         impl #impl_generics ::yew::functional::FunctionProvider for #provider_name #ty_generics #where_clause {
             type TProps = #props_type;
 
-            fn run(#provider_props: &Self::TProps) -> ::yew::html::HtmlResult {
+            fn run(#ctx_ident: &mut ::yew::functional::HookContext, #provider_props: &Self::TProps) -> ::yew::html::HtmlResult {
                 #func
 
-                ::yew::html::IntoHtmlResult::into_html_result(#fn_name #fn_generics (#provider_props))
+                ::yew::html::IntoHtmlResult::into_html_result(#fn_name #fn_generics (#ctx_ident, #provider_props))
             }
         }
 

packages/yew-macro/src/hook/body.rs

@@ -0,0 +1,123 @@
+use proc_macro2::Span;
+use proc_macro_error::emit_error;
+use std::sync::{Arc, Mutex};
+use syn::spanned::Spanned;
+use syn::visit_mut::VisitMut;
+use syn::{
+    parse_quote_spanned, visit_mut, Expr, ExprCall, ExprClosure, ExprForLoop, ExprIf, ExprLoop,
+    ExprMatch, ExprWhile, Ident, Item,
+};
+
+#[derive(Debug, Default)]
+pub struct BodyRewriter {
+    branch_lock: Arc<Mutex<()>>,
+}
+
+impl BodyRewriter {
+    fn is_branched(&self) -> bool {
+        self.branch_lock.try_lock().is_err()
+    }
+
+    fn with_branch<F, O>(&mut self, f: F) -> O
+    where
+        F: FnOnce(&mut BodyRewriter) -> O,
+    {
+        let branch_lock = self.branch_lock.clone();
+        let _branched = branch_lock.try_lock();
+        f(self)
+    }
+}
+
+impl VisitMut for BodyRewriter {
+    fn visit_expr_call_mut(&mut self, i: &mut ExprCall) {
+        let ctx_ident = Ident::new("ctx", Span::mixed_site());
+
+        // Only rewrite hook calls.
+        if let Expr::Path(ref m) = &*i.func {
+            if let Some(m) = m.path.segments.last().as_ref().map(|m| &m.ident) {
+                if m.to_string().starts_with("use_") {
+                    if self.is_branched() {
+                        emit_error!(
+                            m,
+                            "hooks cannot be called at this position.";
+                            help = "move hooks to the top-level of your function.";
+                            note = "see: https://yew.rs/docs/next/concepts/function-components/introduction#hooks"
+                        );
+                    } else {
+                        *i = parse_quote_spanned! { i.span() => ::yew::functional::Hook::run(#i, #ctx_ident) };
+                    }
+
+                    return;
+                }
+            }
+        }
+
+        visit_mut::visit_expr_call_mut(self, i);
+    }
+
+    fn visit_expr_closure_mut(&mut self, i: &mut ExprClosure) {
+        self.with_branch(move |m| visit_mut::visit_expr_closure_mut(m, i))
+    }
+
+    fn visit_expr_if_mut(&mut self, i: &mut ExprIf) {
+        for it in &mut i.attrs {
+            visit_mut::visit_attribute_mut(self, it);
+        }
+
+        visit_mut::visit_expr_mut(self, &mut *i.cond);
+
+        self.with_branch(|m| visit_mut::visit_block_mut(m, &mut i.then_branch));
+
+        if let Some(it) = &mut i.else_branch {
+            self.with_branch(|m| visit_mut::visit_expr_mut(m, &mut *(it).1));
+        }
+    }
+
+    fn visit_expr_loop_mut(&mut self, i: &mut ExprLoop) {
+        self.with_branch(|m| visit_mut::visit_expr_loop_mut(m, i));
+    }
+
+    fn visit_expr_for_loop_mut(&mut self, i: &mut ExprForLoop) {
+        for it in &mut i.attrs {
+            visit_mut::visit_attribute_mut(self, it);
+        }
+        if let Some(it) = &mut i.label {
+            visit_mut::visit_label_mut(self, it);
+        }
+        visit_mut::visit_pat_mut(self, &mut i.pat);
+        visit_mut::visit_expr_mut(self, &mut *i.expr);
+
+        self.with_branch(|m| visit_mut::visit_block_mut(m, &mut i.body));
+    }
+
+    fn visit_expr_match_mut(&mut self, i: &mut ExprMatch) {
+        for it in &mut i.attrs {
+            visit_mut::visit_attribute_mut(self, it);
+        }
+
+        visit_mut::visit_expr_mut(self, &mut *i.expr);
+
+        self.with_branch(|m| {
+            for it in &mut i.arms {
+                visit_mut::visit_arm_mut(m, it);
+            }
+        });
+    }
+
+    fn visit_expr_while_mut(&mut self, i: &mut ExprWhile) {
+        for it in &mut i.attrs {
+            visit_mut::visit_attribute_mut(self, it);
+        }
+        if let Some(it) = &mut i.label {
+            visit_mut::visit_label_mut(self, it);
+        }
+
+        self.with_branch(|m| visit_mut::visit_expr_mut(m, &mut i.cond));
+        self.with_branch(|m| visit_mut::visit_block_mut(m, &mut i.body));
+    }
+
+    fn visit_item_mut(&mut self, _i: &mut Item) {
+        // We don't do anything for items.
+        // for components / hooks in other components / hooks, apply the attribute again.
+    }
+}

packages/yew-macro/src/hook/lifetime.rs

@@ -0,0 +1,121 @@
+use proc_macro2::Span;
+use std::sync::{Arc, Mutex};
+use syn::visit_mut::{self, VisitMut};
+use syn::{
+    GenericArgument, Lifetime, ParenthesizedGenericArguments, Receiver, TypeBareFn, TypeImplTrait,
+    TypeParamBound, TypeReference,
+};
+
+// borrowed from the awesome async-trait crate.
+pub struct CollectLifetimes {
+    pub elided: Vec<Lifetime>,
+    pub explicit: Vec<Lifetime>,
+    pub name: &'static str,
+    pub default_span: Span,
+
+    pub impl_trait_lock: Arc<Mutex<()>>,
+    pub impl_fn_lock: Arc<Mutex<()>>,
+}
+
+impl CollectLifetimes {
+    pub fn new(name: &'static str, default_span: Span) -> Self {
+        CollectLifetimes {
+            elided: Vec::new(),
+            explicit: Vec::new(),
+            name,
+            default_span,
+
+            impl_trait_lock: Arc::default(),
+            impl_fn_lock: Arc::default(),
+        }
+    }
+
+    fn is_impl_trait(&self) -> bool {
+        self.impl_trait_lock.try_lock().is_err()
+    }
+
+    fn is_impl_fn(&self) -> bool {
+        self.impl_fn_lock.try_lock().is_err()
+    }
+
+    fn visit_opt_lifetime(&mut self, lifetime: &mut Option<Lifetime>) {
+        match lifetime {
+            None => *lifetime = Some(self.next_lifetime(None)),
+            Some(lifetime) => self.visit_lifetime(lifetime),
+        }
+    }
+
+    fn visit_lifetime(&mut self, lifetime: &mut Lifetime) {
+        if lifetime.ident == "_" {
+            *lifetime = self.next_lifetime(lifetime.span());
+        } else {
+            self.explicit.push(lifetime.clone());
+        }
+    }
+
+    fn next_lifetime<S: Into<Option<Span>>>(&mut self, span: S) -> Lifetime {
+        let name = format!("{}{}", self.name, self.elided.len());
+        let span = span.into().unwrap_or(self.default_span);
+        let life = Lifetime::new(&name, span);
+        self.elided.push(life.clone());
+        life
+    }
+}
+
+impl VisitMut for CollectLifetimes {
+    fn visit_receiver_mut(&mut self, arg: &mut Receiver) {
+        if let Some((_, lifetime)) = &mut arg.reference {
+            self.visit_opt_lifetime(lifetime);
+        }
+    }
+
+    fn visit_type_reference_mut(&mut self, ty: &mut TypeReference) {
+        // We don't rewrite references in the impl FnOnce(&arg) or fn(&arg)
+        if self.is_impl_fn() {
+            return;
+        }
+
+        self.visit_opt_lifetime(&mut ty.lifetime);
+        visit_mut::visit_type_reference_mut(self, ty);
+    }
+
+    fn visit_generic_argument_mut(&mut self, gen: &mut GenericArgument) {
+        // We don't rewrite types in the impl FnOnce(&arg) -> Type<'_>
+        if self.is_impl_fn() {
+            return;
+        }
+
+        if let GenericArgument::Lifetime(lifetime) = gen {
+            self.visit_lifetime(lifetime);
+        }
+        visit_mut::visit_generic_argument_mut(self, gen);
+    }
+
+    fn visit_type_impl_trait_mut(&mut self, impl_trait: &mut TypeImplTrait) {
+        let impl_trait_lock = self.impl_trait_lock.clone();
+        let _locked = impl_trait_lock.try_lock();
+
+        impl_trait
+            .bounds
+            .insert(0, TypeParamBound::Lifetime(self.next_lifetime(None)));
+
+        visit_mut::visit_type_impl_trait_mut(self, impl_trait);
+    }
+
+    fn visit_parenthesized_generic_arguments_mut(
+        &mut self,
+        generic_args: &mut ParenthesizedGenericArguments,
+    ) {
+        let impl_fn_lock = self.impl_fn_lock.clone();
+        let _maybe_locked = self.is_impl_trait().then(|| impl_fn_lock.try_lock());
+
+        visit_mut::visit_parenthesized_generic_arguments_mut(self, generic_args);
+    }
+
+    fn visit_type_bare_fn_mut(&mut self, i: &mut TypeBareFn) {
+        let impl_fn_lock = self.impl_fn_lock.clone();
+        let _locked = impl_fn_lock.try_lock();
+
+        visit_mut::visit_type_bare_fn_mut(self, i);
+    }
+}