1 use std::iter::Iterator;
2 use std::ops::Index;
3
4 use crate::builder::OsStr;
5 use crate::Arg;
6 use crate::INTERNAL_ERROR_MSG;
7
8 #[derive(PartialEq, Eq, Debug, Clone)]
9 pub(crate) struct Key {
10 key: KeyType,
11 index: usize,
12 }
13
14 #[derive(Default, PartialEq, Eq, Debug, Clone)]
15 pub(crate) struct MKeyMap {
16 /// All of the arguments.
17 args: Vec<Arg>,
18
19 // Cache part:
20 /// Will be set after `_build()`.
21 keys: Vec<Key>,
22 }
23
24 #[derive(Debug, PartialEq, Eq, Hash, Clone)]
25 pub(crate) enum KeyType {
26 Short(char),
27 Long(OsStr),
28 Position(usize),
29 }
30
31 impl KeyType {
is_position(&self) -> bool32 pub(crate) fn is_position(&self) -> bool {
33 matches!(self, KeyType::Position(_))
34 }
35 }
36
37 impl PartialEq<usize> for KeyType {
eq(&self, rhs: &usize) -> bool38 fn eq(&self, rhs: &usize) -> bool {
39 match self {
40 KeyType::Position(x) => x == rhs,
41 _ => false,
42 }
43 }
44 }
45
46 impl PartialEq<&str> for KeyType {
eq(&self, rhs: &&str) -> bool47 fn eq(&self, rhs: &&str) -> bool {
48 match self {
49 KeyType::Long(l) => l == rhs,
50 _ => false,
51 }
52 }
53 }
54
55 impl PartialEq<str> for KeyType {
eq(&self, rhs: &str) -> bool56 fn eq(&self, rhs: &str) -> bool {
57 match self {
58 KeyType::Long(l) => l == rhs,
59 _ => false,
60 }
61 }
62 }
63
64 impl PartialEq<OsStr> for KeyType {
eq(&self, rhs: &OsStr) -> bool65 fn eq(&self, rhs: &OsStr) -> bool {
66 match self {
67 KeyType::Long(l) => l == rhs,
68 _ => false,
69 }
70 }
71 }
72
73 impl PartialEq<char> for KeyType {
eq(&self, rhs: &char) -> bool74 fn eq(&self, rhs: &char) -> bool {
75 match self {
76 KeyType::Short(c) => c == rhs,
77 _ => false,
78 }
79 }
80 }
81
82 impl MKeyMap {
83 /// If any arg has corresponding key in this map, we can search the key with
84 /// u64(for positional argument), char(for short flag), &str and OsString
85 /// (for long flag)
contains<K>(&self, key: K) -> bool where KeyType: PartialEq<K>,86 pub(crate) fn contains<K>(&self, key: K) -> bool
87 where
88 KeyType: PartialEq<K>,
89 {
90 self.keys.iter().any(|x| x.key == key)
91 }
92
93 /// Push an argument in the map.
push(&mut self, new_arg: Arg)94 pub(crate) fn push(&mut self, new_arg: Arg) {
95 self.args.push(new_arg);
96 }
97
98 /// Find the arg have corresponding key in this map, we can search the key
99 /// with u64(for positional argument), char(for short flag), &str and
100 /// OsString (for long flag)
get<K: ?Sized>(&self, key: &K) -> Option<&Arg> where KeyType: PartialEq<K>,101 pub(crate) fn get<K: ?Sized>(&self, key: &K) -> Option<&Arg>
102 where
103 KeyType: PartialEq<K>,
104 {
105 self.keys
106 .iter()
107 .find(|k| &k.key == key)
108 .map(|k| &self.args[k.index])
109 }
110
111 /// Return iterators of all keys.
keys(&self) -> impl Iterator<Item = &KeyType>112 pub(crate) fn keys(&self) -> impl Iterator<Item = &KeyType> {
113 self.keys.iter().map(|x| &x.key)
114 }
115
116 /// Return iterators of all args.
args(&self) -> impl Iterator<Item = &Arg>117 pub(crate) fn args(&self) -> impl Iterator<Item = &Arg> {
118 self.args.iter()
119 }
120
121 /// Return mutable iterators of all args.
args_mut(&mut self) -> impl Iterator<Item = &mut Arg>122 pub(crate) fn args_mut(&mut self) -> impl Iterator<Item = &mut Arg> {
123 self.args.iter_mut()
124 }
125
126 /// We need a lazy build here since some we may change args after creating
127 /// the map, you can checkout who uses `args_mut`.
_build(&mut self)128 pub(crate) fn _build(&mut self) {
129 // There will be at least as many keys as args, so that is a good starting point
130 self.keys.reserve(self.args.len());
131 for (i, arg) in self.args.iter().enumerate() {
132 append_keys(&mut self.keys, arg, i);
133 }
134 }
135
136 /// Remove an arg in the graph by Id, usually used by `mut_arg`. Return
137 /// `Some(arg)` if removed.
remove_by_name(&mut self, name: &str) -> Option<Arg>138 pub(crate) fn remove_by_name(&mut self, name: &str) -> Option<Arg> {
139 self.args
140 .iter()
141 .position(|arg| arg.id == name)
142 // since it's a cold function, using this wouldn't hurt much
143 .map(|i| self.args.remove(i))
144 }
145 }
146
147 impl Index<&'_ KeyType> for MKeyMap {
148 type Output = Arg;
149
index(&self, key: &KeyType) -> &Self::Output150 fn index(&self, key: &KeyType) -> &Self::Output {
151 self.get(key).expect(INTERNAL_ERROR_MSG)
152 }
153 }
154
155 /// Generate key types for an specific Arg.
append_keys(keys: &mut Vec<Key>, arg: &Arg, index: usize)156 fn append_keys(keys: &mut Vec<Key>, arg: &Arg, index: usize) {
157 if let Some(pos_index) = arg.index {
158 let key = KeyType::Position(pos_index);
159 keys.push(Key { key, index });
160 } else {
161 if let Some(short) = arg.short {
162 let key = KeyType::Short(short);
163 keys.push(Key { key, index });
164 }
165 if let Some(long) = arg.long.clone() {
166 let key = KeyType::Long(long.into());
167 keys.push(Key { key, index });
168 }
169
170 for (short, _) in arg.short_aliases.iter() {
171 let key = KeyType::Short(*short);
172 keys.push(Key { key, index });
173 }
174 for (long, _) in arg.aliases.iter() {
175 let key = KeyType::Long(long.into());
176 keys.push(Key { key, index });
177 }
178 }
179 }
180