Implement thrust-macros that expands requires and ensures attrs

Author: coord-e

.gitignore

@@ -1 +1 @@
-/target
+target/

Cargo.lock

@@ -1,3 +1,7 @@
+[workspace]
+members = [".", "thrust-macros"]
+resolver = "2"
+
 [package]
 name = "thrust"
 version = "0.1.0"

Cargo.toml

@@ -66,8 +66,25 @@ impl Callbacks for CompilerCalls {
     }
 }
 
+fn thrust_macros_path() -> Option<std::path::PathBuf> {
+    let exe = std::env::current_exe().ok()?;
+    let dir = exe.parent()?;
+    for ext in ["so", "dylib", "dll"] {
+        let candidate = dir.join(format!("libthrust_macros.{ext}"));
+        if candidate.exists() {
+            return Some(candidate);
+        }
+    }
+    None
+}
+
 pub fn main() {
-    let args = std::env::args().collect::<Vec<_>>();
+    let mut args = std::env::args().collect::<Vec<_>>();
+
+    if let Some(path) = thrust_macros_path() {
+        args.push("--extern".to_owned());
+        args.push(format!("thrust_macros={}", path.display()));
+    }
 
     use tracing_subscriber::{filter::EnvFilter, prelude::*};
     tracing_subscriber::registry()

src/main.rs

@@ -0,0 +1,13 @@
+[package]
+name = "thrust-macros"
+version = "0.1.0"
+authors = ["coord_e <me@coord-e.com>"]
+edition = "2021"
+
+[lib]
+proc-macro = true
+
+[dependencies]
+proc-macro2 = "1"
+quote = "=1.0.36"
+syn = { version = "2", features = ["full"] }

thrust-macros/Cargo.lock

@@ -0,0 +1,246 @@
+use proc_macro::TokenStream;
+use proc_macro2::TokenStream as TokenStream2;
+use quote::{format_ident, quote, ToTokens};
+use syn::{
+    parse_macro_input, FnArg, GenericParam, Generics, ItemFn, Pat, ReturnType, Type,
+    TypeParamBound, WherePredicate,
+};
+
+/// Transforms a function annotated with `#[requires(req)]` and `#[ensures(ens)]` into
+/// an internal multi-function representation for the Thrust analysis backend.
+///
+/// This macro must be paired with `#[ensures(...)]` on the same function.
+///
+/// # Generated output
+///
+/// For `#[requires(req)] #[ensures(ens)] fn f<T>(a: A) -> R`:
+///
+/// - The original function `f` (unchanged, minus the two thrust attrs)
+/// - `_thrust_requires_f<T>(a: <A as Model>::Ty) -> bool` with body `req`
+/// - `_thrust_ensures_f<T>(result: <R as Model>::Ty, a: <A as Model>::Ty) -> bool` with body `ens`
+/// - `_thrust_extern_spec_f<T>(a: A) -> R` referencing the above via path statements
+#[proc_macro_attribute]
+pub fn requires(attr: TokenStream, item: TokenStream) -> TokenStream {
+    let req_expr = TokenStream2::from(attr);
+    let mut func = parse_macro_input!(item as ItemFn);
+
+    // Find and extract the paired #[ensures(...)] attribute
+    let ensures_pos = func.attrs.iter().position(|a| {
+        a.path()
+            .segments
+            .last()
+            .map_or(false, |s| s.ident == "ensures")
+    });
+
+    let ens_expr = match ensures_pos {
+        Some(idx) => {
+            let attr = func.attrs.remove(idx);
+            match attr.parse_args::<TokenStream2>() {
+                Ok(ts) => ts,
+                Err(e) => return e.to_compile_error().into(),
+            }
+        }
+        None => {
+            return syn::Error::new_spanned(
+                &func.sig.ident,
+                "#[requires] must be paired with #[ensures]",
+            )
+            .to_compile_error()
+            .into();
+        }
+    };
+
+    // Strip any stray requires attrs (defensive — they'd cause infinite recursion)
+    func.attrs.retain(|a| {
+        a.path()
+            .segments
+            .last()
+            .map_or(true, |s| s.ident != "requires")
+    });
+
+    transform(func, req_expr, ens_expr).unwrap_or_else(|e| e.to_compile_error().into())
+}
+
+/// Pass-through no-op. When `#[ensures]` is consumed by `#[requires]`, this is never
+/// reached. It exists so that `#[ensures]` is a recognized proc macro attribute and does
+/// not produce an "unknown attribute" error if applied in isolation.
+#[proc_macro_attribute]
+pub fn ensures(_attr: TokenStream, item: TokenStream) -> TokenStream {
+    item
+}
+
+fn transform(
+    func: ItemFn,
+    req_expr: TokenStream2,
+    ens_expr: TokenStream2,
+) -> syn::Result<TokenStream> {
+    let name = &func.sig.ident;
+    let requires_name = format_ident!("_thrust_requires_{}", name);
+    let ensures_name = format_ident!("_thrust_ensures_{}", name);
+    let extern_spec_name = format_ident!("_thrust_extern_spec_{}", name);
+
+    let generics = &func.sig.generics;
+
+    // <T, U, 'a, ...> for function definitions (params only, no where clause)
+    let def_generics = generic_params_tokens(generics);
+
+    // ::<T, U, 'a, ...> for turbofish (empty if no params)
+    let turbofish = generic_turbofish(generics);
+
+    // Additional where predicates: T: thrust_models::Model for qualifying type params
+    let model_preds = model_where_predicates(generics);
+
+    // Full where clause combining original predicates and model predicates
+    let extended_where = extended_where_clause(generics, &model_preds);
+
+    // Parameters with types replaced by <T as thrust_models::Model>::Ty
+    let model_ty_params = fn_params_with_model_ty(&func.sig.inputs);
+
+    // Argument idents for function call expressions
+    let call_args = fn_arg_idents(&func.sig.inputs);
+
+    // Return type for requires/ensures helper (wrapped in ::Ty)
+    let ret_model_ty: Type = match &func.sig.output {
+        ReturnType::Default => syn::parse_quote!(<() as thrust_models::Model>::Ty),
+        ReturnType::Type(_, ty) => syn::parse_quote!(<#ty as thrust_models::Model>::Ty),
+    };
+
+    let orig_inputs = &func.sig.inputs;
+    let orig_output = &func.sig.output;
+
+    let output = quote! {
+        #func
+
+        #[allow(unused_variables)]
+        #[thrust::formula_fn]
+        fn #requires_name #def_generics(#model_ty_params) -> bool #extended_where {
+            #req_expr
+        }
+
+        #[allow(unused_variables)]
+        #[thrust::formula_fn]
+        fn #ensures_name #def_generics(result: #ret_model_ty, #model_ty_params) -> bool #extended_where {
+            #ens_expr
+        }
+
+        #[thrust::extern_spec_fn]
+        #[allow(path_statements)]
+        fn #extern_spec_name #def_generics(#orig_inputs) #orig_output #extended_where {
+            #[thrust::requires_path]
+            #requires_name #turbofish;
+
+            #[thrust::ensures_path]
+            #ensures_name #turbofish;
+
+            #name #turbofish(#call_args)
+        }
+    };
+
+    Ok(output.into())
+}
+
+/// Returns `<T: Bound, U, 'a>` — the generic param list for function definitions,
+/// without a where clause.
+fn generic_params_tokens(generics: &Generics) -> TokenStream2 {
+    if generics.params.is_empty() {
+        return quote!();
+    }
+    let params = &generics.params;
+    quote!(<#params>)
+}
+
+/// Returns `::<T, U, 'a>` for turbofish use, or nothing if no generic params.
+fn generic_turbofish(generics: &Generics) -> TokenStream2 {
+    if generics.params.is_empty() {
+        return quote!();
+    }
+    let args: Vec<TokenStream2> = generics
+        .params
+        .iter()
+        .map(|p| match p {
+            GenericParam::Type(tp) => tp.ident.to_token_stream(),
+            GenericParam::Lifetime(lp) => lp.lifetime.to_token_stream(),
+            GenericParam::Const(cp) => cp.ident.to_token_stream(),
+        })
+        .collect();
+    quote!(::<#(#args),*>)
+}
+
+/// Returns `T: thrust_models::Model` predicates for every type param that does not
+/// already carry an `Fn`, `FnOnce`, or `FnMut` bound.
+fn model_where_predicates(generics: &Generics) -> Vec<WherePredicate> {
+    generics
+        .params
+        .iter()
+        .filter_map(|p| {
+            let GenericParam::Type(tp) = p else {
+                return None;
+            };
+            let has_fn_bound = tp.bounds.iter().any(|b| {
+                let TypeParamBound::Trait(tb) = b else {
+                    return false;
+                };
+                tb.path.segments.last().map_or(false, |s| {
+                    matches!(s.ident.to_string().as_str(), "Fn" | "FnOnce" | "FnMut")
+                })
+            });
+            if has_fn_bound {
+                return None;
+            }
+            let ident = &tp.ident;
+            Some(syn::parse_quote!(#ident: thrust_models::Model))
+        })
+        .collect()
+}
+
+/// Builds `where <original predicates>, <model predicates>`.
+/// Returns an empty token stream when both sets are empty.
+fn extended_where_clause(generics: &Generics, model_preds: &[WherePredicate]) -> TokenStream2 {
+    let existing: Vec<&WherePredicate> = generics
+        .where_clause
+        .as_ref()
+        .map(|wc| wc.predicates.iter().collect())
+        .unwrap_or_default();
+
+    if existing.is_empty() && model_preds.is_empty() {
+        return quote!();
+    }
+
+    quote! { where #(#existing,)* #(#model_preds),* }
+}
+
+/// Maps each typed function parameter `x: T` to `x: <T as thrust_models::Model>::Ty`.
+/// Receiver (`self`) parameters are skipped.
+fn fn_params_with_model_ty(
+    inputs: &syn::punctuated::Punctuated<FnArg, syn::token::Comma>,
+) -> TokenStream2 {
+    let params: Vec<TokenStream2> = inputs
+        .iter()
+        .filter_map(|arg| {
+            let FnArg::Typed(pt) = arg else { return None };
+            let pat = &pt.pat;
+            let ty = &pt.ty;
+            let model_ty: Type = syn::parse_quote!(<#ty as thrust_models::Model>::Ty);
+            Some(quote!(#pat: #model_ty))
+        })
+        .collect();
+    quote!(#(#params),*)
+}
+
+/// Extracts argument identifiers from function parameters for use in a call expression.
+fn fn_arg_idents(inputs: &syn::punctuated::Punctuated<FnArg, syn::token::Comma>) -> TokenStream2 {
+    let args: Vec<TokenStream2> = inputs
+        .iter()
+        .filter_map(|arg| match arg {
+            FnArg::Typed(pt) => {
+                if let Pat::Ident(pi) = &*pt.pat {
+                    Some(pi.ident.to_token_stream())
+                } else {
+                    None
+                }
+            }
+            FnArg::Receiver(_) => Some(quote!(self)),
+        })
+        .collect();
+    quote!(#(#args),*)
+}