From ee9aa492a61d3d58dd6f0fb8382ce8b4dddb00cb Mon Sep 17 00:00:00 2001
From: clearloop <tianyi.gc@gmail.com>
Date: Sat, 30 Nov 2024 17:01:07 +0700
Subject: [PATCH] feat(zink): introduce proc-macro assert

---
 examples/revert.rs         |  7 +++++--
 zink/codegen/src/lib.rs    | 12 +++++------
 zink/codegen/src/revert.rs | 43 +++++++++++++++++++++++++++++++++++++-
 3 files changed, 52 insertions(+), 10 deletions(-)

diff --git a/examples/revert.rs b/examples/revert.rs
index e26e9a53f..620e1197f 100644
--- a/examples/revert.rs
+++ b/examples/revert.rs
@@ -9,13 +9,13 @@ fn main() {}
 
 /// check if the passing address is owner
 #[zink::external]
-pub fn run_revert() {
+pub fn revert() {
     zink::revert!("revert works")
 }
 
 /// check if the passing address is owner
 #[zink::external]
-pub fn run_assert() {
+pub fn assert() {
     zink::assert!(false, "assert works");
 }
 
@@ -26,5 +26,8 @@ fn test_revert() -> anyhow::Result<()> {
 
     let info = contract.execute(["revert()".as_bytes()])?;
     assert_eq!(info.revert, Some("revert works".into()));
+
+    let info = contract.execute(["assert()".as_bytes()])?;
+    assert_eq!(info.revert, Some("assert works".into()));
     Ok(())
 }
diff --git a/zink/codegen/src/lib.rs b/zink/codegen/src/lib.rs
index a98387685..3194ab4e9 100644
--- a/zink/codegen/src/lib.rs
+++ b/zink/codegen/src/lib.rs
@@ -4,7 +4,9 @@
 extern crate proc_macro;
 
 use proc_macro::TokenStream;
-use syn::{parse_macro_input, Attribute, DeriveInput, ItemFn, ItemStruct, LitStr};
+use proc_macro2::Span;
+use quote::ToTokens;
+use syn::{parse_macro_input, Attribute, DeriveInput, Expr, ItemFn, ItemStruct, LitStr};
 
 mod event;
 mod revert;
@@ -27,12 +29,8 @@ pub fn revert(input: TokenStream) -> TokenStream {
 /// message only support raw string.
 #[proc_macro]
 pub fn assert(input: TokenStream) -> TokenStream {
-    todo!(
-        r#"
-1. split input into condition and message then do the following
-2. if not cond then revert!(message);
-"#
-    )
+    let input = parse_macro_input!(input as revert::AssertInput);
+    revert::parse_assert(input)
 }
 
 /// Event logging interface
diff --git a/zink/codegen/src/revert.rs b/zink/codegen/src/revert.rs
index e03076cdb..e8fc1305c 100644
--- a/zink/codegen/src/revert.rs
+++ b/zink/codegen/src/revert.rs
@@ -3,7 +3,10 @@
 use proc_macro::TokenStream;
 use proc_macro2::{Literal, Span};
 use quote::{quote, ToTokens};
-use syn::{Ident, LitStr};
+use syn::{
+    parse::{Parse, ParseStream},
+    parse2, Expr, Ident, LitStr, Token,
+};
 
 /// Revert with message
 pub fn parse(input: LitStr) -> TokenStream {
@@ -34,3 +37,41 @@ pub fn parse(input: LitStr) -> TokenStream {
     }
     .into()
 }
+
+/// Parse assert macro
+pub fn parse_assert(input: AssertInput) -> TokenStream {
+    let cond = input.cond;
+    let revert: Expr = syn::parse2(
+        parse(
+            input
+                .message
+                .unwrap_or(LitStr::new("unknown error", Span::call_site())),
+        )
+        .into(),
+    )
+    .expect("Invalid revert message");
+
+    quote! {
+        if !#cond {
+            #revert
+        }
+    }
+    .into()
+}
+
+/// Assert input
+pub struct AssertInput {
+    pub cond: Expr,
+    pub comma: Token![,],
+    pub message: Option<LitStr>,
+}
+
+impl Parse for AssertInput {
+    fn parse(input: ParseStream) -> syn::Result<Self> {
+        Ok(AssertInput {
+            cond: input.parse()?,
+            comma: input.parse()?,
+            message: input.parse()?,
+        })
+    }
+}