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 varient 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),
/// }
/// ```
#[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()),
        }
    }
}