feat: add support for doc comments

rust/candid/src/pretty/candid.rs

@@ -5,6 +5,8 @@ use crate::types::{
 };
 use pretty::RcDoc;
 
+const DOC_COMMENT_LINE_PREFIX: &str = "/// ";
+
 static KEYWORDS: [&str; 30] = [
     "import",
     "service",
@@ -141,7 +143,9 @@ pub(crate) fn pp_field(field: &TypeField, is_variant: bool) -> RcDoc {
     } else {
         kwd(" :").append(pp_ty(&field.typ))
     };
-    pp_label(&field.label).append(ty_doc)
+    pp_doc_comment(field.doc_comment.as_ref())
+        .append(pp_label(&field.label))
+        .append(ty_doc)
 }
 
 fn pp_fields(fs: &[TypeField], is_variant: bool) -> RcDoc {
@@ -189,22 +193,32 @@ pub fn pp_modes(modes: &[FuncMode]) -> RcDoc {
 
 fn pp_service(serv: &[Binding]) -> RcDoc {
     let doc = concat(
-        serv.iter().map(|Binding { id, typ }| {
-            let func_doc = match typ {
-                IDLType::FuncT(ref f) => pp_function(f),
-                IDLType::VarT(_) => pp_ty(typ),
-                _ => unreachable!(),
-            };
-            pp_text(id).append(kwd(" :")).append(func_doc)
-        }),
+        serv.iter().map(
+            |Binding {
+                 id,
+                 typ,
+                 doc_comment,
+             }| {
+                let func_doc = match typ {
+                    IDLType::FuncT(ref f) => pp_function(f),
+                    IDLType::VarT(_) => pp_ty(typ),
+                    _ => unreachable!(),
+                };
+                pp_doc_comment(doc_comment.as_ref())
+                    .append(pp_text(id))
+                    .append(kwd(" :"))
+                    .append(func_doc)
+            },
+        ),
         ";",
     );
     enclose_space("{", doc, "}")
 }
 
 fn pp_defs(env: &IDLMergedProg) -> RcDoc {
-    lines(env.get_types().iter().map(|(id, typ)| {
-        kwd("type")
+    lines(env.get_types().iter().map(|(id, typ, doc_comment)| {
+        pp_doc_comment(*doc_comment)
+            .append(kwd("type"))
             .append(ident(id))
             .append(kwd("="))
             .append(pp_ty(typ))
@@ -220,6 +234,20 @@ fn pp_actor(ty: &IDLType) -> RcDoc {
     }
 }
 
+fn pp_doc_comment(comment_lines: Option<&Vec<String>>) -> RcDoc {
+    let mut doc_comment = RcDoc::nil();
+    if let Some(comment_lines) = comment_lines {
+        for line in comment_lines {
+            doc_comment = doc_comment.append(
+                RcDoc::text(DOC_COMMENT_LINE_PREFIX)
+                    .append(line)
+                    .append(RcDoc::hardline()),
+            );
+        }
+    }
+    doc_comment
+}
+
 pub fn pp_init_args<'a>(env: &'a IDLMergedProg, args: &'a [IDLArgType]) -> RcDoc<'a> {
     pp_defs(env).append(pp_args(args))
 }
@@ -228,7 +256,9 @@ pub fn compile(env: &IDLMergedProg) -> String {
         None => pp_defs(env).pretty(LINE_WIDTH).to_string(),
         Some(actor) => {
             let defs = pp_defs(env);
-            let actor = kwd("service :").append(pp_actor(actor));
+            let actor = pp_doc_comment(actor.doc_comment.as_ref())
+                .append(kwd("service :"))
+                .append(pp_actor(&actor.typ));
             let doc = defs.append(actor);
             doc.pretty(LINE_WIDTH).to_string()
         }

rust/candid/src/types/syntax.rs

@@ -1,4 +1,4 @@
-use std::fmt;
+use std::{collections::HashMap, fmt};
 
 use crate::types::{ArgType, Field, FuncMode, Function, Label, Type, TypeInner};
 
@@ -183,6 +183,7 @@ impl IDLArgType {
 pub struct TypeField {
     pub label: Label,
     pub typ: IDLType,
+    pub doc_comment: Option<Vec<String>>,
 }
 
 impl From<TypeField> for Field {
@@ -205,12 +206,19 @@ pub enum Dec {
 pub struct Binding {
     pub id: String,
     pub typ: IDLType,
+    pub doc_comment: Option<Vec<String>>,
 }
 
-#[derive(Debug, Default, Clone)]
+#[derive(Debug, Clone)]
+pub struct IDLActorType {
+    pub typ: IDLType,
+    pub doc_comment: Option<Vec<String>>,
+}
+
+#[derive(Debug, Clone)]
 pub struct IDLProg {
     pub decs: Vec<Dec>,
-    pub actor: Option<IDLType>,
+    pub actor: Option<IDLActorType>,
 }
 
 #[derive(Debug)]
@@ -223,7 +231,7 @@ pub struct IDLInitArgs {
 #[derive(Debug, Default)]
 pub struct IDLMergedProg {
     types: Vec<Binding>,
-    pub actor: Option<IDLType>,
+    pub actor: Option<IDLActorType>,
 }
 
 impl From<IDLProg> for IDLMergedProg {
@@ -235,6 +243,15 @@ impl From<IDLProg> for IDLMergedProg {
     }
 }
 
+impl From<IDLActorType> for IDLMergedProg {
+    fn from(other_actor: IDLActorType) -> Self {
+        Self {
+            types: vec![],
+            actor: Some(other_actor),
+        }
+    }
+}
+
 impl From<Vec<Binding>> for IDLMergedProg {
     fn from(bindings: Vec<Binding>) -> Self {
         Self {
@@ -260,14 +277,25 @@ impl IDLMergedProg {
         self.types.extend(types);
     }
 
-    pub fn set_actor(&mut self, other: Option<IDLType>) {
+    pub fn set_actor(&mut self, other: Option<IDLActorType>) {
         self.actor = other;
     }
 
+    pub fn set_comments_in_actor(&mut self, doc_comments: &HashMap<String, Vec<String>>) {
+        self.actor = self.actor.as_ref().map(|t| IDLActorType {
+            typ: set_comments_in_type(&t.typ, doc_comments),
+            doc_comment: t.doc_comment.clone(),
+        });
+    }
+
     pub fn find_type(&self, id: &str) -> Result<&IDLType, String> {
+        self.find_binding(id).map(|b| &b.typ)
+    }
+
+    pub fn find_binding(&self, id: &str) -> Result<&Binding, String> {
         self.types
             .iter()
-            .find_map(|t| if t.id == id { Some(&t.typ) } else { None })
+            .find(|t| t.id == id)
             .ok_or(format!("Type identifier not found: {id}"))
     }
 
@@ -279,8 +307,11 @@ impl IDLMergedProg {
         }
     }
 
-    pub fn get_types(&self) -> Vec<(&str, &IDLType)> {
-        self.types.iter().map(|t| (t.id.as_str(), &t.typ)).collect()
+    pub fn get_types(&self) -> Vec<(&str, &IDLType, Option<&Vec<String>>)> {
+        self.types
+            .iter()
+            .map(|t| (t.id.as_str(), &t.typ, t.doc_comment.as_ref()))
+            .collect()
     }
 
     pub fn get_bindings(&self) -> Vec<Binding> {
@@ -328,3 +359,72 @@ impl IDLMergedProg {
         Err(format!("cannot find method {id}"))
     }
 }
+
+fn set_comments_in_type(t: &IDLType, doc_comments: &HashMap<String, Vec<String>>) -> IDLType {
+    match t {
+        IDLType::PrimT(prim) => IDLType::PrimT(prim.clone()),
+        IDLType::VarT(id) => IDLType::VarT(id.clone()),
+        IDLType::FuncT(func) => IDLType::FuncT(FuncType {
+            modes: func.modes.clone(),
+            args: func
+                .args
+                .iter()
+                .map(|a| IDLArgType {
+                    typ: set_comments_in_type(&a.typ, doc_comments),
+                    name: a.name.clone(),
+                })
+                .collect(),
+            rets: func
+                .rets
+                .iter()
+                .map(|r| set_comments_in_type(r, doc_comments))
+                .collect(),
+        }),
+        IDLType::OptT(t) => IDLType::OptT(Box::new(set_comments_in_type(t, doc_comments))),
+        IDLType::VecT(t) => IDLType::VecT(Box::new(set_comments_in_type(t, doc_comments))),
+        IDLType::RecordT(fields) => {
+            let fields = fields
+                .iter()
+                .map(|f| TypeField {
+                    label: f.label.clone(),
+                    typ: set_comments_in_type(&f.typ, doc_comments),
+                    doc_comment: doc_comments.get(&f.label.to_string()).cloned(),
+                })
+                .collect();
+            IDLType::RecordT(fields)
+        }
+        IDLType::ServT(methods) => {
+            let methods = methods
+                .iter()
+                .map(|m| Binding {
+                    id: m.id.clone(),
+                    typ: set_comments_in_type(&m.typ, doc_comments),
+                    doc_comment: doc_comments.get(&m.id).cloned(),
+                })
+                .collect();
+            IDLType::ServT(methods)
+        }
+        IDLType::VariantT(fields) => {
+            let fields = fields
+                .iter()
+                .map(|f| TypeField {
+                    label: f.label.clone(),
+                    typ: set_comments_in_type(&f.typ, doc_comments),
+                    doc_comment: doc_comments.get(&f.label.to_string()).cloned(),
+                })
+                .collect();
+            IDLType::VariantT(fields)
+        }
+        IDLType::ClassT(args, t) => {
+            let args = args
+                .iter()
+                .map(|a| IDLArgType {
+                    typ: set_comments_in_type(&a.typ, doc_comments),
+                    name: a.name.clone(),
+                })
+                .collect();
+            IDLType::ClassT(args, Box::new(set_comments_in_type(t, doc_comments)))
+        }
+        IDLType::PrincipalT => IDLType::PrincipalT,
+    }
+}

rust/candid/src/types/type_env.rs

@@ -240,6 +240,7 @@ impl TypeEnv {
         TypeField {
             label: field.id.as_ref().clone(),
             typ: self.as_idl_type(&field.ty),
+            doc_comment: None,
         }
     }
 
@@ -256,6 +257,7 @@ impl TypeEnv {
             .map(|(id, t)| Binding {
                 id: id.clone(),
                 typ: self.as_idl_type(t),
+                doc_comment: None,
             })
             .collect()
     }

rust/candid/tests/types.rs

@@ -156,10 +156,12 @@ fn test_struct() {
             TypeField {
                 label: Label::Named("bar".to_string()),
                 typ: IDLType::PrimT(PrimType::Bool),
+                doc_comment: None,
             },
             TypeField {
                 label: Label::Named("foo".to_string()),
                 typ: IDLType::PrimT(PrimType::Int),
+                doc_comment: None,
             },
         ])
     );
@@ -181,10 +183,12 @@ fn test_struct() {
             TypeField {
                 label: Label::Named("g1".to_string()),
                 typ: IDLType::PrimT(PrimType::Int32),
+                doc_comment: None,
             },
             TypeField {
                 label: Label::Named("g2".to_string()),
                 typ: IDLType::PrimT(PrimType::Bool),
+                doc_comment: None,
             },
         ])
     );
@@ -205,10 +209,12 @@ fn test_struct() {
             TypeField {
                 label: Label::Named("head".to_string()),
                 typ: IDLType::PrimT(PrimType::Int32),
+                doc_comment: None,
             },
             TypeField {
                 label: Label::Named("tail".to_string()),
                 typ: IDLType::OptT(Box::new(IDLType::VarT("List".to_string()))),
+                doc_comment: None,
             },
         ])
     );
@@ -229,10 +235,12 @@ fn test_struct() {
             TypeField {
                 label: Label::Named("head".to_string()),
                 typ: IDLType::PrimT(PrimType::Int32),
+                doc_comment: None,
             },
             TypeField {
                 label: Label::Named("tail".to_string()),
                 typ: IDLType::OptT(Box::new(IDLType::VarT("GenericList".to_string()))),
+                doc_comment: None,
             },
         ])
     );
@@ -250,10 +258,12 @@ fn test_struct() {
             TypeField {
                 label: Label::Named("head".to_string()),
                 typ: IDLType::PrimT(PrimType::Int32),
+                doc_comment: None,
             },
             TypeField {
                 label: Label::Named("tail".to_string()),
                 typ: IDLType::OptT(Box::new(IDLType::VarT("GenericList".to_string()))),
+                doc_comment: None,
             },
         ])
     );
@@ -292,33 +302,41 @@ fn test_variant() {
                     TypeField {
                         label: Label::Id(0),
                         typ: IDLType::PrimT(PrimType::Bool),
+                        doc_comment: None,
                     },
                     TypeField {
                         label: Label::Id(1),
                         typ: IDLType::PrimT(PrimType::Int32),
+                        doc_comment: None,
                     },
                 ]),
+                doc_comment: None,
             },
             TypeField {
                 label: Label::Named("Baz".to_string()),
                 typ: IDLType::RecordT(vec![
                     TypeField {
                         label: Label::Named("a".to_string()),
                         typ: IDLType::PrimT(PrimType::Int32),
+                        doc_comment: None,
                     },
                     TypeField {
                         label: Label::Named("b".to_string()),
                         typ: IDLType::PrimT(PrimType::Nat32),
+                        doc_comment: None,
                     },
                 ]),
+                doc_comment: None,
             },
             TypeField {
                 label: Label::Named("Foo".to_string()),
                 typ: IDLType::PrimT(PrimType::Null),
+                doc_comment: None,
             },
             TypeField {
                 label: Label::Named("Newtype".to_string()),
                 typ: IDLType::PrimT(PrimType::Bool),
+                doc_comment: None,
             },
         ])
     );