feat(tvix/eval): Add docstrings as documentation for builtins r/5269

Add a new `documentation: Option<&'static str>` field to Builtin, and
populate it in the `#[builtins]` macro with the docstring of the builtin
function, if any.

Change-Id: Ic68fdf9b314d15a780731974234e2ae43f6a44b0
Reviewed-on: https://cl.tvl.fyi/c/depot/+/7205
Tested-by: BuildkiteCI
Reviewed-by: tazjin <tazjin@tvl.su>

5 files changed, 54 insertions, 2 deletions

diff --git a/tvix/eval/builtin-macros/src/lib.rs b/tvix/eval/builtin-macros/src/lib.rs
index 8276e421c5..e815749ec6 100644
--- a/tvix/eval/builtin-macros/src/lib.rs
+++ b/tvix/eval/builtin-macros/src/lib.rs
@@ -2,11 +2,12 @@ extern crate proc_macro;
 
 use proc_macro::TokenStream;
 use proc_macro2::Span;
-use quote::{quote_spanned, ToTokens};
+use quote::{quote, quote_spanned, ToTokens};
 use syn::parse::Parse;
 use syn::spanned::Spanned;
 use syn::{
-    parse_macro_input, parse_quote, FnArg, Ident, Item, ItemMod, LitStr, Pat, PatIdent, PatType,
+    parse2, parse_macro_input, parse_quote, Attribute, FnArg, Ident, Item, ItemMod, LitStr, Pat,
+    PatIdent, PatType, Token,
 };
 
 struct BuiltinArgs {
@@ -21,6 +22,35 @@ impl Parse for BuiltinArgs {
     }
 }
 
+fn extract_docstring(attrs: &[Attribute]) -> Option<LitStr> {
+    // Rust docstrings are transparently written pre-macro expansion into an attribute that looks
+    // like:
+    //
+    // #[doc = "docstring here"]
+
+    #[allow(dead_code)]
+    #[derive(Debug)]
+    struct Docstring {
+        eq: Token![=],
+        doc: LitStr,
+    }
+
+    impl Parse for Docstring {
+        fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
+            Ok(Self {
+                eq: input.parse()?,
+                doc: input.parse()?,
+            })
+        }
+    }
+
+    attrs
+        .iter()
+        .filter(|attr| attr.path.get_ident().into_iter().any(|id| id == "doc"))
+        .find_map(|attr| parse2::<Docstring>(attr.tokens.clone()).ok())
+        .map(|docstring| docstring.doc)
+}
+
 /// Mark the annotated module as a module for defining Nix builtins.
 ///
 /// A function `fn builtins() -> Vec<Builtin>` will be defined within the annotated module,
@@ -144,10 +174,16 @@ pub fn builtins(_args: TokenStream, item: TokenStream) -> TokenStream {
                 let mut reversed_args = args.clone();
                 reversed_args.reverse();
 
+                let docstring = match extract_docstring(&f.attrs) {
+                    Some(docs) => quote!(Some(#docs)),
+                    None => quote!(None),
+                };
+
                 builtins.push(quote_spanned! { builtin_attr.span() => {
                     crate::internal::Builtin::new(
                         #name,
                         &[#(#builtin_arguments),*],
+                        #docstring,
                         |mut args: Vec<crate::Value>, vm: &mut crate::internal::VM| {
                             #(let #reversed_args = args.pop().unwrap();)*
                             #fn_name(vm, #(#args),*)
diff --git a/tvix/eval/builtin-macros/tests/tests.rs b/tvix/eval/builtin-macros/tests/tests.rs
index d07020f69b..b270594b0f 100644
--- a/tvix/eval/builtin-macros/tests/tests.rs
+++ b/tvix/eval/builtin-macros/tests/tests.rs
@@ -7,6 +7,7 @@ mod builtins {
     use tvix_eval::internal::VM;
     use tvix_eval::{ErrorKind, Value};
 
+    /// Test docstring
     #[builtin("identity")]
     pub fn builtin_identity(_vm: &mut VM, x: Value) -> Result<Value, ErrorKind> {
         Ok(x)
@@ -22,4 +23,7 @@ mod builtins {
 fn builtins() {
     let builtins = builtins::builtins();
     assert_eq!(builtins.len(), 2);
+
+    let identity = builtins.iter().find(|b| b.name() == "identity").unwrap();
+    assert_eq!(identity.documentation(), Some(" Test docstring"));
 }
diff --git a/tvix/eval/src/builtins/impure.rs b/tvix/eval/src/builtins/impure.rs
index d4e7c41703..f28609fa76 100644
--- a/tvix/eval/src/builtins/impure.rs
+++ b/tvix/eval/src/builtins/impure.rs
@@ -115,6 +115,7 @@ pub fn builtins_import(globals: &Weak<GlobalsMap>, source: SourceCode) -> Builti
             strict: true,
             name: "path",
         }],
+        None,
         move |mut args: Vec<Value>, vm: &mut VM| {
             let mut path = super::coerce_value_to_path(&args.pop().unwrap(), vm)?;
             if path.is_dir() {
diff --git a/tvix/eval/src/builtins/mod.rs b/tvix/eval/src/builtins/mod.rs
index 9b6e29db0a..41fa4d4d99 100644
--- a/tvix/eval/src/builtins/mod.rs
+++ b/tvix/eval/src/builtins/mod.rs
@@ -918,6 +918,7 @@ fn placeholders() -> Vec<Builtin> {
                     name: "value",
                 },
             ],
+            None,
             |mut args: Vec<Value>, vm: &mut VM| {
                 vm.emit_warning(WarningKind::NotImplemented("builtins.addErrorContext"));
                 Ok(args.pop().unwrap())
@@ -929,6 +930,7 @@ fn placeholders() -> Vec<Builtin> {
                 strict: true,
                 name: "s",
             }],
+            None,
             |mut args: Vec<Value>, vm: &mut VM| {
                 vm.emit_warning(WarningKind::NotImplemented(
                     "builtins.unsafeDiscardStringContext",
@@ -942,6 +944,7 @@ fn placeholders() -> Vec<Builtin> {
                 strict: true,
                 name: "attrs",
             }],
+            None,
             |args: Vec<Value>, vm: &mut VM| {
                 vm.emit_warning(WarningKind::NotImplemented("builtins.derivation"));
 
diff --git a/tvix/eval/src/value/builtin.rs b/tvix/eval/src/value/builtin.rs
index 944efd85eb..bb14265181 100644
--- a/tvix/eval/src/value/builtin.rs
+++ b/tvix/eval/src/value/builtin.rs
@@ -50,6 +50,8 @@ pub struct Builtin {
     name: &'static str,
     /// Array of arguments to the builtin.
     arguments: &'static [BuiltinArgument],
+    /// Optional documentation for the builtin.
+    documentation: Option<&'static str>,
     func: Rc<dyn BuiltinFn>,
 
     /// Partially applied function arguments.
@@ -60,11 +62,13 @@ impl Builtin {
     pub fn new<F: BuiltinFn + 'static>(
         name: &'static str,
         arguments: &'static [BuiltinArgument],
+        documentation: Option<&'static str>,
         func: F,
     ) -> Self {
         Builtin {
             name,
             arguments,
+            documentation,
             func: Rc::new(func),
             partials: vec![],
         }
@@ -74,6 +78,10 @@ impl Builtin {
         self.name
     }
 
+    pub fn documentation(&self) -> Option<&'static str> {
+        self.documentation
+    }
+
     /// Apply an additional argument to the builtin, which will either
     /// lead to execution of the function or to returning a partial
     /// builtin.