Attribute Macros on Items (#528)

Co-authored-by: John-John Tedro <[email protected]>

crates/rune-macros/src/lib.rs

@@ -22,6 +22,7 @@
 //!
 //! This is part of the [Rune Language](https://rune-rs.github.io).
 
+use ::quote::format_ident;
 use syn::{Generics, Path};
 
 extern crate proc_macro;
@@ -79,7 +80,23 @@ pub fn macro_(
     let attrs = syn::parse_macro_input!(attrs with crate::macro_::Config::parse);
     let macro_ = syn::parse_macro_input!(item with crate::macro_::Macro::parse);
 
-    let output = match macro_.expand(attrs) {
+    let output = match macro_.expand(attrs, format_ident!("function")) {
+        Ok(output) => output,
+        Err(e) => return proc_macro::TokenStream::from(e.to_compile_error()),
+    };
+
+    output.into()
+}
+
+#[proc_macro_attribute]
+pub fn attribute_macro(
+    attrs: proc_macro::TokenStream,
+    item: proc_macro::TokenStream,
+) -> proc_macro::TokenStream {
+    let attrs = syn::parse_macro_input!(attrs with crate::macro_::Config::parse);
+    let macro_ = syn::parse_macro_input!(item with crate::macro_::Macro::parse);
+
+    let output = match macro_.expand(attrs, format_ident!("attribute")) {
         Ok(output) => output,
         Err(e) => return proc_macro::TokenStream::from(e.to_compile_error()),
     };

crates/rune-macros/src/macro_.rs

@@ -1,4 +1,4 @@
-use proc_macro2::{Span, TokenStream};
+use proc_macro2::{Ident, Span, TokenStream};
 use quote::{quote, ToTokens};
 use syn::parse::ParseStream;
 use syn::punctuated::Punctuated;
@@ -109,7 +109,7 @@ impl Macro {
     }
 
     /// Expand the function declaration.
-    pub(crate) fn expand(self, attrs: Config) -> Result<TokenStream, Error> {
+    pub(crate) fn expand(self, attrs: Config, macro_kind: Ident) -> Result<TokenStream, Error> {
         let real_fn_path = {
             let mut segments = Punctuated::default();
 
@@ -181,7 +181,7 @@ impl Macro {
             #[automatically_derived]
             #meta_vis fn #meta_fn() -> rune::__private::MacroMetaData {
                 rune::__private::MacroMetaData {
-                    kind: rune::__private::MacroMetaKind::function(#meta_name, #real_fn_path),
+                    kind: rune::__private::MacroMetaKind::#macro_kind(#meta_name, #real_fn_path),
                     name: #name_string,
                     docs: &#docs[..],
                 }

crates/rune-modules/src/experiments.rs

@@ -45,7 +45,7 @@ fn passthrough(_: &mut MacroContext<'_>, stream: &TokenStream) -> compile::Resul
 /// Implementation for the `make_function!` macro.
 #[rune::macro_]
 fn make_function(ctx: &mut MacroContext<'_>, stream: &TokenStream) -> compile::Result<TokenStream> {
-    let mut parser = Parser::from_token_stream(stream, ctx.stream_span());
+    let mut parser = Parser::from_token_stream(stream, ctx.input_span());
 
     let ident = parser.parse::<ast::Ident>()?;
     let _ = parser.parse::<T![=>]>()?;

crates/rune-modules/src/experiments/stringy_math_macro.rs

@@ -9,7 +9,7 @@ pub(crate) fn stringy_math(
     ctx: &mut MacroContext<'_>,
     stream: &TokenStream,
 ) -> compile::Result<TokenStream> {
-    let mut parser = Parser::from_token_stream(stream, ctx.stream_span());
+    let mut parser = Parser::from_token_stream(stream, ctx.input_span());
 
     let mut output = quote!(0);
 

crates/rune/src/ast.rs

@@ -16,7 +16,7 @@
 //!
 //! #[rune::macro_]
 //! fn ident_to_string(ctx: &mut MacroContext<'_>, stream: &TokenStream) -> compile::Result<TokenStream> {
-//!     let mut p = Parser::from_token_stream(stream, ctx.stream_span());
+//!     let mut p = Parser::from_token_stream(stream, ctx.input_span());
 //!     let ident = p.parse_all::<ast::Ident>()?;
 //!     let ident = ctx.resolve(ident)?.to_owned();
 //!     let string = ctx.lit(&ident);
@@ -148,6 +148,7 @@ mod lit_number;
 mod lit_str;
 mod local;
 mod macro_call;
+mod macro_utils;
 mod pat;
 mod path;
 mod prelude;
@@ -213,6 +214,7 @@ pub use self::lit_number::LitNumber;
 pub use self::lit_str::LitStr;
 pub use self::local::Local;
 pub use self::macro_call::MacroCall;
+pub use self::macro_utils::{EqValue, Group};
 pub use self::pat::{
     Pat, PatBinding, PatIgnore, PatLit, PatObject, PatPath, PatRest, PatTuple, PatVec,
 };

crates/rune/src/ast/attribute.rs

@@ -60,6 +60,13 @@ pub struct Attribute {
     /// The `]` character
     pub close: T![']'],
 }
+impl Attribute {
+    pub(crate) fn input_span(&self) -> Span {
+        self.input
+            .option_span()
+            .unwrap_or_else(|| self.close.span.head())
+    }
+}
 
 impl Parse for Attribute {
     fn parse(p: &mut Parser<'_>) -> Result<Self> {

crates/rune/src/ast/item.rs

@@ -2,6 +2,8 @@ use core::mem::take;
 
 use crate::ast::prelude::*;
 
+use super::Attribute;
+
 /// A declaration.
 #[derive(Debug, Clone, PartialEq, Eq, ToTokens, Spanned)]
 #[non_exhaustive]
@@ -26,7 +28,7 @@ pub enum Item {
 }
 
 impl Item {
-    /// Test if the item has any attributes
+    /// Get the item's attributes
     pub(crate) fn attributes(&self) -> &[ast::Attribute] {
         match self {
             Self::Use(item) => &item.attributes,
@@ -39,6 +41,19 @@ impl Item {
             Self::MacroCall(item) => &item.attributes,
         }
     }
+    /// Get the item's attributes mutably
+    pub(crate) fn attributes_mut(&mut self) -> &mut Vec<ast::Attribute> {
+        match self {
+            Self::Use(item) => &mut item.attributes,
+            Self::Fn(item) => &mut item.attributes,
+            Self::Enum(item) => &mut item.attributes,
+            Self::Struct(item) => &mut item.attributes,
+            Self::Impl(item) => &mut item.attributes,
+            Self::Mod(item) => &mut item.attributes,
+            Self::Const(item) => &mut item.attributes,
+            Self::MacroCall(item) => &mut item.attributes,
+        }
+    }
 
     /// Indicates if the declaration needs a semi-colon or not.
     pub(crate) fn needs_semi_colon(&self) -> bool {
@@ -155,6 +170,17 @@ impl Item {
 
         Ok(item)
     }
+
+    /// Removes the first attribute in the item list and returns it if present.
+    pub(crate) fn remove_first_attribute(&mut self) -> Option<Attribute> {
+        let attributes = self.attributes_mut();
+
+        if !attributes.is_empty() {
+            return Some(attributes.remove(0));
+        }
+
+        None
+    }
 }
 
 impl Parse for Item {

crates/rune/src/ast/macro_call.rs

@@ -29,7 +29,7 @@ pub struct MacroCall {
     pub open: ast::Token,
     /// The tokens provided to the macro.
     #[rune(optional)]
-    pub stream: TokenStream,
+    pub input: TokenStream,
     /// Closing token.
     pub close: ast::Token,
 }
@@ -40,9 +40,9 @@ impl MacroCall {
         !matches!(self.close.kind, K!['}'])
     }
 
-    /// The span of the token stream.
-    pub(crate) fn stream_span(&self) -> Span {
-        if let Some(span) = self.stream.option_span() {
+    /// The span of the input token stream.
+    pub(crate) fn input_span(&self) -> Span {
+        if let Some(span) = self.input.option_span() {
             span
         } else {
             self.open.span.tail()
@@ -106,7 +106,7 @@ impl MacroCall {
             bang,
             path,
             open,
-            stream: TokenStream::from(stream),
+            input: TokenStream::from(stream),
             close,
         })
     }

crates/rune/src/ast/macro_utils.rs

@@ -0,0 +1,77 @@
+use super::prelude::*;
+use super::Eq;
+
+/// An `= ...` e.g. inside an attribute `#[doc = ...]`.
+///
+/// To get unparsed tokens use `EqValue<TokenStream>`.
+#[derive(Debug, Clone, PartialEq, Eq, ToTokens, Parse, Spanned)]
+pub struct EqValue<T> {
+    /// The `=` token.
+    pub eq: Eq,
+    /// The remainder.
+    pub value: T,
+}
+
+/// Parses `[{( ... )}]` ensuring that the delimiter is balanced.
+#[derive(Debug, Clone, PartialEq, Eq, ToTokens, Spanned)]
+pub struct Group {
+    /// The opening delimiter.
+    pub open: ast::Token,
+    /// The content between the delimiters.
+    pub content: TokenStream,
+    /// The closing delimiter.
+    pub close: ast::Token,
+}
+
+impl Parse for Group {
+    fn parse(parser: &mut Parser) -> compile::Result<Self> {
+        let mut level = 1;
+        let open = parser.next()?;
+
+        let delim = match open.kind {
+            ast::Kind::Open(delim) => delim,
+            _ => {
+                return Err(compile::Error::expected(open, Expectation::OpenDelimiter));
+            }
+        };
+
+        let close;
+
+        let mut stream = Vec::new();
+
+        loop {
+            let token = parser.next()?;
+
+            match token.kind {
+                ast::Kind::Open(..) => level += 1,
+                ast::Kind::Close(actual) => {
+                    level -= 1;
+
+                    if level == 0 {
+                        if actual != delim {
+                            return Err(compile::Error::new(
+                                open,
+                                ParseErrorKind::ExpectedMacroCloseDelimiter {
+                                    actual: token.kind,
+                                    expected: ast::Kind::Close(delim),
+                                },
+                            ));
+                        }
+
+                        close = token;
+                        break;
+                    }
+                }
+                _ => (),
+            }
+
+            stream.push(token);
+        }
+
+        Ok(Self {
+            open,
+            content: TokenStream::from(stream),
+            close,
+        })
+    }
+}

crates/rune/src/compile/context.rs

@@ -9,12 +9,12 @@ use crate::compile::meta;
 use crate::compile::Docs;
 use crate::compile::{ComponentRef, ContextError, IntoComponent, Item, ItemBuf, MetaInfo, Names};
 use crate::module::{
-    Fields, Function, InternalEnum, Module, ModuleAssociated, ModuleConstant, ModuleFunction,
-    ModuleMacro, ModuleType, TypeSpecification, UnitType,
+    Fields, Function, InternalEnum, Module, ModuleAssociated, ModuleAttributeMacro, ModuleConstant,
+    ModuleFunction, ModuleMacro, ModuleType, TypeSpecification, UnitType,
 };
 use crate::runtime::{
-    ConstValue, FunctionHandler, MacroHandler, Protocol, RuntimeContext, StaticType, TypeCheck,
-    TypeInfo, TypeOf, VariantRtti,
+    AttributeMacroHandler, ConstValue, FunctionHandler, MacroHandler, Protocol, RuntimeContext,
+    StaticType, TypeCheck, TypeInfo, TypeOf, VariantRtti,
 };
 use crate::Hash;
 
@@ -92,6 +92,8 @@ pub struct Context {
     associated: HashMap<Hash, Vec<Hash>>,
     /// Registered native macro handlers.
     macros: HashMap<Hash, Arc<MacroHandler>>,
+    /// Registered native attribute macro handlers.
+    attribute_macros: HashMap<Hash, Arc<AttributeMacroHandler>>,
     /// Registered types.
     types: HashMap<Hash, ContextType>,
     /// Registered internal enums.
@@ -217,6 +219,10 @@ impl Context {
             self.install_macro(module, m)?;
         }
 
+        for m in &module.attribute_macros {
+            self.install_attribute_macro(module, m)?;
+        }
+
         for m in &module.constants {
             self.install_constant(module, m)?;
         }
@@ -331,6 +337,11 @@ impl Context {
         self.macros.get(&hash)
     }
 
+    /// Lookup the given attribute macro handler.
+    pub(crate) fn lookup_attribute_macro(&self, hash: Hash) -> Option<&Arc<AttributeMacroHandler>> {
+        self.attribute_macros.get(&hash)
+    }
+
     /// Look up the type check implementation for the specified type hash.
     pub(crate) fn type_check_for(&self, hash: Hash) -> Option<TypeCheck> {
         let ty = self.types.get(&hash)?;
@@ -593,7 +604,7 @@ impl Context {
         Ok(())
     }
 
-    /// Install a function and check for duplicates.
+    /// Install a macro and check for duplicates.
     fn install_macro(&mut self, module: &Module, m: &ModuleMacro) -> Result<(), ContextError> {
         let item = module.item.join(&m.item);
         let hash = Hash::type_hash(&item);
@@ -610,6 +621,27 @@ impl Context {
         Ok(())
     }
 
+    /// Install an attribute macro and check for duplicates.
+    fn install_attribute_macro(
+        &mut self,
+        module: &Module,
+        m: &ModuleAttributeMacro,
+    ) -> Result<(), ContextError> {
+        let item = module.item.join(&m.item);
+        let hash = Hash::type_hash(&item);
+        self.attribute_macros.insert(hash, m.handler.clone());
+
+        self.install_meta(ContextMeta {
+            hash,
+            item: Some(item),
+            kind: meta::Kind::AttributeMacro,
+            #[cfg(feature = "doc")]
+            docs: m.docs.clone(),
+        })?;
+
+        Ok(())
+    }
+
     /// Install a constant and check for duplicates.
     fn install_constant(
         &mut self,

crates/rune/src/compile/meta.rs

@@ -120,6 +120,7 @@ impl Meta {
             Kind::ConstFn { .. } => None,
             Kind::Import { .. } => None,
             Kind::Macro => None,
+            Kind::AttributeMacro => None,
             Kind::Module => None,
         }
     }
@@ -173,6 +174,8 @@ pub enum Kind {
     },
     /// A macro item.
     Macro,
+    /// An attribute macro item.
+    AttributeMacro,
     /// A function declaration.
     Function {
         /// Native signature for this function.