11use std:: collections:: HashMap ;
22
33use crate :: pretty:: utils:: * ;
4- use crate :: types:: { Field , FuncMode , Function , Label , SharedLabel , Type , TypeEnv , TypeInner } ;
4+ use crate :: types:: {
5+ Field , FieldDoc , FuncMode , Function , Label , SharedLabel , Type , TypeDoc , TypeDocs , TypeEnv ,
6+ TypeInner ,
7+ } ;
58use pretty:: RcDoc ;
69
710static KEYWORDS : [ & str ; 30 ] = [
@@ -125,6 +128,12 @@ pub fn pp_docs<'a>(docs: &'a [String]) -> RcDoc<'a> {
125128 lines ( docs. iter ( ) . map ( |line| RcDoc :: text ( "// " ) . append ( line) ) )
126129}
127130
131+ fn maybe_pp_docs < ' a > ( docs : Option < & ' a [ String ] > ) -> RcDoc < ' a > {
132+ docs. filter ( |docs| !docs. is_empty ( ) )
133+ . map ( pp_docs)
134+ . unwrap_or_else ( RcDoc :: nil)
135+ }
136+
128137/// This function is kept for backward compatibility.
129138///
130139/// It is recommended to use [`pp_label_raw`] instead, which accepts a [`Label`].
@@ -153,6 +162,56 @@ fn pp_fields(fs: &[Field], is_variant: bool) -> RcDoc<'_> {
153162 enclose_space ( "{" , concat ( fields, ";" ) , "}" )
154163}
155164
165+ fn pp_field_with_doc < ' a > (
166+ field : & ' a Field ,
167+ is_variant : bool ,
168+ doc : Option < & ' a FieldDoc > ,
169+ ) -> RcDoc < ' a > {
170+ let docs = maybe_pp_docs ( doc. map ( |doc| doc. docs . as_slice ( ) ) ) ;
171+ let ty_doc = if is_variant && * field. ty == TypeInner :: Null {
172+ RcDoc :: nil ( )
173+ } else {
174+ kwd ( " :" ) . append ( pp_ty_with_doc (
175+ & field. ty ,
176+ doc. and_then ( |doc| doc. ty . as_deref ( ) ) ,
177+ ) )
178+ } ;
179+ docs. append ( pp_label_raw ( & field. id ) ) . append ( ty_doc)
180+ }
181+
182+ fn pp_fields_with_doc < ' a > (
183+ fs : & ' a [ Field ] ,
184+ is_variant : bool ,
185+ doc : Option < & ' a TypeDoc > ,
186+ ) -> RcDoc < ' a > {
187+ let fields = fs. iter ( ) . map ( |field| {
188+ let field_doc = doc. and_then ( |doc| doc. fields . get ( & field. id . get_id ( ) ) ) ;
189+ pp_field_with_doc ( field, is_variant, field_doc)
190+ } ) ;
191+ enclose_space ( "{" , concat ( fields, ";" ) , "}" )
192+ }
193+
194+ fn has_field_docs ( doc : Option < & TypeDoc > ) -> bool {
195+ doc. map ( |doc| doc. fields . values ( ) . any ( |field| !field. is_empty ( ) ) )
196+ . unwrap_or ( false )
197+ }
198+
199+ fn pp_ty_with_doc < ' a > ( ty : & ' a Type , doc : Option < & ' a TypeDoc > ) -> RcDoc < ' a > {
200+ use TypeInner :: * ;
201+ match ty. as_ref ( ) {
202+ Record ( ref fs) => {
203+ if ty. is_tuple ( ) && !has_field_docs ( doc) {
204+ let tuple = concat ( fs. iter ( ) . map ( |f| pp_ty_with_doc ( & f. ty , None ) ) , ";" ) ;
205+ kwd ( "record" ) . append ( enclose_space ( "{" , tuple, "}" ) )
206+ } else {
207+ kwd ( "record" ) . append ( pp_fields_with_doc ( fs, false , doc) )
208+ }
209+ }
210+ Variant ( ref fs) => kwd ( "variant" ) . append ( pp_fields_with_doc ( fs, true , doc) ) ,
211+ ty => pp_ty_inner ( ty) ,
212+ }
213+ }
214+
156215pub fn pp_function ( func : & Function ) -> RcDoc < ' _ > {
157216 let args = pp_args ( & func. args ) ;
158217 let rets = pp_rets ( & func. rets ) ;
@@ -204,7 +263,7 @@ fn pp_service<'a>(serv: &'a [(String, Type)], docs: Option<&'a DocComments>) ->
204263 enclose_space ( "{" , doc, "}" )
205264}
206265
207- fn pp_defs ( env : & TypeEnv ) -> RcDoc < ' _ > {
266+ fn pp_defs_plain ( env : & TypeEnv ) -> RcDoc < ' _ > {
208267 lines ( env. 0 . iter ( ) . map ( |( id, ty) | {
209268 kwd ( "type" )
210269 . append ( ident ( id) )
@@ -214,6 +273,18 @@ fn pp_defs(env: &TypeEnv) -> RcDoc<'_> {
214273 } ) )
215274}
216275
276+ fn pp_defs < ' a > ( env : & ' a TypeEnv , docs : & ' a DocComments ) -> RcDoc < ' a > {
277+ lines ( env. 0 . iter ( ) . map ( |( id, ty) | {
278+ let type_doc = docs. lookup_type_def ( id) ;
279+ maybe_pp_docs ( type_doc. map ( |doc| doc. docs . as_slice ( ) ) )
280+ . append ( kwd ( "type" ) )
281+ . append ( ident ( id) )
282+ . append ( kwd ( "=" ) )
283+ . append ( pp_ty_with_doc ( ty, type_doc) )
284+ . append ( ";" )
285+ } ) )
286+ }
287+
217288fn pp_class < ' a > ( args : & ' a [ Type ] , t : & ' a Type , docs : Option < & ' a DocComments > ) -> RcDoc < ' a > {
218289 let doc = pp_args ( args) . append ( " ->" ) . append ( RcDoc :: space ( ) ) ;
219290 match t. as_ref ( ) {
@@ -234,13 +305,14 @@ fn pp_actor<'a>(ty: &'a Type, docs: &'a DocComments) -> RcDoc<'a> {
234305
235306/// Pretty-prints the initialization arguments for a Candid actor.
236307pub fn pp_init_args < ' a > ( env : & ' a TypeEnv , args : & ' a [ Type ] ) -> RcDoc < ' a > {
237- pp_defs ( env) . append ( pp_args ( args) )
308+ pp_defs_plain ( env) . append ( pp_args ( args) )
238309}
239310
240311/// Collects doc comments that can be passed to the [compile_with_docs] function.
241312#[ derive( Default , Debug ) ]
242313pub struct DocComments {
243314 service_methods : HashMap < String , Vec < String > > ,
315+ type_defs : HashMap < String , TypeDoc > ,
244316}
245317
246318impl DocComments {
@@ -255,6 +327,22 @@ impl DocComments {
255327 pub fn lookup_service_method ( & self , method : & str ) -> Option < & Vec < String > > {
256328 self . service_methods . get ( method)
257329 }
330+
331+ pub fn add_type_def ( & mut self , name : String , doc : TypeDoc ) {
332+ if !doc. is_empty ( ) {
333+ self . type_defs . insert ( name, doc) ;
334+ }
335+ }
336+
337+ pub fn lookup_type_def ( & self , name : & str ) -> Option < & TypeDoc > {
338+ self . type_defs . get ( name)
339+ }
340+
341+ pub fn extend_types ( & mut self , docs : TypeDocs ) {
342+ for ( name, doc) in docs. named {
343+ self . add_type_def ( name, doc) ;
344+ }
345+ }
258346}
259347
260348pub fn compile ( env : & TypeEnv , actor : & Option < Type > ) -> String {
@@ -281,9 +369,9 @@ pub fn compile(env: &TypeEnv, actor: &Option<Type>) -> String {
281369/// ```
282370pub fn compile_with_docs ( env : & TypeEnv , actor : & Option < Type > , docs : & DocComments ) -> String {
283371 match actor {
284- None => pp_defs ( env) . pretty ( LINE_WIDTH ) . to_string ( ) ,
372+ None => pp_defs ( env, docs ) . pretty ( LINE_WIDTH ) . to_string ( ) ,
285373 Some ( actor) => {
286- let defs = pp_defs ( env) ;
374+ let defs = pp_defs ( env, docs ) ;
287375 let actor = kwd ( "service :" ) . append ( pp_actor ( actor, docs) ) ;
288376 let doc = defs. append ( actor) ;
289377 doc. pretty ( LINE_WIDTH ) . to_string ( )
0 commit comments