riptc 0.1.2

//! This is the main core of the compiler driver that builds up an internal representation of
//! rust trees into units from axum, serde, and inertia, that we care about.
//!
//! Because we don't want to scan nor do we care about everything in a crate, and on top of that we
//! have to do some unconventional things like scan for axum route strings and match on absolute paths
//! of types, we define all of this logic and the elevated types in here. In essense, this crate receives
//! all inputs from a visitor, determines if they're a relevant item (for instance, `.route` being called
//! on an axum router), and then parses that unit (for instance, extracting the route string, the method,
//! and the handler fn) into a format suitable for codegen.
//!
//! Most of the parsing here, for simplicity sake, is designed to be very _flat_. To illustrate what I mean by this,
//! one thing that we do when we encounter a route handler is scan the signature to look for the `Json` extractor.
//! Once we do that "base" step, we don't actually process the body of the handler and instead, we just write
//! the handler with some blank state in it. Afterwards, more expressions will be fed into this that may be props,
//! which we then just write into whatever handler we have already detected.

mod analysis;
mod structures;

pub use structures::*;

use std::collections::{HashMap, HashSet};

use either::Either;
use itertools::Itertools;
use rustc_ast::LitKind;
use rustc_hir::{
    self, Expr, ExprKind, HirId,
    def_id::DefId,
    definitions::DefPath,
    intravisit::{Visitor, nested_filter::NestedFilter},
};
use rustc_middle::ty::{AdtDef, GenericArg, Ty as MTy, TyCtxt, TyKind as MTyKind};
use rustc_span::{Span, Symbol, source_map::Spanned};

use crate::namespace::NamespacedPath;

/// This is the struct that can be though of as the in-between of the Rust syntax tree, and the typescript syntax tree.
/// Here, you will find all of the derived information from the Rust code that is directly relevant to ript functionality,
/// stripping away all other code that we don't care about.
pub struct RiptAnalyzer<'tcx> {
    tcx: TyCtxt<'tcx>,

    /// this is the router that axum uses to parse its route strings.
    /// here, we create one that acts as effectively a simulation that we can feed
    /// axum route paths into to parse out the path information and find collisions
    matchit_router: parking_lot::RwLock<matchit::Router<()>>,

    axum_routes: HashMap<DefId, structures::AxumRoute<'tcx>>,
    axum_map_response_fns: HashSet<DefId>,
    bridge_marked_items: HashMap<DefId, BridgeMarkedItem<'tcx>>,
    inertia_prop_destinations: HashMap<HirId, Either<HirId, structures::InertiaPropDestination>>,
    inertia_props: HashMap<HirId, structures::InertiaProp<'tcx>>,
}

impl<'tcx> RiptAnalyzer<'tcx> {
    pub fn new(tcx: TyCtxt<'tcx>) -> Self {
        Self {
            tcx,
            matchit_router: parking_lot::RwLock::new(matchit::Router::new()),
            axum_routes: Default::default(),
            axum_map_response_fns: Default::default(),
            inertia_prop_destinations: Default::default(),
            inertia_props: Default::default(),
            bridge_marked_items: Default::default(),
        }
    }

    /// All inertia share props which should be defined globally and should be reloadable
    pub fn inertia_share_props(&self) -> impl Iterator<Item = &structures::InertiaProp<'tcx>> {
        self.inertia_props.iter().filter_map(|(id, p)| {
            let dest = self.inertia_prop_destination(*id);
            matches!(dest, InertiaPropDestination::Share).then(|| p)
        })
    }

    pub fn namespaced_path(&self, id: DefId) -> NamespacedPath {
        NamespacedPath::new(self.tcx, id)
    }

    /// All inertia flash props which should be readable globally but not reloadable
    pub fn inertia_flash_props(&self) -> impl Iterator<Item = &structures::InertiaProp<'tcx>> {
        self.inertia_props.values().filter(|p| p.flash_prop())
    }

    /// From an expression, try to extract a method call receiver, args, and span
    fn try_method_call(
        &self,
        ex: Expr<'tcx>,
    ) -> Option<(Expr<'tcx>, impl Iterator<Item = Expr<'tcx>>, Span)> {
        match ex.kind {
            ExprKind::MethodCall(_path_seg, rx, args, span) => {
                Some((*rx, args.iter().copied(), span))
            }
            _ => None,
        }
    }

    /// Given a type dependent def path, attempt to assert that it equals what you are expecting.
    /// Returns an `Option` for easy try operator chaining with pre-validation steps of an expression
    fn try_assert_type_dependent_def_path(&self, hir_id: HirId, expecting: &str) -> Option<()> {
        let dp = self
            .tcx
            .typeck(hir_id.owner)
            .type_dependent_def_id(hir_id)?;

        let dp = self.tcx.def_path(dp);

        (self.fmt_def_path(&dp) == expecting).then_some(())
    }

    fn def_path_and<T>(&self, dp: &DefPath, expecting: &str, and: T) -> Option<T> {
        (self.fmt_def_path(dp) == expecting).then_some(and)
    }

    fn get_def_path_and<T>(
        &self,
        id: DefId,
        expecting: &str,
        and: impl FnOnce() -> Option<T>,
    ) -> Option<T> {
        let dp = self.tcx.def_path(id);
        (self.def_path_and(&dp, expecting, and)?)()
    }

    /// Formats a def path into a string that makes it easy to write out and maintain path checks.
    /// For instance, a given def path of a method will turn into the string:
    /// `crate_name::mod_name::maybe_another_mod_name::{impl#<impl_block_number>}::method_name`
    fn fmt_def_path(&self, dp: &DefPath) -> String {
        let dp = format!(
            "{}{}",
            self.tcx.crate_name(dp.krate),
            dp.to_string_no_crate_verbose()
        );

        dp
    }

    pub fn bridge_marked_items(&self) -> impl Iterator<Item = &structures::BridgeMarkedItem<'tcx>> {
        self.bridge_marked_items.values()
    }

    /// Returns all axum routes that are not inertia routes, aka they do not use the inertia extractor
    pub fn axum_routes(&self) -> impl Iterator<Item = &structures::AxumRoute<'tcx>> {
        self.axum_routes.values().filter(|r| !r.inertia())
    }

    /// Returns all axum routes that are inertia routes, aka they use the inertia extractor
    pub fn inertia_axum_routes(&self) -> impl Iterator<Item = &structures::AxumRoute<'tcx>> {
        self.axum_routes.values().filter(|r| r.inertia())
    }

    pub fn axum_route_path(&self, id: DefId) -> Vec<AxumRoutePathSegment<'tcx>> {
        self.axum_routes
            .get(&id)
            .expect("expected axum route")
            .route_path(self)
            .collect_vec()
    }

    pub fn component_for_inertia_prop(&self, id: HirId) -> Option<Symbol> {
        let prop = self
            .inertia_props
            .get(&id)
            .expect("inertia prop not found in prop destinations");

        let dest = self.inertia_prop_destination(prop.id);

        match dest {
            InertiaPropDestination::Render(component) => Some(component),
            _ => None,
        }
    }

    #[track_caller]
    fn inertia_prop_destination(&self, id: HirId) -> structures::InertiaPropDestination {
        let prop = self
            .inertia_props
            .get(&id)
            .expect("inertia prop not found in prop destinations");

        let mut current = prop.rx_id;
        while let Some(Either::Left(rx_id)) = self.inertia_prop_destinations.get(&current) {
            current = *rx_id;
        }

        let Some(Either::Right(destination)) = self.inertia_prop_destinations.get(&current) else {
            prop.emit_fatal_chain_broken(self.tcx.dcx())
        };

        *destination
    }

    fn fn_params(&self, id: DefId) -> impl Iterator<Item = &MTy<'tcx>> {
        self.tcx
            .fn_sig(id)
            .skip_binder()
            .inputs()
            .iter()
            .map(|i| i.skip_binder())
    }

    fn typeof_expr(&self, ex: Expr<'tcx>) -> MTy<'tcx> {
        let typeck_result = self.tcx.typeck(ex.hir_id.owner.def_id);
        typeck_result.expr_ty(&ex)
    }

    pub fn inertia_props_for_axum_route(
        &self,
        route: DefId,
    ) -> impl Iterator<Item = &structures::InertiaProp<'tcx>> {
        let props_by_component = self
            .inertia_props
            .iter()
            .filter(|(_id, p)| p.axum_handler == route)
            .filter(|(_id, p)| !p.flash_prop())
            .filter_map(|(id, p)| match self.inertia_prop_destination(*id) {
                // we only care about props who's destination is a component, because that indicates
                // that the prop is local to this route. `.back` props are flashed and are effectively global,
                // so those are queried elsewhere
                InertiaPropDestination::Render(component) => Some((component, p)),
                _ => None,
            })
            .into_group_map();

        // as a part of locating the props, we want to scan all of the props we have found and assert that
        // each "grouping" of props (i.e. if you have two calls to `.render` based on some condition) have the
        // same prop keys so we know they don't diverge from one another

        let mut all_prop_names: Option<HashSet<&str>> = None;
        for (component, props) in &props_by_component {
            let props_by_name: HashMap<_, _> = props.iter().map(|p| (p.name(), p)).collect();
            let names_in_this_group: HashSet<_> = props_by_name.keys().copied().collect();

            match all_prop_names.as_ref() {
                Some(known) => {
                    props_by_name
                        .keys()
                        .copied()
                        .collect::<HashSet<_>>()
                        .difference(known)
                        .for_each(|divergent| {
                            let prop = props_by_name.get(divergent).unwrap();
                            prop.emit_err_diverged(self.tcx.dcx(), component.as_str());
                        });
                }
                None => {
                    // otherwise record this set as the canonical set
                    all_prop_names = Some(names_in_this_group);
                }
            }
        }

        props_by_component.into_values().flatten()
    }
}

/// Common operations on types defined in the `rustc_middle` (i.e. post `typeck`) for asserting
/// they are a given shape
pub trait MTyAnalysisExt<'tcx> {
    /// Given a type or a tuple, peel back the tuple and split into a set of types
    fn peel_tuple(self) -> Vec<MTy<'tcx>>;

    /// Extract an adt returning the adt def and its generic arguments
    fn extract_adt(self) -> Option<(AdtDef<'tcx>, impl Iterator<Item = GenericArg<'tcx>> + 'tcx)>;
}

impl<'tcx> MTyAnalysisExt<'tcx> for MTy<'tcx> {
    fn peel_tuple(self) -> Vec<MTy<'tcx>> {
        let MTyKind::Tuple(tys) = self.kind() else {
            return vec![self];
        };

        tys.iter().collect()
    }

    fn extract_adt(self) -> Option<(AdtDef<'tcx>, impl Iterator<Item = GenericArg<'tcx>> + 'tcx)> {
        let MTyKind::Adt(adt, genargs) = self.kind() else {
            return None;
        };

        Some((*adt, genargs.into_iter()))
    }
}

/// Common operations on exprs for asserting they are a given shape
pub trait ExprAnalysisExt<'tcx> {
    fn extract_lit_str(self) -> Option<Symbol>;

    /// Given an expr that you expect to be a function call with a single argument,
    /// extract the def path of the function and its argument expression
    fn extract_call_with_single_arg(self, ra: &RiptAnalyzer<'tcx>)
    -> Option<(DefPath, Expr<'tcx>)>;

    /// Expect the expression to resolve to a path, then typecheck the path and extract
    /// the def id of the thing the path points to
    fn extract_qpath_did(self, ra: &RiptAnalyzer<'tcx>) -> Option<DefId>;
}

impl<'tcx> ExprAnalysisExt<'tcx> for &Expr<'tcx> {
    fn extract_lit_str(self) -> Option<Symbol> {
        let rustc_hir::ExprKind::Lit(Spanned {
            node: LitKind::Str(sym, _cooked_or_raw),
            ..
        }) = self.kind
        else {
            return None;
        };

        Some(*sym)
    }

    fn extract_call_with_single_arg(
        self,
        ra: &RiptAnalyzer<'tcx>,
    ) -> Option<(DefPath, Expr<'tcx>)> {
        let rustc_hir::ExprKind::Call(ex, &[first_arg]) = self.kind else {
            return None;
        };

        let rustc_hir::ExprKind::Path(qpath) = ex.kind else {
            return None;
        };

        let typeck_results = ra.tcx.typeck(first_arg.hir_id.owner);
        let qpath_res = typeck_results.qpath_res(&qpath, first_arg.hir_id);
        let dp = ra.tcx.def_path(qpath_res.def_id());

        Some((dp, first_arg))
    }

    fn extract_qpath_did(self, ra: &RiptAnalyzer<'tcx>) -> Option<DefId> {
        let rustc_hir::ExprKind::Path(qpath) = self.kind else {
            return None;
        };

        let typeck_results = ra.tcx.typeck(self.hir_id.owner);
        let qpath_res = typeck_results.qpath_res(&qpath, self.hir_id);

        Some(qpath_res.def_id())
    }
}

impl<'tcx> Visitor<'tcx> for RiptAnalyzer<'tcx> {
    type NestedFilter = VisitorFilter;

    fn visit_item(&mut self, i: &'tcx rustc_hir::Item<'tcx>) -> Self::Result {
        if let Some(bmi) = analysis::try_bridge_marked_item(self, i) {
            self.bridge_marked_items.insert(bmi.id, bmi);
        }

        rustc_hir::intravisit::walk_item(self, i)
    }

    fn visit_expr(&mut self, ex: &'tcx Expr<'tcx>) -> Self::Result {
        if let Some((dest_id, prop_dest)) = analysis::try_inertia_prop_dest_chain_definer(self, *ex)
        {
            self.inertia_prop_destinations
                .insert(dest_id, Either::Right(prop_dest));
        }

        if let Some(prop) = analysis::try_inertia_prop_from_chain(self, *ex) {
            // after resolving a prop, we want to include this props own receiver
            // in the destination chain so that it can be walked
            self.inertia_prop_destinations
                .insert(ex.hir_id, Either::Left(prop.rx_id));
            self.inertia_props.insert(ex.hir_id, prop);
        }

        if let Some(prop) = analysis::try_inertia_flash_prop_from_chain(self, *ex) {
            self.inertia_prop_destinations
                .insert(ex.hir_id, Either::Left(prop.rx_id));
            self.inertia_props.insert(ex.hir_id, prop);
        }

        if let Some(route) = analysis::try_axum_route_from_route_method_call(self, *ex) {
            self.axum_routes.insert(route.id, route);
        }

        if let Some(mw) = analysis::try_axum_share_middleware_registration(self, *ex) {
            self.axum_map_response_fns.insert(mw);
        }

        rustc_hir::intravisit::walk_expr(self, ex)
    }

    // don't remove this, its required for the nested filter to work
    fn maybe_tcx(&mut self) -> Self::MaybeTyCtxt {
        self.tcx
    }
}

// doc hidden because this is only public since it has to be to be used as an associated type
// of the visitor
//
// the only purpose of this is to enable the visitor to visit deeply and recursively instead of stopping
// on top-level items
#[doc(hidden)]
pub struct VisitorFilter;

impl<'tcx> NestedFilter<'tcx> for VisitorFilter {
    type MaybeTyCtxt = TyCtxt<'tcx>;

    const INTER: bool = true;
    const INTRA: bool = true;
}