From afb5cb1ab4b0f7078d23d3e74363f964c97c80de Mon Sep 17 00:00:00 2001
From: N3xed <dominik.gschwind99@gmail.com>
Date: Wed, 5 Jan 2022 22:34:13 +0100
Subject: [PATCH] Flash args

---
 cargo-idf/src/flash.rs      |  73 ++--------------------
 cargo-idf/src/flash/opts.rs | 119 ++++++++++++++++++++++++++++++++++++
 2 files changed, 123 insertions(+), 69 deletions(-)
 create mode 100644 cargo-idf/src/flash/opts.rs

diff --git a/cargo-idf/src/flash.rs b/cargo-idf/src/flash.rs
index 60ca2cd..8d03bbc 100644
--- a/cargo-idf/src/flash.rs
+++ b/cargo-idf/src/flash.rs
@@ -1,71 +1,6 @@
-use std::ffi::OsStr;
-use std::path::PathBuf;
-use std::str::FromStr;
+pub mod opts;
+pub use opts::FlashOpts;
 
-use anyhow::bail;
-use clap::{AppSettings, ArgEnum, Args};
-use strum::{Display, EnumString};
-
-use crate::build;
-
-#[derive(Args)]
-#[clap(global_setting = AppSettings::DisableVersionFlag)]
-pub struct FlashOpts {
-    /// Which bootloader to flash [possible values: esp-idf, none, <file>]
-    ///
-    /// - `esp-idf` will flash the bootloader compiled locally from the esp-idf.
-    /// - `none` prevents flashing a bootloader.
-    /// - `<file>` will flash the user provided binary file if it exists.
-    #[clap(
-        long,
-        default_value_t = Bootloader::EspIdf,
-        parse(try_from_os_str = Bootloader::try_from_os_str),
-        verbatim_doc_comment
-    )]
-    bootloader: Bootloader,
-
-    /// How to flash the binary
-    #[clap(long, arg_enum, default_value_t = Mode::Esptool)]
-    mode: Mode,
-
-    #[clap(flatten)]
-    build_opts: build::BuildOpts,
-}
-
-#[derive(Debug, ArgEnum, Clone, Copy)]
-pub enum Mode {
-    Esptool,
-    Dfu,
-    Uf2,
-}
-
-#[derive(Debug, Clone, EnumString, Display)]
-#[strum(serialize_all = "kebab-case")]
-pub enum Bootloader {
-    EspIdf,
-    None,
-    #[strum(default)]
-    #[strum(to_string = "<file>")]
-    File(PathBuf),
-}
-
-impl Bootloader {
-    pub fn try_from_os_str(arg: &OsStr) -> Result<Bootloader, anyhow::Error> {
-        let val = if let Some(arg) = arg.to_str() {
-            Bootloader::from_str(arg).unwrap()
-        } else {
-            Bootloader::File(arg.into())
-        };
-
-        if let Bootloader::File(ref path) = val {
-            if !path.is_file() {
-                bail!("'{}' is not a file", path.display())
-            }
-        }
-        Ok(val)
-    }
-}
-
-pub fn run(opts: FlashOpts) -> anyhow::Result<()> {
-    Ok(())
+pub fn run(_opts: FlashOpts) -> anyhow::Result<()> {
+    unimplemented!()
 }
diff --git a/cargo-idf/src/flash/opts.rs b/cargo-idf/src/flash/opts.rs
new file mode 100644
index 0000000..7e47dbc
--- /dev/null
+++ b/cargo-idf/src/flash/opts.rs
@@ -0,0 +1,119 @@
+use std::ffi::{OsStr, OsString};
+use std::path::PathBuf;
+use std::str::FromStr;
+
+use anyhow::bail;
+use clap::{AppSettings, ArgEnum, Args};
+use embuild::utils::OsStrExt;
+use strum::{Display, EnumDiscriminants, EnumString};
+
+use crate::build;
+
+#[derive(Args)]
+#[clap(global_setting = AppSettings::DisableVersionFlag)]
+pub struct FlashOpts {
+    /// One or more images to flash [possible values: all, bootloader, partition-table,
+    /// app, <partition name> <file>, <address> <file>]
+    ///
+    /// - `all`: flash the whole project (bootloader, partition-table, app)
+    /// - `bootloader`: flash bootloader (see `--bootloader` option)
+    /// - `partition-table`: flash the partition table (see `--partition-table` option)
+    /// - `app`: flash the app
+    /// - `<partition name> <file>`: flash <file> at the address of <partition name>
+    /// - `<address> <file>`: flash <file> at <address>
+    #[clap(
+        parse(from_os_str = ImageArg::from_os_str),
+        validator_os = ImageArg::parse_validator(),
+        verbatim_doc_comment,
+        default_value = "all"
+    )]
+    images: Vec<ImageArg>,
+
+    /// The bootloader binary file to use instead of the default
+    #[clap(long, parse(from_os_str), value_name = "file")]
+    bootloader: Option<PathBuf>,
+
+    /// The partition table `.csv` file to use instead of the default
+    #[clap(long, parse(from_os_str), value_name = "file")]
+    partition_table: Option<PathBuf>,
+
+    #[clap(flatten)]
+    build_opts: build::BuildOpts,
+}
+
+#[derive(Debug, ArgEnum, Clone, Copy)]
+pub enum Mode {
+    Esptool,
+    Dfu,
+    Uf2,
+}
+
+#[derive(Debug, Clone, Display, EnumDiscriminants)]
+#[strum_discriminants(name(ImageArgKind))]
+pub enum ImageArg {
+    #[strum(to_string = "<name>")]
+    Name(ImageName),
+    #[strum(to_string = "<address>")]
+    Address(usize),
+    #[strum(to_string = "<partition name or file>")]
+    PartitionOrFile(OsString),
+    Partition(String),
+    File(PathBuf),
+}
+
+impl Default for ImageArgKind {
+    fn default() -> Self {
+        ImageArgKind::Name
+    }
+}
+
+#[derive(Debug, EnumString, Display, Clone, Copy)]
+#[strum(serialize_all = "kebab-case")]
+pub enum ImageName {
+    All,
+    Bootloader,
+    PartitionTable,
+    App,
+}
+
+impl ImageArg {
+    fn from_os_str(arg: &OsStr) -> ImageArg {
+        if let Some(arg) = arg.to_str() {
+            if let Ok(name) = ImageName::from_str(arg) {
+                return ImageArg::Name(name);
+            } else if let Ok(address) = arg.parse::<usize>() {
+                return ImageArg::Address(address);
+            }
+        }
+        ImageArg::PartitionOrFile(arg.to_owned())
+    }
+
+    fn parse_validator() -> impl FnMut(&OsStr) -> Result<(), anyhow::Error> {
+        let mut previous = ImageArgKind::default();
+        move |arg| {
+            let next = Self::from_os_str(arg);
+            previous = Self::parse(previous, next)?.into();
+            Ok(())
+        }
+    }
+
+    /// Parses image arg with a given previous arg kind.
+    ///
+    /// Never returns [`ImageArg::PartitionOrFile`].
+    pub fn parse(last: ImageArgKind, next: ImageArg) -> Result<ImageArg, anyhow::Error> {
+        use ImageArgKind::*;
+        let result = match (last, next) {
+            (Name | File, ImageArg::PartitionOrFile(file)) => {
+                ImageArg::Partition(file.try_to_str()?.into())
+            }
+            (Partition | Address, ImageArg::PartitionOrFile(file)) => ImageArg::File(file.into()),
+            (Name | File, val @ ImageArg::Name(_) | val @ ImageArg::Address(_)) => val,
+            (Partition | Address, _) => bail!("expected <file>"),
+
+            (_, ImageArg::File(_) | ImageArg::Partition(_)) | (PartitionOrFile, _) => {
+                unreachable!("invalid state")
+            }
+        };
+        Ok(result)
+    }
+}