1 files changed, 110 insertions, 26 deletions

diff --git a/rust/format/src/main.rs b/rust/format/src/main.rs
index 9c167eb..0ec7d3a 100644
--- a/rust/format/src/main.rs
+++ b/rust/format/src/main.rs
@@ -3,48 +3,127 @@ use std::{env::args, fs::File, io::Read};
 mod format_layer;
 
 fn main() {
-    let mut keymap_h = String::new();
+    let mode = args().nth(1).expect("Exists");
 
-    File::open(args().skip(1).last().expect("Exists"))
-        .expect("Should work")
-        .read_to_string(&mut keymap_h)
-        .expect("Failed to read keymap_h");
+    let mut file_to_format = String::new();
+    {
+        File::open(args().skip(1).last().expect("Exists"))
+            .expect("Should work")
+            .read_to_string(&mut file_to_format)
+            .expect("Failed to read file to format");
 
-    let mut lines: Vec<_> = keymap_h.lines().collect();
-    if lines[0] == "// clang-format off" {
-        lines.remove(0);
-    }
-    if lines.last().expect("Exists") == &"// clang-format on" {
-        lines.remove(lines.len() - 1);
+        let mut lines: Vec<_> = file_to_format.lines().collect();
+        if lines[0] == "// clang-format off" {
+            lines.remove(0);
+        }
+        if lines.last().expect("Exists") == &"// clang-format on" {
+            lines.remove(lines.len() - 1);
+        }
+
+        file_to_format = lines.join("\n");
     }
 
-    keymap_h = lines.join("\n");
+    let output = match mode.as_str() {
+        "keymap.h" => {
+            let mut c_column_max = [0; 14];
 
-    let mut c_column_max = [0; 14];
+            // The first run calculates the correct c_column_max values, and the second one
+            // actually formats.
+            let out = format_keymap_h(&mut c_column_max, file_to_format);
+            format_keymap_h(&mut c_column_max, out)
+        }
+        "ledmap.h" => {
+            let mut c_column_max = [0; 14];
 
-    let out = calculate(&mut c_column_max, keymap_h);
-    let output = calculate(&mut c_column_max, out.join("\n"));
+            // We also need two passes here.
+            let out = format_ledmap_h(&mut c_column_max, file_to_format);
+            format_ledmap_h(&mut c_column_max, out)
+        }
+        other => unreachable!("Invalid mode: '{other}'"),
+    };
+    print!("// clang-format off\n{}// clang-format on\n", output);
+}
+
+fn format_ledmap_h(c_column_max: &mut [usize; 14], ledmap_h: String) -> String {
+    let mut output: String = String::new();
+    let mut c_layer: Vec<String> = vec![];
+
+    ledmap_h
+        .lines()
+        .filter(|line| !line.trim().is_empty())
+        .for_each(|line| {
+            let line = line.trim();
+            let first_char = line.chars().next().unwrap();
+            let second_char = line.chars().nth(1).unwrap();
+
+            if line.is_empty() {
+                return;
+            }
+
+            match (first_char, second_char) {
+                ('/', '/') => {
+                    // Comment, don't touch.
+                    output.push_str(line);
+                    output.push('\n');
+                }
+                ('c', 'o') | ('}', ';') => {
+                    // Start or end of the ledmap def, leave it alone
+                    output.push_str(line);
+                    output.push('\n');
+                }
+                ('[', '0'..='9') | ('/', '*') => {
+                    // Start of a new layer
+                    assert!(
+                        c_layer.is_empty(),
+                        "A new layer cannot start, when the current layer is not empty."
+                    );
+                    output.push_str(&format!("  {}", line));
+                    output.push('\n');
+                }
+                ('}', ',') => {
+                    // End of a layer
+                    output.push_str(&format_layer::format_layer(
+                        c_layer.join("\n"),
+                        c_column_max,
+                    ));
+                    c_layer.clear();
 
-    print!(
-        "// clang-format off\n{}\n// clang-format on",
-        output.join("\n")
-    );
+                    output.push_str(&format!("  {}", line));
+                    output.push('\n');
+                }
+                _ => {
+                    // We are in a layer
+                    c_layer.push(line.to_owned());
+                }
+            }
+        });
+    output
 }
 
-fn calculate(c_column_max: &mut [usize; 14], keymap_h: String) -> Vec<String> {
-    let mut output: Vec<String> = vec![];
+fn format_keymap_h(c_column_max: &mut [usize; 14], keymap_h: String) -> String {
+    let mut output: String = String::new();
     let mut c_layer: Vec<String> = vec![];
     keymap_h
         .lines()
         .filter(|line| !line.trim().is_empty())
         .for_each(|line| {
-            let first_char = line.trim().chars().take(1).last().unwrap();
             let line = line.trim();
+            let first_char = line.chars().take(1).last().unwrap();
+
+            if line.is_empty() {
+                return;
+            }
 
             match first_char {
+                '/' => {
+                    // Comment, don't touch
+                    output.push_str(line);
+                    output.push('\n');
+                }
                 'c' | '}' => {
                     // Start or end of the kemap def, leave it alone
-                    output.push(line.to_owned());
+                    output.push_str(line);
+                    output.push('\n');
                 }
                 '[' => {
                     // Start of a new layer
@@ -52,14 +131,19 @@ fn calculate(c_column_max: &mut [usize; 14], keymap_h: String) -> Vec<String> {
                         c_layer.is_empty(),
                         "A new layer cannot start, when the current layer is not empty."
                     );
-                    output.push(format!("  {}", line));
+                    output.push_str(&format!("  {}", line));
+                    output.push('\n');
                 }
                 ')' => {
                     // End of a layer
-                    output.push(format_layer::format_layer(c_layer.join("\n"), c_column_max));
+                    output.push_str(&format_layer::format_layer(
+                        c_layer.join("\n"),
+                        c_column_max,
+                    ));
                     c_layer.clear();
 
-                    output.push(format!("  {}", line));
+                    output.push_str(&format!("  {}", line));
+                    output.push('\n');
                 }
                 _ => {
                     // We are in a layer