thrust-macros: split lib.rs into context, invariant, and spec modules

thrust-macros/src/context.rs

@@ -0,0 +1,36 @@
+//! Expansion of `#[thrust_macros::context]`.
+//!
+//! Stamps each method in an `impl`/`trait` block with the enclosing header so
+//! method-level `requires`/`ensures` can recover the outer generics.
+
+use proc_macro::TokenStream;
+use quote::ToTokens as _;
+
+use crate::fn_outer_item::FnOuterItem;
+
+pub fn expand(item: TokenStream) -> TokenStream {
+    let mut outer_item = syn::parse_macro_input!(item as FnOuterItem);
+    let outer_header = outer_item.clone().into_header_only();
+    match &mut outer_item {
+        FnOuterItem::ItemImpl(item_impl) => {
+            for item in &mut item_impl.items {
+                let syn::ImplItem::Fn(item) = item else {
+                    continue;
+                };
+                item.attrs
+                    .push(syn::parse_quote!(#[thrust::_outer_context(#outer_header)]));
+            }
+        }
+        FnOuterItem::ItemTrait(item_trait) => {
+            for item in &mut item_trait.items {
+                let syn::TraitItem::Fn(item) = item else {
+                    continue;
+                };
+                item.attrs
+                    .push(syn::parse_quote!(#[thrust::_outer_context(#outer_header)]));
+            }
+        }
+    }
+
+    outer_item.into_token_stream().into()
+}

thrust-macros/src/fn_outer_item.rs

@@ -0,0 +1,57 @@
+/// An `impl` or `trait` header carried by the `#[thrust::_outer_context(..)]`
+/// attribute so a method can recover its enclosing generics.
+#[derive(Debug, Clone)]
+pub enum FnOuterItem {
+    ItemImpl(syn::ItemImpl),
+    ItemTrait(syn::ItemTrait),
+}
+
+impl syn::parse::Parse for FnOuterItem {
+    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
+        use syn::parse::discouraged::Speculative as _;
+
+        let fork = input.fork();
+        if let Ok(item_impl) = fork.parse::<syn::ItemImpl>() {
+            input.advance_to(&fork);
+            return Ok(Self::ItemImpl(item_impl));
+        }
+
+        let fork = input.fork();
+        if let Ok(item_trait) = fork.parse::<syn::ItemTrait>() {
+            input.advance_to(&fork);
+            return Ok(Self::ItemTrait(item_trait));
+        }
+
+        Err(input.error("expected an impl block or a trait definition"))
+    }
+}
+
+impl quote::ToTokens for FnOuterItem {
+    fn to_tokens(&self, tokens: &mut proc_macro2::TokenStream) {
+        match self {
+            FnOuterItem::ItemImpl(item_impl) => item_impl.to_tokens(tokens),
+            FnOuterItem::ItemTrait(item_trait) => item_trait.to_tokens(tokens),
+        }
+    }
+}
+
+impl FnOuterItem {
+    pub fn into_header_only(mut self) -> Self {
+        match &mut self {
+            FnOuterItem::ItemImpl(item_impl) => {
+                item_impl.items.clear();
+            }
+            FnOuterItem::ItemTrait(item_trait) => {
+                item_trait.items.clear();
+            }
+        }
+        self
+    }
+
+    pub fn generics(&self) -> &syn::Generics {
+        match self {
+            FnOuterItem::ItemImpl(item_impl) => &item_impl.generics,
+            FnOuterItem::ItemTrait(item_trait) => &item_trait.generics,
+        }
+    }
+}