Skip to main content

zsh/extensions/
stringsort.rs

1//! String manipulation and sorting for zshrs
2
3use std::cmp::Ordering;
4
5/// Duplicate a string (equivalent to dupstring/ztrdup in C)
6#[inline]
7pub fn dupstring(s: &str) -> String {
8    s.to_string()
9}
10
11/// Duplicate a string with a specified length
12pub fn dupstring_wlen(s: &str, len: usize) -> String {
13    if len >= s.len() {
14        s.to_string()
15    } else {
16        s[..len].to_string()
17    }
18}
19
20/// Concatenate three strings
21pub fn tricat(s1: &str, s2: &str, s3: &str) -> String {
22    let mut result = String::with_capacity(s1.len() + s2.len() + s3.len());
23    result.push_str(s1);
24    result.push_str(s2);
25    result.push_str(s3);
26    result
27}
28
29/// Concatenate two strings
30pub fn bicat(s1: &str, s2: &str) -> String {
31    let mut result = String::with_capacity(s1.len() + s2.len());
32    result.push_str(s1);
33    result.push_str(s2);
34    result
35}
36
37/// Duplicate a prefix of a string
38pub fn dupstrpfx(s: &str, len: usize) -> String {
39    dupstring_wlen(s, len)
40}
41
42/// Append a string to another, returning the result
43pub fn appstr(base: &str, append: &str) -> String {
44    bicat(base, append)
45}
46
47/// Get pointer to the last character of a string
48pub fn strend(s: &str) -> Option<char> {
49    s.chars().last()
50}
51
52/// Sort flags
53pub mod sort_flags {
54    /// `SORTIT_BACKWARDS` constant.
55    pub const SORTIT_BACKWARDS: u32 = 1;
56    /// `SORTIT_NUMERICALLY` constant.
57    pub const SORTIT_NUMERICALLY: u32 = 2;
58    /// `SORTIT_NUMERICALLY_SIGNED` constant.
59    pub const SORTIT_NUMERICALLY_SIGNED: u32 = 4;
60    /// `SORTIT_IGNORING_CASE` constant.
61    pub const SORTIT_IGNORING_CASE: u32 = 8;
62    /// `SORTIT_IGNORING_BACKSLASHES` constant.
63    pub const SORTIT_IGNORING_BACKSLASHES: u32 = 16;
64}
65
66/// Compare two strings with various options
67pub fn zstrcmp(a: &str, b: &str, flags: u32) -> Ordering {
68    let ignore_case = flags & sort_flags::SORTIT_IGNORING_CASE != 0;
69    let ignore_backslash = flags & sort_flags::SORTIT_IGNORING_BACKSLASHES != 0;
70    let numeric = flags & sort_flags::SORTIT_NUMERICALLY != 0;
71    let numeric_signed = flags & sort_flags::SORTIT_NUMERICALLY_SIGNED != 0;
72
73    // Prepare strings for comparison
74    let (a_cmp, b_cmp): (std::borrow::Cow<str>, std::borrow::Cow<str>) = if ignore_case {
75        (
76            std::borrow::Cow::Owned(a.to_lowercase()),
77            std::borrow::Cow::Owned(b.to_lowercase()),
78        )
79    } else {
80        (std::borrow::Cow::Borrowed(a), std::borrow::Cow::Borrowed(b))
81    };
82
83    let (a_final, b_final): (std::borrow::Cow<str>, std::borrow::Cow<str>) = if ignore_backslash {
84        (
85            std::borrow::Cow::Owned(a_cmp.replace('\\', "")),
86            std::borrow::Cow::Owned(b_cmp.replace('\\', "")),
87        )
88    } else {
89        (a_cmp, b_cmp)
90    };
91
92    if numeric || numeric_signed {
93        numeric_compare(&a_final, &b_final, numeric_signed)
94    } else {
95        a_final.cmp(&b_final)
96    }
97}
98
99/// Numeric-aware string comparison
100fn numeric_compare(a: &str, b: &str, signed: bool) -> Ordering {
101    let mut a_chars = a.chars().peekable();
102    let mut b_chars = b.chars().peekable();
103
104    loop {
105        let a_next = a_chars.peek().copied();
106        let b_next = b_chars.peek().copied();
107
108        match (a_next, b_next) {
109            (None, None) => return Ordering::Equal,
110            (None, Some(_)) => return Ordering::Less,
111            (Some(_), None) => return Ordering::Greater,
112            (Some(ac), Some(bc)) => {
113                // Check if we're at the start of a number
114                let a_is_digit = ac.is_ascii_digit();
115                let b_is_digit = bc.is_ascii_digit();
116                let a_is_neg = signed
117                    && ac == '-'
118                    && a_chars
119                        .clone()
120                        .nth(1)
121                        .map(|c| c.is_ascii_digit())
122                        .unwrap_or(false);
123                let b_is_neg = signed
124                    && bc == '-'
125                    && b_chars
126                        .clone()
127                        .nth(1)
128                        .map(|c| c.is_ascii_digit())
129                        .unwrap_or(false);
130
131                if a_is_digit || b_is_digit || a_is_neg || b_is_neg {
132                    // Extract and compare numbers
133                    let a_num = extract_number(&mut a_chars, signed);
134                    let b_num = extract_number(&mut b_chars, signed);
135
136                    match a_num.cmp(&b_num) {
137                        Ordering::Equal => continue,
138                        other => return other,
139                    }
140                } else {
141                    // Regular character comparison
142                    a_chars.next();
143                    b_chars.next();
144                    match ac.cmp(&bc) {
145                        Ordering::Equal => continue,
146                        other => return other,
147                    }
148                }
149            }
150        }
151    }
152}
153
154/// Extract a number from a character iterator
155fn extract_number<I: Iterator<Item = char>>(
156    chars: &mut std::iter::Peekable<I>,
157    signed: bool,
158) -> i64 {
159    let mut negative = false;
160    let mut num: i64 = 0;
161    let mut has_digit = false;
162
163    // Check for sign
164    if signed {
165        if let Some(&'-') = chars.peek() {
166            chars.next();
167            negative = true;
168        } else if let Some(&'+') = chars.peek() {
169            chars.next();
170        }
171    }
172
173    // Skip leading zeros
174    while let Some(&'0') = chars.peek() {
175        chars.next();
176        has_digit = true;
177    }
178
179    // Collect digits
180    while let Some(&c) = chars.peek() {
181        if c.is_ascii_digit() {
182            has_digit = true;
183            num = num
184                .saturating_mul(10)
185                .saturating_add((c as i64) - ('0' as i64));
186            chars.next();
187        } else {
188            break;
189        }
190    }
191
192    if !has_digit {
193        return 0;
194    }
195
196    if negative {
197        -num
198    } else {
199        num
200    }
201}
202
203/// Sort an array of strings with various options
204pub fn strmetasort(array: &mut [String], flags: u32) {
205    if array.len() < 2 {
206        return;
207    }
208
209    let backwards = flags & sort_flags::SORTIT_BACKWARDS != 0;
210
211    array.sort_by(|a, b| {
212        let cmp = zstrcmp(a, b, flags);
213        if backwards {
214            cmp.reverse()
215        } else {
216            cmp
217        }
218    });
219}
220
221/// Sort string slices with various options
222pub fn sort_strings(array: &mut [&str], flags: u32) {
223    if array.len() < 2 {
224        return;
225    }
226
227    let backwards = flags & sort_flags::SORTIT_BACKWARDS != 0;
228
229    array.sort_by(|a, b| {
230        let cmp = zstrcmp(a, b, flags);
231        if backwards {
232            cmp.reverse()
233        } else {
234            cmp
235        }
236    });
237}
238
239/// Natural sort comparison (numbers sorted numerically within strings)
240pub fn natural_cmp(a: &str, b: &str) -> Ordering {
241    zstrcmp(a, b, sort_flags::SORTIT_NUMERICALLY)
242}
243
244/// Case-insensitive comparison
245pub fn strcasecmp(a: &str, b: &str) -> Ordering {
246    a.to_lowercase().cmp(&b.to_lowercase())
247}
248
249/// Find first occurrence of substring
250pub fn strstr(haystack: &str, needle: &str) -> Option<usize> {
251    haystack.find(needle)
252}
253
254/// Check if string starts with prefix
255pub fn strprefix(s: &str, prefix: &str) -> bool {
256    s.starts_with(prefix)
257}
258
259/// Check if string ends with suffix
260pub fn strsuffix(s: &str, suffix: &str) -> bool {
261    s.ends_with(suffix)
262}
263
264/// Join strings with a separator
265pub fn strjoin<I, S>(iter: I, sep: &str) -> String
266where
267    I: IntoIterator<Item = S>,
268    S: AsRef<str>,
269{
270    iter.into_iter()
271        .map(|s| s.as_ref().to_string())
272        .collect::<Vec<_>>()
273        .join(sep)
274}
275
276/// Split string by separator
277pub fn strsplit(s: &str, sep: char) -> Vec<&str> {
278    s.split(sep).collect()
279}
280
281/// Trim whitespace from both ends
282pub fn strtrim(s: &str) -> &str {
283    s.trim()
284}
285
286/// Convert string to lowercase
287pub fn strlower(s: &str) -> String {
288    s.to_lowercase()
289}
290
291/// Convert string to uppercase
292pub fn strupper(s: &str) -> String {
293    s.to_uppercase()
294}
295
296#[cfg(test)]
297mod tests {
298    use super::*;
299
300    #[test]
301    fn test_dupstring() {
302        let _g = crate::test_util::global_state_lock();
303        assert_eq!(dupstring("hello"), "hello");
304        assert_eq!(dupstring(""), "");
305    }
306
307    #[test]
308    fn test_dupstring_wlen() {
309        let _g = crate::test_util::global_state_lock();
310        assert_eq!(dupstring_wlen("hello", 3), "hel");
311        assert_eq!(dupstring_wlen("hi", 10), "hi");
312    }
313
314    #[test]
315    fn test_tricat() {
316        let _g = crate::test_util::global_state_lock();
317        assert_eq!(tricat("a", "b", "c"), "abc");
318        assert_eq!(tricat("hello", " ", "world"), "hello world");
319    }
320
321    #[test]
322    fn test_bicat() {
323        let _g = crate::test_util::global_state_lock();
324        assert_eq!(bicat("hello", "world"), "helloworld");
325    }
326
327    #[test]
328    fn test_strend() {
329        let _g = crate::test_util::global_state_lock();
330        assert_eq!(strend("hello"), Some('o'));
331        assert_eq!(strend(""), None);
332    }
333
334    #[test]
335    fn test_zstrcmp_basic() {
336        let _g = crate::test_util::global_state_lock();
337        assert_eq!(zstrcmp("abc", "abc", 0), Ordering::Equal);
338        assert_eq!(zstrcmp("abc", "abd", 0), Ordering::Less);
339        assert_eq!(zstrcmp("abd", "abc", 0), Ordering::Greater);
340    }
341
342    #[test]
343    fn test_zstrcmp_case_insensitive() {
344        let _g = crate::test_util::global_state_lock();
345        let flags = sort_flags::SORTIT_IGNORING_CASE;
346        assert_eq!(zstrcmp("ABC", "abc", flags), Ordering::Equal);
347        assert_eq!(zstrcmp("ABC", "ABD", flags), Ordering::Less);
348    }
349
350    #[test]
351    fn test_zstrcmp_ignore_backslash() {
352        let _g = crate::test_util::global_state_lock();
353        let flags = sort_flags::SORTIT_IGNORING_BACKSLASHES;
354        assert_eq!(zstrcmp("a\\bc", "abc", flags), Ordering::Equal);
355    }
356
357    #[test]
358    fn test_zstrcmp_numeric() {
359        let _g = crate::test_util::global_state_lock();
360        let flags = sort_flags::SORTIT_NUMERICALLY;
361        assert_eq!(zstrcmp("file2", "file10", flags), Ordering::Less);
362        assert_eq!(zstrcmp("file10", "file2", flags), Ordering::Greater);
363        assert_eq!(zstrcmp("file10", "file10", flags), Ordering::Equal);
364    }
365
366    #[test]
367    fn test_zstrcmp_numeric_signed() {
368        let _g = crate::test_util::global_state_lock();
369        let flags = sort_flags::SORTIT_NUMERICALLY_SIGNED;
370        assert_eq!(zstrcmp("-5", "3", flags), Ordering::Less);
371        assert_eq!(zstrcmp("-10", "-2", flags), Ordering::Less);
372    }
373
374    #[test]
375    fn test_strmetasort() {
376        let _g = crate::test_util::global_state_lock();
377        let mut arr = vec![
378            "file10".to_string(),
379            "file2".to_string(),
380            "file1".to_string(),
381        ];
382        strmetasort(&mut arr, sort_flags::SORTIT_NUMERICALLY);
383        assert_eq!(arr, vec!["file1", "file2", "file10"]);
384    }
385
386    #[test]
387    fn test_strmetasort_backwards() {
388        let _g = crate::test_util::global_state_lock();
389        let mut arr = vec!["a".to_string(), "c".to_string(), "b".to_string()];
390        strmetasort(&mut arr, sort_flags::SORTIT_BACKWARDS);
391        assert_eq!(arr, vec!["c", "b", "a"]);
392    }
393
394    #[test]
395    fn test_natural_cmp() {
396        let _g = crate::test_util::global_state_lock();
397        assert_eq!(natural_cmp("item2", "item10"), Ordering::Less);
398    }
399
400    #[test]
401    fn test_strcasecmp() {
402        let _g = crate::test_util::global_state_lock();
403        assert_eq!(strcasecmp("Hello", "HELLO"), Ordering::Equal);
404        assert_eq!(strcasecmp("abc", "ABD"), Ordering::Less);
405    }
406
407    #[test]
408    fn test_strstr() {
409        let _g = crate::test_util::global_state_lock();
410        assert_eq!(strstr("hello world", "world"), Some(6));
411        assert_eq!(strstr("hello", "xyz"), None);
412    }
413
414    #[test]
415    fn test_strprefix_suffix() {
416        let _g = crate::test_util::global_state_lock();
417        assert!(strprefix("hello", "hel"));
418        assert!(!strprefix("hello", "ell"));
419        assert!(strsuffix("hello", "llo"));
420        assert!(!strsuffix("hello", "ell"));
421    }
422
423    #[test]
424    fn test_strjoin() {
425        let _g = crate::test_util::global_state_lock();
426        assert_eq!(strjoin(["a", "b", "c"], ","), "a,b,c");
427        assert_eq!(strjoin(Vec::<&str>::new(), ","), "");
428    }
429
430    #[test]
431    fn test_strsplit() {
432        let _g = crate::test_util::global_state_lock();
433        assert_eq!(strsplit("a,b,c", ','), vec!["a", "b", "c"]);
434    }
435
436    #[test]
437    fn test_strtrim() {
438        let _g = crate::test_util::global_state_lock();
439        assert_eq!(strtrim("  hello  "), "hello");
440    }
441
442    #[test]
443    fn test_case_conversion() {
444        let _g = crate::test_util::global_state_lock();
445        assert_eq!(strlower("HeLLo"), "hello");
446        assert_eq!(strupper("HeLLo"), "HELLO");
447    }
448}