cubecl_common/
format.rs

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

cubecl_common/
format.rs
