tinymist-query 0.15.4-rc1

//! Analysis of function signatures.

use std::cell::RefCell;

use itertools::Either;
use tinymist_analysis::{ArgInfo, ArgsInfo, PartialSignature, func_signature};
use tinymist_derive::BindTyCtx;
use tinymist_std::hash::FxHashSet;

use super::{Definition, SharedContext, prelude::*};
use crate::analysis::PostTypeChecker;
use crate::docs::{DocText, UntypedDefDocs, UntypedSignatureDocs, UntypedVarDocs};
use crate::syntax::{DeclExpr, classify_def_loosely};
use crate::ty::{
    BoundCheckContext, BoundChecker, DocSource, DynTypeBounds, ParamAttrs, ParamTy, SigTy,
    SigWithTy, TyCtx, TyCtxMut, TypeInfo, TypeVar,
};

pub use tinymist_analysis::{PrimarySignature, Signature};

/// The language object that the signature is being analyzed for.
#[derive(Debug, Clone)]
pub enum SignatureTarget {
    /// A static node without knowing the function at runtime.
    Def(Option<Source>, Definition),
    /// A static node without knowing the function at runtime.
    SyntaxFast(Source, Span),
    /// A static node without knowing the function at runtime.
    Syntax(Source, Span),
    /// A function that is known at runtime.
    Runtime(Func),
    /// A function that is known at runtime.
    Convert(Func),
}

impl SignatureTarget {
    /// Returns the span of the callee node.
    pub fn span(&self) -> Span {
        match self {
            SignatureTarget::Def(_, def) => def.decl.span(),
            SignatureTarget::SyntaxFast(_, span) | SignatureTarget::Syntax(_, span) => *span,
            SignatureTarget::Runtime(func) | SignatureTarget::Convert(func) => func.span(),
        }
    }
}

#[typst_macros::time(span = callee_node.span())]
pub(crate) fn analyze_signature(
    ctx: &Arc<SharedContext>,
    callee_node: SignatureTarget,
) -> Option<Signature> {
    ctx.compute_signature(callee_node.clone(), move |ctx| {
        crate::log_debug_ct!("analyzing signature for {callee_node:?}");
        analyze_type_signature(ctx, &callee_node)
            .or_else(|| analyze_dyn_signature(ctx, &callee_node))
    })
}

#[typst_macros::time(span = callee_node.span())]
fn analyze_type_signature(
    ctx: &Arc<SharedContext>,
    callee_node: &SignatureTarget,
) -> Option<Signature> {
    let (type_info, ty) = match callee_node {
        SignatureTarget::Convert(..) => return None,
        SignatureTarget::SyntaxFast(source, span) | SignatureTarget::Syntax(source, span) => {
            let type_info = ctx.type_check(source);
            let ty = type_info.type_of_span(*span)?;
            Some((type_info, ty))
        }
        SignatureTarget::Def(source, def) => {
            let span = def.decl.span();
            let type_info = ctx.type_check(source.as_ref()?);
            let ty = type_info.type_of_span(span)?;
            Some((type_info, ty))
        }
        SignatureTarget::Runtime(func) => {
            let source = ctx.source_by_id(func.span().id()?).ok()?;
            let node = source.find(func.span())?;
            let def = classify_def_loosely(node.parent()?.clone())?;
            let type_info = ctx.type_check(&source);
            let ty = type_info.type_of_span(def.name()?.span())?;
            Some((type_info, ty))
        }
    }?;

    sig_of_type(ctx, &type_info, ty)
}

pub(crate) fn sig_of_type(
    ctx: &Arc<SharedContext>,
    type_info: &TypeInfo,
    ty: Ty,
) -> Option<Signature> {
    // todo multiple sources
    let mut srcs = ty.sources();
    srcs.sort();
    crate::log_debug_ct!("check type signature of ty: {ty:?} => {srcs:?}");
    let type_var = srcs.into_iter().next()?;
    match type_var {
        DocSource::Var(v) => {
            let mut ty_ctx =
                SignatureTypeContext::new(PostTypeChecker::new(ctx.clone(), type_info));
            let sig_ty = Ty::Func(ty.sig_repr(true, &mut ty_ctx)?);
            let sig_ty = type_info.simplify(sig_ty, false);
            let Ty::Func(sig_ty) = sig_ty else {
                static WARN_ONCE: std::sync::Once = std::sync::Once::new();
                WARN_ONCE.call_once(|| {
                    // todo: seems like a bug
                    log::warn!("expected function type, got {sig_ty:?}");
                });
                return None;
            };

            // todo: this will affect inlay hint: _var_with
            ty_ctx.reset_cycle_guard();
            let (var_with, docstring) = match type_info.var_docs.get(&v.def).map(|x| x.as_ref()) {
                Some(UntypedDefDocs::Function(sig)) => (vec![], Either::Left(sig.as_ref())),
                Some(UntypedDefDocs::Variable(docs)) => find_alias_stack(&mut ty_ctx, &v, docs)?,
                _ => return None,
            };

            let docstring = match docstring {
                Either::Left(docstring) => docstring,
                Either::Right(func) => return Some(wind_stack(var_with, ctx.type_of_func(func))),
            };

            let mut param_specs = Vec::new();
            let mut has_fill_or_size_or_stroke = false;
            let mut _broken = false;

            if docstring.pos.len() != sig_ty.positional_params().len() {
                static WARN_ONCE: std::sync::Once = std::sync::Once::new();
                WARN_ONCE.call_once(|| {
                    // todo: seems like a bug
                    log::warn!("positional params mismatch: {docstring:#?} != {sig_ty:#?}");
                });
                return None;
            }

            for (doc, ty) in docstring.pos.iter().zip(sig_ty.positional_params()) {
                let default = doc.default.clone();
                let ty = ty.clone();

                let name = doc.name.clone();
                if matches!(name.as_ref(), "fill" | "stroke" | "size") {
                    has_fill_or_size_or_stroke = true;
                }

                param_specs.push(Interned::new(ParamTy {
                    name,
                    docs: Some(DocText::plain(doc.docs.clone())),
                    default,
                    ty,
                    attrs: ParamAttrs::positional(),
                }));
            }

            for (name, ty) in sig_ty.named_params() {
                let docstring = docstring.named.get(name);
                let default = Some(
                    docstring
                        .and_then(|doc| doc.default.clone())
                        .unwrap_or_else(|| "unknown".into()),
                );
                let ty = ty.clone();

                if matches!(name.as_ref(), "fill" | "stroke" | "size") {
                    has_fill_or_size_or_stroke = true;
                }

                param_specs.push(Interned::new(ParamTy {
                    name: name.clone(),
                    docs: docstring.map(|doc| DocText::plain(doc.docs.clone())),
                    default,
                    ty,
                    attrs: ParamAttrs::named(),
                }));
            }

            if let Some(doc) = docstring.rest.as_ref() {
                let default = doc.default.clone();

                param_specs.push(Interned::new(ParamTy {
                    name: doc.name.clone(),
                    docs: Some(DocText::plain(doc.docs.clone())),
                    default,
                    ty: sig_ty.rest_param().cloned().unwrap_or(Ty::Any),
                    attrs: ParamAttrs::variadic(),
                }));
            }

            let sig = Signature::Primary(Arc::new(PrimarySignature {
                docs: Some(DocText::plain(docstring.docs.clone())),
                param_specs,
                has_fill_or_size_or_stroke,
                sig_ty,
                _broken,
            }));
            Some(wind_stack(var_with, sig))
        }
        src @ (DocSource::Builtin(..) | DocSource::Ins(..)) => {
            Some(ctx.type_of_func(src.as_func()?))
        }
    }
}

/// A type context for signature rendering that memoizes `(TypeVar, polarity)`
/// visits within a single traversal so recursive bounds are only expanded once.
struct SignatureTypeContext<'a> {
    inner: PostTypeChecker<'a>,
    visited: RefCell<FxHashSet<(DeclExpr, bool)>>,
}

impl<'a> SignatureTypeContext<'a> {
    fn new(inner: PostTypeChecker<'a>) -> Self {
        Self {
            inner,
            visited: RefCell::new(FxHashSet::default()),
        }
    }

    fn reset_cycle_guard(&self) {
        self.visited.borrow_mut().clear();
    }
}

impl TyCtx for SignatureTypeContext<'_> {
    fn global_bounds(&self, var: &Interned<TypeVar>, pol: bool) -> Option<DynTypeBounds> {
        if !self.visited.borrow_mut().insert((var.def.clone(), pol)) {
            return None;
        }

        self.inner.global_bounds(var, pol)
    }

    fn local_bind_of(&self, var: &Interned<TypeVar>) -> Option<Ty> {
        self.inner.local_bind_of(var)
    }
}

impl TyCtxMut for SignatureTypeContext<'_> {
    type Snap = <TypeInfo as TyCtxMut>::Snap;

    fn start_scope(&mut self) -> Self::Snap {
        self.inner.start_scope()
    }

    fn end_scope(&mut self, snap: Self::Snap) {
        self.inner.end_scope(snap)
    }

    fn bind_local(&mut self, var: &Interned<TypeVar>, ty: Ty) {
        self.inner.bind_local(var, ty);
    }

    fn type_of_func(&mut self, func: &Func) -> Option<Interned<SigTy>> {
        self.inner.type_of_func(func)
    }

    fn type_of_value(&mut self, val: &Value) -> Ty {
        self.inner.type_of_value(val)
    }

    fn check_module_item(&mut self, module: TypstFileId, key: &StrRef) -> Option<Ty> {
        self.inner.check_module_item(module, key)
    }
}

fn wind_stack(var_with: Vec<WithElem>, sig: Signature) -> Signature {
    if var_with.is_empty() {
        return sig;
    }

    let (primary, mut base_args) = match sig {
        Signature::Primary(primary) => (primary, eco_vec![]),
        Signature::Partial(partial) => (partial.signature.clone(), partial.with_stack.clone()),
    };

    let mut accepting = primary.pos().iter().skip(base_args.len());

    // Ignoring docs at the moment
    for (_d, w) in var_with {
        if let Some(w) = w {
            let mut items = eco_vec![];
            for pos in w.with.positional_params() {
                let Some(arg) = accepting.next() else {
                    break;
                };
                items.push(ArgInfo {
                    name: Some(arg.name.clone()),
                    term: Some(pos.clone()),
                });
            }
            // todo: ignored spread arguments
            if !items.is_empty() {
                base_args.push(ArgsInfo { items });
            }
        }
    }

    Signature::Partial(Arc::new(PartialSignature {
        signature: primary,
        with_stack: base_args,
    }))
}

type WithElem<'a> = (&'a UntypedVarDocs, Option<Interned<SigWithTy>>);

fn find_alias_stack<'a>(
    ctx: &'a mut SignatureTypeContext,
    var: &Interned<TypeVar>,
    docs: &'a UntypedVarDocs,
) -> Option<(Vec<WithElem<'a>>, Either<&'a UntypedSignatureDocs, Func>)> {
    let mut checker = AliasStackChecker {
        ctx,
        stack: vec![(docs, None)],
        res: None,
        checking_with: true,
    };
    Ty::Var(var.clone()).bounds(true, &mut checker);

    checker.res.map(|res| (checker.stack, res))
}

#[derive(BindTyCtx)]
#[bind(ctx)]
struct AliasStackChecker<'a, 'b> {
    ctx: &'a mut SignatureTypeContext<'b>,
    stack: Vec<WithElem<'a>>,
    res: Option<Either<&'a UntypedSignatureDocs, Func>>,
    checking_with: bool,
}

impl BoundChecker for AliasStackChecker<'_, '_> {
    fn check_var(&mut self, u: &Interned<TypeVar>, pol: bool, ctx: &mut BoundCheckContext) {
        crate::log_debug_ct!("collecting var {u:?} {pol:?}");
        if self.res.is_some() {
            return;
        }

        if self.checking_with {
            ctx.check_var_rec(u, pol, self);
            return;
        }

        let docs = self.ctx.inner.info.var_docs.get(&u.def).map(|x| x.as_ref());

        crate::log_debug_ct!("collecting var {u:?} {pol:?} => {docs:?}");
        // todo: bind builtin functions
        match docs {
            Some(UntypedDefDocs::Function(sig)) => {
                self.res = Some(Either::Left(sig));
            }
            Some(UntypedDefDocs::Variable(docs)) => {
                self.checking_with = true;
                self.stack.push((docs, None));
                ctx.check_var_rec(u, pol, self);
                self.stack.pop();
                self.checking_with = false;
            }
            _ => {}
        }
    }

    fn collect(&mut self, ty: &Ty, pol: bool) {
        if self.res.is_some() {
            return;
        }

        match (self.checking_with, ty) {
            (true, Ty::With(w)) => {
                crate::log_debug_ct!("collecting with {ty:?} {pol:?}");
                self.stack.last_mut().unwrap().1 = Some(w.clone());
                self.checking_with = false;
                w.sig.bounds(pol, self);
                self.checking_with = true;
            }
            (false, ty) => {
                if let Some(src) = ty.as_source() {
                    match src {
                        DocSource::Var(u) => {
                            let mut ctx = BoundCheckContext::default();
                            self.check_var(&u, pol, &mut ctx);
                        }
                        src @ (DocSource::Builtin(..) | DocSource::Ins(..)) => {
                            if let Some(func) = src.as_func() {
                                self.res = Some(Either::Right(func));
                            }
                        }
                    }
                }
            }
            _ => {}
        }
    }
}

#[typst_macros::time(span = callee_node.span())]
fn analyze_dyn_signature(
    ctx: &Arc<SharedContext>,
    callee_node: &SignatureTarget,
) -> Option<Signature> {
    let func = match callee_node {
        SignatureTarget::Def(_source, def) => def.value()?.to_func()?,
        SignatureTarget::SyntaxFast(..) => return None,
        SignatureTarget::Syntax(source, span) => {
            let def = ctx.def_of_span(source, *span)?;
            def.value()?.to_func()?
        }
        SignatureTarget::Convert(func) | SignatureTarget::Runtime(func) => func.clone(),
    };

    Some(func_signature(func))
}