riptc 0.1.2

//! This module contains the code for taking the analyzer
//! state and emitting fully valid typescript output.

mod exports;
mod preamble;
mod type_conversion;

use std::{
    collections::{HashMap, HashSet},
    io,
    rc::Rc,
};

use bon::{Builder, builder};
use itertools::Itertools;
use rustc_hir::def_id::DefId;
use rustc_middle::ty::TyCtxt;
use swc_atoms::Atom;
use swc_common::{DUMMY_SP, FileName, SourceMap, comments::SingleThreadedComments};
use swc_ecma_ast::{
    Decl, EsVersion, ExportDecl, Module, ModuleDecl, ModuleItem, TplElement, TsType, TsTypeElement,
};
use swc_ecma_codegen::{Emitter, text_writer::JsWriter};

const PREAMBLE: &str = r#"// Auto-generated by RIPT CLI.
// Do not modify this file directly.
// For more information, see https://github.com/kindness-ai/ript

"#;

use crate::{
    analyzer::{self, AxumRoutePathSegment, RiptAnalyzer},
    callbacks::ConfigStrExt,
    namespace::{NamespaceNode, NamespacedPath},
    swc_utils,
};

#[derive(Builder)]
pub struct RiptCodegen<'tcx, 'swc> {
    analyzer: RiptAnalyzer<'tcx>,
    tcx: TyCtxt<'tcx>,
    swc_emitter: Emitter<'swc, JsWriter<'swc, &'swc mut dyn io::Write>, SourceMap>,
    #[builder(default)]
    namespace_root: parking_lot::RwLock<NamespaceNode<ModuleItem>>,
    #[builder(default = Module {
        span: DUMMY_SP,
        body: vec![],
        shebang: None,
    })]
    module: Module,
}

impl<'tcx> RiptCodegen<'tcx, '_> {
    pub fn emit_all(mut self) -> io::Result<()> {
        exports::insert_flash_props(&self);
        exports::insert_share_props(&self);
        exports::insert_bridge_marked_items(&self, self.analyzer.bridge_marked_items());
        let inertia_axum_routes = self.analyzer.inertia_axum_routes().collect_vec();

        let axum_routes = self.analyzer.axum_routes().collect_vec();

        exports::insert_axum_route_types(&self, axum_routes.iter().copied());
        exports::insert_inertia_route_href_builders(&self, inertia_axum_routes.iter().copied());
        exports::insert_inertia_route_definitions_type(&self, inertia_axum_routes.iter().copied());

        let preamble = preamble::wrap_preamble(&mut self.swc_emitter)?;
        self.swc_emitter.wr.preamble(&preamble)?;

        let child_items = node_to_module_items(self.namespace_root.into_inner());

        self.module.body.extend(child_items);
        self.swc_emitter.emit_module(&self.module)
    }

    fn insert_node(&self, path: NamespacedPath, item: ModuleItem) {
        self.namespace_root
            .write()
            .insert()
            .root_path(path)
            .item(item)
            .call()
    }

    fn insert_ts_type(&self, path: NamespacedPath, ts_type: TsType) {
        let module_item = ModuleItem::ModuleDecl(ModuleDecl::ExportDecl(ExportDecl {
            span: DUMMY_SP,
            decl: Decl::TsTypeAlias(Box::new(swc_ecma_ast::TsTypeAliasDecl {
                span: DUMMY_SP,
                declare: false,
                id: swc_utils::ident(path.item_name().as_str()),
                type_params: None,
                type_ann: Box::new(ts_type),
            })),
        }));

        self.insert_node(path, module_item)
    }

    /// Convert any given rust type into an swc type
    fn as_swc(
        &self,
        ty: rustc_middle::ty::Ty<'tcx>,
        span: rustc_span::Span,
    ) -> swc_ecma_ast::TsType {
        type_conversion::convert_ty(self, &span, ty, 0)
    }

    fn axum_route_path_template_literal(
        &self,
        route_id: DefId,
    ) -> (swc_ecma_ast::Tpl, impl Iterator<Item = (Atom, TsType)>) {
        let ra = &self.analyzer;
        let segments = ra.axum_route_path(route_id);

        let mut tpl_exprs = vec![];
        let mut tpl_quasis = vec![];

        let mut params = vec![];

        for segment in segments {
            match segment {
                AxumRoutePathSegment::Static(sym) => {
                    tpl_quasis.push(TplElement {
                        span: DUMMY_SP,
                        cooked: None,
                        raw: swc_ecma_ast::Str::from_tpl_raw(sym.as_str()),
                        tail: false,
                    });
                }
                AxumRoutePathSegment::Dynamic(sym, ty) => {
                    let sym_atom = swc_atoms::Atom::from(sym.as_str());

                    tpl_exprs.push(Box::new(swc_ecma_ast::Expr::Ident(swc_ecma_ast::Ident {
                        span: DUMMY_SP,
                        sym: swc_atoms::Atom::from(sym.as_str()),
                        optional: false,
                        ctxt: swc_common::SyntaxContext::empty(),
                    })));

                    params.push((sym_atom, self.as_swc(ty, Default::default())));
                }
            }
        }

        if !tpl_quasis.is_empty() {
            tpl_quasis.last_mut().unwrap().tail = true;
        }

        (
            swc_ecma_ast::Tpl {
                span: DUMMY_SP,
                exprs: tpl_exprs,
                quasis: tpl_quasis,
            },
            params.into_iter(),
        )
    }

    /// When we have something like share / flash props, they are effectively global and cannot necessarily be restricted to a single
    /// given page or route. Because of this, we have to be able to flatten flash / share props with the same name into a single type.
    fn inertia_prop_tys_flattened_into_unions<'a>(
        &'a self,
        props: impl Iterator<Item = &'a analyzer::InertiaProp<'tcx>>,
    ) -> Vec<TsTypeElement> {
        let mut omittable_map: HashSet<&str> = HashSet::default();
        let mut seen_tys: HashMap<&str, HashSet<TsType>> = HashMap::default();
        for prop in props {
            let prop_ty = self.as_swc(prop.ty(), prop.span());

            let prop_ty = if prop.lazy() || prop.deferred() {
                swc_utils::union_type(vec![prop_ty, swc_utils::null_type()].into_iter())
            } else {
                prop_ty
            };

            // if we have `T` | `null`, we want to insert those separately
            let prop_tys = swc_utils::split_union_type(Box::new(prop_ty));

            seen_tys
                .entry(prop.name())
                .or_default()
                .extend(prop_tys.map(|p| *p));

            if prop.flash_prop() {
                omittable_map.insert(prop.name());
            }
        }

        seen_tys
            .into_iter()
            .map(|(name, ty)| {
                let optional = omittable_map.contains(name);

                swc_utils::object_member_type_element()
                    .key(name)
                    .type_ann(swc_utils::union_type(ty.into_iter()))
                    .optional(optional)
                    .build()
            })
            .collect_vec()
    }
}

#[builder]
pub fn new_swc_emitter<'swc>(
    w: &'swc mut dyn io::Write,
    comments: &'swc SingleThreadedComments,
) -> io::Result<Emitter<'swc, JsWriter<'swc, &'swc mut dyn io::Write>, SourceMap>> {
    // hack to make comments work without an actual source map
    let cm = Rc::new(SourceMap::default());
    let _fm = cm.new_source_file(
        Rc::new(FileName::Custom("synthesized.ts".into())),
        "".into(),
    );

    let mut jsw = JsWriter::new(Rc::clone(&cm), "\n", w, None);
    jsw.preamble(PREAMBLE)?;

    Ok(Emitter {
        cfg: swc_ecma_codegen::Config::default().with_target(EsVersion::latest()),
        cm,
        comments: Some(comments),
        wr: jsw,
    })
}

/// Recursively convert a `NamespaceNode` into a list of `ModuleItem`s.
/// Each child becomes a `namespace <child_name> { ... }` statement,
/// and each type becomes `export type <typeName> = ...`.
fn node_to_module_items(node: NamespaceNode<ModuleItem>) -> Vec<ModuleItem> {
    let mut items = Vec::new();

    let mut sorted_children: Vec<_> = node.children.into_iter().collect();
    sorted_children.sort_by(|a, b| {
        *(&a.0
            .as_str()
            .namespace_name()
            .cmp(&b.0.as_str().namespace_name()))
    });

    for (child_name, child_node) in sorted_children {
        let child_items = node_to_module_items(child_node);

        let ns_decl =
            swc_utils::nested_namespace_decl(child_name.as_str().namespace_name(), child_items);
        items.push(ModuleItem::ModuleDecl(ns_decl));
    }

    for (_type_name, item) in node.items {
        items.push(item);
    }

    items
}