Add support for docstrings.

lambdaclass · Jan 14, 2025 · abe73e5 · abe73e5
1 parent 6fe75de
commit abe73e5
Show file tree

Hide file tree

Showing 4 changed files with 77 additions and 8 deletions.
diff --git a/src/ast/common.rs b/src/ast/common.rs
@@ -20,7 +20,8 @@ impl From<Span> for Range<usize> {
 
 #[derive(Clone, Debug, Eq, Hash, PartialEq, PartialOrd, Ord)]
 pub struct DocString {
-    contents: String,
+    pub contents: Vec<String>,
+    pub span: Span,
 }
 
 /// Identifiers, without a path or generics.

diff --git a/src/grammar.lalrpop b/src/grammar.lalrpop
@@ -43,6 +43,7 @@ extern {
     "boolean" => Token::Boolean(<bool>),
 
     // Other
+    "docstring" => Token::DocString(<String>),
 
     "(" => Token::LeftParen,
     ")" => Token::RightParen,
@@ -142,6 +143,13 @@ SemiColonSeparated<T>: Vec<T> = {
 
 // -- Common
 
+DocString: ast::common::DocString = {
+  <lo:@L> <contents:"docstring"+> <hi:@R> => ast::common::DocString {
+    contents,
+    span: ast::common::Span::new(lo, hi),
+  },
+};
+
 Ident: ast::common::Ident = {
   <lo:@L> <name:"identifier"> <hi:@R> => ast::common::Ident {
     name,
@@ -223,9 +231,9 @@ ExternalModule: ast::common::Ident = {
 }
 
 Module: ast::modules::Module = {
-  <lo:@L> "mod" <name:Ident> "{" <external_modules:List<ExternalModule>?> <imports:ImportList?> <contents:ModuleItems?> "}" <hi:@R> => {
+  <lo:@L> <doc_string:DocString?> "mod" <name:Ident> "{" <external_modules:List<ExternalModule>?> <imports:ImportList?> <contents:ModuleItems?> "}" <hi:@R> => {
     ast::modules::Module {
-      doc_string: None,
+      doc_string,
       imports: imports.unwrap_or_else(Vec::new),
       external_modules: external_modules.unwrap_or(vec![]),
       name,
@@ -293,10 +301,10 @@ ModuleDefItem: ast::modules::ModuleDefItem = {
 // Constants
 
 ConstantDef: ast::constants::ConstantDef = {
-  <is_pub:"pub"?> "const" <name:Ident> ":" <type_spec:TypeDescriptor> "=" <exp:Expression> => {
+  <doc_string:DocString?> <is_pub:"pub"?> "const" <name:Ident> ":" <type_spec:TypeDescriptor> "=" <exp:Expression> => {
     ast::constants::ConstantDef {
       decl: ast::constants::ConstantDecl {
-        doc_string: None,
+        doc_string,
         name,
         is_pub: is_pub.is_some(),
         r#type: type_spec
@@ -338,11 +346,11 @@ Attribute: ast::common::Attribute = {
 }
 
 FunctionDecl: ast::functions::FunctionDecl = {
-  <lo:@L> <attributes:List<Attribute>?> <is_pub:"pub"?> <is_extern:"extern"?>
+  <lo:@L> <doc_string:DocString?> <attributes:List<Attribute>?> <is_pub:"pub"?> <is_extern:"extern"?>
       "fn" <name:Ident> <generic_params:GenericParams?> "(" <params:Comma<Param>> ")"
         <ret_type:FunctionRetType?> <hi:@R> =>
     ast::functions::FunctionDecl {
-        doc_string: None,
+        doc_string,
         generic_params: generic_params.unwrap_or(vec![]),
         attributes: attributes.unwrap_or(vec![]),
         name,

diff --git a/src/parser/mod.rs b/src/parser/mod.rs
@@ -42,6 +42,7 @@ pub fn parse_ast<'db>(db: &'db dyn salsa::Database, source: ProgramSource<'db>)
 #[cfg(test)]
 mod tests {
     use super::{grammar, lexer::Lexer};
+    use crate::ast;
 
     #[test]
     fn parse_simple_program() {
@@ -358,4 +359,60 @@ mod ModuleName {
         let parser = grammar::ProgramParser::new();
         parser.parse(lexer).unwrap();
     }
+
+    #[test]
+    fn parse_doc_strings() {
+        let source = r##"
+/// Documentation for `MyMod`.
+mod MyMod {
+    /// The PI number.
+    ///
+    /// Note: Rounded to an integer.
+    const ROUNDED_PI: u8 = 3;
+
+    /// Increment a number.
+    pub fn my_func(x: u64) -> u64 {
+        return x + 1;
+    }
+}
+"##;
+        let lexer = Lexer::new(source);
+        let parser = grammar::ProgramParser::new();
+        let module = parser.parse(lexer).unwrap();
+
+        let const_item = match &module.modules[0].contents[0] {
+            ast::modules::ModuleDefItem::Constant(x) => x,
+            _ => unreachable!(),
+        };
+        let fn_item = match &module.modules[0].contents[1] {
+            ast::modules::ModuleDefItem::Function(x) => x,
+            _ => unreachable!(),
+        };
+
+        assert_eq!(
+            module.modules[0].doc_string,
+            Some(ast::common::DocString {
+                contents: vec![" Documentation for `MyMod`.".to_string()],
+                span: ast::common::Span::new(1, 31),
+            }),
+        );
+        assert_eq!(
+            const_item.decl.doc_string,
+            Some(ast::common::DocString {
+                contents: vec![
+                    " The PI number.".to_string(),
+                    "".to_string(),
+                    " Note: Rounded to an integer.".to_string(),
+                ],
+                span: ast::common::Span::new(48, 111),
+            }),
+        );
+        assert_eq!(
+            fn_item.decl.doc_string,
+            Some(ast::common::DocString {
+                contents: vec![" Increment a number.".to_string()],
+                span: ast::common::Span::new(147, 170),
+            }),
+        );
+    }
 }
diff --git a/src/parser/tokens.rs b/src/parser/tokens.rs
@@ -21,7 +21,7 @@ impl From<Infallible> for LexingError {
 }
 
 #[derive(Logos, logos_display::Debug, logos_display::Display, PartialEq, Clone)]
-#[logos(error = LexingError, skip r"[ \t\n\f]+", skip r"//[^\n]*", skip r"/\*(?:[^*]|\*[^/])*\*/")]
+#[logos(error = LexingError, skip r"[ \t\n\f]+", skip r"//[^/][^\n]*", skip r"/\*(?:[^*]|\*[^/])*\*/")]
 pub enum Token {
     #[token("let")]
     KeywordLet,
@@ -154,4 +154,7 @@ pub enum Token {
     Ampersand,
     #[token("|")]
     OperatorBitwiseOr,
+
+    #[regex("///[^\n]*", |lex| lex.slice().strip_prefix("///").unwrap().to_string())]
+    DocString(String),
 }