From adf9cce3bf9e49775def3245ca503b124a52ed66 Mon Sep 17 00:00:00 2001
From: Kajetan Puchalski <kajetan.puchalski@tuta.io>
Date: Wed, 14 Aug 2024 01:58:21 +0100
Subject: [PATCH] driver: Print emitted asm with syntax highlighting

---
 Cargo.lock | 31 +++++++++++++++++++++++++++++++
 Cargo.toml |  2 ++
 src/lib.rs | 39 ++++++++++++++++++++++++++++++++++++++-
 3 files changed, 71 insertions(+), 1 deletion(-)

diff --git a/Cargo.lock b/Cargo.lock
index e2fa12d..6cfde9f 100644
--- a/Cargo.lock
+++ b/Cargo.lock
@@ -87,12 +87,14 @@ dependencies = [
  "colog",
  "display_tree",
  "env_logger",
+ "lliw",
  "llvm-sys",
  "log",
  "regex",
  "serial_test",
  "strum",
  "strum_macros",
+ "synoptic",
 ]
 
 [[package]]
@@ -307,6 +309,12 @@ version = "2.1.0"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "9a3a5bfb195931eeb336b2a7b4d761daec841b97f947d34394601737a7bba5e4"
 
+[[package]]
+name = "if_chain"
+version = "1.0.2"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "cb56e1aa765b4b4f3aadfab769793b7087bb03a4ea4920644a6d238e2df5b9ed"
+
 [[package]]
 name = "is_terminal_polyfill"
 version = "1.70.1"
@@ -325,6 +333,12 @@ version = "0.2.155"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "97b3888a4aecf77e811145cadf6eef5901f4782c53886191b2f693f24761847c"
 
+[[package]]
+name = "lliw"
+version = "0.2.0"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "2d502c8bcc35a4f7ca9a7ffb7ac27b15ba30b1b92c2d69a1e4437e2635d73af7"
+
 [[package]]
 name = "llvm-sys"
 version = "180.0.0"
@@ -611,12 +625,29 @@ dependencies = [
  "unicode-ident",
 ]
 
+[[package]]
+name = "synoptic"
+version = "2.0.4"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "347d7a7719cafb0dc9079d7d042dbaef32e10a8abf327585675148f068fcc053"
+dependencies = [
+ "if_chain",
+ "regex",
+ "unicode-width",
+]
+
 [[package]]
 name = "unicode-ident"
 version = "1.0.12"
 source = "registry+https://github.com/rust-lang/crates.io-index"
 checksum = "3354b9ac3fae1ff6755cb6db53683adb661634f67557942dea4facebec0fee4b"
 
+[[package]]
+name = "unicode-width"
+version = "0.1.13"
+source = "registry+https://github.com/rust-lang/crates.io-index"
+checksum = "0336d538f7abc86d282a4189614dfaa90810dfc2c6f6427eaf88e16311dd225d"
+
 [[package]]
 name = "utf8parse"
 version = "0.2.2"
diff --git a/Cargo.toml b/Cargo.toml
index f5ec7c0..d12233f 100644
--- a/Cargo.toml
+++ b/Cargo.toml
@@ -23,3 +23,5 @@ strum = { version = "0.26.3", features = ["derive"] }
 strum_macros = "0.26.4"
 llvm-sys = { version = "180", features = ["prefer-dynamic"], optional = true }
 cfg-if = "1.0.0"
+synoptic = "2.0.4"
+lliw = "0.2.0"
diff --git a/src/lib.rs b/src/lib.rs
index 494b2ca..a3c776b 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -1,10 +1,12 @@
 #![feature(let_chains)]
 
 use display_tree::format_tree;
+use lliw::Fg;
 use std::fs;
 use std::io::{Error, Write};
 use std::process::Command;
 use strum::EnumIs;
+use synoptic;
 
 pub mod ast;
 pub mod codegen;
@@ -138,7 +140,10 @@ impl Driver {
     fn emit(&self, code: codegen::Program) -> Result<String, Error> {
         let output_path = format!("{}.s", self.name);
         let asm = code.emit();
-        log::debug!("Emitted assembly:\n\n{}", asm);
+
+        if log::log_enabled!(log::Level::Debug) {
+            Driver::print_asm_with_highlight(&asm);
+        }
 
         let mut file = fs::File::create(&output_path)?;
         file.write_all(asm.as_bytes())?;
@@ -146,6 +151,38 @@ impl Driver {
         Ok(output_path)
     }
 
+    fn print_asm_with_highlight(asm: &str) {
+        fn colour(name: &str) -> Fg {
+            match name {
+                "comment" => Fg::LightBlack,
+                "digit" => Fg::Purple,
+                "string" => Fg::Green,
+                "keyword" => Fg::Yellow,
+                "function" => Fg::Red,
+                _ => panic!("unknown token name"),
+            }
+        }
+
+        let mut highlight = synoptic::from_extension("asm", 4).unwrap();
+        let highlighted_asm = asm.split('\n').map(|line| line.to_string()).collect();
+        highlight.run(&highlighted_asm);
+
+        for (line_number, line) in highlighted_asm.iter().enumerate() {
+            print!("{}\t", line_number + 1);
+            // Line returns tokens for the corresponding line
+            for token in highlight.line(line_number, &line) {
+                // Tokens can either require highlighting or not require highlighting
+                match token {
+                    synoptic::TokOpt::Some(text, kind) => {
+                        print!("{}{text}{}", colour(&kind), Fg::Reset)
+                    }
+                    synoptic::TokOpt::None(text) => print!("{text}"),
+                }
+            }
+            println!();
+        }
+    }
+
     fn assemble(&self, path: &str) {
         let _ = Command::new("gcc")
             .args([path, "-masm=intel", "-o", &self.name])