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 { |
32 | pub(crate) fn is_position(&self) -> bool { |
33 | matches!(self, KeyType::Position(_)) |
34 | } |
35 | } |
36 | |
37 | impl PartialEq<usize> for KeyType { |
38 | fn eq(&self, rhs: &usize) -> bool { |
39 | match self { |
40 | KeyType::Position(x: &usize) => x == rhs, |
41 | _ => false, |
42 | } |
43 | } |
44 | } |
45 | |
46 | impl PartialEq<&str> for KeyType { |
47 | fn eq(&self, rhs: &&str) -> bool { |
48 | match self { |
49 | KeyType::Long(l: &OsStr) => l == rhs, |
50 | _ => false, |
51 | } |
52 | } |
53 | } |
54 | |
55 | impl PartialEq<str> for KeyType { |
56 | fn eq(&self, rhs: &str) -> bool { |
57 | match self { |
58 | KeyType::Long(l: &OsStr) => l == rhs, |
59 | _ => false, |
60 | } |
61 | } |
62 | } |
63 | |
64 | impl PartialEq<OsStr> for KeyType { |
65 | fn eq(&self, rhs: &OsStr) -> bool { |
66 | match self { |
67 | KeyType::Long(l: &OsStr) => l == rhs, |
68 | _ => false, |
69 | } |
70 | } |
71 | } |
72 | |
73 | impl PartialEq<char> for KeyType { |
74 | fn eq(&self, rhs: &char) -> bool { |
75 | match self { |
76 | KeyType::Short(c: &char) => 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) |
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. |
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) |
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. |
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. |
117 | pub(crate) fn args(&self) -> impl Iterator<Item = &Arg> { |
118 | self.args.iter() |
119 | } |
120 | |
121 | /// Return mutable iterators of all args. |
122 | pub(crate) fn args_mut(&mut self) -> impl Iterator<Item = &mut Arg> { |
123 | self.args.iter_mut() |
124 | } |
125 | |
126 | /// Mutate every argument. |
127 | pub(crate) fn mut_args<F>(&mut self, f: F) |
128 | where |
129 | F: FnMut(Arg) -> Arg, |
130 | { |
131 | let mut args = std::mem::take(&mut self.args); |
132 | self.args.extend(args.drain(..).map(f)); |
133 | } |
134 | |
135 | /// We need a lazy build here since some we may change args after creating |
136 | /// the map, you can checkout who uses `args_mut`. |
137 | pub(crate) fn _build(&mut self) { |
138 | // There will be at least as many keys as args, so that is a good starting point |
139 | self.keys.reserve(self.args.len()); |
140 | for (i, arg) in self.args.iter().enumerate() { |
141 | append_keys(&mut self.keys, arg, i); |
142 | } |
143 | } |
144 | |
145 | /// Remove an arg in the graph by Id, usually used by `mut_arg`. Return |
146 | /// `Some(arg)` if removed. |
147 | pub(crate) fn remove_by_name(&mut self, name: &str) -> Option<Arg> { |
148 | self.args |
149 | .iter() |
150 | .position(|arg| arg.id == name) |
151 | // since it's a cold function, using this wouldn't hurt much |
152 | .map(|i| self.args.remove(i)) |
153 | } |
154 | } |
155 | |
156 | impl Index<&'_ KeyType> for MKeyMap { |
157 | type Output = Arg; |
158 | |
159 | fn index(&self, key: &KeyType) -> &Self::Output { |
160 | self.get(key).expect(INTERNAL_ERROR_MSG) |
161 | } |
162 | } |
163 | |
164 | /// Generate key types for an specific Arg. |
165 | fn append_keys(keys: &mut Vec<Key>, arg: &Arg, index: usize) { |
166 | if let Some(pos_index: usize) = arg.index { |
167 | let key: KeyType = KeyType::Position(pos_index); |
168 | keys.push(Key { key, index }); |
169 | } else { |
170 | if let Some(short: char) = arg.short { |
171 | let key: KeyType = KeyType::Short(short); |
172 | keys.push(Key { key, index }); |
173 | } |
174 | if let Some(long: Str) = arg.long.clone() { |
175 | let key: KeyType = KeyType::Long(long.into()); |
176 | keys.push(Key { key, index }); |
177 | } |
178 | |
179 | for (short: &char, _) in arg.short_aliases.iter() { |
180 | let key: KeyType = KeyType::Short(*short); |
181 | keys.push(Key { key, index }); |
182 | } |
183 | for (long: &Str, _) in arg.aliases.iter() { |
184 | let key: KeyType = KeyType::Long(long.into()); |
185 | keys.push(Key { key, index }); |
186 | } |
187 | } |
188 | } |
189 | |