Revert "Simplify PR"

This reverts commit 3d21790.

tools/calyx-ffi-macro/src/lib.rs

@@ -239,3 +239,180 @@ pub fn calyx_ffi(attrs: TokenStream, item: TokenStream) -> TokenStream {
     }
     .into()
 }
+
+/// Extracts tests marked with `#[calyx_ffi_test]` and generates wrapper `#[test]`s and modules.
+#[derive(Default)]
+struct CalyxFFITestModuleVisitor {
+    pub wrappers: Vec<syn::ItemFn>,
+    pub tests: Vec<syn::ItemFn>,
+}
+
+impl syn::visit::Visit<'_> for CalyxFFITestModuleVisitor {
+    fn visit_item_fn(&mut self, i: &syn::ItemFn) {
+        let has_calyx_ffi_test = i
+            .attrs
+            .iter()
+            .any(|attr| attr.path().is_ident("calyx_ffi_test"));
+        if has_calyx_ffi_test {
+            let fn_name = &i.sig.ident;
+            let dut_type = get_ffi_test_dut_type(i).expect(
+                "calyx_ffi_test should enforce that the type is well-formed",
+            );
+
+            self.wrappers.push(syn::parse_quote! {
+                pub(crate) unsafe fn #fn_name(ffi: &mut CalyxFFI) {
+                    let dut = ffi.new_comp::<#dut_type>();
+                    let dut_ref = &mut *dut.borrow_mut();
+                    let dut_pointer = dut_ref as *mut dyn CalyxFFIComponent as *mut _ as *mut #dut_type;
+                    let dut_concrete: &mut #dut_type = &mut *dut_pointer;
+                    super::#fn_name(dut_concrete);
+                }
+            });
+            self.tests.push(syn::parse_quote! {
+                #[test]
+                pub(crate) fn #fn_name() {
+                    let mut ffi = CalyxFFI::new();
+                    unsafe {
+                        super::calyx_ffi_generated_wrappers::#fn_name(&mut ffi);
+                    }
+                }
+            });
+        }
+    }
+}
+
+#[proc_macro_attribute]
+pub fn calyx_ffi_tests(args: TokenStream, item: TokenStream) -> TokenStream {
+    if !args.is_empty() {
+        return util::compile_error(
+            &args.into_iter().next().unwrap().span().into(),
+            "#[calyx_ffi_tests] takes no arguments".into(),
+        );
+    }
+
+    let mut module = parse_macro_input!(item as syn::ItemMod);
+    let module_name = &module.ident;
+
+    let mut visitor = CalyxFFITestModuleVisitor::default();
+    syn::visit::visit_item_mod(&mut visitor, &module);
+    let wrappers = visitor.wrappers;
+    let tests = visitor.tests;
+
+    let test_names = wrappers.iter().map(|test| test.sig.ident.clone());
+    let generated_wrappers = quote! {
+        pub(crate) mod calyx_ffi_generated_wrappers {
+            use super::*;
+
+            pub(crate) const CALYX_FFI_TESTS: &'static [unsafe fn(&mut CalyxFFI) -> ()] = &[
+                #(#test_names),*
+            ];
+
+            #(#wrappers)*
+        }
+    };
+    let generated_wrappers_item: syn::Item =
+        syn::parse2(generated_wrappers).unwrap();
+
+    let generated_tests = quote! {
+        pub(crate) mod calyx_ffi_generated_tests {
+            use super::*;
+
+            #(#tests)*
+        }
+    };
+    let generated_tests_item: syn::Item = syn::parse2(generated_tests).unwrap();
+
+    let items_to_add = vec![generated_wrappers_item, generated_tests_item];
+    if let Some((_, ref mut items)) = module.content {
+        items.extend(items_to_add);
+    } else {
+        module.content = Some((syn::token::Brace::default(), items_to_add));
+    }
+
+    quote! {
+        #module
+
+        pub mod calyx_ffi_generated_top {
+            use super::*;
+
+            pub unsafe fn run_tests(ffi: &mut CalyxFFI) {
+                for test in #module_name::calyx_ffi_generated_wrappers::CALYX_FFI_TESTS {
+                    test(ffi);
+                }
+            }
+        }
+    }
+    .into()
+}
+
+#[proc_macro_attribute]
+pub fn calyx_ffi_test(args: TokenStream, item: TokenStream) -> TokenStream {
+    if !args.is_empty() {
+        return util::compile_error(
+            &args.into_iter().next().unwrap().span().into(),
+            "#[calyx_ffi_test] takes no arguments".into(),
+        );
+    }
+
+    let mut func = parse_macro_input!(item as syn::ItemFn);
+    let dut_type = get_ffi_test_dut_type(&func);
+    let Ok(dut_type) = dut_type else {
+        return dut_type.err().unwrap();
+    };
+
+    let check_trait_impl = quote! {
+        {
+            fn assert_is_calyx_ffi_component<T: CalyxFFIComponent>() {}
+            assert_is_calyx_ffi_component::<#dut_type>();
+        }
+    };
+
+    let check_trait_impl_stmts: syn::Block = syn::parse2(check_trait_impl)
+        .expect("Failed to parse check_trait_impl as a block");
+
+    let new_stmts: Vec<syn::Stmt> = check_trait_impl_stmts
+        .stmts
+        .iter()
+        .chain(func.block.stmts.iter())
+        .cloned()
+        .collect();
+
+    let new_block = syn::Block {
+        brace_token: func.block.brace_token,
+        stmts: new_stmts,
+    };
+    func.block = Box::new(new_block);
+
+    quote! {
+        #func
+    }
+    .into()
+}
+
+fn get_ffi_test_dut_type(
+    func: &syn::ItemFn,
+) -> Result<&syn::Type, TokenStream> {
+    let inputs: Vec<&syn::FnArg> = func.sig.inputs.iter().collect();
+
+    let bad_sig_msg = "#[calyx_ffi_test] tests must take exactly one argument, namely, a mutable reference to the DUT".into();
+
+    if inputs.len() != 1 {
+        return Err(util::compile_error(&func.span(), bad_sig_msg));
+    }
+    let input = inputs.first().unwrap();
+
+    let syn::FnArg::Typed(pat_ty) = input else {
+        return Err(util::compile_error(&func.span(), bad_sig_msg));
+    };
+
+    let syn::Type::Reference(syn::TypeReference {
+        mutability: Some(syn::token::Mut { span: _ }),
+        ref elem,
+        ..
+    }) = *pat_ty.ty
+    else {
+        return Err(util::compile_error(&func.span(), bad_sig_msg));
+    };
+
+    Ok(elem)
+}

tools/calyx-ffi/src/prelude.rs

@@ -2,7 +2,7 @@ pub use super::{
     interface::{CalyxFFI, CalyxFFIComponent, CalyxFFIComponentRef},
     value::Value,
 };
-pub use calyx_ffi_macro::calyx_ffi;
+pub use calyx_ffi_macro::{calyx_ffi, calyx_ffi_test, calyx_ffi_tests};
 pub use calyx_ir;
 pub use interp;
 pub use paste;