Implement modules' extra, to get function & module comments in docs.

crates/project/src/lib.rs

@@ -108,7 +108,10 @@ where
         self.read_source_files()?;
 
         let destination = destination.unwrap_or(self.root.join("doc"));
-        let parsed_modules = self.parse_sources()?;
+        let mut parsed_modules = self.parse_sources()?;
+        for (_, module) in parsed_modules.iter_mut() {
+            module.attach_doc_and_module_comments();
+        }
         let checked_modules = self.type_check(parsed_modules)?;
         self.event_listener.handle_event(Event::GeneratingDocFiles {
             output_path: destination.clone(),
@@ -229,7 +232,7 @@ where
         } in self.sources.drain(0..)
         {
             match aiken_lang::parser::module(&code, kind) {
-                Ok((mut ast, _)) => {
+                Ok((mut ast, extra)) => {
                     // Store the name
                     ast.name = name.clone();
 
@@ -239,6 +242,7 @@ where
                         code,
                         name,
                         path,
+                        extra,
                         package: self.config.name.clone(),
                     };
 
@@ -286,6 +290,7 @@ where
                 path,
                 code,
                 kind,
+                extra,
                 // TODO: come back and figure out where to use this
                 package: _package,
                 ast,
@@ -324,7 +329,7 @@ where
                     name.clone(),
                     CheckedModule {
                         kind,
-                        // extra,
+                        extra,
                         name,
                         code,
                         ast,

crates/project/src/module.rs

@@ -4,7 +4,10 @@ use std::{
     path::PathBuf,
 };
 
-use aiken_lang::ast::{Definition, ModuleKind, TypedModule, UntypedModule};
+use aiken_lang::{
+    ast::{DataType, Definition, ModuleKind, TypedModule, UntypedModule},
+    parser::extra::{comments_before, Comment, ModuleExtra},
+};
 use petgraph::{algo, graph::NodeIndex, Direction, Graph};
 
 use crate::error::Error;
@@ -17,7 +20,7 @@ pub struct ParsedModule {
     pub kind: ModuleKind,
     pub package: String,
     pub ast: UntypedModule,
-    // extra: ModuleExtra,
+    pub extra: ModuleExtra,
 }
 
 impl ParsedModule {
@@ -33,6 +36,55 @@ 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>);
@@ -168,7 +220,7 @@ pub struct CheckedModule {
     pub input_path: PathBuf,
     pub kind: ModuleKind,
     pub ast: TypedModule,
-    // pub extra: ModuleExtra,
+    pub extra: ModuleExtra,
 }
 
 impl CheckedModule {