From dbad3d9bdee55da21568543708996bd0b1d55696 Mon Sep 17 00:00:00 2001 From: ts0yu <120932697+ts0yu@users.noreply.github.com> Date: Wed, 28 Aug 2024 09:15:11 +0000 Subject: [PATCH] feat: simplify inspector --- Cargo.lock | 9 +- src/engine/inspector.rs | 247 +++++++++++----------------------------- src/lib.rs | 1 - 3 files changed, 68 insertions(+), 189 deletions(-) diff --git a/Cargo.lock b/Cargo.lock index 456efbf..50ee68d 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -681,18 +681,15 @@ dependencies = [ "alloy-contract", "alloy-sol-macro", "alloy-sol-types", - "csv", - "plotly", - "rand", - "rand_distr", - "serde", - "serde_json", "async-trait", "clap", + "csv", "plotly", "rand", "rand_distr", "rug", + "serde", + "serde_json", "tokio", ] diff --git a/src/engine/inspector.rs b/src/engine/inspector.rs index 87369ae..912f583 100644 --- a/src/engine/inspector.rs +++ b/src/engine/inspector.rs @@ -1,13 +1,6 @@ -use std::collections::HashMap; -use plotly::{Plot, Scatter}; - -use std::fs::File; -use std::io::prelude::*; -use std::path::Path; use std::fs::OpenOptions; -use std::{error::Error, io, process}; -use serde::{Serialize, Deserialize}; -use serde_json::{Value, json}; + +use serde::{Deserialize, Serialize}; /// Trait allowing custom behavior to be defined for logging and inspecting values. pub trait Inspector { @@ -18,185 +11,109 @@ pub trait Inspector { fn inspect(&self, step: usize) -> Option; /// Save the inspector state. - fn save(&self, save_type: Option); -} - -pub enum SaveData { - ToCsv(String), - ToNewCsv(String), - ToJson(String), - ToNewJson(String), + fn save(&self); } +/// Type that allows for logging indexed values to files on disc. #[derive(Serialize, Deserialize, Debug, Clone)] pub struct LogMessage { - id: usize, // Id of the log - name: String, // Name of the log - data: String, // The data being stored inside the log + /// Index of the log message. + pub id: usize, + + /// Key of the log message. + pub name: String, + + /// Data of the log message. + pub data: String, } impl LogMessage { + /// Public constructor function for a new [`LogMessage`]. pub fn new(name: String, data: String) -> Self { - LogMessage { - id: 0, - name: name, - data: data, - } - } - - pub fn update_counter(&mut self, counter: usize) { - self.id = counter; + Self { id: 0, name, data } } } #[derive(Debug)] +/// Custom implementation of an [`Inspector`] for logging values to a file (CSV or JSON). pub struct Logger { - pub values: Vec, // Stores all the messages - pub counter: usize, // Used for keeping position of the id tag + values: Vec, + counter: usize, + file_path: String, + format: LogFormat, +} + +#[derive(Debug)] +/// Enum to specify the logging format. +enum LogFormat { + Csv, + Json, } impl Logger { - pub fn new() -> Self { - Logger { + /// Public constructor function for a new [`Logger`] for CSV format. + pub fn new_csv(file_path: String) -> Self { + Self { values: Vec::new(), counter: 0, + file_path, + format: LogFormat::Csv, } } - fn append_to_csv(record: LogMessage, path: &Path) -> Result<(), csv::Error> { - let file = OpenOptions::new() - .append(true) - .create(true) - .open(path)?; - - let mut writer = csv::Writer::from_writer(file); - - writer.serialize(( - record.id, - record.name, - record.data - ))?; - writer.flush()?; - - Ok(()) - } - - fn create_csv(file_location: &String) -> Result<(), csv::Error> { - let file = OpenOptions::new() - .write(true) - .create(true) - .truncate(true) - .open(file_location)?; - - let mut writer = csv::Writer::from_writer(file); - writer.write_record(&["id", "name", "data"])?; - writer.flush()?; - - Ok(()) + /// Public constructor function for a new [`Logger`] for JSON format. + pub fn new_json(file_path: String) -> Self { + Self { + values: Vec::new(), + counter: 0, + file_path, + format: LogFormat::Json, + } } - fn create_json(file_location: &String) -> Result<(), Box> { - let file = OpenOptions::new() - .write(true) - .create(true) - .truncate(true) - .open(file_location)?; - - let empty_array = json!([]); - serde_json::to_writer(file, &empty_array)?; - - Ok(()) - } - - fn append_to_json(record: LogMessage, path: &Path) -> Result<(), Box> { - let mut records = Self::read_json_file(&path.to_string_lossy().to_string())?; - - records.push(record); - + /// Append a log message to the appropriate file format. + fn append_to_file(&self, record: &LogMessage) -> Result<(), Box> { let file = OpenOptions::new() - .write(true) - .truncate(true) - .open(path)?; - - serde_json::to_writer(file, &records)?; - - Ok(()) - } - - fn read_json_file(file_loc: &String) -> Result, Box> { - let file = File::open(file_loc)?; - let records: Vec = serde_json::from_reader(file)?; - Ok(records) - } - - /* - fn read_csv_file(file_loc: &String) -> Result, csv::Error> { - let mut records: Vec = vec![]; - let file = OpenOptions::new() - .write(true) + .append(true) .create(true) - .truncate(true) - .open(file_loc)?; + .open(&self.file_path)?; - let mut reader = csv::Reader::from_reader(file); - - for record in reader.deserialize() { - // println!("Line"); - let record: Record = record?; - records.push(record); + match self.format { + LogFormat::Csv => { + let mut writer = csv::Writer::from_writer(file); + writer.serialize((record.id, &record.name, &record.data))?; + writer.flush()?; + } + LogFormat::Json => { + let mut records: Vec = + serde_json::from_reader(&file).unwrap_or_default(); + records.push(record.clone()); + serde_json::to_writer(file, &records)?; + } } - Ok(records) + Ok(()) } - */ } impl Inspector for Logger { fn log(&mut self, mut value: LogMessage) { - value.update_counter(self.counter); + value.id = self.counter; self.counter += 1; - self.values.push(value); + self.values.push(value.clone()); + + if let Err(e) = self.append_to_file(&value) { + eprintln!("Failed to append to file: {}", e); + } } fn inspect(&self, step: usize) -> Option { - Some(self.values[step].clone()) + self.values.get(step).cloned() } - fn save(&self, save_type: Option) { - let mut file_loc = String::new(); - match save_type.unwrap() { - SaveData::ToNewCsv(file_location) => { - Logger::create_csv(&file_location); - for log in self.values.clone() { - let file_path = Path::new(&file_location); - Logger::append_to_csv(log, &file_path); - } - }, - SaveData::ToCsv(file_location) => { - for log in self.values.clone() { - let file_path = Path::new(&file_location); - Logger::append_to_csv(log, &file_path); - } - }, - - SaveData::ToNewJson(file_location) => { - Logger::create_json(&file_location); - for log in self.values.clone() { - let file_path = Path::new(&file_location); - Logger::append_to_json(log, &file_path); - } - }, - - SaveData::ToJson(file_location) => { - for log in self.values.clone() { - let file_path = Path::new(&file_location); - Logger::append_to_json(log, &file_path); - } - } - } - } + fn save(&self) {} } -/// No-op implementation of an [`Inspector`] for custom usecases. +/// No-op implementation of an [`Inspector`] for custom use cases. pub struct EmptyInspector; impl Inspector for EmptyInspector { @@ -206,37 +123,3 @@ impl Inspector for EmptyInspector { fn log(&mut self, _value: f64) {} fn save(&self) {} } - -// pub struct Plotter { -// values: Vec, -// } - -// impl Plotter { -// pub fn new() -> Self { -// Plotter { -// values: Vec::new(), -// } -// } -// } - -// impl Inspector for Plotter { -// fn log(&mut self, value: f64) { -// self.values.push(value); -// } - -// fn inspect(&self, step: usize) -> Option { -// Some(self.values[step]) -// } - -// fn save(&self, save_type: Option) { -// let mut plot = Plot::new(); - -// let timesteps: Vec = (0..self.values.len()).collect(); -// let values: Vec = timesteps.iter().map(|v| v).collect(); - -// let trace = Scatter::new(timesteps, values).mode(plotly::common::Mode::Markers); - -// plot.add_trace(trace); -// plot.show(); -// } -// } diff --git a/src/lib.rs b/src/lib.rs index 60ef5ee..17a56bf 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -28,7 +28,6 @@ use alloy::{ }; use crate::{engine::inspector::Inspector, types::PoolManager::PoolKey}; -use crate::engine::inspector::SaveData; /// Provider type that includes all necessary fillers to execute transactions on an [`Anvil`] node. pub type AnvilProvider = FillProvider<