feat(fmt): lsp find references (#4934)

* adds support for LSP references
* Updated LSP docs links

* In the process of adding more granularity for references, fixed a bug for completions (from engines) so we're now capable of offering completions based on partial text :)
* map should be offered when we aren't in a specific part of an attribute
* test that mssql default map doesn't show up inside specific sections of `@default`

---------

Co-authored-by: Flavian Desverne <[email protected]>

prisma-fmt/src/lib.rs

@@ -7,6 +7,7 @@ mod merge_schemas;
 mod native;
 mod offsets;
 mod preview;
+mod references;
 mod schema_file_input;
 mod text_document_completion;
 mod validate;
@@ -84,13 +85,30 @@ pub fn code_actions(schema_files: String, params: &str) -> String {
 
     let Ok(input) = serde_json::from_str::<SchemaFileInput>(&schema_files) else {
         warn!("Failed to parse schema file input");
-        return serde_json::to_string(&text_document_completion::empty_completion_list()).unwrap();
+        return serde_json::to_string(&code_actions::empty_code_actions()).unwrap();
     };
 
     let actions = code_actions::available_actions(input.into(), params);
     serde_json::to_string(&actions).unwrap()
 }
 
+pub fn references(schema_files: String, params: &str) -> String {
+    let params: lsp_types::ReferenceParams = if let Ok(params) = serde_json::from_str(params) {
+        params
+    } else {
+        warn!("Failed to parse params to references() as ReferenceParams.");
+        return serde_json::to_string(&references::empty_references()).unwrap();
+    };
+
+    let Ok(input) = serde_json::from_str::<SchemaFileInput>(&schema_files) else {
+        warn!("Failed to parse schema file input");
+        return serde_json::to_string(&references::empty_references()).unwrap();
+    };
+
+    let references = references::references(input.into(), params);
+    serde_json::to_string(&references).unwrap()
+}
+
 /// The two parameters are:
 /// - The [`SchemaFileInput`] to reformat, as a string.
 /// - An LSP

prisma-fmt/src/references.rs

@@ -0,0 +1,245 @@
+use log::*;
+use lsp_types::{Location, ReferenceParams, Url};
+use psl::{
+    diagnostics::FileId,
+    error_tolerant_parse_configuration,
+    parser_database::ParserDatabase,
+    schema_ast::ast::{
+        AttributePosition, CompositeTypePosition, EnumPosition, Field, FieldId, FieldPosition, FieldType, Identifier,
+        ModelId, ModelPosition, SchemaPosition, SourcePosition, Top, WithIdentifier, WithName,
+    },
+    Diagnostics, SourceFile,
+};
+
+use crate::{offsets::position_to_offset, span_to_range, LSPContext};
+
+pub(super) type ReferencesContext<'a> = LSPContext<'a, ReferenceParams>;
+
+pub(crate) fn empty_references() -> Vec<Location> {
+    Vec::new()
+}
+
+fn empty_identifiers<'ast>() -> impl Iterator<Item = &'ast Identifier> {
+    std::iter::empty()
+}
+
+pub(crate) fn references(schema_files: Vec<(String, SourceFile)>, params: ReferenceParams) -> Vec<Location> {
+    let (_, config, _) = error_tolerant_parse_configuration(&schema_files);
+
+    let db = {
+        let mut diag = Diagnostics::new();
+        ParserDatabase::new(&schema_files, &mut diag)
+    };
+
+    let Some(initiating_file_id) = db.file_id(params.text_document_position.text_document.uri.as_str()) else {
+        warn!("Initating file name is not found in the schema");
+        return empty_references();
+    };
+
+    let initiating_doc = db.source(initiating_file_id);
+
+    let position = if let Some(pos) = position_to_offset(&params.text_document_position.position, initiating_doc) {
+        pos
+    } else {
+        warn!("Received a position outside of the document boundaries in ReferenceParams");
+        return empty_references();
+    };
+
+    let target_position = db.ast(initiating_file_id).find_at_position(position);
+
+    let ctx = ReferencesContext {
+        db: &db,
+        config: &config,
+        initiating_file_id,
+        params: &params,
+    };
+
+    reference_locations_for_target(ctx, target_position)
+}
+
+fn reference_locations_for_target(ctx: ReferencesContext<'_>, target: SchemaPosition) -> Vec<Location> {
+    let identifiers: Vec<&Identifier> = match target {
+        // Blocks
+        SchemaPosition::Model(model_id, ModelPosition::Name(name)) => {
+            let model = ctx.db.walk((ctx.initiating_file_id, model_id));
+
+            std::iter::once(model.ast_model().identifier())
+                .chain(find_where_used_as_field_type(&ctx, name))
+                .collect()
+        }
+        SchemaPosition::Enum(enum_id, EnumPosition::Name(name)) => {
+            let enm = ctx.db.walk((ctx.initiating_file_id, enum_id));
+
+            std::iter::once(enm.ast_enum().identifier())
+                .chain(find_where_used_as_field_type(&ctx, name))
+                .collect()
+        }
+        SchemaPosition::CompositeType(composite_id, CompositeTypePosition::Name(name)) => {
+            let ct = ctx.db.walk((ctx.initiating_file_id, composite_id));
+
+            std::iter::once(ct.ast_composite_type().identifier())
+                .chain(find_where_used_as_field_type(&ctx, name))
+                .collect()
+        }
+        SchemaPosition::DataSource(_, SourcePosition::Name(name)) => find_where_used_as_ds_name(&ctx, name)
+            .into_iter()
+            .chain(find_where_used_for_native_type(&ctx, name))
+            .collect(),
+
+        // Fields
+        SchemaPosition::Model(_, ModelPosition::Field(_, FieldPosition::Type(r#type)))
+        | SchemaPosition::CompositeType(_, CompositeTypePosition::Field(_, FieldPosition::Type(r#type))) => {
+            find_where_used_as_top_name(&ctx, r#type)
+                .into_iter()
+                .chain(find_where_used_as_field_type(&ctx, r#type))
+                .collect()
+        }
+
+        // Attributes
+        SchemaPosition::Model(
+            model_id,
+            ModelPosition::Field(
+                field_id,
+                FieldPosition::Attribute(_, _, AttributePosition::ArgumentValue(arg_name, arg_value)),
+            ),
+        ) => match arg_name {
+            Some("fields") => find_where_used_as_field_name(&ctx, arg_value.as_str(), model_id, ctx.initiating_file_id)
+                .into_iter()
+                .collect(),
+            Some("references") => {
+                let field = &ctx.db.ast(ctx.initiating_file_id)[model_id][field_id];
+                let referenced_model = field.field_type.name();
+
+                let Some(ref_model_id) = ctx.db.find_model(referenced_model) else {
+                    warn!("Could not find model with name: {}", referenced_model);
+                    return empty_references();
+                };
+
+                find_where_used_as_field_name(&ctx, arg_value.as_str(), ref_model_id.id.1, ref_model_id.id.0)
+                    .into_iter()
+                    .collect()
+            }
+            _ => vec![],
+        },
+
+        // ? This might make more sense to add as a definition rather than a reference
+        SchemaPosition::Model(_, ModelPosition::Field(_, FieldPosition::Attribute(name, _, _)))
+        | SchemaPosition::CompositeType(_, CompositeTypePosition::Field(_, FieldPosition::Attribute(name, _, _))) => {
+            match ctx.datasource().map(|ds| &ds.name) {
+                Some(ds_name) if name.contains(ds_name) => find_where_used_as_ds_name(&ctx, ds_name)
+                    .into_iter()
+                    .chain(find_where_used_for_native_type(&ctx, ds_name))
+                    .collect(),
+                _ => vec![],
+            }
+        }
+
+        SchemaPosition::Model(
+            model_id,
+            ModelPosition::ModelAttribute(_attr_name, _, AttributePosition::ArgumentValue(_, arg_val)),
+        ) => find_where_used_as_field_name(&ctx, arg_val.as_str(), model_id, ctx.initiating_file_id)
+            .into_iter()
+            .collect(),
+
+        _ => vec![],
+    };
+
+    identifiers
+        .iter()
+        .filter_map(|ident| ident_to_location(ident, &ctx))
+        .collect()
+}
+
+fn find_where_used_as_field_name<'ast>(
+    ctx: &'ast ReferencesContext<'_>,
+    name: &str,
+    model_id: ModelId,
+    file_id: FileId,
+) -> Option<&'ast Identifier> {
+    let model = ctx.db.walk((file_id, model_id));
+
+    match model.scalar_fields().find(|field| field.name() == name) {
+        Some(field) => Some(field.ast_field().identifier()),
+        None => None,
+    }
+}
+
+fn find_where_used_for_native_type<'ast>(
+    ctx: &ReferencesContext<'ast>,
+    name: &'ast str,
+) -> impl Iterator<Item = &'ast Identifier> {
+    fn find_native_type_locations<'ast>(
+        name: &'ast str,
+        fields: impl Iterator<Item = (FieldId, &'ast Field)> + 'ast,
+    ) -> Box<dyn Iterator<Item = &'ast Identifier> + 'ast> {
+        Box::new(fields.filter_map(move |field| {
+            field
+                .1
+                .attributes
+                .iter()
+                .find(|attr| extract_ds_from_native_type(attr.name()) == Some(name))
+                .map(|attr| attr.identifier())
+        }))
+    }
+
+    ctx.db.walk_tops().flat_map(move |top| match top.ast_top() {
+        Top::CompositeType(composite_type) => find_native_type_locations(name, composite_type.iter_fields()),
+        Top::Model(model) => find_native_type_locations(name, model.iter_fields()),
+
+        Top::Enum(_) | Top::Source(_) | Top::Generator(_) => Box::new(empty_identifiers()),
+    })
+}
+
+fn find_where_used_as_field_type<'ast>(
+    ctx: &'ast ReferencesContext<'_>,
+    name: &'ast str,
+) -> impl Iterator<Item = &'ast Identifier> {
+    fn get_relevent_identifiers<'a>(
+        fields: impl Iterator<Item = (FieldId, &'a Field)>,
+        name: &str,
+    ) -> Vec<&'a Identifier> {
+        fields
+            .filter_map(|(_id, field)| match &field.field_type {
+                FieldType::Supported(id) if id.name == name => Some(id),
+                _ => None,
+            })
+            .collect()
+    }
+
+    ctx.db.walk_tops().flat_map(|top| match top.ast_top() {
+        Top::Model(model) => get_relevent_identifiers(model.iter_fields(), name),
+        Top::CompositeType(composite_type) => get_relevent_identifiers(composite_type.iter_fields(), name),
+        // * Cannot contain field types
+        Top::Enum(_) | Top::Source(_) | Top::Generator(_) => vec![],
+    })
+}
+
+fn find_where_used_as_top_name<'ast>(ctx: &'ast ReferencesContext<'_>, name: &'ast str) -> Option<&'ast Identifier> {
+    ctx.db.find_top(name).map(|top| top.ast_top().identifier())
+}
+
+fn find_where_used_as_ds_name<'ast>(ctx: &'ast ReferencesContext<'_>, name: &'ast str) -> Option<&'ast Identifier> {
+    ctx.db
+        .find_source(name)
+        .map(|source| ctx.db.ast(source.0)[source.1].identifier())
+}
+
+fn extract_ds_from_native_type(attr_name: &str) -> Option<&str> {
+    attr_name.split('.').next()
+}
+
+fn ident_to_location<'ast>(id: &'ast Identifier, ctx: &'ast ReferencesContext<'_>) -> Option<Location> {
+    let file_id = id.span.file_id;
+
+    let source = ctx.db.source(file_id);
+    let range = span_to_range(id.span, source);
+    let file_name = ctx.db.file_name(file_id);
+
+    let uri = if let Ok(uri) = Url::parse(file_name) {
+        uri
+    } else {
+        return None;
+    };
+
+    Some(Location { uri, range })
+}

prisma-fmt/src/text_document_completion.rs

@@ -5,7 +5,8 @@ use lsp_types::*;
 use psl::{
     diagnostics::Span,
     error_tolerant_parse_configuration,
-    parser_database::{ast, ParserDatabase, SourceFile},
+    parser_database::{ast, ParserDatabase, ReferentialAction, SourceFile},
+    schema_ast::ast::AttributePosition,
     Diagnostics, PreviewFeature,
 };
 
@@ -85,19 +86,46 @@ fn push_ast_completions(ctx: CompletionContext<'_>, completion_list: &mut Comple
         .relation_mode()
         .unwrap_or_else(|| ctx.connector().default_relation_mode());
 
-    match ctx.db.ast(ctx.initiating_file_id).find_at_position(position) {
+    let find_at_position = ctx.db.ast(ctx.initiating_file_id).find_at_position(position);
+
+    fn push_referential_action(completion_list: &mut CompletionList, referential_action: ReferentialAction) {
+        completion_list.items.push(CompletionItem {
+            label: referential_action.as_str().to_owned(),
+            kind: Some(CompletionItemKind::ENUM),
+            // what is the difference between detail and documentation?
+            detail: Some(referential_action.documentation().to_owned()),
+            ..Default::default()
+        });
+    }
+
+    match find_at_position {
         ast::SchemaPosition::Model(
             _model_id,
-            ast::ModelPosition::Field(_, ast::FieldPosition::Attribute("relation", _, Some(attr_name))),
+            ast::ModelPosition::Field(
+                _,
+                ast::FieldPosition::Attribute("relation", _, AttributePosition::Argument(attr_name)),
+            ),
         ) if attr_name == "onDelete" || attr_name == "onUpdate" => {
             for referential_action in ctx.connector().referential_actions(&relation_mode).iter() {
-                completion_list.items.push(CompletionItem {
-                    label: referential_action.as_str().to_owned(),
-                    kind: Some(CompletionItemKind::ENUM),
-                    // what is the difference between detail and documentation?
-                    detail: Some(referential_action.documentation().to_owned()),
-                    ..Default::default()
-                });
+                push_referential_action(completion_list, referential_action);
+            }
+        }
+
+        ast::SchemaPosition::Model(
+            _model_id,
+            ast::ModelPosition::Field(
+                _,
+                ast::FieldPosition::Attribute("relation", _, AttributePosition::ArgumentValue(attr_name, value)),
+            ),
+        ) => {
+            if let Some(attr_name) = attr_name {
+                if attr_name == "onDelete" || attr_name == "onUpdate" {
+                    ctx.connector()
+                        .referential_actions(&relation_mode)
+                        .iter()
+                        .filter(|ref_action| ref_action.to_string().starts_with(&value))
+                        .for_each(|ref_action| push_referential_action(completion_list, ref_action));
+                }
             }
         }
 

prisma-fmt/tests/references/mod.rs

@@ -0,0 +1,2 @@
+mod test_api;
+mod tests;

prisma-fmt/tests/references/scenarios/composite_type_as_type/a.prisma

@@ -0,0 +1,14 @@
+generator client {
+  provider        = "prisma-client-js"
+  previewFeatures = ["prismaSchemaFolder", "views"]
+}
+
+datasource db {
+  provider = "mongodb"
+  url      = env("DATABASE_URL")
+}
+
+type Address {
+  city     String
+  postCode String
+}

prisma-fmt/tests/references/scenarios/composite_type_as_type/b.prisma

@@ -0,0 +1,5 @@
+model User {
+  id       String @id @map("_id")
+  authorId String
+  address  Add<|>ress
+}