Print app size when using any image format, cleanup/use new features (#281)

* Use the new #[default] enum variant attribute

* Add support for printing app size using any image format

Author: jessebraham

espflash/src/cli/mod.rs

@@ -17,6 +17,7 @@ use std::{
 use clap::{builder::ArgPredicate, Args};
 use comfy_table::{modifiers, presets::UTF8_FULL, Attribute, Cell, Color, Table};
 use esp_idf_part::{DataType, Partition, PartitionTable};
+use indicatif::HumanCount;
 use log::{debug, info};
 use miette::{IntoDiagnostic, Result, WrapErr};
 use serialport::{SerialPortType, UsbPortInfo};
@@ -296,6 +297,8 @@ pub fn save_elf_as_image(
             flash_freq,
         )?;
 
+        display_image_size(image.app_size(), image.part_size());
+
         let mut file = fs::OpenOptions::new()
             .write(true)
             .truncate(true)
@@ -317,13 +320,13 @@ pub fn save_elf_as_image(
             // Take flash_size as input parameter, if None, use default value of 4Mb
             let padding_bytes = vec![
                 0xffu8;
-                flash_size.unwrap_or(FlashSize::Flash4Mb).size() as usize
+                flash_size.unwrap_or_default().size() as usize
                     - file.metadata().into_diagnostic()?.len() as usize
             ];
             file.write_all(&padding_bytes).into_diagnostic()?;
         }
     } else {
-        let flash_image = chip.into_target().get_flash_image(
+        let image = chip.into_target().get_flash_image(
             &image,
             None,
             None,
@@ -333,8 +336,10 @@ pub fn save_elf_as_image(
             flash_size,
             flash_freq,
         )?;
-        let parts: Vec<_> = flash_image.ota_segments().collect();
 
+        display_image_size(image.app_size(), image.part_size());
+
+        let parts = image.ota_segments().collect::<Vec<_>>();
         match parts.as_slice() {
             [single] => fs::write(&image_path, &single.data).into_diagnostic()?,
             parts => {
@@ -349,6 +354,20 @@ pub fn save_elf_as_image(
     Ok(())
 }
 
+pub(crate) fn display_image_size(app_size: u32, part_size: Option<u32>) {
+    if let Some(part_size) = part_size {
+        let percent = app_size as f32 / part_size as f32 * 100.0;
+        println!(
+            "App/part. size:    {}/{} bytes, {:.2}%",
+            HumanCount(app_size as u64),
+            HumanCount(part_size as u64),
+            percent
+        );
+    } else {
+        println!("App size:          {} bytes", HumanCount(app_size as u64));
+    }
+}
+
 /// Write an ELF image to a target device's flash
 pub fn flash_elf_image(
     flasher: &mut Flasher,
@@ -425,7 +444,8 @@ pub fn erase_partitions(
     }
 
     // Look for any data partitions with specific data subtype
-    // There might be multiple partition of the same subtype, e.g. when using multiple FAT partitions
+    // There might be multiple partition of the same subtype, e.g. when using
+    // multiple FAT partitions
     if let Some(partition_types) = erase_data_parts {
         for ty in partition_types {
             for part in partition_table.partitions() {

espflash/src/flasher/mod.rs

@@ -38,7 +38,7 @@ const FLASH_SECTORS_PER_BLOCK: usize = FLASH_SECTOR_SIZE / FLASH_BLOCK_SIZE;
 /// Supported flash frequencies
 ///
 /// Note that not all frequencies are supported by each target device.
-#[derive(Debug, Clone, Copy, Hash, PartialEq, Eq, Display, EnumVariantNames)]
+#[derive(Debug, Default, Clone, Copy, Hash, PartialEq, Eq, Display, EnumVariantNames)]
 #[repr(u8)]
 pub enum FlashFrequency {
     /// 12 MHz
@@ -63,6 +63,7 @@ pub enum FlashFrequency {
     #[strum(serialize = "30M")]
     Flash30M,
     /// 40 MHz
+    #[default]
     #[strum(serialize = "40M")]
     Flash40M,
     /// 48 MHz
@@ -100,14 +101,15 @@ impl FromStr for FlashFrequency {
 }
 
 /// Supported flash modes
-#[derive(Copy, Clone, Debug, EnumVariantNames)]
+#[derive(Copy, Clone, Debug, Default, EnumVariantNames)]
 #[strum(serialize_all = "lowercase")]
 pub enum FlashMode {
     /// Quad I/O (4 pins used for address & data)
     Qio,
     /// Quad Output (4 pins used for data)
     Qout,
     /// Dual I/O (2 pins used for address & data)
+    #[default]
     Dio,
     /// Dual Output (2 pins used for data)
     Dout,
@@ -132,7 +134,7 @@ impl FromStr for FlashMode {
 /// Supported flash sizes
 ///
 /// Note that not all sizes are supported by each target device.
-#[derive(Clone, Copy, Debug, Eq, PartialEq, Display, EnumVariantNames)]
+#[derive(Clone, Copy, Debug, Default, Eq, PartialEq, Display, EnumVariantNames)]
 #[repr(u8)]
 pub enum FlashSize {
     /// 256 KB
@@ -148,6 +150,7 @@ pub enum FlashSize {
     #[strum(serialize = "2M")]
     Flash2Mb = 0x15,
     /// 4 MB
+    #[default]
     #[strum(serialize = "4M")]
     Flash4Mb = 0x16,
     /// 8 MB
@@ -472,7 +475,7 @@ impl Flasher {
             Ok(size) => size,
             Err(_) => {
                 warn!(
-                    "Could not detect flash size (FlashID=0x{:02X}, SizeID=0x{:02X}), defaulting to 4MB\n",
+                    "Could not detect flash size (FlashID=0x{:02X}, SizeID=0x{:02X}), defaulting to 4MB",
                     flash_id,
                     size_id
                 );
@@ -678,7 +681,7 @@ impl Flasher {
         let mut target = self.chip.flash_target(self.spi_params, self.use_stub);
         target.begin(&mut self.connection).flashing()?;
 
-        let flash_image = self.chip.into_target().get_flash_image(
+        let image = self.chip.into_target().get_flash_image(
             &image,
             bootloader,
             partition_table,
@@ -691,7 +694,10 @@ impl Flasher {
             flash_freq,
         )?;
 
-        for segment in flash_image.flash_segments() {
+        #[cfg(feature = "cli")]
+        crate::cli::display_image_size(image.app_size(), image.part_size());
+
+        for segment in image.flash_segments() {
             target
                 .write_segment(&mut self.connection, segment)
                 .flashing()?;

espflash/src/image_format/direct_boot.rs

@@ -56,6 +56,14 @@ impl<'a> ImageFormat<'a> for DirectBootFormat<'a> {
     {
         Box::new(once(self.segment.borrow()))
     }
+
+    fn app_size(&self) -> u32 {
+        self.segment.data.len() as u32
+    }
+
+    fn part_size(&self) -> Option<u32> {
+        None
+    }
 }
 
 #[cfg(test)]

espflash/src/image_format/esp8266.rs

@@ -17,6 +17,7 @@ use crate::{
 pub struct Esp8266Format<'a> {
     irom_data: Option<RomSegment<'a>>,
     flash_segment: RomSegment<'a>,
+    app_size: u32,
 }
 
 impl<'a> Esp8266Format<'a> {
@@ -37,9 +38,9 @@ impl<'a> Esp8266Format<'a> {
         );
 
         // Common header
-        let flash_mode = flash_mode.unwrap_or(FlashMode::Dio) as u8;
-        let flash_freq = flash_freq.unwrap_or(FlashFrequency::Flash40M);
-        let flash_size = flash_size.unwrap_or(FlashSize::Flash4Mb);
+        let flash_mode = flash_mode.unwrap_or_default() as u8;
+        let flash_freq = flash_freq.unwrap_or_default();
+        let flash_size = flash_size.unwrap_or_default();
         let flash_config =
             encode_flash_size(flash_size)? + encode_flash_frequency(Chip::Esp8266, flash_freq)?;
         let segment_count = image.ram_segments(Chip::Esp8266).count() as u8;
@@ -87,9 +88,16 @@ impl<'a> Esp8266Format<'a> {
             data: Cow::Owned(common_data),
         };
 
+        let app_size = irom_data
+            .clone()
+            .map(|d| d.data.len() as u32)
+            .unwrap_or_default()
+            + flash_segment.data.len() as u32;
+
         Ok(Self {
             irom_data,
             flash_segment,
+            app_size,
         })
     }
 }
@@ -118,6 +126,14 @@ impl<'a> ImageFormat<'a> for Esp8266Format<'a> {
                 .chain(once(self.flash_segment.borrow())),
         )
     }
+
+    fn app_size(&self) -> u32 {
+        self.app_size
+    }
+
+    fn part_size(&self) -> Option<u32> {
+        None
+    }
 }
 
 fn merge_rom_segments<'a>(

espflash/src/image_format/idf_bootloader.rs

@@ -1,8 +1,7 @@
 use std::{borrow::Cow, io::Write, iter::once};
 
 use bytemuck::{bytes_of, from_bytes, Pod, Zeroable};
-use esp_idf_part::{Partition, PartitionTable, Type};
-use indicatif::HumanCount;
+use esp_idf_part::{PartitionTable, Type};
 use sha2::{Digest, Sha256};
 
 use super::{
@@ -39,6 +38,8 @@ pub struct IdfBootloaderFormat<'a> {
     bootloader: Cow<'a, [u8]>,
     partition_table: PartitionTable,
     flash_segment: RomSegment<'a>,
+    app_size: u32,
+    part_size: u32,
 }
 
 impl<'a> IdfBootloaderFormat<'a> {
@@ -179,7 +180,14 @@ impl<'a> IdfBootloaderFormat<'a> {
             .or_else(|| partition_table.find_by_type(Type::App))
             .unwrap();
 
-        check_partition_stats(factory_partition, &data)?;
+        let app_size = data.len() as u32;
+        let part_size = factory_partition.size();
+
+        // The size of the application must not exceed the size of the factory
+        // partition.
+        if app_size as f32 / part_size as f32 > 1.0 {
+            return Err(Error::ElfTooBig);
+        }
 
         let flash_segment = RomSegment {
             addr: factory_partition.offset(),
@@ -191,6 +199,8 @@ impl<'a> IdfBootloaderFormat<'a> {
             bootloader,
             partition_table,
             flash_segment,
+            app_size,
+            part_size,
         })
     }
 }
@@ -219,6 +229,14 @@ impl<'a> ImageFormat<'a> for IdfBootloaderFormat<'a> {
     {
         Box::new(once(self.flash_segment.borrow()))
     }
+
+    fn app_size(&self) -> u32 {
+        self.app_size
+    }
+
+    fn part_size(&self) -> Option<u32> {
+        Some(self.part_size)
+    }
 }
 
 fn encode_flash_size(size: FlashSize) -> Result<u8, FlashDetectError> {
@@ -237,22 +255,6 @@ fn encode_flash_size(size: FlashSize) -> Result<u8, FlashDetectError> {
     }
 }
 
-fn check_partition_stats(part: &Partition, data: &Vec<u8>) -> Result<(), Error> {
-    let perc = data.len() as f32 / part.size() as f32 * 100.0;
-    println!(
-        "App/part. size:    {}/{} bytes, {:.2}%",
-        HumanCount(data.len() as u64),
-        HumanCount(part.size() as u64),
-        perc
-    );
-
-    if perc > 100.0 {
-        return Err(Error::ElfTooBig);
-    }
-
-    Ok(())
-}
-
 /// Actual alignment (in data bytes) required for a segment header: positioned
 /// so that after we write the next 8 byte header, file_offs % IROM_ALIGN ==
 /// segment.addr % IROM_ALIGN

espflash/src/image_format/mod.rs

@@ -55,6 +55,13 @@ pub trait ImageFormat<'a>: Send {
     fn ota_segments<'b>(&'b self) -> Box<dyn Iterator<Item = RomSegment<'b>> + 'b>
     where
         'a: 'b;
+
+    /// The size of the application binary
+    fn app_size(&self) -> u32;
+
+    /// If applicable, the size of the application partition (if it can be
+    /// determined)
+    fn part_size(&self) -> Option<u32>;
 }
 
 /// All supported firmware image formats

espflash/src/targets/flash_target/esp8266.rs

@@ -10,6 +10,7 @@ use crate::{
 };
 
 /// Applications running from an ESP8266's flash
+#[derive(Default)]
 pub struct Esp8266Target;
 
 impl Esp8266Target {
@@ -18,12 +19,6 @@ impl Esp8266Target {
     }
 }
 
-impl Default for Esp8266Target {
-    fn default() -> Self {
-        Self::new()
-    }
-}
-
 impl FlashTarget for Esp8266Target {
     fn begin(&mut self, connection: &mut Connection) -> Result<(), Error> {
         connection.command(Command::FlashBegin {