From fbd148016a5ce88038fb141e8a7205e9329829a7 Mon Sep 17 00:00:00 2001
From: Rodgers <cptrodgers@gmail.com>
Date: Fri, 25 Oct 2019 15:01:23 +0700
Subject: [PATCH] Get retry times

---
 .gitignore        |  1 +
 examples/retry.rs |  7 +++++++
 src/lib.rs        | 31 ++++++++++++++++++++++++++++++-
 3 files changed, 38 insertions(+), 1 deletion(-)
 create mode 100644 examples/retry.rs

diff --git a/.gitignore b/.gitignore
index 2f3e84f..c315266 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,3 +2,4 @@
 **/*.rs.bk
 Cargo.lock
 *.iml
+.idea
diff --git a/examples/retry.rs b/examples/retry.rs
new file mode 100644
index 0000000..1a2b5ab
--- /dev/null
+++ b/examples/retry.rs
@@ -0,0 +1,7 @@
+#[fort::root(retry = "2")]
+fn main() {
+    loop {
+        println!("Undying main!");
+        panic!("Error")
+    }
+}
diff --git a/src/lib.rs b/src/lib.rs
index 471622f..3fc9ffe 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -11,6 +11,7 @@ extern crate proc_macro;
 
 use proc_macro::TokenStream;
 use quote::{quote, quote_spanned};
+use syn::{NestedMeta, AttributeArgs, Lit, LitInt, Meta, MetaNameValue};
 
 /// Supplies bastion runtime to given `main`
 ///
@@ -24,8 +25,9 @@ use quote::{quote, quote_spanned};
 /// ```
 #[cfg(not(test))]
 #[proc_macro_attribute]
-pub fn root(_attr: TokenStream, item: TokenStream) -> TokenStream {
+pub fn root(attr: TokenStream, item: TokenStream) -> TokenStream {
     let input = syn::parse_macro_input!(item as syn::ItemFn);
+    let attr_args = syn::parse_macro_input!(attr as syn::AttributeArgs);
 
     let ret = &input.sig.output;
     let args = input.sig.inputs.iter();
@@ -40,6 +42,12 @@ pub fn root(_attr: TokenStream, item: TokenStream) -> TokenStream {
         return TokenStream::from(tokens);
     }
 
+    let retry_time = if let Some(retry_meta) = get_meta(&attr_args, "retry") {
+        parse_retry(&retry_meta)
+    } else {
+        0
+    };
+
     let result = quote! {
         use bastion::prelude::*;
 
@@ -62,3 +70,24 @@ pub fn root(_attr: TokenStream, item: TokenStream) -> TokenStream {
 
     result.into()
 }
+
+fn parse_retry(name_value: &MetaNameValue) -> usize {
+    if let Lit::Int(lit_int) = &name_value.lit {
+        lit_int.base10_parse::<usize>().unwrap()
+    } else {
+        0
+    }
+}
+
+fn get_meta(attr_args: &AttributeArgs, ident: &str) -> Option<MetaNameValue> {
+    let mut result = None;
+    for attr_args in attr_args {
+        if let NestedMeta::Meta(Meta::NameValue(name_value)) = attr_args {
+            if name_value.path.is_ident(ident) {
+                result = Some(name_value.clone())
+            }
+        }
+    }
+
+    result
+}