dioxus-mirror/packages/router-macro/src/lib.rs

extern crate proc_macro;

use layout::Layout;
use nest::{Nest, NestId};
use proc_macro::TokenStream;
use quote::{__private::Span, format_ident, quote, ToTokens};
use redirect::Redirect;
use route::{Route, RouteType};
use segment::RouteSegment;
use syn::{parse::ParseStream, parse_macro_input, Ident, Token, Type};

use proc_macro2::TokenStream as TokenStream2;

use crate::{layout::LayoutId, route_tree::RouteTree};

mod layout;
mod nest;
mod query;
mod redirect;
mod route;
mod route_tree;
mod segment;

/// Derives the Routable trait for an enum of routes
///
/// Each variant must:
/// 1. Be struct-like with {}'s
/// 2. Contain all of the dynamic parameters of the current and nested routes
/// 3. Have a #[route("route")] attribute
///
/// Route Segments:
/// 1. Static Segments: "/static"
/// 2. Dynamic Segments: "/:dynamic" (where dynamic has a type that is FromStr in all child Variants)
/// 3. Catch all Segments: "/:...segments" (where segments has a type that is FromSegments in all child Variants)
/// 4. Query Segments: "/?:query" (where query has a type that is FromQuery in all child Variants)
///
/// Routes are matched:
/// 1. By there specificity this order: Query Routes ("/?:query"), Static Routes ("/route"), Dynamic Routes ("/:route"), Catch All Routes ("/:...route")
/// 2. By the order they are defined in the enum
///
/// All features:
/// ```rust, skip
/// #[rustfmt::skip]
/// #[derive(Clone, Debug, PartialEq, Routable)]
/// enum Route {
///     // Define routes with the route macro. If the name of the component is not the same as the variant, you can specify it as the second parameter and the props type as the third
///     #[route("/", IndexComponent, ComponentProps)]
///     Index {},
///     // Nests with parameters have types taken from child routes
///     // Everything inside the nest has the added parameter `user_id: usize`
///     #[nest("/user/:user_id")]
///         // All children of layouts will be rendered inside the Outlet in the layout component
///         // Creates a Layout UserFrame that has the parameter `user_id: usize`
///         #[layout(UserFrame)]
///             // If there is a component with the name Route1 and props with the name Route1Props, you do not need to pass in the component and type
///             #[route("/:dynamic?:query")]
///             Route1 {
///                 // The type is taken from the first instance of the dynamic parameter
///                 user_id: usize,
///                 dynamic: usize,
///                 query: String,
///                 extra: String,
///             },
///             #[route("/hello_world")]
///             // You can opt out of the layout by using the `!` prefix
///             #[layout(!UserFrame)]
///             Route2 { user_id: usize },
///         // End layouts with #[end_layout]
///         #[end_layout]
///     // End nests with #[end_nest]
///     #[end_nest]
///     // Redirects take a path and a function that takes the parameters from the path and returns a new route
///     #[redirect("/:id/user", |id: usize| Route::Route3 { dynamic: id.to_string()})]
///     #[route("/:dynamic")]
///     Route3 { dynamic: String },
///     #[child]
///     NestedRoute(NestedRoute),
/// }
/// ```
///
/// # `#[route("path", component, props)]`
///
/// The `#[route]` attribute is used to define a route. It takes up to 3 parameters:
/// - `path`: The path to the enum variant (relative to the parent nest)
/// - (optional) `component`: The component to render when the route is matched. If not specified, the name of the variant is used
/// - (optional) `props`: The props type for the component. If not specified, the name of the variant with `Props` appended is used
///
/// Routes are the most basic attribute. They allow you to define a route and the component to render when the route is matched. The component must take all dynamic parameters of the route and all parent nests.
/// The next variant will be tied to the component. If you link to that variant, the component will be rendered.
///
/// ```rust, skip
/// #[derive(Clone, Debug, PartialEq, Routable)]
/// enum Route {
///     // Define routes that renders the IndexComponent that takes the IndexProps
///     // The Index component will be rendered when the route is matched (e.g. when the user navigates to /)
///     #[route("/", Index, IndexProps)]
///     Index {},
/// }
/// ```
///
/// # `#[redirect("path", function)]`
///
/// The `#[redirect]` attribute is used to define a redirect. It takes 2 parameters:
/// - `path`: The path to the enum variant (relative to the parent nest)
/// - `function`: A function that takes the parameters from the path and returns a new route
///
/// ```rust, skip
/// #[derive(Clone, Debug, PartialEq, Routable)]
/// enum Route {
///     // Redirects the /:id route to the Index route
///     #[redirect("/:id", |_: usize| Route::Index {})]
///     #[route("/", Index, IndexProps)]
///     Index {},
/// }
/// ```
///
/// Redirects allow you to redirect a route to another route. The function must take all dynamic parameters of the route and all parent nests.
///
/// # `#[nest("path")]`
///
/// The `#[nest]` attribute is used to define a nest. It takes 1 parameter:
/// - `path`: The path to the nest (relative to the parent nest)
///
/// Nests effect all nests, routes and redirects defined until the next `#[end_nest]` attribute. All children of nests are relative to the nest route and must include all dynamic parameters of the nest.
///
/// ```rust, skip
/// #[derive(Clone, Debug, PartialEq, Routable)]
/// enum Route {
///     // Nests all child routes in the /blog route
///     #[nest("/blog")]
///         // This is at /blog/:id
///         #[redirect("/:id", |_: usize| Route::Index {})]
///         // This is at /blog
///         #[route("/", Index, IndexProps)]
///         Index {},
/// }
/// ```
///
/// # `#[end_nest]`
///
/// The `#[end_nest]` attribute is used to end a nest. It takes no parameters.
///
/// ```rust, skip
/// #[derive(Clone, Debug, PartialEq, Routable)]
/// enum Route {
///     #[nest("/blog")]
///         // This is at /blog/:id
///         #[redirect("/:id", |_: usize| Route::Index {})]
///         // This is at /blog
///         #[route("/", Index, IndexProps)]
///         Index {},
///     // Ends the nest
///     #[end_nest]
///     // This is at /
///     #[route("/")]
///     Home {},
/// }
/// ```
///
/// # `#[layout(component)]`
///
/// The `#[layout]` attribute is used to define a layout. It takes 2 parameters:
/// - `component`: The component to render when the route is matched. If not specified, the name of the variant is used
/// - (optional) `props`: The props type for the component. If not specified, the name of the variant with `Props` appended is used
///
/// The layout component allows you to wrap all children of the layout in a component. The child routes are rendered in the Outlet of the layout component. The layout component must take all dynamic parameters of the nests it is nested in.
///
/// ```rust, skip
/// #[derive(Clone, Debug, PartialEq, Routable)]
/// enum Route {
///     #[layout(BlogFrame)]
///         #[redirect("/:id", |_: usize| Route::Index {})]
///         // Index will be rendered in the Outlet of the BlogFrame component
///         #[route("/", Index, IndexProps)]
///         Index {},
/// }
/// ```
///
/// # `#[end_layout]`
///
/// The `#[end_layout]` attribute is used to end a layout. It takes no parameters.
///
/// ```rust, skip
/// #[derive(Clone, Debug, PartialEq, Routable)]
/// enum Route {
///     #[layout(BlogFrame)]
///         #[redirect("/:id", |_: usize| Route::Index {})]
///         // Index will be rendered in the Outlet of the BlogFrame component
///         #[route("/", Index, IndexProps)]
///         Index {},
///     // Ends the layout
///     #[end_layout]
///     // This will be rendered standalone
///     #[route("/")]
///     Home {},
/// }
/// ```
#[proc_macro_derive(
    Routable,
    attributes(route, nest, end_nest, layout, end_layout, redirect, child)
)]
pub fn routable(input: TokenStream) -> TokenStream {
    let routes_enum = parse_macro_input!(input as syn::ItemEnum);

    let route_enum = match RouteEnum::parse(routes_enum) {
        Ok(route_enum) => route_enum,
        Err(err) => return err.to_compile_error().into(),
    };

    let error_type = route_enum.error_type();
    let parse_impl = route_enum.parse_impl();
    let display_impl = route_enum.impl_display();
    let routable_impl = route_enum.routable_impl();
    let name = &route_enum.name;
    let vis = &route_enum.vis;

    quote! {
        #vis fn Outlet(cx: dioxus::prelude::Scope) -> dioxus::prelude::Element {
            dioxus_router::prelude::GenericOutlet::<#name>(cx)
        }

        #vis fn Router(cx: dioxus::prelude::Scope<dioxus_router::prelude::GenericRouterProps<#name>>) -> dioxus::prelude::Element {
            dioxus_router::prelude::GenericRouter(cx)
        }

        #vis fn Link<'a>(cx: dioxus::prelude::Scope<'a, dioxus_router::prelude::GenericLinkProps<'a, #name>>) -> dioxus::prelude::Element<'a> {
            dioxus_router::prelude::GenericLink(cx)
        }

        #vis fn GoBackButton<'a>(cx: dioxus::prelude::Scope<'a, dioxus_router::prelude::GenericHistoryButtonProps<'a>>) -> dioxus::prelude::Element<'a> {
            dioxus_router::prelude::GenericGoBackButton::<#name>(cx)
        }

        #vis fn GoForwardButton<'a>(cx: dioxus::prelude::Scope<'a, dioxus_router::prelude::GenericHistoryButtonProps<'a>>) -> dioxus::prelude::Element<'a> {
            dioxus_router::prelude::GenericGoForwardButton::<#name>(cx)
        }

        #vis fn use_route(cx: &dioxus::prelude::ScopeState) -> Option<#name> {
            dioxus_router::prelude::use_generic_route(cx)
        }

        #vis fn use_navigator(cx: &dioxus::prelude::ScopeState) -> &dioxus_router::prelude::GenericNavigator<#name> {
            dioxus_router::prelude::use_generic_navigator(cx)
        }

        #error_type

        #display_impl

        #routable_impl

        #parse_impl
    }
    .into()
}

struct RouteEnum {
    vis: syn::Visibility,
    name: Ident,
    redirects: Vec<Redirect>,
    routes: Vec<Route>,
    nests: Vec<Nest>,
    layouts: Vec<Layout>,
    site_map: Vec<SiteMapSegment>,
}

impl RouteEnum {
    fn parse(data: syn::ItemEnum) -> syn::Result<Self> {
        let name = &data.ident;
        let vis = &data.vis;

        let mut site_map = Vec::new();
        let mut site_map_stack: Vec<Vec<SiteMapSegment>> = Vec::new();

        let mut routes = Vec::new();

        let mut redirects = Vec::new();

        let mut layouts: Vec<Layout> = Vec::new();
        let mut layout_stack = Vec::new();

        let mut nests = Vec::new();
        let mut nest_stack = Vec::new();

        for variant in &data.variants {
            let mut excluded = Vec::new();
            // Apply the any nesting attributes in order
            for attr in &variant.attrs {
                if attr.path.is_ident("nest") {
                    let mut children_routes = Vec::new();
                    {
                        // add all of the variants of the enum to the children_routes until we hit an end_nest
                        let mut level = 0;
                        'o: for variant in &data.variants {
                            children_routes.push(variant.fields.clone());
                            for attr in &variant.attrs {
                                if attr.path.is_ident("nest") {
                                    level += 1;
                                } else if attr.path.is_ident("end_nest") {
                                    level -= 1;
                                    if level < 0 {
                                        break 'o;
                                    }
                                }
                            }
                        }
                    }

                    let nest_index = nests.len();

                    let parser = |input: ParseStream| {
                        Nest::parse(
                            input,
                            children_routes
                                .iter()
                                .filter_map(|f: &syn::Fields| match f {
                                    syn::Fields::Named(fields) => Some(fields.clone()),
                                    _ => None,
                                })
                                .collect(),
                            nest_index,
                        )
                    };
                    let nest = attr.parse_args_with(parser)?;

                    // add the current segment to the site map stack
                    let segments: Vec<_> = nest
                        .segments
                        .iter()
                        .map(|seg| {
                            let segment_type = seg.into();
                            SiteMapSegment {
                                segment_type,
                                children: Vec::new(),
                            }
                        })
                        .collect();
                    if !segments.is_empty() {
                        site_map_stack.push(segments);
                    }

                    nests.push(nest);
                    nest_stack.push(NestId(nest_index));
                } else if attr.path.is_ident("end_nest") {
                    nest_stack.pop();
                    // pop the current nest segment off the stack and add it to the parent or the site map
                    if let Some(segment) = site_map_stack.pop() {
                        let children = site_map_stack
                            .last_mut()
                            .map(|seg| &mut seg.last_mut().unwrap().children)
                            .unwrap_or(&mut site_map);

                        // Turn the list of segments in the segments stack into a tree
                        let mut iter = segment.into_iter().rev();
                        let mut current = iter.next().unwrap();
                        for mut segment in iter {
                            segment.children.push(current);
                            current = segment;
                        }

                        children.push(current);
                    }
                } else if attr.path.is_ident("layout") {
                    let parser = |input: ParseStream| {
                        let bang: Option<Token![!]> = input.parse().ok();
                        let exclude = bang.is_some();
                        Ok((
                            exclude,
                            Layout::parse(input, nest_stack.iter().rev().cloned().collect())?,
                        ))
                    };
                    let (exclude, layout): (bool, Layout) = attr.parse_args_with(parser)?;

                    if exclude {
                        let Some(layout_index) =
                            layouts.iter().position(|l| l.comp == layout.comp) else {
                                return Err(syn::Error::new(
                                    Span::call_site(),
                                    "Attempted to exclude a layout that does not exist",
                                ));
                            };
                        excluded.push(LayoutId(layout_index));
                    } else {
                        let layout_index = layouts.len();
                        layouts.push(layout);
                        layout_stack.push(LayoutId(layout_index));
                    }
                } else if attr.path.is_ident("end_layout") {
                    layout_stack.pop();
                } else if attr.path.is_ident("redirect") {
                    let parser = |input: ParseStream| {
                        Redirect::parse(
                            input,
                            nest_stack.iter().rev().cloned().collect(),
                            redirects.len(),
                        )
                    };
                    let redirect = attr.parse_args_with(parser)?;
                    redirects.push(redirect);
                }
            }

            let mut active_nests = nest_stack.clone();
            active_nests.reverse();
            let mut active_layouts = layout_stack.clone();
            active_layouts.retain(|&id| !excluded.contains(&id));

            let route = Route::parse(active_nests, active_layouts, variant.clone())?;

            // add the route to the site map
            let mut segment = SiteMapSegment::new(&route.segments);
            if let RouteType::Child(child) = &route.ty {
                let new_segment = SiteMapSegment {
                    segment_type: SegmentType::Child(child.ty.clone()),
                    children: Vec::new(),
                };
                match &mut segment {
                    Some(segment) => {
                        fn set_last_child_to(
                            segment: &mut SiteMapSegment,
                            new_segment: SiteMapSegment,
                        ) {
                            if let Some(last) = segment.children.last_mut() {
                                set_last_child_to(last, new_segment);
                            } else {
                                segment.children = vec![new_segment];
                            }
                        }
                        set_last_child_to(segment, new_segment);
                    }
                    None => {
                        segment = Some(new_segment);
                    }
                }
            }

            if let Some(segment) = segment {
                let parent = site_map_stack.last_mut();
                let children = match parent {
                    Some(parent) => &mut parent.last_mut().unwrap().children,
                    None => &mut site_map,
                };
                children.push(segment);
            }

            routes.push(route);
        }

        let myself = Self {
            vis: vis.clone(),
            name: name.clone(),
            routes,
            redirects,
            nests,
            layouts,
            site_map,
        };

        Ok(myself)
    }

    fn impl_display(&self) -> TokenStream2 {
        let mut display_match = Vec::new();

        for route in &self.routes {
            display_match.push(route.display_match(&self.nests));
        }

        let name = &self.name;

        quote! {
            impl std::fmt::Display for #name {
                fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                    #[allow(unused)]
                    match self {
                        #(#display_match)*
                    }
                    Ok(())
                }
            }
        }
    }

    fn parse_impl(&self) -> TokenStream2 {
        let tree = RouteTree::new(&self.routes, &self.nests, &self.redirects);
        let name = &self.name;

        let error_name = format_ident!("{}MatchError", self.name);
        let tokens = tree.roots.iter().map(|&id| {
            let route = tree.get(id).unwrap();
            route.to_tokens(&self.nests, &tree, self.name.clone(), error_name.clone())
        });

        quote! {
            impl<'a> core::convert::TryFrom<&'a str> for #name {
                type Error = <Self as std::str::FromStr>::Err;

                fn try_from(s: &'a str) -> Result<Self, Self::Error> {
                    s.parse()
                }
            }

            impl std::str::FromStr for #name {
                type Err = dioxus_router::routable::RouteParseError<#error_name>;

                fn from_str(s: &str) -> Result<Self, Self::Err> {
                    let route = s;
                    let (route, query) = route.split_once('?').unwrap_or((route, ""));
                    let mut segments = route.split('/');
                    // skip the first empty segment
                    if s.starts_with('/') {
                        segments.next();
                    }
                    let mut errors = Vec::new();

                    #(#tokens)*

                    Err(dioxus_router::routable::RouteParseError {
                        attempted_routes: errors,
                    })
                }
            }
        }
    }

    fn error_name(&self) -> Ident {
        Ident::new(&(self.name.to_string() + "MatchError"), Span::call_site())
    }

    fn error_type(&self) -> TokenStream2 {
        let match_error_name = self.error_name();

        let mut type_defs = Vec::new();
        let mut error_variants = Vec::new();
        let mut display_match = Vec::new();

        for route in &self.routes {
            let route_name = &route.route_name;

            let error_name = route.error_ident();
            let route_str = &route.route;

            error_variants.push(quote! { #route_name(#error_name) });
            display_match.push(quote! { Self::#route_name(err) => write!(f, "Route '{}' ('{}') did not match:\n{}", stringify!(#route_name), #route_str, err)? });
            type_defs.push(route.error_type());
        }

        for nest in &self.nests {
            let error_variant = nest.error_variant();
            let error_name = nest.error_ident();
            let route_str = &nest.route;

            error_variants.push(quote! { #error_variant(#error_name) });
            display_match.push(quote! { Self::#error_variant(err) => write!(f, "Nest '{}' ('{}') did not match:\n{}", stringify!(#error_name), #route_str, err)? });
            type_defs.push(nest.error_type());
        }

        for redirect in &self.redirects {
            let error_variant = redirect.error_variant();
            let error_name = redirect.error_ident();
            let route_str = &redirect.route;

            error_variants.push(quote! { #error_variant(#error_name) });
            display_match.push(quote! { Self::#error_variant(err) => write!(f, "Redirect '{}' ('{}') did not match:\n{}", stringify!(#error_name), #route_str, err)? });
            type_defs.push(redirect.error_type());
        }

        quote! {
            #(#type_defs)*

            #[allow(non_camel_case_types)]
            #[derive(Debug, PartialEq)]
            pub enum #match_error_name {
                #(#error_variants),*
            }

            impl std::fmt::Display for #match_error_name {
                fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
                    match self {
                        #(#display_match),*
                    }
                    Ok(())
                }
            }
        }
    }

    fn routable_impl(&self) -> TokenStream2 {
        let name = &self.name;
        let site_map = &self.site_map;

        let mut matches = Vec::new();

        // Collect all routes matches
        for route in &self.routes {
            matches.push(route.routable_match(&self.layouts, &self.nests));
        }

        quote! {
            impl dioxus_router::routable::Routable for #name where Self: Clone {
                const SITE_MAP: &'static [dioxus_router::routable::SiteMapSegment] = &[
                    #(#site_map,)*
                ];

                fn render<'a>(&self, cx: &'a dioxus::prelude::ScopeState, level: usize) -> dioxus::prelude::Element<'a> {
                    let myself = self.clone();
                    match (level, myself) {
                        #(#matches)*
                        _ => None
                    }
                }
            }
        }
    }
}

struct SiteMapSegment {
    pub segment_type: SegmentType,
    pub children: Vec<SiteMapSegment>,
}

impl SiteMapSegment {
    fn new(segments: &[RouteSegment]) -> Option<Self> {
        let mut current = None;
        // walk backwards through the new segments, adding children as we go
        for segment in segments.iter().rev() {
            let segment_type = segment.into();
            let mut segment = SiteMapSegment {
                segment_type,
                children: Vec::new(),
            };
            // if we have a current segment, add it as a child
            if let Some(current) = current.take() {
                segment.children.push(current)
            }
            current = Some(segment);
        }
        current
    }
}

impl ToTokens for SiteMapSegment {
    fn to_tokens(&self, tokens: &mut TokenStream2) {
        let segment_type = &self.segment_type;
        let children = if let SegmentType::Child(ty) = &self.segment_type {
            quote! { #ty::SITE_MAP }
        } else {
            let children = self
                .children
                .iter()
                .map(|child| child.to_token_stream())
                .collect::<Vec<_>>();
            quote! {
                &[
                    #(#children,)*
                ]
            }
        };

        tokens.extend(quote! {
            dioxus_router::routable::SiteMapSegment {
                segment_type: #segment_type,
                children: #children,
            }
        });
    }
}

enum SegmentType {
    Static(String),
    Dynamic(String),
    CatchAll(String),
    Child(Type),
}

impl ToTokens for SegmentType {
    fn to_tokens(&self, tokens: &mut TokenStream2) {
        match self {
            SegmentType::Static(s) => {
                tokens.extend(quote! { dioxus_router::routable::SegmentType::Static(#s) })
            }
            SegmentType::Dynamic(s) => {
                tokens.extend(quote! { dioxus_router::routable::SegmentType::Dynamic(#s) })
            }
            SegmentType::CatchAll(s) => {
                tokens.extend(quote! { dioxus_router::routable::SegmentType::CatchAll(#s) })
            }
            SegmentType::Child(_) => {
                tokens.extend(quote! { dioxus_router::routable::SegmentType::Child })
            }
        }
    }
}

impl<'a> From<&'a RouteSegment> for SegmentType {
    fn from(value: &'a RouteSegment) -> Self {
        match value {
            segment::RouteSegment::Static(s) => SegmentType::Static(s.to_string()),
            segment::RouteSegment::Dynamic(s, _) => SegmentType::Dynamic(s.to_string()),
            segment::RouteSegment::CatchAll(s, _) => SegmentType::CatchAll(s.to_string()),
        }
    }
}