fix: support verus_spec in verusdoc

Co-authored-by: Xinyi Wan <64517311+rikosellic@users.noreply.github.com>

Author: Marsman1996

source/builtin_macros/src/attr_rewrite.rs

@@ -36,6 +36,7 @@
 /// - Refer to `examples/syntax_attr.rs`.
 use proc_macro2::TokenStream;
 use quote::{ToTokens, quote, quote_spanned};
+use syn::parse::Parser;
 use syn::visit_mut::VisitMut;
 use syn::{Expr, Item, ItemConst, parse2, spanned::Spanned};
 
@@ -487,9 +488,88 @@ pub(crate) fn rewrite_verus_spec_on_fun_or_loop(
             // To avoid misuse of the unverified function,
             // we add `requires false` and thus prevent verified function to use it.
             // Allow unverified code to use the function without changing in/output.
+            let mut rustdoc_attrs: Vec<syn::Attribute> = vec![];
+            if crate::rustdoc::env_rustdoc() {
+                let mut verus_fun: verus_syn::ItemFn = syn_to_verus_syn(fun.clone());
+                verus_fun.sig.spec = spec_attr.spec.clone();
+
+                // Set return variable name
+                if let Some((verus_syn::Pat::Ident(pat_ident), _)) = &spec_attr.ret_pat {
+                    if let verus_syn::ReturnType::Type(_, _, opt_name, _) =
+                        &mut verus_fun.sig.output
+                    {
+                        *opt_name = Some(Box::new((
+                            verus_syn::token::Paren::default(),
+                            verus_syn::Pat::Ident(pat_ident.clone()),
+                            verus_syn::Token![:](pat_ident.span()),
+                        )));
+                    }
+                }
+
+                crate::rustdoc::process_item_fn(&mut verus_fun);
+
+                for attr in &verus_fun.attrs {
+                    if attr.path().is_ident("doc")
+                        && attr.to_token_stream().to_string().contains("verusdoc_special_attr")
+                    {
+                        if let Ok(doc_attrs) =
+                            syn::Attribute::parse_outer.parse(attr.to_token_stream().into())
+                        {
+                            rustdoc_attrs.extend(doc_attrs);
+                        }
+                    }
+                }
+            }
+
             if let Some(with) = &spec_attr.spec.with {
-                let extra_funs = rewrite_unverified_func(&mut fun, with.with.span(), erase);
+                let mut extra_funs = rewrite_unverified_func(&mut fun, with.with.span(), erase);
+
+                if crate::rustdoc::env_rustdoc() {
+                    if let Some(unverified_fun) = extra_funs.last_mut() {
+                        unverified_fun.attrs.extend(rustdoc_attrs.clone());
+                    }
+                    fun.attrs.push(crate::syntax::mk_rust_attr_syn(
+                        with.with.span(),
+                        "doc",
+                        quote! {hidden},
+                    ));
+                }
                 extra_funs.iter().for_each(|f| f.to_tokens(&mut new_stream));
+            } else if crate::rustdoc::env_rustdoc() {
+                fun.attrs.extend(rustdoc_attrs);
+            }
+
+            // Inject doc attribute in rustdoc mode
+            if crate::rustdoc::env_rustdoc() {
+                let mut verus_fun: verus_syn::ItemFn = syn_to_verus_syn(fun.clone());
+                verus_fun.sig.spec = spec_attr.spec.clone();
+
+                // Set return variable name
+                if let Some((verus_syn::Pat::Ident(pat_ident), _)) = &spec_attr.ret_pat {
+                    if let verus_syn::ReturnType::Type(_, _, opt_name, _) =
+                        &mut verus_fun.sig.output
+                    {
+                        *opt_name = Some(Box::new((
+                            verus_syn::token::Paren::default(),
+                            verus_syn::Pat::Ident(pat_ident.clone()),
+                            verus_syn::Token![:](pat_ident.span()),
+                        )));
+                    }
+                }
+
+                crate::rustdoc::process_item_fn(&mut verus_fun);
+
+                for attr in &verus_fun.attrs {
+                    if attr.path().is_ident("doc")
+                        && attr.to_token_stream().to_string().contains("verusdoc_special_attr")
+                    {
+                        if let Ok(doc_attrs) =
+                            syn::Attribute::parse_outer.parse(attr.to_token_stream().into())
+                        {
+                            fun.attrs.extend(doc_attrs);
+                        }
+                    }
+                }
             }
 
             // Update function signature based on verus_spec.
@@ -820,6 +900,13 @@ fn rewrite_unverified_func(
         Some(syn::token::Semi { spans: [span] }),
     );
     unverified_fun.attrs_mut().push(mk_verus_attr_syn(span, quote! { external_body }));
+    if !crate::rustdoc::env_rustdoc() {
+        unverified_fun.attrs_mut().push(crate::syntax::mk_rust_attr_syn(
+            span,
+            "doc",
+            quote! {hidden},
+        ));
+    }
     if let Some(block) = unverified_fun.block_mut() {
         // For an unverified function, if it is in keep mode,
         // we erase the function body to avoid using

source/builtin_macros/src/rustdoc.rs

@@ -29,13 +29,14 @@
 // some data explaining the function mode, param modes, and return mode.
 
 use proc_macro2::Span;
+use quote::ToTokens;
 use std::iter::FromIterator;
 use verus_syn::punctuated::Punctuated;
 use verus_syn::spanned::Spanned;
 use verus_syn::token;
 use verus_syn::{
-    AssumeSpecification, AttrStyle, Attribute, Block, Expr, ExprBlock, ExprPath, FnMode, Ident,
-    ImplItemFn, ItemFn, Pat, PatIdent, Path, PathArguments, PathSegment, Publish, QSelf,
+    AssumeSpecification, AttrStyle, Attribute, Block, Expr, ExprBlock, ExprPath, FnArg, FnMode,
+    Ident, ImplItemFn, ItemFn, Pat, PatIdent, Path, PathArguments, PathSegment, Publish, QSelf,
     ReturnType, Signature, TraitItemFn, Type, TypeGroup, TypePath,
 };
 
@@ -103,6 +104,10 @@ fn attr_for_sig(
 
     v.push(encoded_sig_info(sig));
 
+    if let Some(with_spec) = &sig.spec.with {
+        v.push(encoded_str("with", &format_with_spec(with_spec)));
+    }
+
     match &sig.spec.requires {
         Some(es) => {
             for expr in es.exprs.exprs.iter() {
@@ -330,6 +335,89 @@ fn encoded_str(kind: &str, data: &str) -> String {
     "```rust\n// verusdoc_special_attr ".to_string() + kind + "\n" + data + "\n```"
 }
 
+fn format_with_spec(with_spec: &verus_syn::WithSpecOnFn) -> String {
+    let mut lines: Vec<String> = vec![];
+
+    let inputs = format_fn_args(&with_spec.inputs);
+    for input in inputs {
+        let input = normalize_ws(input.trim());
+        lines.push(format!("    {input},"));
+    }
+
+    if let Some((_, outputs)) = &with_spec.outputs {
+        lines.push("    ->".to_string());
+        let outputs = format_pat_types(outputs);
+        for output in outputs {
+            let output = normalize_ws(output.trim());
+            lines.push(format!("    {output},"));
+        }
+    }
+
+    lines.join("\n")
+}
+
+fn format_pat_types(outputs: &Punctuated<verus_syn::PatType, verus_syn::Token![,]>) -> Vec<String> {
+    if outputs.is_empty() {
+        return vec![];
+    }
+
+    outputs.iter().map(format_pat_type).collect()
+}
+
+fn format_fn_args(inputs: &Punctuated<FnArg, verus_syn::Token![,]>) -> Vec<String> {
+    if inputs.is_empty() {
+        return vec![];
+    }
+    inputs.iter().map(format_fn_arg).collect()
+}
+
+fn format_fn_arg(arg: &FnArg) -> String {
+    let tracked = if arg.tracked.is_some() { "tracked " } else { "" };
+    match &arg.kind {
+        verus_syn::FnArgKind::Receiver(receiver) => {
+            let s = normalize_ws(&receiver.to_token_stream().to_string());
+            format!("{tracked}{s}")
+        }
+        verus_syn::FnArgKind::Typed(pt) => {
+            let s = format_pat_type(pt);
+            format!("{tracked}{s}")
+        }
+    }
+}
+
+fn format_pat_type(pt: &verus_syn::PatType) -> String {
+    let pat = normalize_ws(&pt.pat.to_token_stream().to_string());
+    let ty = tighten_type_spacing(&normalize_ws(&pt.ty.to_token_stream().to_string()));
+    format!("{pat}: {ty}")
+}
+
+fn normalize_ws(s: &str) -> String {
+    s.split_whitespace().collect::<Vec<&str>>().join(" ")
+}
+
+fn tighten_type_spacing(s: &str) -> String {
+    let bytes = s.as_bytes();
+    let mut out = String::with_capacity(bytes.len());
+    let mut i = 0usize;
+    while i < bytes.len() {
+        let b = bytes[i];
+        if b == b' ' {
+            let prev = if i > 0 { bytes[i - 1] } else { 0 };
+            let next = if i + 1 < bytes.len() { bytes[i + 1] } else { 0 };
+            let remove =
+                matches!(prev, b'<' | b'(' | b',' | b':') || matches!(next, b'>' | b')' | b',');
+            if !remove {
+                out.push(' ');
+            }
+            i += 1;
+            continue;
+        }
+        out.push(b as char);
+        i += 1;
+    }
+    out
+}
+
 /// Create an attr that looks like #[doc = "doc_str"]
 fn doc_attr_from_string(doc_str: &str, span: Span) -> Attribute {
     let path = Path {

source/verusdoc/src/main.rs

@@ -33,7 +33,7 @@ enum VerusDocAttr {
 }
 
 // Types of spec clauses we handle.
-static SPEC_NAMES: [&str; 4] = ["requires", "ensures", "recommends", "body"];
+static SPEC_NAMES: [&str; 5] = ["with", "requires", "ensures", "recommends", "body"];
 
 fn main() {
     // Manipulate the auto-generated files in `doc/` to clean them up to make
@@ -256,14 +256,26 @@ fn update_docblock(
     }
 
     for spec_name in SPEC_NAMES.iter() {
-        let code_blocks: Vec<NodeRef> = attrs
+        let mut code_blocks: Vec<NodeRef> = attrs
             .iter()
             .filter_map(|a| match a {
                 VerusDocAttr::Specification(s, nr) if s == spec_name => Some(nr.clone()),
                 _ => None,
             })
             .collect();
 
+        // De-duplicate identical spec blocks (can happen with #[verus_spec] + rustdoc pass)
+        let mut seen: Vec<String> = Vec::new();
+        code_blocks.retain(|nr| {
+            let text = nr.text_contents();
+            if seen.iter().any(|t| t == &text) {
+                false
+            } else {
+                seen.push(text);
+                true
+            }
+        });
+
         let is_body = spec_name == &"body";
 
         if code_blocks.len() > 0 && !is_body {
@@ -289,22 +301,29 @@ fn update_docblock(
     }
 
     // Add mode info to the signature
-
-    for attr in attrs.iter() {
-        match attr {
-            VerusDocAttr::ModeInfo(doc_mode_info) => {
-                update_sig_info(
-                    docblock_elem,
-                    UpdateSigMode::DocSigInfo(doc_mode_info),
-                    opt_trait_info,
-                );
-                break;
-            }
-            VerusDocAttr::BroadcastGroup => {
-                update_sig_info(docblock_elem, UpdateSigMode::BroadcastGroup, opt_trait_info);
-                break;
+    // If there are multiple ModeInfo attrs (caused by the `#[verus_spec]` macro expansion),
+    // choose the one with a non-empty `ret_name`.
+    let mode_infos: Vec<&DocSigInfo> = attrs
+        .iter()
+        .filter_map(|a| match a {
+            VerusDocAttr::ModeInfo(info) => Some(info),
+            _ => None,
+        })
+        .collect();
+    let info_to_use =
+        mode_infos.iter().find(|info| !info.ret_name.is_empty()).or_else(|| mode_infos.first());
+
+    if let Some(info) = info_to_use {
+        update_sig_info(docblock_elem, UpdateSigMode::DocSigInfo(info), opt_trait_info);
+    } else {
+        for attr in attrs.iter() {
+            match attr {
+                VerusDocAttr::BroadcastGroup => {
+                    update_sig_info(docblock_elem, UpdateSigMode::BroadcastGroup, opt_trait_info);
+                    break;
+                }
+                _ => {}
             }
-            _ => {}
         }
     }
 }