feat(lsp): hover and goto definition

crates/aiken-lang/src/ast.rs

@@ -66,6 +66,14 @@ impl UntypedModule {
     }
 }
 
+impl TypedModule {
+    pub fn find_node(&self, byte_index: usize) -> Option<Located<'_>> {
+        self.definitions
+            .iter()
+            .find_map(|definition| definition.find_node(byte_index))
+    }
+}
+
 pub type TypedFunction = Function<Arc<Type>, TypedExpr>;
 pub type UntypedFunction = Function<(), UntypedExpr>;
 
@@ -318,6 +326,48 @@ impl<A, B, C> Definition<A, B, C> {
     }
 }
 
+impl TypedDefinition {
+    pub fn find_node(&self, byte_index: usize) -> Option<Located<'_>> {
+        // Note that the fn span covers the function head, not
+        // the entire statement.
+        if let Definition::Fn(Function { body, .. })
+        | Definition::Validator(Validator {
+            fun: Function { body, .. },
+            ..
+        })
+        | Definition::Test(Function { body, .. }) = self
+        {
+            if let Some(expression) = body.find_node(byte_index) {
+                return Some(Located::Expression(expression));
+            }
+        }
+
+        if self.location().contains(byte_index) {
+            Some(Located::Definition(self))
+        } else {
+            None
+        }
+    }
+}
+
+#[derive(Debug, Clone, PartialEq)]
+pub enum Located<'a> {
+    Expression(&'a TypedExpr),
+    Definition(&'a TypedDefinition),
+}
+
+impl<'a> Located<'a> {
+    pub fn definition_location(&self) -> Option<DefinitionLocation<'_>> {
+        match self {
+            Self::Expression(expression) => expression.definition_location(),
+            Self::Definition(definition) => Some(DefinitionLocation {
+                module: None,
+                span: definition.location(),
+            }),
+        }
+    }
+}
+
 #[derive(Debug, PartialEq, Eq, Clone)]
 pub struct DefinitionLocation<'module> {
     pub module: Option<&'module str>,
@@ -379,6 +429,12 @@ impl CallArg<UntypedExpr> {
     }
 }
 
+impl TypedCallArg {
+    pub fn find_node(&self, byte_index: usize) -> Option<&TypedExpr> {
+        self.value.find_node(byte_index)
+    }
+}
+
 #[derive(Debug, Clone, PartialEq)]
 pub struct RecordConstructor<T> {
     pub location: Span,
@@ -816,6 +872,10 @@ impl TypedClause {
             end: self.then.location().end,
         }
     }
+
+    pub fn find_node(&self, byte_index: usize) -> Option<&TypedExpr> {
+        self.then.find_node(byte_index)
+    }
 }
 
 pub type UntypedClauseGuard = ClauseGuard<()>;
@@ -946,7 +1006,7 @@ pub struct IfBranch<Expr> {
     pub location: Span,
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub struct TypedRecordUpdateArg {
     pub label: String,
     pub location: Span,
@@ -954,6 +1014,12 @@ pub struct TypedRecordUpdateArg {
     pub index: usize,
 }
 
+impl TypedRecordUpdateArg {
+    pub fn find_node(&self, byte_index: usize) -> Option<&TypedExpr> {
+        self.value.find_node(byte_index)
+    }
+}
+
 #[derive(Debug, Clone, PartialEq)]
 pub struct UntypedRecordUpdateArg {
     pub label: String,
@@ -1023,6 +1089,10 @@ impl Span {
             end: self.end().max(other.end()),
         }
     }
+
+    pub fn contains(&self, byte_index: usize) -> bool {
+        byte_index >= self.start && byte_index < self.end
+    }
 }
 
 impl fmt::Debug for Span {

crates/aiken-lang/src/expr.rs

@@ -12,7 +12,7 @@ use crate::{
     tipo::{ModuleValueConstructor, PatternConstructor, Type, ValueConstructor},
 };
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub enum TypedExpr {
     Int {
         location: Span,
@@ -307,6 +307,96 @@ impl TypedExpr {
             | Self::RecordUpdate { location, .. } => *location,
         }
     }
+
+    // This could be optimised in places to exit early if the first of a series
+    // of expressions is after the byte index.
+    pub fn find_node(&self, byte_index: usize) -> Option<&Self> {
+        if !self.location().contains(byte_index) {
+            return None;
+        }
+
+        match self {
+            TypedExpr::ErrorTerm { .. }
+            | TypedExpr::Var { .. }
+            | TypedExpr::Int { .. }
+            | TypedExpr::String { .. }
+            | TypedExpr::ByteArray { .. }
+            | TypedExpr::ModuleSelect { .. } => Some(self),
+
+            TypedExpr::Trace { text, then, .. } => text
+                .find_node(byte_index)
+                .or_else(|| then.find_node(byte_index))
+                .or(Some(self)),
+
+            TypedExpr::Pipeline { expressions, .. } | TypedExpr::Sequence { expressions, .. } => {
+                expressions.iter().find_map(|e| e.find_node(byte_index))
+            }
+
+            TypedExpr::Fn { body, .. } => body.find_node(byte_index).or(Some(self)),
+
+            TypedExpr::Tuple {
+                elems: elements, ..
+            }
+            | TypedExpr::List { elements, .. } => elements
+                .iter()
+                .find_map(|e| e.find_node(byte_index))
+                .or(Some(self)),
+
+            TypedExpr::Call { fun, args, .. } => args
+                .iter()
+                .find_map(|arg| arg.find_node(byte_index))
+                .or_else(|| fun.find_node(byte_index))
+                .or(Some(self)),
+
+            TypedExpr::BinOp { left, right, .. } => left
+                .find_node(byte_index)
+                .or_else(|| right.find_node(byte_index)),
+
+            TypedExpr::Assignment { value, .. } => value.find_node(byte_index),
+
+            TypedExpr::When {
+                subjects, clauses, ..
+            } => subjects
+                .iter()
+                .find_map(|subject| subject.find_node(byte_index))
+                .or_else(|| {
+                    clauses
+                        .iter()
+                        .find_map(|clause| clause.find_node(byte_index))
+                })
+                .or(Some(self)),
+
+            TypedExpr::RecordAccess {
+                record: expression, ..
+            }
+            | TypedExpr::TupleIndex {
+                tuple: expression, ..
+            } => expression.find_node(byte_index).or(Some(self)),
+
+            TypedExpr::RecordUpdate { spread, args, .. } => args
+                .iter()
+                .find_map(|arg| arg.find_node(byte_index))
+                .or_else(|| spread.find_node(byte_index))
+                .or(Some(self)),
+
+            TypedExpr::If {
+                branches,
+                final_else,
+                ..
+            } => branches
+                .iter()
+                .find_map(|branch| {
+                    branch
+                        .condition
+                        .find_node(byte_index)
+                        .or_else(|| branch.body.find_node(byte_index))
+                })
+                .or_else(|| final_else.find_node(byte_index))
+                .or(Some(self)),
+
+            TypedExpr::UnOp { value, .. } => value.find_node(byte_index).or(Some(self)),
+        }
+    }
 }
 
 #[derive(Debug, Clone, PartialEq)]

crates/aiken-lang/src/tipo.rs

@@ -715,7 +715,7 @@ pub enum PatternConstructor {
     },
 }
 
-#[derive(Debug, Clone)]
+#[derive(Debug, Clone, PartialEq)]
 pub enum ModuleValueConstructor {
     Record {
         name: String,