diff --git a/cairo-vm-cli/src/main.rs b/cairo-vm-cli/src/main.rs index ac816e9fbc..46797ccba3 100644 --- a/cairo-vm-cli/src/main.rs +++ b/cairo-vm-cli/src/main.rs @@ -1,13 +1,12 @@ #![deny(warnings)] #![forbid(unsafe_code)] -use bincode::enc::write::Writer; use cairo_vm::cairo_run::{self, EncodeTraceError}; use cairo_vm::hint_processor::builtin_hint_processor::builtin_hint_processor_definition::BuiltinHintProcessor; use cairo_vm::vm::errors::cairo_run_errors::CairoRunError; use cairo_vm::vm::errors::trace_errors::TraceError; use cairo_vm::vm::errors::vm_errors::VirtualMachineError; +use cairo_vm::write::BufWriter; use clap::{Parser, ValueHint}; -use std::io::{self, Write}; use std::path::PathBuf; use thiserror::Error; @@ -70,39 +69,6 @@ enum Error { Trace(#[from] TraceError), } -struct FileWriter { - buf_writer: io::BufWriter, - bytes_written: usize, -} - -impl Writer for FileWriter { - fn write(&mut self, bytes: &[u8]) -> Result<(), bincode::error::EncodeError> { - self.buf_writer - .write_all(bytes) - .map_err(|e| bincode::error::EncodeError::Io { - inner: e, - index: self.bytes_written, - })?; - - self.bytes_written += bytes.len(); - - Ok(()) - } -} - -impl FileWriter { - fn new(buf_writer: io::BufWriter) -> Self { - Self { - buf_writer, - bytes_written: 0, - } - } - - fn flush(&mut self) -> io::Result<()> { - self.buf_writer.flush() - } -} - fn run(args: impl Iterator) -> Result<(), Error> { let args = Args::try_parse_from(args); let args = match args { @@ -144,8 +110,7 @@ fn run(args: impl Iterator) -> Result<(), Error> { let relocated_trace = vm.get_relocated_trace()?; let trace_file = std::fs::File::create(trace_path)?; - let mut trace_writer = - FileWriter::new(io::BufWriter::with_capacity(3 * 1024 * 1024, trace_file)); + let mut trace_writer = BufWriter::with_capacity(3 * 1024 * 1024, trace_file); cairo_run::write_encoded_trace(relocated_trace, &mut trace_writer)?; trace_writer.flush()?; @@ -153,8 +118,7 @@ fn run(args: impl Iterator) -> Result<(), Error> { if let Some(memory_path) = args.memory_file { let memory_file = std::fs::File::create(memory_path)?; - let mut memory_writer = - FileWriter::new(io::BufWriter::with_capacity(5 * 1024 * 1024, memory_file)); + let mut memory_writer = BufWriter::with_capacity(5 * 1024 * 1024, memory_file); cairo_run::write_encoded_memory(&cairo_runner.relocated_memory, &mut memory_writer)?; memory_writer.flush()?; diff --git a/src/cairo_run.rs b/src/cairo_run.rs index 4976c829d6..e7ee97254b 100644 --- a/src/cairo_run.rs +++ b/src/cairo_run.rs @@ -7,9 +7,9 @@ use crate::{ security::verify_secure_runner, vm_core::VirtualMachine, }, + write::Writer, }; -use bincode::enc::write::Writer; use felt::Felt252; #[cfg(feature = "std")] diff --git a/src/lib.rs b/src/lib.rs index c6387994c4..2ed5961bfe 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -48,6 +48,91 @@ mod stdlib { pub use crate::without_std::*; } +pub mod write { + pub use bincode::enc::write::SliceWriter; + pub use bincode::enc::write::Writer; + pub use bincode::error::EncodeError; + + pub struct VecWriter { + buf: Vec, + } + + impl VecWriter { + pub fn new() -> Self { + Self { buf: Vec::new() } + } + + pub fn with_capacity(cap: usize) -> Self { + Self { + buf: Vec::with_capacity(cap), + } + } + + pub fn as_slice(&self) -> &[u8] { + &self.buf + } + } + + impl Writer for VecWriter { + fn write(&mut self, bytes: &[u8]) -> Result<(), bincode::error::EncodeError> { + self.buf.extend_from_slice(bytes); + Ok(()) + } + } + + impl Into> for VecWriter { + fn into(self) -> Vec { + self.buf + } + } + + #[cfg(feature = "std")] + use std::io::Write; + + #[cfg(feature = "std")] + pub struct BufWriter { + buf_writer: std::io::BufWriter, + bytes_written: usize, + } + + #[cfg(feature = "std")] + impl Writer for BufWriter { + fn write(&mut self, bytes: &[u8]) -> Result<(), EncodeError> { + self.buf_writer + .write_all(bytes) + .map_err(|e| EncodeError::Io { + inner: e, + index: self.bytes_written, + })?; + + self.bytes_written += bytes.len(); + + Ok(()) + } + } + + #[cfg(feature = "std")] + impl BufWriter { + pub fn new(writer: T) -> Self { + Self { + buf_writer: std::io::BufWriter::new(writer), + bytes_written: 0, + } + } + + pub fn with_capacity(cap: usize, writer: T) -> Self { + Self { + buf_writer: std::io::BufWriter::with_capacity(cap, writer), + bytes_written: 0, + } + } + + pub fn flush(&mut self) -> std::io::Result<()> { + self.buf_writer.flush() + } + } +} + pub extern crate felt; pub mod cairo_run; pub mod hint_processor;