From 85b49cfc7420795569c45c8301ee1fa2ba141a35 Mon Sep 17 00:00:00 2001
From: Edward Shen <code@eddie.sh>
Date: Tue, 1 Oct 2024 20:53:32 -0700
Subject: [PATCH] Work-in-progress

---
 Cargo.lock              |  1 +
 teto_macros/Cargo.toml  |  1 +
 teto_macros/src/lib.rs  | 27 +++++++++++++++------------
 teto_macros/tests/ui.rs |  2 +-
 4 files changed, 18 insertions(+), 13 deletions(-)

diff --git a/Cargo.lock b/Cargo.lock
index 1b0283b..97c7ebf 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -138,6 +138,7 @@ dependencies = [
 name = "teto_macros"
 version = "0.1.0"
 dependencies = [
+ "proc-macro2",
  "quote",
  "syn",
  "teto",
diff --git a/teto_macros/Cargo.toml b/teto_macros/Cargo.toml
index d608cc8..e43d8ac 100644
--- a/teto_macros/Cargo.toml
+++ b/teto_macros/Cargo.toml
@@ -10,6 +10,7 @@ proc-macro = true
 [dependencies]
 syn = { version = "2.0.77", features = ["full", "extra-traits"] }
 quote = "1.0.37"
+proc-macro2 = "1.0.86"
 
 [dev-dependencies]
 teto = { version = "*", path = "../teto", features = ["macros"] }
diff --git a/teto_macros/src/lib.rs b/teto_macros/src/lib.rs
index 3300ad3..61dd3ae 100644
--- a/teto_macros/src/lib.rs
+++ b/teto_macros/src/lib.rs
@@ -6,21 +6,13 @@ use quote::quote;
 use syn::{FnArg, ItemFn};
 
 #[proc_macro_attribute]
-pub fn test(_args: TokenStream, tokens: TokenStream) -> TokenStream {
+pub fn test(args: TokenStream, tokens: TokenStream) -> TokenStream {
     let ItemFn {
         attrs,
         vis,
         mut sig,
         block,
-    } = match syn::parse(tokens) {
-        Ok(inner) => inner,
-        Err(e) => {
-            return TokenStream::from(
-                syn::Error::new(e.span(), "teto::test can only be called on functions")
-                    .to_compile_error(),
-            )
-        }
-    };
+    } = parse_token(args, tokens)?;
 
     let inputs = std::mem::take(&mut sig.inputs);
     // Since this macro replaces #[test], we can assume that the only arg (if
@@ -31,8 +23,8 @@ pub fn test(_args: TokenStream, tokens: TokenStream) -> TokenStream {
     });
     let inner_fn_invocation = if let Some(test_token_name) = test_token_name {
         quote! {
-            let token = const { #test_token_name::__private_teto_macros_only__create_teto_test_token };
-            inner(token())
+            let token = <#test_token_name>::__private_teto_macros_only__create_teto_test_token();
+            inner(token)
         }
     } else {
         quote! {
@@ -56,6 +48,17 @@ pub fn test(_args: TokenStream, tokens: TokenStream) -> TokenStream {
     };
 
     TokenStream::from(wrapped)
+}q
+
+fn parse_token(args: TokenStream, tokens: TokenStream) -> syn::Result<ItemFn> {
+    if !args.is_empty() {
+        return Err(syn::Error::new_spanned(
+            proc_macro2::TokenStream::from(args),
+            "teto macros do not accept arguments",
+        ));
+    }
+    syn::parse(tokens)
+        .map_err(|e| syn::Error::new(e.span(), "teto::test can only be called on functions"))
 }
 
 // #[proc_macro_attribute]
diff --git a/teto_macros/tests/ui.rs b/teto_macros/tests/ui.rs
index 578b927..0e69c46 100644
--- a/teto_macros/tests/ui.rs
+++ b/teto_macros/tests/ui.rs
@@ -1,6 +1,6 @@
 #[test]
 fn ui() {
     let t = trybuild::TestCases::new();
-    t.pass("tests/ui/pass/*.rs");
     t.compile_fail("tests/ui/compile_fail/*.rs");
+    t.pass("tests/ui/pass/*.rs");
 }