xash3d_protocol/
color.rs

1// SPDX-License-Identifier: LGPL-3.0-only
2// SPDX-FileCopyrightText: 2023 Denis Drakhnia <numas13@gmail.com>
3
4//! Color codes for strings.
5
6use std::borrow::Cow;
7
8/// Color codes `^digit`.
9#[derive(Copy, Clone, Debug, PartialEq, Eq)]
10pub enum Color {
11    /// Black is coded as `^0`.
12    Black,
13    /// Red is coded as `^1`.
14    Red,
15    /// Green is coded as `^2`.
16    Green,
17    /// Yellow is coded as `^3`.
18    Yellow,
19    /// Blue is coded as `^4`.
20    Blue,
21    /// Cyan is coded as `^5`.
22    Cyan,
23    /// Magenta is coded as `^6`.
24    Magenta,
25    /// White is coded as `^7`.
26    White,
27}
28
29impl TryFrom<&str> for Color {
30    type Error = ();
31
32    fn try_from(value: &str) -> Result<Self, Self::Error> {
33        Ok(match value {
34            "^0" => Self::Black,
35            "^1" => Self::Red,
36            "^2" => Self::Green,
37            "^3" => Self::Yellow,
38            "^4" => Self::Blue,
39            "^5" => Self::Cyan,
40            "^6" => Self::Magenta,
41            "^7" => Self::White,
42            _ => return Err(()),
43        })
44    }
45}
46
47/// Test if string starts with color code.
48///
49/// # Examples
50/// ```rust
51/// # use xash3d_protocol::color::is_color_code;
52/// assert_eq!(is_color_code("hello"), false);
53/// assert_eq!(is_color_code("^4blue ocean"), true);
54/// ```
55#[inline]
56pub fn is_color_code(s: &str) -> bool {
57    matches!(s.as_bytes(), [b'^', c, ..] if c.is_ascii_digit())
58}
59
60/// Trim color codes from a start of string.
61///
62/// # Examples
63///
64/// ```rust
65/// # use xash3d_protocol::color::trim_start_color;
66/// assert_eq!(trim_start_color("hello"), ("", "hello"));
67/// assert_eq!(trim_start_color("^1red apple"), ("^1", "red apple"));
68/// assert_eq!(trim_start_color("^1^2^3yellow roof"), ("^3", "yellow roof"));
69/// ```
70#[inline]
71pub fn trim_start_color(s: &str) -> (&str, &str) {
72    let mut n = 0;
73    while is_color_code(&s[n..]) {
74        n += 2;
75    }
76    if n > 0 {
77        (&s[n - 2..n], &s[n..])
78    } else {
79        s.split_at(0)
80    }
81}
82
83/// Iterator for colored parts of a string.
84///
85/// # Examples
86///
87/// ```rust
88/// # use xash3d_protocol::color::ColorIter;
89/// let colored = "^1red flower^7 and ^2green grass";
90/// let mut iter = ColorIter::new(colored);
91/// assert_eq!(iter.next(), Some(("^1", "red flower")));
92/// assert_eq!(iter.next(), Some(("^7", " and ")));
93/// assert_eq!(iter.next(), Some(("^2", "green grass")));
94/// assert_eq!(iter.next(), None);
95/// ```
96pub struct ColorIter<'a> {
97    inner: &'a str,
98}
99
100impl<'a> ColorIter<'a> {
101    /// Creates a new `ColorIter`.
102    pub fn new(inner: &'a str) -> Self {
103        Self { inner }
104    }
105}
106
107impl<'a> Iterator for ColorIter<'a> {
108    type Item = (&'a str, &'a str);
109
110    fn next(&mut self) -> Option<Self::Item> {
111        if self.inner.is_empty() {
112            return None;
113        }
114        let (color, tail) = trim_start_color(self.inner);
115        let offset = tail
116            .char_indices()
117            .map(|(i, _)| i)
118            .find(|&i| is_color_code(&tail[i..]))
119            .unwrap_or(tail.len());
120        let (head, tail) = tail.split_at(offset);
121        self.inner = tail;
122        Some((color, head))
123    }
124}
125
126/// Trim color codes from a string.
127///
128/// # Examples
129///
130/// ```rust
131/// # use xash3d_protocol::color::trim_color;
132/// assert_eq!(trim_color("^1no^7 ^2colors^7"), "no colors");
133/// ```
134pub fn trim_color(s: &str) -> Cow<'_, str> {
135    let (_, s) = trim_start_color(s);
136    if !s.chars().any(|c| c == '^') {
137        return Cow::Borrowed(s);
138    }
139
140    let mut out = String::with_capacity(s.len());
141    for (_, s) in ColorIter::new(s) {
142        out.push_str(s);
143    }
144
145    Cow::Owned(out)
146}
147
148#[cfg(test)]
149mod tests {
150    use super::*;
151
152    #[test]
153    fn trim_start_colors() {
154        assert_eq!(trim_start_color("foo^2bar"), ("", "foo^2bar"));
155        assert_eq!(trim_start_color("^foo^2bar"), ("", "^foo^2bar"));
156        assert_eq!(trim_start_color("^1foo^2bar"), ("^1", "foo^2bar"));
157        assert_eq!(trim_start_color("^1^2^3foo^2bar"), ("^3", "foo^2bar"));
158    }
159
160    #[test]
161    fn trim_colors() {
162        assert_eq!(trim_color("foo^2bar"), "foobar");
163        assert_eq!(trim_color("^1foo^2bar^3"), "foobar");
164        assert_eq!(trim_color("^1foo^2bar^3"), "foobar");
165        assert_eq!(trim_color("^1foo^bar^3"), "foo^bar");
166        assert_eq!(trim_color("^1foo^2bar^"), "foobar^");
167        assert_eq!(trim_color("^foo^bar^"), "^foo^bar^");
168        assert_eq!(trim_color("\u{fe0f}^1foo^bar^"), "\u{fe0f}foo^bar^");
169        assert_eq!(
170            trim_color("^1^2^3foo\u{fe0f}^2^\u{fe0f}^bar^"),
171            "foo\u{fe0f}^\u{fe0f}^bar^"
172        );
173    }
174}