add wasmcov macro

near · Jan 16, 2024 · c382ec6 · c382ec6
1 parent 1343059
commit c382ec6
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 0 deletions.
diff --git a/near-sdk-macros/src/core_impl/code_generator/impl_item_method_info.rs b/near-sdk-macros/src/core_impl/code_generator/impl_item_method_info.rs
@@ -52,11 +52,13 @@ impl ImplItemMethodInfo {
     fn void_return_body_tokens(&self) -> TokenStream2 {
         let contract_init = self.contract_init_tokens();
         let method_invocation = self.method_invocation_tokens();
+        let wasmcov = self.wasmcov_tokens();
         let contract_ser = self.contract_ser_tokens();
 
         quote! {
             #contract_init
             #method_invocation;
+            #wasmcov
             #contract_ser
         }
     }
@@ -67,12 +69,14 @@ impl ImplItemMethodInfo {
         let contract_ser = self.contract_ser_tokens();
         let value_ser = self.value_ser_tokens();
         let value_return = self.value_return_tokens();
+        let wasmcov = self.wasmcov_tokens();
 
         quote! {
             #contract_init
             #method_invocation_with_return
             #value_ser
             #value_return
+            #wasmcov
             #contract_ser
         }
     }
@@ -84,6 +88,7 @@ impl ImplItemMethodInfo {
         let value_ser = self.value_ser_tokens();
         let value_return = self.value_return_tokens();
         let result_identifier = self.result_identifier();
+        let wasmcov = self.wasmcov_tokens();
 
         quote! {
             #contract_init
@@ -92,6 +97,7 @@ impl ImplItemMethodInfo {
                 ::std::result::Result::Ok(#result_identifier) => {
                     #value_ser
                     #value_return
+                    #wasmcov
                     #contract_ser
                 }
                 ::std::result::Result::Err(err) => ::near_sdk::FunctionError::panic(&err)
@@ -247,6 +253,17 @@ impl ImplItemMethodInfo {
         }
     }
 
+    fn wasmcov_tokens(&self) -> TokenStream2 {
+        quote! {
+            let mut coverage = vec![];
+            unsafe {
+                // Note that this function is not thread-safe! Use a lock if needed.
+                wasmcov::capture_coverage(&mut coverage).unwrap();
+            };
+            ::near_sdk::env::log_str();
+        }
+    }
+
     fn contract_ser_tokens(&self) -> TokenStream2 {
         use MethodKind::*;
 

diff --git a/near-sdk/Cargo.toml b/near-sdk/Cargo.toml
@@ -43,6 +43,7 @@ near-vm-logic = { version = "0.17", optional = true }
 near-primitives-core = { version = "0.17", optional = true }
 near-primitives = { version = "0.17", optional = true }
 near-crypto = { version = "0.17", optional = true }
+wasmcov = { version = "0.1.0" }
 
 [dev-dependencies]
 rand = "0.8.4"

diff --git a/near-sdk/src/lib.rs b/near-sdk/src/lib.rs
@@ -67,3 +67,5 @@ pub use bs58;
 pub use schemars;
 pub use serde;
 pub use serde_json;
+#[doc(hidden)]
+pub use wasmcov;