From e12c36f5878c4f8c050f1e1dff3281f875f8a637 Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Piotr=20Miko=C5=82ajczyk?= <piomiko41@gmail.com>
Date: Tue, 15 Oct 2024 13:18:04 +0200
Subject: [PATCH] Use spans

---
 src/logging.rs         | 13 +++----------
 src/main.rs            | 19 +++++++++----------
 src/scenario_config.rs |  6 +++---
 3 files changed, 15 insertions(+), 23 deletions(-)

diff --git a/src/logging.rs b/src/logging.rs
index 88dd238..4433e99 100644
--- a/src/logging.rs
+++ b/src/logging.rs
@@ -1,14 +1,7 @@
-use std::{io, time::Instant};
+use std::io;
 
-use tracing::{info, level_filters::LevelFilter};
-use tracing_subscriber::{EnvFilter, fmt::format::FmtSpan};
-
-pub fn measure<T, Action: FnOnce() -> T>(action_name: &str, action: Action) -> T {
-    let start = Instant::now();
-    let result = action();
-    info!("{action_name}: {:?}", start.elapsed());
-    result
-}
+use tracing::level_filters::LevelFilter;
+use tracing_subscriber::{fmt::format::FmtSpan, EnvFilter};
 
 fn get_filter() -> EnvFilter {
     EnvFilter::builder()
diff --git a/src/main.rs b/src/main.rs
index 0854d29..31bc5c9 100644
--- a/src/main.rs
+++ b/src/main.rs
@@ -1,10 +1,10 @@
 use scenario_config::ScenarioConfig;
+use tracing::info_span;
 
 use crate::{
     folding::{prepare_folding, verify_folding, FoldingSchemeExt, HyperNovaFolding, NovaFolding},
-    logging::init_logging,
+    logging::{init_logging},
 };
-use crate::logging::measure;
 
 mod circuit;
 mod logging;
@@ -22,17 +22,16 @@ fn scenario<FS: FoldingSchemeExt>(
     // ============== FOLDING PREPARATION ==========================================================
 
     let start_state = config.start_ivc_state.clone();
-    let (mut folding, folding_vp) = measure("Prepare folding", || {
-        prepare_folding::<FS>(&config.circuit, start_state.clone(), rng)
-    });
+    let (mut folding, folding_vp) = info_span!("Prepare folding")
+        .in_scope(|| prepare_folding::<FS>(&config.circuit, start_state.clone(), rng));
+
+    let input = info_span!("Transform input")
+        .in_scope(|| folding.transform_inputs(config.input().to_vec(), start_state, &mut *rng));
 
     // ============== FOLDING ======================================================================
 
-    let input = measure("Transform input", || {
-        folding.transform_inputs(config.input().to_vec(), start_state.clone(), &mut *rng)
-    });
     for (i, step_input) in input.into_iter().enumerate() {
-        measure(&format!("Prove_step {i}"), || {
+        info_span!("Folding step", step = i).in_scope(|| {
             folding
                 .prove_step(
                     &mut *rng,
@@ -45,7 +44,7 @@ fn scenario<FS: FoldingSchemeExt>(
 
     // ============== FOLDING VERIFICATION =========================================================
 
-    measure("Folding verification", || {
+    info_span!("Folding verification").in_scope(|| {
         verify_folding(
             &folding,
             folding_vp,
diff --git a/src/scenario_config.rs b/src/scenario_config.rs
index 0387d63..a90f7a5 100644
--- a/src/scenario_config.rs
+++ b/src/scenario_config.rs
@@ -1,9 +1,9 @@
 use ark_bn254::Fr;
 use num_traits::Zero;
 use sonobe::frontend::circom::CircomFCircuit;
+use tracing::info_span;
 
 use crate::{circuit::create_circuit, input::prepare_input};
-use crate::logging::measure;
 
 #[derive(Clone)]
 pub struct ScenarioConfig {
@@ -18,8 +18,8 @@ impl ScenarioConfig {
         Self {
             num_inputs: 6,
             start_ivc_state: vec![Fr::zero(); 2],
-            circuit: measure("Prepare circuit", create_circuit),
-            input: measure("Prepare input", prepare_input),
+            circuit: info_span!("Prepare circuit").in_scope(create_circuit),
+            input: info_span!("Prepare input").in_scope(prepare_input),
         }
     }