From aa7bdb536db2a0021550f54f781cd5ada51f81e8 Mon Sep 17 00:00:00 2001
From: Putta Khunchalee <ultimaweapon@outlook.com>
Date: Sat, 17 Aug 2024 00:44:21 +0700
Subject: [PATCH] Uses syn to parse integration test body

---
 testing-macros/src/lib.rs  |  6 ++--
 testing-macros/src/qemu.rs | 73 +++++++-------------------------------
 2 files changed, 16 insertions(+), 63 deletions(-)

diff --git a/testing-macros/src/lib.rs b/testing-macros/src/lib.rs
index b93a4f7..6356471 100644
--- a/testing-macros/src/lib.rs
+++ b/testing-macros/src/lib.rs
@@ -1,6 +1,6 @@
 use self::qemu::parse_qemu_attribute;
 use proc_macro::TokenStream;
-use syn::Error;
+use syn::{parse_macro_input, Error, ItemFn};
 
 mod qemu;
 
@@ -10,7 +10,9 @@ mod qemu;
 /// must put everything that are needed within this function.
 #[proc_macro_attribute]
 pub fn qemu(_: TokenStream, item: TokenStream) -> TokenStream {
-    parse_qemu_attribute(item.into())
+    let item = parse_macro_input!(item as ItemFn);
+
+    parse_qemu_attribute(item)
         .unwrap_or_else(Error::into_compile_error)
         .into()
 }
diff --git a/testing-macros/src/qemu.rs b/testing-macros/src/qemu.rs
index 4616dd4..f756329 100644
--- a/testing-macros/src/qemu.rs
+++ b/testing-macros/src/qemu.rs
@@ -1,61 +1,21 @@
-use proc_macro2::{Delimiter, Group, Span, TokenStream, TokenTree};
-use quote::quote_spanned;
+use proc_macro2::{Span, TokenStream};
+use quote::ToTokens;
 use serde::{Deserialize, Serialize};
 use std::collections::HashMap;
 use std::env::VarError;
 use std::fmt::Write;
 use std::fs::create_dir_all;
 use std::path::{Path, PathBuf};
-use syn::Error;
-
-pub fn parse_qemu_attribute(item: TokenStream) -> Result<TokenStream, Error> {
-    // Get the function body.
-    let mut output = Vec::new();
-    let mut state = State::Start;
-    let mut name = None;
-    let mut body = None;
-
-    for t in item {
-        match &state {
-            State::Start => match &t {
-                TokenTree::Ident(i) if i == "fn" => state = State::Fn,
-                _ => {}
-            },
-            State::Fn => match &t {
-                TokenTree::Ident(i) => {
-                    name = Some(i.to_string());
-                    state = State::FnName;
-                }
-                _ => unreachable!(),
-            },
-            State::FnName => match &t {
-                TokenTree::Group(g) if g.delimiter() == Delimiter::Parenthesis => {
-                    state = State::Params
-                }
-                _ => {}
-            },
-            State::Params => match t {
-                TokenTree::Group(g) if g.delimiter() == Delimiter::Brace => {
-                    body = Some(g);
-                    break;
-                }
-                _ => {}
-            },
-        }
-
-        output.push(t);
-    }
+use syn::{parse_quote, Error, ItemFn};
 
+pub fn parse_qemu_attribute(mut item: ItemFn) -> Result<TokenStream, Error> {
     // Get path for the integration test data.
-    let item = body.unwrap();
     let root = match std::env::var("CARGO_TARGET_TMPDIR") {
         Ok(v) => PathBuf::from(v),
         Err(e) => match e {
             VarError::NotPresent => {
                 // We are not running by the integration test, keep the original function body.
-                output.push(TokenTree::Group(item));
-
-                return Ok(TokenStream::from_iter(output));
+                return Ok(item.into_token_stream());
             }
             VarError::NotUnicode(_) => {
                 return Err(Error::new(
@@ -67,21 +27,19 @@ pub fn parse_qemu_attribute(item: TokenStream) -> Result<TokenStream, Error> {
     };
 
     // Generate a test project.
-    let name = name.unwrap();
-    let span = item.span();
+    let name = item.sig.ident.to_string();
+    let body = item.block.brace_token.span.join().source_text().unwrap();
 
-    generate_test(root.join("project"), &name, &span.source_text().unwrap())?;
+    generate_test(root.join("project"), &name, &body)?;
 
     // Construct a new body.
     let root = root.to_str().unwrap();
-    let body = Group::new(
-        Delimiter::Brace,
-        quote_spanned!(span=> ::zfi_testing::run_qemu_test(::std::path::Path::new(#root));),
-    );
 
-    output.push(TokenTree::Group(body));
+    item.block = Box::new(parse_quote!({
+        ::zfi_testing::run_qemu_test(::std::path::Path::new(#root));
+    }));
 
-    Ok(TokenStream::from_iter(output))
+    Ok(item.into_token_stream())
 }
 
 fn generate_test<P: AsRef<Path>>(dir: P, name: &str, body: &str) -> Result<(), Error> {
@@ -223,13 +181,6 @@ fn generate_test<P: AsRef<Path>>(dir: P, name: &str, body: &str) -> Result<(), E
     Ok(())
 }
 
-enum State {
-    Start,
-    Fn,
-    FnName,
-    Params,
-}
-
 #[derive(Serialize, Deserialize)]
 struct Cargo {
     package: Option<Package>,