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

crates/lang/src/ast.rs

@@ -356,6 +356,12 @@ pub struct RecordConstructor<T> {
     pub sugar: bool,
 }
 
+impl<A> RecordConstructor<A> {
+    pub fn put_doc(&mut self, new_doc: String) {
+        self.doc = Some(new_doc);
+    }
+}
+
 #[derive(Debug, Clone, PartialEq)]
 pub struct RecordConstructorArg<T> {
     pub label: Option<String>,
@@ -366,6 +372,12 @@ pub struct RecordConstructorArg<T> {
     pub doc: Option<String>,
 }
 
+impl<T: PartialEq> RecordConstructorArg<T> {
+    pub fn put_doc(&mut self, new_doc: String) {
+        self.doc = Some(new_doc);
+    }
+}
+
 pub type TypedArg = Arg<Arc<Type>>;
 pub type UntypedArg = Arg<()>;
 

crates/lang/src/parser.rs

@@ -33,22 +33,18 @@ pub fn module(
     let tokens = tokens.into_iter().filter(|(token, span)| match token {
         Token::ModuleComment => {
             extra.module_comments.push(*span);
-
             false
         }
         Token::DocComment => {
             extra.doc_comments.push(*span);
-
             false
         }
         Token::Comment => {
             extra.comments.push(*span);
-
             false
         }
         Token::EmptyLine => {
             extra.empty_lines.push(span.start);
-
             false
         }
         Token::NewLine => false,

crates/lang/src/parser/extra.rs

@@ -1,4 +1,5 @@
 use crate::ast::Span;
+use std::iter::Peekable;
 
 #[derive(Debug, PartialEq, Eq, Default, Clone)]
 pub struct ModuleExtra {
@@ -33,3 +34,22 @@ impl<'a> From<(&Span, &'a str)> for Comment<'a> {
         }
     }
 }
+
+pub fn comments_before<'a>(
+    comment_spans: &mut Peekable<impl Iterator<Item = &'a Span>>,
+    byte: usize,
+    src: &'a str,
+) -> Vec<&'a str> {
+    let mut comments = vec![];
+    while let Some(Span { start, .. }) = comment_spans.peek() {
+        if start <= &byte {
+            let comment = comment_spans
+                .next()
+                .expect("Comment before accessing next span");
+            comments.push(Comment::from((comment, src)).content)
+        } else {
+            break;
+        }
+    }
+    comments
+}

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 {