feat(lsp): hover and goto definition

Cargo.lock

@@ -60,7 +60,7 @@ dependencies = [
  "clap",
  "hex",
  "ignore",
- "indoc",
+ "indoc 1.0.8",
  "miette",
  "owo-colors",
  "pallas-addresses",
@@ -81,7 +81,7 @@ dependencies = [
  "chumsky",
  "hex",
  "indexmap",
- "indoc",
+ "indoc 1.0.8",
  "itertools",
  "miette",
  "ordinal",
@@ -100,6 +100,8 @@ dependencies = [
  "aiken-lang",
  "aiken-project",
  "crossbeam-channel",
+ "indoc 2.0.0",
+ "itertools",
  "lsp-server",
  "lsp-types",
  "miette",
@@ -108,6 +110,7 @@ dependencies = [
  "thiserror",
  "tracing",
  "url",
+ "urlencoding",
 ]
 
 [[package]]
@@ -1082,6 +1085,12 @@ version = "1.0.8"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "da2d6f23ffea9d7e76c53eee25dfb67bcd8fde7f1198b0855350698c9f07c780"
 
+[[package]]
+name = "indoc"
+version = "2.0.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "6fe2b9d82064e8a0226fddb3547f37f28eaa46d0fc210e275d835f08cf3b76a7"
+
 [[package]]
 name = "instant"
 version = "0.1.12"
@@ -2670,6 +2679,12 @@ dependencies = [
  "serde",
 ]
 
+[[package]]
+name = "urlencoding"
+version = "2.1.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "e8db7427f936968176eaa7cdf81b7f98b980b18495ec28f1b5791ac3bfe3eea9"
+
 [[package]]
 name = "utf8parse"
 version = "0.2.0"

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,

crates/aiken-lsp/Cargo.toml

@@ -13,6 +13,8 @@ authors = ["Lucas Rosa <x@rvcas.dev>"]
 aiken-lang = { path = '../aiken-lang', version = "0.0.28" }
 aiken-project = { path = '../aiken-project', version = "0.0.28" }
 crossbeam-channel = "0.5.6"
+indoc = "2.0.0"
+itertools = "0.10.5"
 lsp-server = "0.6.0"
 lsp-types = "0.93.2"
 miette = "5.4.1"
@@ -21,3 +23,4 @@ serde_json = "1.0.87"
 thiserror = "1.0.37"
 tracing = "0.1.37"
 url = "2.3.1"
+urlencoding = "2.1.2"

crates/aiken-lsp/src/lib.rs

@@ -2,10 +2,6 @@ use std::env;
 
 use aiken_project::{config::Config, paths};
 use lsp_server::Connection;
-use lsp_types::{
-    OneOf, SaveOptions, ServerCapabilities, TextDocumentSyncCapability, TextDocumentSyncKind,
-    TextDocumentSyncOptions, TextDocumentSyncSaveOptions,
-};
 
 mod cast;
 pub mod error;
@@ -53,21 +49,32 @@ pub fn start() -> Result<(), Error> {
     Ok(())
 }
 
-fn capabilities() -> ServerCapabilities {
-    ServerCapabilities {
-        text_document_sync: Some(TextDocumentSyncCapability::Options(
-            TextDocumentSyncOptions {
+fn capabilities() -> lsp_types::ServerCapabilities {
+    lsp_types::ServerCapabilities {
+        completion_provider: Some(lsp_types::CompletionOptions {
+            resolve_provider: None,
+            trigger_characters: Some(vec![".".into(), " ".into()]),
+            all_commit_characters: None,
+            work_done_progress_options: lsp_types::WorkDoneProgressOptions {
+                work_done_progress: None,
+            },
+        }),
+        document_formatting_provider: Some(lsp_types::OneOf::Left(true)),
+        definition_provider: Some(lsp_types::OneOf::Left(true)),
+        hover_provider: Some(lsp_types::HoverProviderCapability::Simple(true)),
+        text_document_sync: Some(lsp_types::TextDocumentSyncCapability::Options(
+            lsp_types::TextDocumentSyncOptions {
                 open_close: None,
-                change: Some(TextDocumentSyncKind::FULL),
+                change: Some(lsp_types::TextDocumentSyncKind::FULL),
                 will_save: None,
                 will_save_wait_until: None,
-                save: Some(TextDocumentSyncSaveOptions::SaveOptions(SaveOptions {
+                save: Some(lsp_types::TextDocumentSyncSaveOptions::SaveOptions(
+                    lsp_types::SaveOptions {
                         include_text: Some(false),
-                })),
                     },
                 )),
-        // definition_provider: Some(OneOf::Left(true)),
-        document_formatting_provider: Some(OneOf::Left(true)),
+            },
+        )),
         ..Default::default()
     }
 }

crates/aiken-lsp/src/server.rs

@@ -4,18 +4,24 @@ use std::{
     path::{Path, PathBuf},
 };
 
-use aiken_lang::{ast::ModuleKind, parser};
+use aiken_lang::{
+    ast::{Located, ModuleKind, Span},
+    parser,
+    tipo::pretty::Printer,
+};
 use aiken_project::{
     config,
     error::{Error as ProjectError, GetSource},
+    module::CheckedModule,
 };
+use indoc::formatdoc;
 use lsp_server::{Connection, Message};
 use lsp_types::{
     notification::{
         DidChangeTextDocument, DidSaveTextDocument, Notification, Progress, PublishDiagnostics,
         ShowMessage,
     },
-    request::{Formatting, Request, WorkDoneProgressCreate},
+    request::{Formatting, GotoDefinition, HoverRequest, Request, WorkDoneProgressCreate},
     DocumentFormattingParams, InitializeParams, TextEdit,
 };
 use miette::Diagnostic;
@@ -25,7 +31,8 @@ use crate::{
     error::Error as ServerError,
     line_numbers::LineNumbers,
     utils::{
-        path_to_uri, text_edit_replace, COMPILING_PROGRESS_TOKEN, CREATE_COMPILING_PROGRESS_TOKEN,
+        path_to_uri, span_to_lsp_range, text_edit_replace, uri_to_module_name,
+        COMPILING_PROGRESS_TOKEN, CREATE_COMPILING_PROGRESS_TOKEN,
     },
 };
 
@@ -230,12 +237,138 @@ impl Server {
                     }
                 }
             }
+            HoverRequest::METHOD => {
+                let params = cast_request::<HoverRequest>(request)?;
+
+                let opt_hover = self.hover(params)?;
+
+                Ok(lsp_server::Response {
+                    id,
+                    error: None,
+                    result: Some(serde_json::to_value(opt_hover)?),
+                })
+            }
+
+            GotoDefinition::METHOD => {
+                let params = cast_request::<GotoDefinition>(request)?;
+
+                let location = self.goto_definition(params)?;
+
+                Ok(lsp_server::Response {
+                    id,
+                    error: None,
+                    result: Some(serde_json::to_value(location)?),
+                })
+            }
+
             unsupported => Err(ServerError::UnsupportedLspRequest {
                 request: unsupported.to_string(),
             }),
         }
     }
 
+    fn goto_definition(
+        &self,
+        params: lsp_types::GotoDefinitionParams,
+    ) -> Result<Option<lsp_types::Location>, ServerError> {
+        let params = params.text_document_position_params;
+
+        let (line_numbers, node) = match self.node_at_position(&params) {
+            Some(location) => location,
+            None => return Ok(None),
+        };
+
+        let location = match node.definition_location() {
+            Some(location) => location,
+            None => return Ok(None),
+        };
+
+        let (uri, line_numbers) = match location.module {
+            None => (params.text_document.uri, &line_numbers),
+            Some(name) => {
+                let module = match self
+                    .compiler
+                    .as_ref()
+                    .and_then(|compiler| compiler.sources.get(name))
+                {
+                    Some(module) => module,
+
+                    None => return Ok(None),
+                };
+
+                let url = url::Url::parse(&format!("file:///{}", &module.path))
+                    .expect("goto definition URL parse");
+
+                (url, &module.line_numbers)
+            }
+        };
+
+        let range = span_to_lsp_range(location.span, line_numbers);
+
+        Ok(Some(lsp_types::Location { uri, range }))
+    }
+
+    fn node_at_position(
+        &self,
+        params: &lsp_types::TextDocumentPositionParams,
+    ) -> Option<(LineNumbers, Located<'_>)> {
+        let module = self.module_for_uri(&params.text_document.uri);
+
+        let module = module?;
+
+        let line_numbers = LineNumbers::new(&module.code);
+
+        let byte_index = line_numbers.byte_index(
+            params.position.line as usize,
+            params.position.character as usize,
+        );
+
+        let node = module.find_node(byte_index);
+
+        let node = node?;
+
+        Some((line_numbers, node))
+    }
+
+    fn module_for_uri(&self, uri: &url::Url) -> Option<&CheckedModule> {
+        self.compiler.as_ref().and_then(|compiler| {
+            let module_name = uri_to_module_name(uri, &self.root).expect("uri to module name");
+
+            compiler.modules.get(&module_name)
+        })
+    }
+
+    fn hover(
+        &self,
+        params: lsp_types::HoverParams,
+    ) -> Result<Option<lsp_types::Hover>, ServerError> {
+        let params = params.text_document_position_params;
+
+        let (line_numbers, found) = match self.node_at_position(&params) {
+            Some(value) => value,
+            None => return Ok(None),
+        };
+
+        let expression = match found {
+            Located::Expression(expression) => expression,
+            Located::Definition(_) => return Ok(None),
+        };
+
+        // Show the type of the hovered node to the user
+        let type_ = Printer::new().pretty_print(expression.tipo().as_ref(), 0);
+
+        let contents = formatdoc! {r#"
+            ```aiken
+            {type_}
+            ```
+        "#};
+
+        Ok(Some(lsp_types::Hover {
+            contents: lsp_types::HoverContents::Scalar(lsp_types::MarkedString::String(contents)),
+            range: Some(span_to_lsp_range(expression.location(), &line_numbers)),
+        }))
+    }
+
     pub fn listen(&mut self, connection: Connection) -> Result<(), ServerError> {
         self.create_compilation_progress_token(&connection)?;
         self.start_watching_aiken_toml(&connection)?;
@@ -366,21 +499,13 @@ impl Server {
 
                 let line_numbers = LineNumbers::new(&src);
 
-                let start = line_numbers.line_and_column_number(labeled_span.inner().offset());
-                let end = line_numbers.line_and_column_number(
-                    labeled_span.inner().offset() + labeled_span.inner().len(),
-                );
-
                 let lsp_diagnostic = lsp_types::Diagnostic {
-                    range: lsp_types::Range::new(
-                        lsp_types::Position {
-                            line: start.line as u32 - 1,
-                            character: start.column as u32 - 1,
-                        },
-                        lsp_types::Position {
-                            line: end.line as u32 - 1,
-                            character: end.column as u32 - 1,
+                    range: span_to_lsp_range(
+                        Span {
+                            start: labeled_span.inner().offset(),
+                            end: labeled_span.inner().offset() + labeled_span.inner().len(),
                         },
+                        &line_numbers,
                     ),
                     severity: Some(severity),
                     code: error

crates/aiken-lsp/src/utils.rs

@@ -1,8 +1,11 @@
-use std::path::PathBuf;
+use std::path::{Path, PathBuf};
 
+use aiken_lang::ast::Span;
+use itertools::Itertools;
 use lsp_types::TextEdit;
+use urlencoding::decode;
 
-use crate::error::Error;
+use crate::{error::Error, line_numbers::LineNumbers};
 
 pub const COMPILING_PROGRESS_TOKEN: &str = "compiling-gleam";
 pub const CREATE_COMPILING_PROGRESS_TOKEN: &str = "create-compiling-progress-token";
@@ -32,3 +35,52 @@ pub fn path_to_uri(path: PathBuf) -> Result<lsp_types::Url, Error> {
 
     Ok(uri)
 }
+
+pub fn span_to_lsp_range(location: Span, line_numbers: &LineNumbers) -> lsp_types::Range {
+    let start = line_numbers.line_and_column_number(location.start);
+    let end = line_numbers.line_and_column_number(location.end);
+
+    lsp_types::Range {
+        start: lsp_types::Position {
+            line: start.line as u32 - 1,
+            character: start.column as u32 - 1,
+        },
+        end: lsp_types::Position {
+            line: end.line as u32 - 1,
+            character: end.column as u32 - 1,
+        },
+    }
+}
+
+pub fn uri_to_module_name(uri: &url::Url, root: &Path) -> Option<String> {
+    let path = if cfg!(target_os = "windows") {
+        let mut uri_path = decode(&uri.path().replace('/', "\\"))
+            .expect("Invalid formatting")
+            .to_string();
+
+        if uri_path.starts_with('\\') {
+            uri_path = uri_path
+                .strip_prefix('\\')
+                .expect("Failed to remove \"\\\" prefix")
+                .to_string();
+        }
+
+        PathBuf::from(uri_path)
+    } else {
+        PathBuf::from(uri.path())
+    };
+
+    let components = path
+        .strip_prefix(root)
+        .ok()?
+        .components()
+        .skip(1)
+        .map(|c| c.as_os_str().to_string_lossy());
+
+    let module_name = Itertools::intersperse(components, "/".into())
+        .collect::<String>()
+        .strip_suffix(".ak")?
+        .to_string();
+
+    Some(module_name)
+}

crates/aiken-project/src/blueprint/validator.rs

@@ -230,7 +230,8 @@ mod test {
             let (mut ast, extra) =
                 parser::module(source_code, kind).expect("Failed to parse module");
             ast.name = name.clone();
-            let mut module = ParsedModule {
+
+            ParsedModule {
                 kind,
                 ast,
                 code: source_code.to_string(),
@@ -238,9 +239,7 @@ mod test {
                 path: PathBuf::new(),
                 extra,
                 package: self.package.to_string(),
-            };
-            module.attach_doc_and_module_comments();
-            module
+            }
         }
 
         fn check(&mut self, module: ParsedModule) -> CheckedModule {
@@ -261,7 +260,7 @@ mod test {
             self.module_types
                 .insert(module.name.clone(), ast.type_info.clone());
 
-            CheckedModule {
+            let mut checked_module = CheckedModule {
                 kind: module.kind,
                 extra: module.extra,
                 name: module.name,
@@ -269,7 +268,11 @@ mod test {
                 package: module.package,
                 input_path: module.path,
                 ast,
-            }
+            };
+
+            checked_module.attach_doc_and_module_comments();
+
+            checked_module
         }
     }
 

crates/aiken-project/src/lib.rs

@@ -467,7 +467,7 @@ where
                     // Store the name
                     ast.name = name.clone();
 
-                    let mut module = ParsedModule {
+                    let module = ParsedModule {
                         kind,
                         ast,
                         code,
@@ -489,8 +489,6 @@ where
                         .into());
                     }
 
-                    module.attach_doc_and_module_comments();
-
                     parsed_modules.insert(module.name.clone(), module);
                 }
                 Err(errs) => {
@@ -561,18 +559,19 @@ where
                 self.module_types
                     .insert(name.clone(), ast.type_info.clone());
 
-                self.checked_modules.insert(
-                    name.clone(),
-                    CheckedModule {
+                let mut checked_module = CheckedModule {
                     kind,
                     extra,
-                        name,
+                    name: name.clone(),
                     code,
                     ast,
                     package,
                     input_path: path,
-                    },
-                );
+                };
+
+                checked_module.attach_doc_and_module_comments();
+
+                self.checked_modules.insert(name, checked_module);
             }
         }
 

crates/aiken-project/src/module.rs

@@ -1,7 +1,7 @@
 use crate::error::Error;
 use aiken_lang::{
     ast::{
-        DataType, Definition, ModuleKind, TypedDataType, TypedFunction, TypedModule,
+        DataType, Definition, Located, ModuleKind, TypedDataType, TypedFunction, TypedModule,
         TypedValidator, UntypedModule,
     },
     builder::{DataTypeKey, FunctionAccessKey},
@@ -41,55 +41,6 @@ impl ParsedModule {
 
         (name, deps)
     }
-
-    pub fn attach_doc_and_module_comments(&mut self) {
-        // Module Comments
-        self.ast.docs = self
-            .extra
-            .module_comments
-            .iter()
-            .map(|span| {
-                Comment::from((span, self.code.as_str()))
-                    .content
-                    .to_string()
-            })
-            .collect();
-
-        // Order definitions to avoid dissociating doc comments from them
-        let mut definitions: Vec<_> = self.ast.definitions.iter_mut().collect();
-        definitions.sort_by(|a, b| a.location().start.cmp(&b.location().start));
-
-        // Doc Comments
-        let mut doc_comments = self.extra.doc_comments.iter().peekable();
-        for def in &mut definitions {
-            let docs: Vec<&str> =
-                comments_before(&mut doc_comments, def.location().start, &self.code);
-            if !docs.is_empty() {
-                let doc = docs.join("\n");
-                def.put_doc(doc);
-            }
-
-            if let Definition::DataType(DataType { constructors, .. }) = def {
-                for constructor in constructors {
-                    let docs: Vec<&str> =
-                        comments_before(&mut doc_comments, constructor.location.start, &self.code);
-                    if !docs.is_empty() {
-                        let doc = docs.join("\n");
-                        constructor.put_doc(doc);
-                    }
-
-                    for argument in constructor.arguments.iter_mut() {
-                        let docs: Vec<&str> =
-                            comments_before(&mut doc_comments, argument.location.start, &self.code);
-                        if !docs.is_empty() {
-                            let doc = docs.join("\n");
-                            argument.put_doc(doc);
-                        }
-                    }
-                }
-            }
-        }
-    }
 }
 
 pub struct ParsedModules(HashMap<String, ParsedModule>);
@@ -223,6 +174,61 @@ pub struct CheckedModule {
     pub extra: ModuleExtra,
 }
 
+impl CheckedModule {
+    pub fn find_node(&self, byte_index: usize) -> Option<Located<'_>> {
+        self.ast.find_node(byte_index)
+    }
+
+    pub fn attach_doc_and_module_comments(&mut self) {
+        // Module Comments
+        self.ast.docs = self
+            .extra
+            .module_comments
+            .iter()
+            .map(|span| {
+                Comment::from((span, self.code.as_str()))
+                    .content
+                    .to_string()
+            })
+            .collect();
+
+        // Order definitions to avoid dissociating doc comments from them
+        let mut definitions: Vec<_> = self.ast.definitions.iter_mut().collect();
+        definitions.sort_by(|a, b| a.location().start.cmp(&b.location().start));
+
+        // Doc Comments
+        let mut doc_comments = self.extra.doc_comments.iter().peekable();
+        for def in &mut definitions {
+            let docs: Vec<&str> =
+                comments_before(&mut doc_comments, def.location().start, &self.code);
+            if !docs.is_empty() {
+                let doc = docs.join("\n");
+                def.put_doc(doc);
+            }
+
+            if let Definition::DataType(DataType { constructors, .. }) = def {
+                for constructor in constructors {
+                    let docs: Vec<&str> =
+                        comments_before(&mut doc_comments, constructor.location.start, &self.code);
+                    if !docs.is_empty() {
+                        let doc = docs.join("\n");
+                        constructor.put_doc(doc);
+                    }
+
+                    for argument in constructor.arguments.iter_mut() {
+                        let docs: Vec<&str> =
+                            comments_before(&mut doc_comments, argument.location.start, &self.code);
+                        if !docs.is_empty() {
+                            let doc = docs.join("\n");
+                            argument.put_doc(doc);
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
 #[derive(Default, Debug, Clone)]
 pub struct CheckedModules(HashMap<String, CheckedModule>);