From eb566e2125e847a3e3efbd2bc15a88a1c454a7df Mon Sep 17 00:00:00 2001
From: Koby Hall <102518238+kobyhallx@users.noreply.github.com>
Date: Thu, 4 Jan 2024 19:27:57 +0100
Subject: [PATCH] feat(lsp): goto trait from trait impl (#3956)

# Description

## Problem\*

Resolves
feat(lsp): goto definition of trait #3725

## Summary\*

Allows to go to Trait definition from Trait Implementation.

## Additional Context



## Documentation\*

Check one:
- [x] No documentation needed.
- [ ] Documentation included in this PR.
- [ ] **[Exceptional Case]** Documentation to be submitted in a separate
PR.

# PR Checklist\*

- [x] I have tested the changes locally.
- [x] I have formatted the changes with [Prettier](https://prettier.io/)
and/or `cargo fmt` on default settings.
---
 compiler/noirc_frontend/src/hir_def/traits.rs |  4 +--
 compiler/noirc_frontend/src/node_interner.rs  | 26 +++++++++++++++++--
 2 files changed, 26 insertions(+), 4 deletions(-)

diff --git a/compiler/noirc_frontend/src/hir_def/traits.rs b/compiler/noirc_frontend/src/hir_def/traits.rs
index ea9c2e2928..7bf5870e1f 100644
--- a/compiler/noirc_frontend/src/hir_def/traits.rs
+++ b/compiler/noirc_frontend/src/hir_def/traits.rs
@@ -6,7 +6,7 @@ use crate::{
     Generics, Ident, NoirFunction, Type, TypeVariable, TypeVariableId,
 };
 use fm::FileId;
-use noirc_errors::Span;
+use noirc_errors::{Location, Span};
 
 #[derive(Clone, Debug, PartialEq, Eq)]
 pub struct TraitFunction {
@@ -56,7 +56,7 @@ pub struct Trait {
 
     pub name: Ident,
     pub generics: Generics,
-    pub span: Span,
+    pub location: Location,
 
     /// When resolving the types of Trait elements, all references to `Self` resolve
     /// to this TypeVariable. Then when we check if the types of trait impl elements
diff --git a/compiler/noirc_frontend/src/node_interner.rs b/compiler/noirc_frontend/src/node_interner.rs
index 9082df1bcd..261b11158c 100644
--- a/compiler/noirc_frontend/src/node_interner.rs
+++ b/compiler/noirc_frontend/src/node_interner.rs
@@ -468,6 +468,7 @@ impl NodeInterner {
     /// The [Location] may not necessarily point to the beginning of the item
     /// so we check if the location's span is contained within the start or end
     /// of each items [Span]
+    #[tracing::instrument(skip(self))]
     pub fn find_location_index(&self, location: Location) -> Option<impl Into<Index>> {
         let mut location_candidate: Option<(&Index, &Location)> = None;
 
@@ -504,7 +505,7 @@ impl NodeInterner {
             id: type_id,
             name: unresolved_trait.trait_def.name.clone(),
             crate_id: unresolved_trait.crate_id,
-            span: unresolved_trait.trait_def.span,
+            location: Location::new(unresolved_trait.trait_def.span, unresolved_trait.file_id),
             generics: vecmap(&unresolved_trait.trait_def.generics, |_| {
                 // Temporary type variable ids before the trait is resolved to its actual ids.
                 // This lets us record how many arguments the type expects so that other types
@@ -1142,6 +1143,7 @@ impl NodeInterner {
     }
 
     /// Adds a trait implementation to the list of known implementations.
+    #[tracing::instrument(skip(self))]
     pub fn add_trait_implementation(
         &mut self,
         object_type: Type,
@@ -1274,7 +1276,9 @@ impl NodeInterner {
     /// Returns the [Location] of the definition of the given Ident found at [Span] of the given [FileId].
     /// Returns [None] when definition is not found.
     pub fn get_definition_location_from(&self, location: Location) -> Option<Location> {
-        self.find_location_index(location).and_then(|index| self.resolve_location(index))
+        self.find_location_index(location)
+            .and_then(|index| self.resolve_location(index))
+            .or_else(|| self.try_resolve_trait_impl_location(location))
     }
 
     /// For a given [Index] we return [Location] to which we resolved to
@@ -1429,6 +1433,24 @@ impl NodeInterner {
     pub(crate) fn ordering_type(&self) -> Type {
         self.ordering_type.clone().expect("Expected ordering_type to be set in the NodeInterner")
     }
+
+    /// Attempts to resolve [Location] of [Trait] based on [Location] of [TraitImpl]
+    /// This is used by LSP to resolve the location of a trait based on the location of a trait impl.
+    ///
+    /// Example:
+    /// impl Foo for Bar { ... } -> trait Foo { ... }
+    fn try_resolve_trait_impl_location(&self, location: Location) -> Option<Location> {
+        self.trait_implementations
+            .iter()
+            .find(|shared_trait_impl| {
+                let trait_impl = shared_trait_impl.borrow();
+                trait_impl.file == location.file && trait_impl.ident.span().contains(&location.span)
+            })
+            .and_then(|shared_trait_impl| {
+                let trait_impl = shared_trait_impl.borrow();
+                self.traits.get(&trait_impl.trait_id).map(|trait_| trait_.location)
+            })
+    }
 }
 
 impl Methods {