Skip to content

Commit d3e3426

Browse files
jessebrahamjessebraham
committed
Extract common functionality for erasing partition tables into a function
1 parent abd8208 commit d3e3426

File tree

3 files changed

+74
-87
lines changed

3 files changed

+74
-87
lines changed

‎cargo-espflash/src/main.rs‎

Lines changed: 7 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
use std::{
2-
collections::HashMap,
32
fs,
43
path::PathBuf,
54
process::{exit, Command, ExitStatus, Stdio},
@@ -9,11 +8,10 @@ use cargo_metadata::Message;
98
use clap::{Args, Parser, Subcommand};
109
use espflash::{
1110
cli::{
12-
self, board_info, clap_enum_variants, config::Config, connect, erase_partition,
11+
self, board_info, clap_enum_variants, config::Config, connect, erase_partitions,
1312
flash_elf_image, monitor::monitor, parse_partition_table, partition_table,
1413
save_elf_as_image, serial_monitor, ConnectArgs, FlashConfigArgs, PartitionTableArgs,
1514
},
16-
error::{MissingPartition, MissingPartitionTable},
1715
image_format::ImageFormatKind,
1816
logging::initialize_logger,
1917
targets::Chip,
@@ -200,46 +198,12 @@ fn flash(
200198
};
201199

202200
if args.flash_args.erase_parts.is_some() || args.flash_args.erase_data_parts.is_some() {
203-
let partition_table = match &partition_table {
204-
Some(partition_table) => partition_table,
205-
None => return Err((MissingPartitionTable {}).into()),
206-
};
207-
208-
// Using a hashmap to deduplicate entries
209-
let mut parts_to_erase = None;
210-
211-
// Look for any part with specific label
212-
if let Some(part_labels) = args.flash_args.erase_parts {
213-
for label in part_labels {
214-
let part = partition_table
215-
.find(label.as_str())
216-
.ok_or(MissingPartition::from(label))?;
217-
parts_to_erase
218-
.get_or_insert(HashMap::new())
219-
.insert(part.offset(), part);
220-
}
221-
}
222-
// Look for any data partitions with specific data subtype
223-
// There might be multiple partition of the same subtype, e.g. when using multiple FAT partitions
224-
if let Some(partition_types) = args.flash_args.erase_data_parts {
225-
for ty in partition_types {
226-
for part in partition_table.partitions() {
227-
if part.ty() == esp_idf_part::Type::Data
228-
&& part.subtype() == esp_idf_part::SubType::Data(ty)
229-
{
230-
parts_to_erase
231-
.get_or_insert(HashMap::new())
232-
.insert(part.offset(), part);
233-
}
234-
}
235-
}
236-
}
237-
238-
if let Some(parts) = parts_to_erase {
239-
parts
240-
.iter()
241-
.try_for_each(|(_, p)| erase_partition(&mut flasher, p))?;
242-
}
201+
erase_partitions(
202+
&mut flasher,
203+
partition_table.clone(),
204+
args.flash_args.erase_parts,
205+
args.flash_args.erase_data_parts,
206+
)?;
243207
}
244208

245209
flash_elf_image(

‎espflash/src/bin/espflash.rs‎

Lines changed: 7 additions & 43 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,4 @@
11
use std::{
2-
collections::HashMap,
32
fs::{self, File},
43
io::Read,
54
num::ParseIntError,
@@ -9,11 +8,10 @@ use std::{
98
use clap::{Args, Parser, Subcommand};
109
use espflash::{
1110
cli::{
12-
self, board_info, clap_enum_variants, config::Config, connect, erase_partition,
11+
self, board_info, clap_enum_variants, config::Config, connect, erase_partitions,
1312
flash_elf_image, monitor::monitor, parse_partition_table, partition_table,
1413
save_elf_as_image, serial_monitor, ConnectArgs, FlashConfigArgs, PartitionTableArgs,
1514
},
16-
error::{MissingPartition, MissingPartitionTable},
1715
image_format::ImageFormatKind,
1816
logging::initialize_logger,
1917
update::check_for_update,
@@ -133,46 +131,12 @@ fn flash(args: FlashArgs, config: &Config) -> Result<()> {
133131
};
134132

135133
if args.flash_args.erase_parts.is_some() || args.flash_args.erase_data_parts.is_some() {
136-
let partition_table = match &partition_table {
137-
Some(partition_table) => partition_table,
138-
None => return Err((MissingPartitionTable {}).into()),
139-
};
140-
141-
// Using a hashmap to deduplicate entries
142-
let mut parts_to_erase = None;
143-
144-
// Look for any part with specific label
145-
if let Some(part_labels) = args.flash_args.erase_parts {
146-
for label in part_labels {
147-
let part = partition_table
148-
.find(label.as_str())
149-
.ok_or(MissingPartition::from(label))?;
150-
parts_to_erase
151-
.get_or_insert(HashMap::new())
152-
.insert(part.offset(), part);
153-
}
154-
}
155-
// Look for any data partitions with specific data subtype
156-
// There might be multiple partition of the same subtype, e.g. when using multiple FAT partitions
157-
if let Some(partition_types) = args.flash_args.erase_data_parts {
158-
for ty in partition_types {
159-
for part in partition_table.partitions() {
160-
if part.ty() == esp_idf_part::Type::Data
161-
&& part.subtype() == esp_idf_part::SubType::Data(ty)
162-
{
163-
parts_to_erase
164-
.get_or_insert(HashMap::new())
165-
.insert(part.offset(), part);
166-
}
167-
}
168-
}
169-
}
170-
171-
if let Some(parts) = parts_to_erase {
172-
parts
173-
.iter()
174-
.try_for_each(|(_, p)| erase_partition(&mut flasher, p))?;
175-
}
134+
erase_partitions(
135+
&mut flasher,
136+
partition_table.clone(),
137+
args.flash_args.erase_parts,
138+
args.flash_args.erase_data_parts,
139+
)?;
176140
}
177141

178142
flash_elf_image(

‎espflash/src/cli/mod.rs‎

Lines changed: 60 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
//! [espflash]: https://crates.io/crates/espflash
99
1010
use std::{
11+
collections::HashMap,
1112
fs,
1213
io::Write,
1314
path::{Path, PathBuf},
@@ -24,6 +25,7 @@ use strum::VariantNames;
2425
use self::{config::Config, monitor::monitor, serial::get_serial_port_info};
2526
use crate::{
2627
elf::ElfFirmwareImage,
28+
error::{MissingPartition, MissingPartitionTable},
2729
flasher::{FlashFrequency, FlashMode, FlashSize, Flasher},
2830
image_format::ImageFormatKind,
2931
interface::Interface,
@@ -377,17 +379,74 @@ pub fn flash_elf_image(
377379
Ok(())
378380
}
379381

382+
/// Parse a [PartitionTable] from the provided path
380383
pub fn parse_partition_table(path: &Path) -> Result<PartitionTable> {
381384
let data = fs::read(path)
382385
.into_diagnostic()
383386
.wrap_err("Failed to open partition table")?;
387+
384388
PartitionTable::try_from(data).into_diagnostic()
385389
}
386390

387-
pub fn erase_partition(flasher: &mut Flasher, part: &Partition) -> Result<()> {
391+
/// Erase one or more partitions by label or [DataType]
392+
pub fn erase_partitions(
393+
flasher: &mut Flasher,
394+
partition_table: Option<PartitionTable>,
395+
erase_parts: Option<Vec<String>>,
396+
erase_data_parts: Option<Vec<DataType>>,
397+
) -> Result<()> {
398+
let partition_table = match &partition_table {
399+
Some(partition_table) => partition_table,
400+
None => return Err((MissingPartitionTable {}).into()),
401+
};
402+
403+
// Using a hashmap to deduplicate entries
404+
let mut parts_to_erase = None;
405+
406+
// Look for any partitions with specific labels
407+
if let Some(part_labels) = erase_parts {
408+
for label in part_labels {
409+
let part = partition_table
410+
.find(label.as_str())
411+
.ok_or(MissingPartition::from(label))?;
412+
413+
parts_to_erase
414+
.get_or_insert(HashMap::new())
415+
.insert(part.offset(), part);
416+
}
417+
}
418+
419+
// Look for any data partitions with specific data subtype
420+
// There might be multiple partition of the same subtype, e.g. when using multiple FAT partitions
421+
if let Some(partition_types) = erase_data_parts {
422+
for ty in partition_types {
423+
for part in partition_table.partitions() {
424+
if part.ty() == esp_idf_part::Type::Data
425+
&& part.subtype() == esp_idf_part::SubType::Data(ty)
426+
{
427+
parts_to_erase
428+
.get_or_insert(HashMap::new())
429+
.insert(part.offset(), part);
430+
}
431+
}
432+
}
433+
}
434+
435+
if let Some(parts) = parts_to_erase {
436+
parts
437+
.iter()
438+
.try_for_each(|(_, p)| erase_partition(flasher, p))?;
439+
}
440+
441+
Ok(())
442+
}
443+
444+
fn erase_partition(flasher: &mut Flasher, part: &Partition) -> Result<()> {
388445
log::info!("Erasing {} ({:?})...", part.name(), part.subtype());
446+
389447
let offset = part.offset();
390448
let size = part.size();
449+
391450
flasher.erase_region(offset, size).into_diagnostic()
392451
}
393452

0 commit comments

Comments
 (0)