Skip to content

Commit

Permalink
Add an optional callback trait which can be implemented and provided …
Browse files Browse the repository at this point in the history 
…to most flashing functions
  • Loading branch information
@jessebraham
jessebraham committed Jan 11, 2023
1 parent 911b28d commit 421a8dc
Show file tree
Hide file tree
Showing 8 changed files with 95 additions and 112 deletions.
6 changes: 3 additions & 3 deletions cargo-espflash/src/main.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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
Expand Down
15 changes: 6 additions & 9 deletions espflash/src/bin/espflash.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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();
Expand Down Expand Up @@ -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(())
}
88 changes: 36 additions & 52 deletions espflash/src/cli/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,6 @@
//! [espflash]: https://crates.io/crates/espflash

use std::{
borrow::Cow,
collections::HashMap,
fs,
io::Write,
Expand All @@ -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};
Expand All @@ -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,
Expand Down Expand Up @@ -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)?;
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -485,6 +468,7 @@ pub fn flash_elf_image(
flash_mode,
flash_size,
flash_freq,
Some(&mut EspflashProgress::default()),
)?;
info!("Flashing has completed!");

Expand Down
55 changes: 24 additions & 31 deletions espflash/src/flasher/mod.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -352,7 +362,7 @@ impl Flasher {
addr: text_addr,
data: Cow::Borrowed(&text),
},
None,
&mut None,
)
.flashing()?;

Expand All @@ -366,7 +376,7 @@ impl Flasher {
addr: data_addr,
data: Cow::Borrowed(&data),
},
None,
&mut None,
)
.flashing()?;

Expand Down Expand Up @@ -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);
Expand All @@ -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()?;
}

Expand All @@ -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)?;

Expand All @@ -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()?;
}

Expand All @@ -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,
Expand All @@ -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(())
Expand All @@ -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,
Expand All @@ -736,6 +728,7 @@ impl Flasher {
flash_mode,
flash_size,
flash_freq,
progress,
)
}

Expand Down
12 changes: 7 additions & 5 deletions espflash/src/targets/flash_target/esp32.rs
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
};

Expand Down Expand Up @@ -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;

Expand Down Expand Up @@ -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;
Expand All @@ -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(())
}

Expand Down
Loading

0 comments on commit 421a8dc

Please sign in to comment.