Refactor and simplify the elf module a bit (#796)

* Eliminate the `ElfFirmwareImage` struct

* Eliminate redundant `RomSegment` struct, rename `CodeSegment` to `Segment`

* Re-export types which need to be public from `elf` module

* Update `CHANGELOG.md`

* Remove unused `segments_with_load_addresses` function

Author: jessebraham

CHANGELOG.md

@@ -23,6 +23,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - The `command`, `elf` and `error` modules are no longer public (#772)
 - `write-bin` now works for files whose lengths are not divisible by 4 (#780, #788)
 - `get_usb_pid` is now `usb_pid` and no longer needlessly returns a `Result` (#795)
+- `CodeSegment` and `RomSegment` have been merged into a single `Segment` struct (#796)
 
 ### Fixed
 
@@ -34,6 +35,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Removed
 
 - Removed the `libudev` feature (#742)
+- The `FirmwareImage` trait no longer includes the `segments_with_load_addresses` function (#796)
 
 ## [3.3.0] - 2025-01-13
 

espflash/src/cli/mod.rs

@@ -26,6 +26,7 @@ use indicatif::{style::ProgressStyle, HumanCount, ProgressBar};
 use log::{debug, info, warn};
 use miette::{IntoDiagnostic, Result, WrapErr};
 use serialport::{FlowControl, SerialPortInfo, SerialPortType, UsbPortInfo};
+use xmas_elf::ElfFile;
 
 use self::{
     config::Config,
@@ -34,8 +35,7 @@ use self::{
 };
 use crate::{
     connection::reset::{ResetAfterOperation, ResetBeforeOperation},
-    elf::ElfFirmwareImage,
-    error::{Error, MissingPartition, MissingPartitionTable},
+    error::{ElfError, Error, MissingPartition, MissingPartitionTable},
     flasher::{
         parse_partition_table,
         FlashData,
@@ -595,14 +595,16 @@ pub fn save_elf_as_image(
     skip_padding: bool,
     xtal_freq: XtalFrequency,
 ) -> Result<()> {
-    let image = ElfFirmwareImage::try_from(elf_data)?;
+    let elf = ElfFile::new(elf_data)
+        .map_err(ElfError::from)
+        .into_diagnostic()?;
 
     if merge {
         // To get a chip revision, the connection is needed
         // For simplicity, the revision None is used
         let image =
             chip.into_target()
-                .get_flash_image(&image, flash_data.clone(), None, xtal_freq)?;
+                .get_flash_image(&elf, flash_data.clone(), None, xtal_freq)?;
 
         display_image_size(image.app_size(), image.part_size());
 
@@ -636,7 +638,7 @@ pub fn save_elf_as_image(
     } else {
         let image = chip
             .into_target()
-            .get_flash_image(&image, flash_data, None, xtal_freq)?;
+            .get_flash_image(&elf, flash_data, None, xtal_freq)?;
 
         display_image_size(image.app_size(), image.part_size());
 

espflash/src/elf.rs

@@ -9,77 +9,45 @@ use std::{
 };
 
 use xmas_elf::{
-    program::Type,
     sections::{SectionData, ShType},
     ElfFile,
 };
 
-use crate::{
-    error::{ElfError, Error},
-    targets::Chip,
-};
+use crate::targets::Chip;
 
 /// Operations for working with firmware images
 pub trait FirmwareImage<'a> {
     /// Firmware image entry point
     fn entry(&self) -> u32;
 
     /// Firmware image segments
-    fn segments(&'a self) -> Box<dyn Iterator<Item = CodeSegment<'a>> + 'a>;
-
-    /// Firmware image segments, with their associated load addresses
-    fn segments_with_load_addresses(&'a self) -> Box<dyn Iterator<Item = CodeSegment<'a>> + 'a>;
+    fn segments(&'a self) -> Box<dyn Iterator<Item = Segment<'a>> + 'a>;
 
     /// Firmware image ROM segments
-    fn rom_segments(&'a self, chip: Chip) -> Box<dyn Iterator<Item = CodeSegment<'a>> + 'a> {
+    fn rom_segments(&'a self, chip: Chip) -> Box<dyn Iterator<Item = Segment<'a>> + 'a> {
         Box::new(
             self.segments()
                 .filter(move |segment| chip.into_target().addr_is_flash(segment.addr)),
         )
     }
 
     /// Firmware image RAM segments
-    fn ram_segments(&'a self, chip: Chip) -> Box<dyn Iterator<Item = CodeSegment<'a>> + 'a> {
+    fn ram_segments(&'a self, chip: Chip) -> Box<dyn Iterator<Item = Segment<'a>> + 'a> {
         Box::new(
             self.segments()
                 .filter(move |segment| !chip.into_target().addr_is_flash(segment.addr)),
         )
     }
 }
 
-/// A firmware image built from an ELF file
-#[derive(Debug)]
-pub struct ElfFirmwareImage<'a> {
-    elf: ElfFile<'a>,
-}
-
-impl<'a> ElfFirmwareImage<'a> {
-    pub fn new(elf: ElfFile<'a>) -> Self {
-        Self { elf }
-    }
-}
-
-impl<'a> TryFrom<&'a [u8]> for ElfFirmwareImage<'a> {
-    type Error = Error;
-
-    fn try_from(value: &'a [u8]) -> Result<Self, Self::Error> {
-        let elf = ElfFile::new(value).map_err(ElfError::from)?;
-
-        let image = ElfFirmwareImage::new(elf);
-
-        Ok(image)
-    }
-}
-
-impl<'a> FirmwareImage<'a> for ElfFirmwareImage<'a> {
+impl<'a> FirmwareImage<'a> for ElfFile<'a> {
     fn entry(&self) -> u32 {
-        self.elf.header.pt2.entry_point() as u32
+        self.header.pt2.entry_point() as u32
     }
 
-    fn segments(&'a self) -> Box<dyn Iterator<Item = CodeSegment<'a>> + 'a> {
+    fn segments(&'a self) -> Box<dyn Iterator<Item = Segment<'a>> + 'a> {
         Box::new(
-            self.elf
-                .section_iter()
+            self.section_iter()
                 .filter(|header| {
                     header.size() > 0
                         && header.get_type() == Ok(ShType::ProgBits)
@@ -88,50 +56,33 @@ impl<'a> FirmwareImage<'a> for ElfFirmwareImage<'a> {
                 })
                 .flat_map(move |header| {
                     let addr = header.address() as u32;
-                    let data = match header.get_data(&self.elf) {
+                    let data = match header.get_data(self) {
                         Ok(SectionData::Undefined(data)) => data,
                         _ => return None,
                     };
-                    Some(CodeSegment::new(addr, data))
-                }),
-        )
-    }
-
-    fn segments_with_load_addresses(&'a self) -> Box<dyn Iterator<Item = CodeSegment<'a>> + 'a> {
-        Box::new(
-            self.elf
-                .program_iter()
-                .filter(|header| {
-                    header.file_size() > 0
-                        && header.get_type() == Ok(Type::Load)
-                        && header.offset() > 0
-                })
-                .flat_map(move |header| {
-                    let addr = header.physical_addr() as u32;
-                    let from = header.offset() as usize;
-                    let to = header.offset() as usize + header.file_size() as usize;
-                    let data = &self.elf.input[from..to];
-                    Some(CodeSegment::new(addr, data))
+                    Some(Segment::new(addr, data))
                 }),
         )
     }
 }
 
-#[derive(Eq, Clone, Default)]
 /// A segment of code from the source ELF
-pub struct CodeSegment<'a> {
+#[derive(Default, Clone, Eq)]
+pub struct Segment<'a> {
     /// Base address of the code segment
     pub addr: u32,
-    data: Cow<'a, [u8]>,
+    /// Segment data
+    pub data: Cow<'a, [u8]>,
 }
 
-impl<'a> CodeSegment<'a> {
+impl<'a> Segment<'a> {
     pub fn new(addr: u32, data: &'a [u8]) -> Self {
-        let mut segment = CodeSegment {
+        let mut segment = Segment {
             addr,
             data: Cow::Borrowed(data),
         };
         segment.pad_align(4);
+
         segment
     }
 
@@ -149,7 +100,7 @@ impl<'a> CodeSegment<'a> {
                     (Cow::Owned(data), Cow::Owned(tail))
                 }
             };
-            let new = CodeSegment {
+            let new = Segment {
                 addr: self.addr,
                 data: head,
             };
@@ -183,28 +134,39 @@ impl<'a> CodeSegment<'a> {
             self.data = Cow::Owned(data);
         }
     }
+
+    /// Borrow the segment for the given lifetime
+    pub fn borrow<'b>(&'b self) -> Segment<'b>
+    where
+        'a: 'b,
+    {
+        Segment {
+            addr: self.addr,
+            data: Cow::Borrowed(self.data.as_ref()),
+        }
+    }
 }
 
-impl AddAssign<&'_ [u8]> for CodeSegment<'_> {
+impl AddAssign<&'_ [u8]> for Segment<'_> {
     fn add_assign(&mut self, rhs: &'_ [u8]) {
         let mut data = take(&mut self.data).into_owned();
         data.extend_from_slice(rhs);
         self.data = Cow::Owned(data);
     }
 }
 
-impl AddAssign<&'_ CodeSegment<'_>> for CodeSegment<'_> {
-    fn add_assign(&mut self, rhs: &'_ CodeSegment<'_>) {
+#[allow(clippy::suspicious_op_assign_impl)]
+impl AddAssign<&'_ Segment<'_>> for Segment<'_> {
+    fn add_assign(&mut self, rhs: &'_ Segment<'_>) {
         let mut data = take(&mut self.data).into_owned();
-        // pad or truncate
-        #[allow(clippy::suspicious_op_assign_impl)]
+        // Pad or truncate:
         data.resize((rhs.addr - self.addr) as usize, 0);
         data.extend_from_slice(rhs.data());
         self.data = Cow::Owned(data);
     }
 }
 
-impl Debug for CodeSegment<'_> {
+impl Debug for Segment<'_> {
     fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
         f.debug_struct("CodeSegment")
             .field("addr", &self.addr)
@@ -213,50 +175,20 @@ impl Debug for CodeSegment<'_> {
     }
 }
 
-impl PartialEq for CodeSegment<'_> {
+impl PartialEq for Segment<'_> {
     fn eq(&self, other: &Self) -> bool {
         self.addr.eq(&other.addr)
     }
 }
 
-impl PartialOrd for CodeSegment<'_> {
+impl PartialOrd for Segment<'_> {
     fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
         Some(self.cmp(other))
     }
 }
 
-impl Ord for CodeSegment<'_> {
+impl Ord for Segment<'_> {
     fn cmp(&self, other: &Self) -> Ordering {
         self.addr.cmp(&other.addr)
     }
 }
-
-/// A segment of data to write to the flash
-#[derive(Debug, Clone)]
-pub struct RomSegment<'a> {
-    /// ROM address at which the segment begins
-    pub addr: u32,
-    /// Segment data
-    pub data: Cow<'a, [u8]>,
-}
-
-impl<'a> RomSegment<'a> {
-    pub fn borrow<'b>(&'b self) -> RomSegment<'b>
-    where
-        'a: 'b,
-    {
-        RomSegment {
-            addr: self.addr,
-            data: Cow::Borrowed(self.data.as_ref()),
-        }
-    }
-}
-
-impl<'a> From<CodeSegment<'a>> for RomSegment<'a> {
-    fn from(segment: CodeSegment<'a>) -> Self {
-        RomSegment {
-            addr: segment.addr,
-            data: segment.data,
-        }
-    }
-}