Add an optional callback trait which can be implemented and provided to most flashing functions

Author: jessebraham

cargo-espflash/src/main.rs

@@ -10,8 +10,8 @@ use espflash::{
     cli::{
         self, board_info, config::Config, connect, erase_partitions, flash_elf_image,
         monitor::monitor, parse_partition_table, partition_table, print_board_info,
-        save_elf_as_image, serial_monitor, ConnectArgs, FlashConfigArgs, MonitorArgs,
-        PartitionTableArgs,
+        save_elf_as_image, serial_monitor, ConnectArgs, EspflashProgress, FlashConfigArgs,
+        MonitorArgs, PartitionTableArgs,
     },
     image_format::ImageFormatKind,
     logging::initialize_logger,
@@ -171,7 +171,7 @@ fn flash(args: FlashArgs, config: &Config) -> Result<()> {
     print_board_info(&mut flasher)?;
 
     if args.flash_args.ram {
-        flasher.load_elf_to_ram(&elf_data)?;
+        flasher.load_elf_to_ram(&elf_data, Some(&mut EspflashProgress::default()))?;
     } else {
         let bootloader = args
             .flash_args

espflash/src/bin/espflash.rs

@@ -8,10 +8,10 @@ use std::{
 use clap::{Args, Parser, Subcommand};
 use espflash::{
     cli::{
-        self, board_info, build_progress_bar_callback, config::Config, connect, erase_partitions,
-        flash_elf_image, monitor::monitor, parse_partition_table, partition_table,
-        print_board_info, progress_bar, save_elf_as_image, serial_monitor, ConnectArgs,
-        FlashConfigArgs, MonitorArgs, PartitionTableArgs,
+        self, board_info, config::Config, connect, erase_partitions, flash_elf_image,
+        monitor::monitor, parse_partition_table, partition_table, print_board_info,
+        save_elf_as_image, serial_monitor, ConnectArgs, EspflashProgress, FlashConfigArgs,
+        MonitorArgs, PartitionTableArgs,
     },
     image_format::ImageFormatKind,
     logging::initialize_logger,
@@ -126,7 +126,7 @@ fn flash(args: FlashArgs, config: &Config) -> Result<()> {
     let elf_data = fs::read(&args.image).into_diagnostic()?;
 
     if args.flash_args.ram {
-        flasher.load_elf_to_ram(&elf_data)?;
+        flasher.load_elf_to_ram(&elf_data, Some(&mut EspflashProgress::default()))?;
     } else {
         let bootloader = args.flash_args.bootloader.as_deref();
         let partition_table = args.flash_args.partition_table.as_deref();
@@ -233,10 +233,7 @@ fn write_bin(args: WriteBinArgs, config: &Config) -> Result<()> {
     let mut buffer = Vec::with_capacity(size.try_into().into_diagnostic()?);
     f.read_to_end(&mut buffer).into_diagnostic()?;
 
-    let progress = progress_bar(format!("segment 0x{:X}", args.addr), None);
-    let progress_cb = Some(build_progress_bar_callback(progress));
-
-    flasher.write_bin_to_flash(args.addr, &buffer, progress_cb)?;
+    flasher.write_bin_to_flash(args.addr, &buffer, Some(&mut EspflashProgress::default()))?;
 
     Ok(())
 }

espflash/src/cli/mod.rs

@@ -8,7 +8,6 @@
 //! [espflash]: https://crates.io/crates/espflash
 
 use std::{
-    borrow::Cow,
     collections::HashMap,
     fs,
     io::Write,
@@ -18,7 +17,7 @@ use std::{
 use clap::Args;
 use comfy_table::{modifiers, presets::UTF8_FULL, Attribute, Cell, Color, Table};
 use esp_idf_part::{DataType, Partition, PartitionTable};
-use indicatif::{style::ProgressStyle, HumanCount, ProgressBar, ProgressDrawTarget};
+use indicatif::{style::ProgressStyle, HumanCount, ProgressBar};
 use log::{debug, info};
 use miette::{IntoDiagnostic, Result, WrapErr};
 use serialport::{SerialPortType, UsbPortInfo};
@@ -27,7 +26,7 @@ use self::{config::Config, monitor::monitor, serial::get_serial_port_info};
 use crate::{
     elf::ElfFirmwareImage,
     error::{MissingPartition, MissingPartitionTable},
-    flasher::{FlashFrequency, FlashMode, FlashSize, Flasher},
+    flasher::{FlashFrequency, FlashMode, FlashSize, Flasher, ProgressCallbacks},
     image_format::ImageFormatKind,
     interface::Interface,
     targets::Chip,
@@ -165,55 +164,6 @@ pub struct MonitorArgs {
     connect_args: ConnectArgs,
 }
 
-/// Create a new [ProgressBar] with some message and styling applied
-pub fn progress_bar<S>(msg: S, len: Option<u64>) -> ProgressBar
-where
-    S: Into<Cow<'static, str>>,
-{
-    // If no length was provided, or the provided length is 0, we will initially
-    // hide the progress bar. This is done to avoid drawing a full-width progress
-    // bar before we've determined the length.
-    let draw_target = match len {
-        Some(len) if len > 0 => ProgressDrawTarget::stderr(),
-        _ => ProgressDrawTarget::hidden(),
-    };
-
-    ProgressBar::with_draw_target(len, draw_target)
-        .with_message(msg)
-        .with_style(
-            ProgressStyle::default_bar()
-                .template("[{elapsed_precise}] [{bar:40}] {pos:>7}/{len:7} {msg}")
-                .unwrap()
-                .progress_chars("=> "),
-        )
-}
-
-/// Create a callback function for the provided [ProgressBar]
-pub fn build_progress_bar_callback(pb: ProgressBar) -> Box<dyn Fn(usize, usize)> {
-    // This is a bit of an odd callback function, as it handles the entire lifecycle
-    // of the progress bar.
-    Box::new(move |current: usize, total: usize| {
-        // If the length has not yet been set, set it and then change the draw target to
-        // make the progress bar visible.
-        match pb.length() {
-            Some(0) | None => {
-                pb.set_length(total as u64);
-                pb.set_draw_target(ProgressDrawTarget::stderr());
-            }
-            _ => {}
-        }
-
-        // Set the new position of the progress bar.
-        pb.set_position(current as u64);
-
-        // If we have reached the end, make sure to finish the progress bar to ensure
-        // proper output on the terminal.
-        if current == total {
-            pb.finish();
-        }
-    })
-}
-
 /// Select a serial port and establish a connection with a target device
 pub fn connect(args: &ConnectArgs, config: &Config) -> Result<Flasher> {
     let port_info = get_serial_port_info(args, config)?;
@@ -453,6 +403,39 @@ pub(crate) fn display_image_size(app_size: u32, part_size: Option<u32>) {
     }
 }
 
+/// Progress callback implementations for use in `cargo-espflash` and `espflash`
+#[derive(Default)]
+pub struct EspflashProgress {
+    pb: Option<ProgressBar>,
+}
+
+impl ProgressCallbacks for EspflashProgress {
+    fn init(&mut self, addr: u32, len: usize) {
+        let pb = ProgressBar::new(len as u64)
+            .with_message(format!("{addr:#X}"))
+            .with_style(
+                ProgressStyle::default_bar()
+                    .template("[{elapsed_precise}] [{bar:40}] {pos:>7}/{len:7} {msg}")
+                    .unwrap()
+                    .progress_chars("=> "),
+            );
+
+        self.pb = Some(pb);
+    }
+
+    fn update(&mut self, current: usize) {
+        if let Some(ref pb) = self.pb {
+            pb.set_position(current as u64);
+        }
+    }
+
+    fn finish(&mut self) {
+        if let Some(ref pb) = self.pb {
+            pb.finish();
+        }
+    }
+}
+
 /// Write an ELF image to a target device's flash
 pub fn flash_elf_image(
     flasher: &mut Flasher,
@@ -485,6 +468,7 @@ pub fn flash_elf_image(
         flash_mode,
         flash_size,
         flash_freq,
+        Some(&mut EspflashProgress::default()),
     )?;
     info!("Flashing has completed!");
 

espflash/src/flasher/mod.rs

@@ -263,6 +263,16 @@ pub struct DeviceInfo {
     pub mac_address: String,
 }
 
+/// Progress update callbacks
+pub trait ProgressCallbacks {
+    /// Initialize some progress report
+    fn init(&mut self, addr: u32, total: usize);
+    /// Update some progress report
+    fn update(&mut self, current: usize);
+    /// Finish some prgoress report
+    fn finish(&mut self);
+}
+
 /// Connect to and flash a target device
 pub struct Flasher {
     /// Connection for flash operations
@@ -352,7 +362,7 @@ impl Flasher {
                     addr: text_addr,
                     data: Cow::Borrowed(&text),
                 },
-                None,
+                &mut None,
             )
             .flashing()?;
 
@@ -366,7 +376,7 @@ impl Flasher {
                     addr: data_addr,
                     data: Cow::Borrowed(&data),
                 },
-                None,
+                &mut None,
             )
             .flashing()?;
 
@@ -605,7 +615,11 @@ impl Flasher {
     /// Load an elf image to ram and execute it
     ///
     /// Note that this will not touch the flash on the device
-    pub fn load_elf_to_ram(&mut self, elf_data: &[u8]) -> Result<(), Error> {
+    pub fn load_elf_to_ram(
+        &mut self,
+        elf_data: &[u8],
+        mut progress: Option<&mut dyn ProgressCallbacks>,
+    ) -> Result<(), Error> {
         let image = ElfFirmwareImage::try_from(elf_data)?;
         if image.rom_segments(self.chip).next().is_some() {
             return Err(Error::ElfNotRamLoadable);
@@ -620,20 +634,8 @@ impl Flasher {
         target.begin(&mut self.connection).flashing()?;
 
         for segment in image.ram_segments(self.chip) {
-            // Only display progress bars when the "cli" feature is enabled.
-            let progress_cb = if cfg!(feature = "cli") {
-                use crate::cli::{build_progress_bar_callback, progress_bar};
-
-                let progress = progress_bar(format!("segment 0x{:X}", segment.addr), None);
-                let progress_cb = build_progress_bar_callback(progress);
-
-                Some(progress_cb)
-            } else {
-                None
-            };
-
             target
-                .write_segment(&mut self.connection, segment.into(), progress_cb)
+                .write_segment(&mut self.connection, segment.into(), &mut progress)
                 .flashing()?;
         }
 
@@ -650,6 +652,7 @@ impl Flasher {
         flash_mode: Option<FlashMode>,
         flash_size: Option<FlashSize>,
         flash_freq: Option<FlashFrequency>,
+        mut progress: Option<&mut dyn ProgressCallbacks>,
     ) -> Result<(), Error> {
         let image = ElfFirmwareImage::try_from(elf_data)?;
 
@@ -676,20 +679,8 @@ impl Flasher {
         crate::cli::display_image_size(image.app_size(), image.part_size());
 
         for segment in image.flash_segments() {
-            // Only display progress bars when the "cli" feature is enabled.
-            let progress_cb = if cfg!(feature = "cli") {
-                use crate::cli::{build_progress_bar_callback, progress_bar};
-
-                let progress = progress_bar(format!("segment 0x{:X}", segment.addr), None);
-                let progress_cb = build_progress_bar_callback(progress);
-
-                Some(progress_cb)
-            } else {
-                None
-            };
-
             target
-                .write_segment(&mut self.connection, segment, progress_cb)
+                .write_segment(&mut self.connection, segment, &mut progress)
                 .flashing()?;
         }
 
@@ -703,7 +694,7 @@ impl Flasher {
         &mut self,
         addr: u32,
         data: &[u8],
-        progress_cb: Option<Box<dyn Fn(usize, usize)>>,
+        mut progress: Option<&mut dyn ProgressCallbacks>,
     ) -> Result<(), Error> {
         let segment = RomSegment {
             addr,
@@ -712,7 +703,7 @@ impl Flasher {
 
         let mut target = self.chip.flash_target(self.spi_params, self.use_stub);
         target.begin(&mut self.connection).flashing()?;
-        target.write_segment(&mut self.connection, segment, progress_cb)?;
+        target.write_segment(&mut self.connection, segment, &mut progress)?;
         target.finish(&mut self.connection, true).flashing()?;
 
         Ok(())
@@ -727,6 +718,7 @@ impl Flasher {
         flash_mode: Option<FlashMode>,
         flash_size: Option<FlashSize>,
         flash_freq: Option<FlashFrequency>,
+        progress: Option<&mut dyn ProgressCallbacks>,
     ) -> Result<(), Error> {
         self.load_elf_to_flash_with_format(
             elf_data,
@@ -736,6 +728,7 @@ impl Flasher {
             flash_mode,
             flash_size,
             flash_freq,
+            progress,
         )
     }
 

espflash/src/targets/flash_target/esp32.rs

@@ -11,7 +11,7 @@ use crate::{
     connection::{Connection, USB_SERIAL_JTAG_PID},
     elf::RomSegment,
     error::Error,
-    flasher::{SpiAttachParams, FLASH_SECTOR_SIZE},
+    flasher::{ProgressCallbacks, SpiAttachParams, FLASH_SECTOR_SIZE},
     targets::Chip,
 };
 
@@ -98,7 +98,7 @@ impl FlashTarget for Esp32Target {
         &mut self,
         connection: &mut Connection,
         segment: RomSegment,
-        progress_cb: Option<Box<dyn Fn(usize, usize)>>,
+        progress: &mut Option<&mut dyn ProgressCallbacks>,
     ) -> Result<(), Error> {
         let addr = segment.addr;
 
@@ -131,6 +131,8 @@ impl FlashTarget for Esp32Target {
         let chunks = compressed.chunks(flash_write_size);
         let num_chunks = chunks.len();
 
+        progress.as_mut().map(|cb| cb.init(addr, num_chunks));
+
         // decode the chunks to see how much data the device will have to save
         let mut decoder = ZlibDecoder::new(Vec::new());
         let mut decoded_size = 0;
@@ -154,11 +156,11 @@ impl FlashTarget for Esp32Target {
                 },
             )?;
 
-            if let Some(ref cb) = progress_cb {
-                cb(i + 1, num_chunks);
-            }
+            progress.as_mut().map(|cb| cb.update(i + 1));
         }
 
+        progress.as_mut().map(|cb| cb.finish());
+
         Ok(())
     }