tinymist-query 0.14.14-rc1

use std::sync::OnceLock;

use tinymist_analysis::adt::interner::Interned;
use typst::syntax::Span;

use crate::{
    StrRef,
    analysis::{Definition, SearchCtx},
    prelude::*,
    syntax::{RefExpr, SyntaxClass, get_index_info},
};

/// The [`textDocument/references`] request is sent from the client to the
/// server to resolve project-wide references for the symbol denoted by the
/// given text document position.
///
/// [`textDocument/references`]: https://microsoft.github.io/language-server-protocol/specification#textDocument_references
#[derive(Debug, Clone)]
pub struct ReferencesRequest {
    /// The path of the document to request for.
    pub path: PathBuf,
    /// The source code position to request for.
    pub position: LspPosition,
}

impl SemanticRequest for ReferencesRequest {
    type Response = Vec<LspLocation>;

    fn request(self, ctx: &mut LocalContext) -> Option<Self::Response> {
        let source = ctx.source_by_path(&self.path).ok()?;
        let syntax = ctx.classify_for_decl(&source, self.position)?;

        let locations = find_references(ctx, &source, syntax)?;

        crate::log_debug_ct!("references: {locations:?}");
        Some(locations)
    }
}

pub(crate) fn find_references(
    ctx: &mut LocalContext,
    source: &Source,
    syntax: SyntaxClass<'_>,
) -> Option<Vec<LspLocation>> {
    let finding_label = match syntax {
        SyntaxClass::VarAccess(..) | SyntaxClass::Callee(..) => false,
        SyntaxClass::Label { .. }
        | SyntaxClass::Ref {
            suffix_colon: false,
            ..
        } => true,
        SyntaxClass::ImportPath(..)
        | SyntaxClass::IncludePath(..)
        | SyntaxClass::Ref {
            suffix_colon: true, ..
        }
        | SyntaxClass::At { node: _ }
        | SyntaxClass::Normal(..) => {
            return None;
        }
    };

    let def = ctx.def_of_syntax(source, syntax)?;

    let worker = ReferencesWorker {
        ctx: ctx.fork_for_search(),
        references: vec![],
        def,
        module_path: OnceLock::new(),
    };

    if finding_label {
        worker.label_root()
    } else {
        // todo: reference of builtin items?
        worker.ident_root()
    }
}

struct ReferencesWorker<'a> {
    ctx: SearchCtx<'a>,
    references: Vec<LspLocation>,
    def: Definition,
    module_path: OnceLock<StrRef>,
}

impl ReferencesWorker<'_> {
    fn label_root(mut self) -> Option<Vec<LspLocation>> {
        for ref_fid in self.ctx.ctx.depended_files() {
            self.file(ref_fid)?;
        }

        Some(self.references)
    }

    fn ident_root(mut self) -> Option<Vec<LspLocation>> {
        self.file(self.def.decl.file_id()?);
        while let Some(ref_fid) = self.ctx.worklist.pop() {
            self.file(ref_fid);
        }

        Some(self.references)
    }

    fn file(&mut self, ref_fid: TypstFileId) -> Option<()> {
        log::debug!("references: file: {ref_fid:?}");

        // todo: find references in data files
        if ref_fid
            .vpath()
            .as_rooted_path()
            .extension()
            .is_none_or(|e| e != "typ")
        {
            return Some(());
        }

        let src = self.ctx.ctx.source_by_id(ref_fid).ok()?;
        let index = get_index_info(&src);
        match self.def.decl.kind() {
            DefKind::Constant | DefKind::Function | DefKind::Struct | DefKind::Variable => {
                if !index.identifiers.contains(self.def.decl.name()) {
                    return Some(());
                }
            }
            DefKind::Module => {
                let ref_by_ident = index.identifiers.contains(self.def.decl.name());
                let ref_by_path = index.paths.contains(self.module_path());
                if !(ref_by_ident || ref_by_path) {
                    return Some(());
                }
            }
            DefKind::Reference => {}
        }

        let ei = self.ctx.ctx.expr_stage(&src);
        let uri = self.ctx.ctx.uri_for_id(ref_fid).ok()?;

        let t = ei.get_refs(self.def.decl.clone());
        self.push_idents(&ei.source, &uri, t);

        if ei.is_exported(&self.def.decl) {
            self.ctx.push_dependents(ref_fid);
        }

        Some(())
    }

    fn push_idents<'b>(
        &mut self,
        src: &Source,
        url: &Url,
        idents: impl Iterator<Item = (&'b Span, &'b Interned<RefExpr>)>,
    ) {
        self.push_ranges(
            src,
            url,
            idents.map(|(span, expr)| {
                let adjust = match expr.decl.as_ref() {
                    Decl::Label(..) => Some((1, -1)),
                    Decl::ContentRef(..) => Some((1, 0)),
                    _ => None,
                };

                (*span, adjust)
            }),
        );
    }

    fn push_ranges(
        &mut self,
        src: &Source,
        url: &Url,
        spans: impl Iterator<Item = (Span, Option<(isize, isize)>)>,
    ) {
        self.references.extend(spans.filter_map(|(span, adjust)| {
            // todo: this is not necessary a name span
            let mut range = src.range(span)?;
            if let Some((start, end)) = adjust {
                range.start = (range.start as isize + start) as usize;
                range.end = (range.end as isize + end) as usize;
            }
            let range = self.ctx.ctx.to_lsp_range(range, src);
            Some(LspLocation {
                uri: url.clone(),
                range,
            })
        }));
    }

    // todo: references of package
    fn module_path(&self) -> &StrRef {
        self.module_path.get_or_init(|| {
            self.def
                .decl
                .file_id()
                .and_then(|fid| {
                    fid.vpath()
                        .as_rooted_path()
                        .file_name()?
                        .to_str()
                        .map(From::from)
                })
                .unwrap_or_default()
        })
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::tests::*;

    #[test]
    fn test() {
        snapshot_testing("references", &|ctx, path| {
            let source = ctx.source_by_path(&path).unwrap();

            let request = ReferencesRequest {
                path: path.clone(),
                position: find_test_position(&source),
            };

            let result = request.request(ctx);
            let mut result = result.map(|v| {
                v.into_iter()
                    .map(|loc| {
                        let fp = file_uri(loc.uri.as_str());
                        format!(
                            "{fp}@{}:{}:{}:{}",
                            loc.range.start.line,
                            loc.range.start.character,
                            loc.range.end.line,
                            loc.range.end.character
                        )
                    })
                    .collect::<Vec<_>>()
            });
            // sort
            if let Some(result) = result.as_mut() {
                result.sort();
            }

            assert_snapshot!(JsonRepr::new_pure(result));
        });
    }
}