diff --git a/examples/revert.rs b/examples/revert.rs index d0ae2d7cf..620e1197f 100644 --- a/examples/revert.rs +++ b/examples/revert.rs @@ -9,10 +9,16 @@ 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 assert() { + zink::assert!(false, "assert works"); +} + #[test] fn test_revert() -> anyhow::Result<()> { use zint::Contract; @@ -20,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 38b8e877d..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; @@ -21,6 +23,16 @@ pub fn revert(input: TokenStream) -> TokenStream { revert::parse(input) } +/// Check and expression and revert with the input message +/// +/// This is similar with the builtin `assert!` in rust, but the revert +/// message only support raw string. +#[proc_macro] +pub fn assert(input: TokenStream) -> TokenStream { + let input = parse_macro_input!(input as revert::AssertInput); + revert::parse_assert(input) +} + /// Event logging interface /// /// ```ignore 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, +} + +impl Parse for AssertInput { + fn parse(input: ParseStream) -> syn::Result { + Ok(AssertInput { + cond: input.parse()?, + comma: input.parse()?, + message: input.parse()?, + }) + } +} diff --git a/zink/src/lib.rs b/zink/src/lib.rs index 3aeee8205..3b26061af 100644 --- a/zink/src/lib.rs +++ b/zink/src/lib.rs @@ -13,7 +13,7 @@ pub mod storage; pub use self::{asm::Asm, event::Event}; pub use storage::{DoubleKeyMapping, Mapping, Storage}; -pub use zink_codegen::{external, revert, storage, Event}; +pub use zink_codegen::{assert, external, revert, storage, Event}; /// Generate a keccak hash of the input (sha3) #[cfg(not(target_family = "wasm"))]