cubecl_common/
format.rs

1use alloc::format;
2use alloc::string::String;
3
4/// Format strings for use in identifiers and types.
5pub fn format_str(string: &str, markers: &[(char, char)], include_space: bool) -> String {
6    let mut result = String::new();
7    let mut depth = 0;
8    let indentation = 4;
9
10    let mut prev = ' ';
11    let mut in_string = false;
12
13    for c in string.chars() {
14        if c == ' ' {
15            if in_string {
16                result.push(c);
17            }
18
19            continue;
20        }
21        if c == '"' {
22            in_string = !in_string;
23        }
24
25        let mut found_marker = false;
26
27        for (start, end) in markers {
28            let (start, end) = (*start, *end);
29
30            if c == start {
31                depth += 1;
32                if prev != ' ' && include_space {
33                    result.push(' ');
34                }
35                result.push(start);
36                result.push('\n');
37                result.push_str(&" ".repeat(indentation * depth));
38                found_marker = true;
39            } else if c == end {
40                depth -= 1;
41                if prev != start {
42                    if prev == ' ' {
43                        result.pop();
44                    }
45                    result.push_str(",\n");
46                    result.push_str(&" ".repeat(indentation * depth));
47                    result.push(end);
48                } else {
49                    for _ in 0..(&" ".repeat(indentation * depth).len()) + 1 + indentation {
50                        result.pop();
51                    }
52                    result.push(end);
53                }
54                found_marker = true;
55            }
56        }
57
58        if found_marker {
59            prev = c;
60            continue;
61        }
62
63        if c == ',' && depth > 0 {
64            if prev == ' ' {
65                result.pop();
66            }
67
68            result.push_str(",\n");
69            result.push_str(&" ".repeat(indentation * depth));
70            continue;
71        }
72
73        if c == ':' && include_space {
74            result.push(c);
75            result.push(' ');
76            prev = ' ';
77        } else {
78            result.push(c);
79            prev = c;
80        }
81    }
82
83    result
84}
85
86/// Format a debug type.
87pub fn format_debug<F: core::fmt::Debug>(string: &F) -> String {
88    let string = format!("{string:?}");
89    format_str(&string, &[('(', ')'), ('[', ']'), ('{', '}')], true)
90}
91
92#[cfg(test)]
93mod tests {
94    use hashbrown::HashMap;
95
96    use super::*;
97
98    #[derive(Debug)]
99    #[allow(unused)]
100    struct Test {
101        map: HashMap<String, u32>,
102    }
103
104    #[test]
105    fn test_format_debug() {
106        let test = Test {
107            map: HashMap::from_iter([("Hey with space".to_string(), 8)]),
108        };
109
110        let formatted = format_debug(&test);
111        let expected = r#"Test {
112    map: {
113        "Hey with space": 8,
114    },
115}"#;
116
117        assert_eq!(expected, formatted);
118    }
119}
cubecl_common/format.rs

cubecl_common/
format.rs
