rrbutani · August 17, 2024 20:25 · rrbutani · Aug 17, 2024
diff --git a/clap_ordered_flag_group_derive.rs b/clap_ordered_flag_group_derive.rs
 use clap::Parser;

 // NOTE: we want to model reporting args as an enum (of which we have multiple
 // copies) instead of a struct of vecs (one per option) because we wish to
 // preserve the order of these args.
 //
 // The clap cookbook entry for `find` provides a reference for how to model
 // these kinds of "order matters" flags:
 // https://docs.rs/clap/4.5.16/clap/_derive/_cookbook/find/index.html
 //
 // We take inspiration from ^ but don't want to drop down to using the
 // procedural APIs (instead of the declarative stuff) for our entire arg parser;
 // instead we describe the reporting args "declaratively" as an enum and
 // then have this macro massages our description into a struct that can derive
 // `Arg` + an adapter that maps back to the enum, ultimately producing an
 // ordered `Vec<{enum}>`.
 macro_rules! enum_to_args_struct {
    (
        $(#[doc = $doc:tt])*
        #[derive($($derives:tt)*)]
        $(#[group($($group_attr:tt)*)])*
        pub enum $name:ident {
            $(
                $(#[doc = $variant_doc:tt])*
                $(#[arg($($variant_arg_attr:tt)*)])*
                $variant_name:ident $(($variant_payload:ty))?
            ),* $(,)?
        } as $args_type_name:ident via $arg_struct_name:ident
    ) => {
        $(#[doc = $doc])*
        #[derive($($derives)*)]
        pub enum $name {
            $(
                $(#[doc = $variant_doc])*
                $variant_name$(($variant_payload))?,
            )*
        }

        #[doc = "Arg struct (deriving [`clap::Args`]) for [`"]
        #[doc = core::stringify!($name)]
        #[doc = "`]."]
        #[derive(clap::Args, $($derives)*)]
        $(#[group($($group_attr)*)])*
        #[allow(non_snake_case)]
        pub(crate) struct $arg_struct_name {
            $(
                $(#[doc = $variant_doc])*
                $(#[arg($($variant_arg_attr)*)])*
                #[cfg_attr(
                    all($($variant_payload, feature = "_null")?),
                    arg(
                        num_args = 0,
                        value_parser = clap::value_parser!(bool),
                        default_missing_value = "true",
                    )
                )]
                // NOTE: we specify `Vec` here to satisify clap's derive proc
                // macros; they "detect" `Vec<_>` args "textually" (i.e. by
                // looking for "Vec" in the field's `syn::Ty`'s last segment)
                $variant_name: Vec<enum_to_args_struct!(
                    @VARIANT_TY: $(($variant_payload))?
                )>,
            )*
        }

        // Define the "adapter" that maps args back to the enum type, preserving
        // order:
        #[doc = "Args type for [`"]
        #[doc = core::stringify!($name)]
        #[doc = "`]; preserves command-line order."]
        #[derive($($derives)*)]
        pub struct $args_type_name(pub Vec<$name>);

        // Arg declaration (as derived for the struct type) is fine; we don't
        // need to make any changes:
        impl clap::Args for $args_type_name {
            fn group_id() -> Option<clap::Id> { $arg_struct_name::group_id() }
            fn augment_args(app: clap::Command) -> clap::Command {
                $arg_struct_name::augment_args(app)
            }
            fn augment_args_for_update(app: clap::Command) -> clap::Command {
                $arg_struct_name::augment_args_for_update(app)
            }
        }

        // Parsing is different however:
        impl clap::FromArgMatches for $args_type_name {
            fn from_arg_matches(m: &clap::ArgMatches) -> Result<Self, clap::Error> {
                Self::from_arg_matches_mut(&mut m.clone())
            }
            fn from_arg_matches_mut(m: &mut clap::ArgMatches) -> Result<Self, clap::Error> {
                let mut this = $args_type_name(Vec::new());
                this.update_from_arg_matches_mut(m)?;
                Ok(this)
            }
            fn update_from_arg_matches(&mut self, m: &clap::ArgMatches) -> Result<(), clap::Error> {
                self.update_from_arg_matches_mut(&mut m.clone())
            }
            #[allow(non_snake_case)]
            fn update_from_arg_matches_mut(&mut self, m: &mut clap::ArgMatches) -> Result<(), clap::Error> {
                // get counts for each arg:
                $(
                    let $variant_name = m
                        .indices_of(core::stringify!($variant_name))
                        .map(|it| it.collect::<Vec<usize>>())
                        .unwrap_or_default();
                )*

                // parse args into the `Vec<{payload}>` struct:
                let mut struct_args = $arg_struct_name::from_arg_matches_mut(m)?;

                // for each arg, assert that the number of payloads we got
                // matches the count from `indices_of`; then insert in order:
                let mut ordered = std::collections::BTreeMap::new();
                $({
                    assert_eq!(
                        $variant_name.len(), struct_args.$variant_name.len(),
                    );
                    // map to enum variant:
                    let it = struct_args.$variant_name.drain(..)
                        .zip($variant_name);
                    for (payload, idx) in it {
                        // if there's an actual payload, this is easy:
                        $(
                            let payload: $variant_payload = payload;
                            let res = $name::$variant_name(payload);

                            #[cfg(any())] // inhibit the no-payload codepath
                        )?
                        // otherwise, `payload` should be a bool that's `true`;
                        // we should check this and then discard the value
                        let res = {
                            let payload: bool = payload;
                            assert!(payload);

                            $name::$variant_name
                        };

                        ordered.insert(idx, res);
                    }

                })*

                // finally, return:
                self.0 = ordered.into_values().collect();
                Ok(())
            }
        }
    };

    // Translate each variant to a struct field; if there is a payload this is
    // straight-forward:
    (@VARIANT_TY: ($ty:ty)) => { $ty };
    // If not, translate as a 0 arg option of type `bool`:
    (@VARIANT_TY: ) => { bool };
 }

 ////////////////////////////////////////////////////////////////////////////////

 // Example:

 enum_to_args_struct! {
    #[derive(Debug, Clone, PartialEq, Eq)]
    #[group(multiple = true, required = false)]
    pub enum ReportingOption {
        /// Print packages that contain the given executable.
        #[arg(long = "pkgs")]
        Packages,

        /// Print available versions of the given package.
        #[arg(long = "vers")]
        Versions,

        /// Print environment variable value.
        #[arg(short = 'g', value_name = "VAR")]
        EnvVarValue(String),

        /// Print executable's full environment.
        #[arg(long = "env")]
        Environment,

        /// Print executable's full path.
        #[arg(long = "path")]
        Path,

        /// Print executable's version.
        #[arg(long = "ver")]
        Version,
    } as ReportingArgs via ReportingArg
 }

 #[derive(clap::Parser, Debug, PartialEq)]
 struct Args {
    #[command(flatten)]
    report_opts: ReportingArgs,
 }

 fn main() {
    assert_eq!(
        dbg!(Args::parse_from([
            "argv0", "--pkgs", "-g", "hello", "--pkgs", "--env", "--ver",
            "--vers", "--vers", "--pkgs", "-g", "yo",
        ])),
        Args { report_opts: ReportingArgs(vec![
            ReportingOption::Packages,
            ReportingOption::EnvVarValue("hello".to_string()),
            ReportingOption::Packages,
            ReportingOption::Environment,
            ReportingOption::Version,
            ReportingOption::Versions,
            ReportingOption::Versions,
            ReportingOption::Packages,
            ReportingOption::EnvVarValue("yo".to_string()),
        ]) },
    );
 }
	use clap::Parser;

	// NOTE: we want to model reporting args as an enum (of which we have multiple
	// copies) instead of a struct of vecs (one per option) because we wish to
	// preserve the order of these args.
	//
	// The clap cookbook entry for `find` provides a reference for how to model
	// these kinds of "order matters" flags:
	// https://docs.rs/clap/4.5.16/clap/_derive/_cookbook/find/index.html
	//
	// We take inspiration from ^ but don't want to drop down to using the
	// procedural APIs (instead of the declarative stuff) for our entire arg parser;
	// instead we describe the reporting args "declaratively" as an enum and
	// then have this macro massages our description into a struct that can derive
	// `Arg` + an adapter that maps back to the enum, ultimately producing an
	// ordered `Vec<{enum}>`.
	macro_rules! enum_to_args_struct {
	(
	$(#[doc = $doc:tt])*
	#[derive($($derives:tt)*)]
	$(#[group($($group_attr:tt))])
	pub enum $name:ident {
	$(
	$(#[doc = $variant_doc:tt])*
	$(#[arg($($variant_arg_attr:tt))])
	$variant_name:ident $(($variant_payload:ty))?
	),* $(,)?
	} as $args_type_name:ident via $arg_struct_name:ident
	) => {
	$(#[doc = $doc])*
	#[derive($($derives)*)]
	pub enum $name {
	$(
	$(#[doc = $variant_doc])*
	$variant_name$(($variant_payload))?,
	)*
	}

	#[doc = "Arg struct (deriving [`clap::Args`]) for [`"]
	#[doc = core::stringify!($name)]
	#[doc = "`]."]
	#[derive(clap::Args, $($derives)*)]
	$(#[group($($group_attr))])
	#[allow(non_snake_case)]
	pub(crate) struct $arg_struct_name {
	$(
	$(#[doc = $variant_doc])*
	$(#[arg($($variant_arg_attr))])
	#[cfg_attr(
	all($($variant_payload, feature = "_null")?),
	arg(
	num_args = 0,
	value_parser = clap::value_parser!(bool),
	default_missing_value = "true",
	)
	)]
	// NOTE: we specify `Vec` here to satisify clap's derive proc
	// macros; they "detect" `Vec<_>` args "textually" (i.e. by
	// looking for "Vec" in the field's `syn::Ty`'s last segment)
	$variant_name: Vec<enum_to_args_struct!(
	@VARIANT_TY: $(($variant_payload))?
	)>,
	)*
	}

	// Define the "adapter" that maps args back to the enum type, preserving
	// order:
	#[doc = "Args type for [`"]
	#[doc = core::stringify!($name)]
	#[doc = "`]; preserves command-line order."]
	#[derive($($derives)*)]
	pub struct $args_type_name(pub Vec<$name>);

	// Arg declaration (as derived for the struct type) is fine; we don't
	// need to make any changes:
	impl clap::Args for $args_type_name {
	fn group_id() -> Option<clap::Id> { $arg_struct_name::group_id() }
	fn augment_args(app: clap::Command) -> clap::Command {
	$arg_struct_name::augment_args(app)
	}
	fn augment_args_for_update(app: clap::Command) -> clap::Command {
	$arg_struct_name::augment_args_for_update(app)
	}
	}

	// Parsing is different however:
	impl clap::FromArgMatches for $args_type_name {
	fn from_arg_matches(m: &clap::ArgMatches) -> Result<Self, clap::Error> {
	Self::from_arg_matches_mut(&mut m.clone())
	}
	fn from_arg_matches_mut(m: &mut clap::ArgMatches) -> Result<Self, clap::Error> {
	let mut this = $args_type_name(Vec::new());
	this.update_from_arg_matches_mut(m)?;
	Ok(this)
	}
	fn update_from_arg_matches(&mut self, m: &clap::ArgMatches) -> Result<(), clap::Error> {
	self.update_from_arg_matches_mut(&mut m.clone())
	}
	#[allow(non_snake_case)]
	fn update_from_arg_matches_mut(&mut self, m: &mut clap::ArgMatches) -> Result<(), clap::Error> {
	// get counts for each arg:
	$(
	let $variant_name = m
	.indices_of(core::stringify!($variant_name))
	.map(\|it\| it.collect::<Vec<usize>>())
	.unwrap_or_default();
	)*

	// parse args into the `Vec<{payload}>` struct:
	let mut struct_args = $arg_struct_name::from_arg_matches_mut(m)?;

	// for each arg, assert that the number of payloads we got
	// matches the count from `indices_of`; then insert in order:
	let mut ordered = std::collections::BTreeMap::new();
	$({
	assert_eq!(
	$variant_name.len(), struct_args.$variant_name.len(),
	);
	// map to enum variant:
	let it = struct_args.$variant_name.drain(..)
	.zip($variant_name);
	for (payload, idx) in it {
	// if there's an actual payload, this is easy:
	$(
	let payload: $variant_payload = payload;
	let res = $name::$variant_name(payload);

	#[cfg(any())] // inhibit the no-payload codepath
	)?
	// otherwise, `payload` should be a bool that's `true`;
	// we should check this and then discard the value
	let res = {
	let payload: bool = payload;
	assert!(payload);

	$name::$variant_name
	};

	ordered.insert(idx, res);
	}

	})*

	// finally, return:
	self.0 = ordered.into_values().collect();
	Ok(())
	}
	}
	};

	// Translate each variant to a struct field; if there is a payload this is
	// straight-forward:
	(@VARIANT_TY: ($ty:ty)) => { $ty };
	// If not, translate as a 0 arg option of type `bool`:
	(@VARIANT_TY: ) => { bool };
	}

	////////////////////////////////////////////////////////////////////////////////

	// Example:

	enum_to_args_struct! {
	#[derive(Debug, Clone, PartialEq, Eq)]
	#[group(multiple = true, required = false)]
	pub enum ReportingOption {
	/// Print packages that contain the given executable.
	#[arg(long = "pkgs")]
	Packages,

	/// Print available versions of the given package.
	#[arg(long = "vers")]
	Versions,

	/// Print environment variable value.
	#[arg(short = 'g', value_name = "VAR")]
	EnvVarValue(String),

	/// Print executable's full environment.
	#[arg(long = "env")]
	Environment,

	/// Print executable's full path.
	#[arg(long = "path")]
	Path,

	/// Print executable's version.
	#[arg(long = "ver")]
	Version,
	} as ReportingArgs via ReportingArg
	}

	#[derive(clap::Parser, Debug, PartialEq)]
	struct Args {
	#[command(flatten)]
	report_opts: ReportingArgs,
	}

	fn main() {
	assert_eq!(
	dbg!(Args::parse_from([
	"argv0", "--pkgs", "-g", "hello", "--pkgs", "--env", "--ver",
	"--vers", "--vers", "--pkgs", "-g", "yo",
	])),
	Args { report_opts: ReportingArgs(vec![
	ReportingOption::Packages,
	ReportingOption::EnvVarValue("hello".to_string()),
	ReportingOption::Packages,
	ReportingOption::Environment,
	ReportingOption::Version,
	ReportingOption::Versions,
	ReportingOption::Versions,
	ReportingOption::Packages,
	ReportingOption::EnvVarValue("yo".to_string()),
	]) },
	);
	}