apollo-language-server 0.7.0

use std::{collections::HashMap, sync::Arc};

use apollo_compiler::Schema;
use apollo_parser::{
    cst::{CstNode, Document},
    SyntaxKind,
};
use ropey::Rope;

use crate::{
    graph::subgraph::Subgraph,
    utils::{
        get_ancestor_node_of_kind::get_ancestor_node_of_kind,
        lsp_range_from_ast_sourcespan::lsp_range_from_ast_sourcespan,
        lsp_range_from_cst_textrange::lsp_range_from_cst_textrange,
    },
};

pub fn goto_definition_at_position(
    uri: &lsp::Url,
    document: &Document,
    schema: &Schema,
    source_text: &Rope,
    position: &lsp::Position,
    subgraphs_by_uri: Option<&HashMap<lsp::Url, Arc<Subgraph>>>,
) -> Option<Vec<lsp::LocationLink>> {
    let char = source_text.try_line_to_byte(position.line as usize).ok()?;
    let offset = char + position.character as usize;
    if document.syntax().text_range().end() < offset.try_into().unwrap() {
        return None;
    }

    goto_definition_named_type(uri, document, schema, source_text, offset, subgraphs_by_uri)
        .or_else(|| {
            goto_definition_directive_application(
                uri,
                document,
                schema,
                source_text,
                offset,
                subgraphs_by_uri,
            )
        })
}

fn goto_definition_named_type(
    uri: &lsp::Url,
    document: &Document,
    schema: &Schema,
    source_text: &Rope,
    offset: usize,
    subgraphs_by_uri: Option<&HashMap<lsp::Url, Arc<Subgraph>>>,
) -> Option<Vec<lsp::LocationLink>> {
    let token = {
        let token_at_offset = document
            .syntax()
            .token_at_offset(offset.try_into().unwrap());
        let left_biased = token_at_offset.clone().left_biased();
        let right_biased = token_at_offset.right_biased();

        left_biased
            .filter(|t| t.kind() != SyntaxKind::WHITESPACE)
            .or_else(|| right_biased.filter(|t| t.kind() != SyntaxKind::WHITESPACE))?
    };
    let parent = token.parent()?;
    let grandparent = parent.parent()?;

    let mut definitions_in_schema = vec![];
    if let Some(subgraphs) = subgraphs_by_uri {
        for subgraph in subgraphs.values() {
            if let Some(type_in_subgraph) = subgraph.schema().types.get(token.text()) {
                if type_in_subgraph.is_built_in()
                    || subgraph.builtins.contains(&token.text().to_string())
                {
                    continue;
                }
                definitions_in_schema.push((type_in_subgraph, subgraph.uri.clone()));
            }
        }
    } else {
        let type_in_schema = schema.types.get(token.text())?;
        definitions_in_schema.push((type_in_schema, uri.clone()));
    }

    if definitions_in_schema.is_empty() {
        return None;
    }

    // When a `NamedType`` is hovered, the offset will point to an identifier
    // token. In order to be sure it's a named type, we need to look up to the
    // token's ancestors to be sure it's a `NamedType`` and not an identifier in
    // a different context.
    (token.kind() == SyntaxKind::IDENT
        && parent.kind() == SyntaxKind::NAME
        && grandparent.kind() == SyntaxKind::NAMED_TYPE)
        .then(|| {
            definitions_in_schema
                .iter()
                .filter_map(|(ty, uri)| {
                    let schema_source_text = subgraphs_by_uri
                        .map(|subgraphs| &subgraphs[uri].source_text)
                        .unwrap_or(source_text);

                    Some(lsp::LocationLink {
                        target_uri: uri.clone(),
                        target_range: lsp_range_from_ast_sourcespan(
                            ty.location()?,
                            schema_source_text,
                        )?,
                        target_selection_range: lsp_range_from_ast_sourcespan(
                            ty.name().location()?,
                            schema_source_text,
                        )?,
                        origin_selection_range: None,
                    })
                })
                .collect::<Vec<_>>()
        })
}

fn goto_definition_directive_application(
    uri: &lsp::Url,
    document: &Document,
    schema: &Schema,
    source_text: &Rope,
    offset: usize,
    subgraphs_by_uri: Option<&HashMap<lsp::Url, Arc<Subgraph>>>,
) -> Option<Vec<lsp::LocationLink>> {
    let token = {
        let token_at_offset = document
            .syntax()
            .token_at_offset(offset.try_into().unwrap());
        let left_biased = token_at_offset.clone().left_biased();
        let right_biased = token_at_offset.right_biased();

        left_biased
            .filter(|t| t.kind() != SyntaxKind::WHITESPACE)
            .or_else(|| right_biased.filter(|t| t.kind() != SyntaxKind::WHITESPACE))?
    }
    .parent()?;

    let directive_node = get_ancestor_node_of_kind(
        &token,
        SyntaxKind::DIRECTIVE,
        Some(vec![SyntaxKind::ARGUMENTS]),
    )?;
    let at_symbol = directive_node.children_with_tokens().next()?;
    let directive_name = directive_node.children().next()?;

    let mut definitions_in_schema = vec![];
    if let Some(subgraphs) = subgraphs_by_uri {
        for subgraph in subgraphs.values() {
            if let Some(directive_in_subgraph) = subgraph
                .schema()
                .directive_definitions
                .get(directive_name.text().to_string().as_str())
            {
                if directive_in_subgraph.is_built_in()
                    || subgraph
                        .builtins
                        .contains(&directive_name.text().to_string())
                {
                    continue;
                }
                definitions_in_schema.push((directive_in_subgraph, subgraph.uri.clone()));
            }
        }
    } else {
        let directive_in_schema = schema
            .directive_definitions
            .get(directive_name.text().to_string().as_str())?;
        definitions_in_schema.push((directive_in_schema, uri.clone()));
    }

    if definitions_in_schema.is_empty() {
        return None;
    }

    Some(
        definitions_in_schema
            .iter()
            .filter_map(|(directive, uri)| {
                let schema_source_text = subgraphs_by_uri
                    .map(|subgraphs| &subgraphs[uri].source_text)
                    .unwrap_or(source_text);
                Some(lsp::LocationLink {
                    target_uri: uri.clone(),
                    origin_selection_range: Some(lsp_range_from_cst_textrange(
                        // Make sure we include the '@' symbol in the origin range
                        directive_name.text_range().cover(at_symbol.text_range()),
                        source_text,
                        None,
                    )),
                    target_range: lsp_range_from_ast_sourcespan(
                        directive.location()?,
                        schema_source_text,
                    )?,
                    target_selection_range: lsp_range_from_ast_sourcespan(
                        directive.name.location()?,
                        schema_source_text,
                    )?,
                })
            })
            .collect::<Vec<_>>(),
    )
}

#[cfg(test)]
mod tests {
    use crate::graph::{supergraph::KnownSubgraphs, Graph, GraphConfig};
    use insta::assert_snapshot;
    use itertools::Itertools;
    use ropey::Rope;
    use tower_lsp::lsp_types::{self as lsp};

    const URI_A: &str = "file:///a.graphql";
    const SOURCE_TEXT_A: &str = r#"
extend schema
  @link(
    url: "https://specs.apollo.dev/federation/v2.8"
    import: ["@key", "@override", "@requires", "@external", "@shareable"]
  )

directive @hello on FIELD_DEFINITION  

type Query {
  user: User
  team: Team
  me: ID!
}

"""
User description
spanning multiple lines
"""
type User {
  id: ID! @deprecated(reason: "Use `me` instead")
  me: ID! 
  name: String @hello
}

type Team {
  id: ID!
  members: [User!]!
} 
  "#;

    const URI_B: &str = "file:///b.graphql";
    const SOURCE_TEXT_B: &str = r#"
extend schema
  @link(
    url: "https://specs.apollo.dev/federation/v2.8"
    import: ["@key", "@override", "@requires", "@external", "@shareable"]
  )

type Query {
  user: User
  me: ID!
}

"""
User description
spanning multiple lines
"""
type User @key(fields: "id") {
  id: ID! @deprecated(reason: "Use `me` instead")
  me: ID!
  name: String
}"#;

    fn get_testing_graph() -> Graph {
        let mut graph = Graph::new(
            lsp::Url::parse(URI_A).unwrap(),
            SOURCE_TEXT_A.to_string(),
            0,
            KnownSubgraphs::default(),
            GraphConfig::default(),
        );
        graph.update(
            lsp::Url::parse(URI_B).unwrap(),
            SOURCE_TEXT_B.to_string(),
            0,
            GraphConfig::default(),
        );
        graph
    }

    fn get_definition_start_end_cases(
        graph: &Graph,
        uri: &str,
        line: u32,
        character: u32,
        length: u32,
    ) -> (Vec<lsp::LocationLink>, Vec<lsp::LocationLink>) {
        let before_start = graph.goto_definition(
            &lsp::Url::parse(uri).unwrap(),
            &lsp::Position {
                line,
                character: character - 1,
            },
        );

        let start = graph.goto_definition(
            &lsp::Url::parse(uri).unwrap(),
            &lsp::Position { line, character },
        );

        let end = graph.goto_definition(
            &lsp::Url::parse(uri).unwrap(),
            &lsp::Position {
                line,
                character: character + length,
            },
        );

        let beyond_end = graph.goto_definition(
            &lsp::Url::parse(uri).unwrap(),
            &lsp::Position {
                line,
                character: character + length + 1,
            },
        );

        assert!(before_start.is_none() || before_start.unwrap().is_empty());
        assert!(beyond_end.is_none() || beyond_end.unwrap().is_empty());

        (
            start.expect("Expected `start` to be Some"),
            end.expect("Expected `end` to be Some"),
        )
    }

    fn build_snapshot(links: &[lsp::LocationLink], source_text: &str) -> String {
        let mut result = String::new();

        let sorted_links = links
            .iter()
            .sorted_by(|a, b| a.target_uri.as_str().cmp(b.target_uri.as_str()))
            .collect::<Vec<_>>();

        if let Some(origin_selection_range) = sorted_links[0].origin_selection_range {
            let source_text = Rope::from_str(source_text);
            let origin_start_offset = source_text
                .try_line_to_byte(origin_selection_range.start.line as usize)
                .expect("Expected start_line to be in bounds")
                + origin_selection_range.start.character as usize;
            let origin_end_offset = source_text
                .try_line_to_byte(origin_selection_range.end.line as usize)
                .expect("Expected start_line to be in bounds")
                + origin_selection_range.end.character as usize;

            let origin_text = source_text.slice(origin_start_offset..origin_end_offset);

            result.push_str(
                format!(
                    "---- Origin selection range ---- \n\nOrigin source text:\n\n```\n{}\n```\n\nOrigin selection range: {:#?}\n\n\n",
                    origin_text,
                    origin_selection_range,
                )
                .as_str(),
            );
        }

        for link in sorted_links {
            let linked_source_text: Rope = if link.target_uri.as_str() == URI_A {
                Rope::from_str(SOURCE_TEXT_A)
            } else {
                Rope::from_str(SOURCE_TEXT_B)
            };
            // get start and end offsets of linked range
            let start_offset = linked_source_text
                .try_line_to_byte(link.target_range.start.line as usize)
                .expect("Expected start_line to be in bounds")
                + link.target_range.start.character as usize;
            let end_offset = linked_source_text
                .try_line_to_byte(link.target_range.end.line as usize)
                .expect("Expected start_line to be in bounds")
                + link.target_range.end.character as usize;

            // get text of linked range
            let text_at_linked_range = linked_source_text.slice(start_offset..end_offset);
            result.push_str(
                format!(
                    "---- Link to range at URI: {:?} ---- \n\nTargeted source text:\n\n```\n{}\n```\n\nTarget range: {:#?}\n\nTarget selection range: {:#?}\n\n\n",
                    link.target_uri.as_str(),
                    text_at_linked_range,
                    link.target_range,
                    link.target_selection_range,
                )
                .as_str(),
            )
        }

        result
    }

    #[test]
    fn test_goto_type_definition_in_multiple_subgraphs() {
        let graph = get_testing_graph();

        // 'User' reference at line 10, character 8 (with length 4)
        let (start, end) = get_definition_start_end_cases(&graph, URI_A, 10, 8, 4);

        let snapshot = build_snapshot(&start, SOURCE_TEXT_A);
        assert_snapshot!(snapshot);
        assert_eq!(start, end);
    }

    #[test]
    fn test_goto_type_definition_in_one_subgraph() {
        let graph = get_testing_graph();

        // 'Team' reference at line 11, character 8 (with length 4)
        let (start, end) = get_definition_start_end_cases(&graph, URI_A, 11, 8, 4);

        let snapshot = build_snapshot(&start, SOURCE_TEXT_A);
        assert_snapshot!(snapshot);
        assert_eq!(start, end);
    }

    #[test]
    fn text_goto_directive_definition() {
        let graph = get_testing_graph();

        // '@hello' reference at line 22, character 15 (with length 6)
        let (start, end) = get_definition_start_end_cases(&graph, URI_A, 22, 15, 6);

        let snapshot = build_snapshot(&start, SOURCE_TEXT_A);
        assert_snapshot!(snapshot);
        assert_eq!(start, end);
    }
}