1 files changed, 44 insertions, 34 deletions

diff --git a/src/render/markdown.rs b/src/render/markdown.rs
index f0cb64a..23d3d59 100644
--- a/src/render/markdown.rs
+++ b/src/render/markdown.rs
@@ -1,4 +1,3 @@
-// use colored::{Color, Colorize};
 use crossterm::style::{Color, Stylize};
 use syntect::highlighting::{Color as SyntectColor, FontStyle, Style, Theme};
 use syntect::parsing::SyntaxSet;
@@ -15,7 +14,7 @@ pub struct MarkdownRender {
     code_color: Color,
     md_syntax: SyntaxReference,
     code_syntax: Option<SyntaxReference>,
-    line_type: LineType,
+    prev_line_type: LineType,
 }
 
 impl MarkdownRender {
@@ -32,7 +31,7 @@ impl MarkdownRender {
             code_color,
             md_syntax,
             code_syntax: None,
-            line_type,
+            prev_line_type: line_type,
         }
     }
 
@@ -43,11 +42,27 @@ impl MarkdownRender {
             .join("\n")
     }
 
-    pub fn render_line(&mut self, line: &str) -> Option<String> {
+    pub fn render_line_stateless(&self, line: &str) -> String {
+        let output = if self.is_code_block() && detect_code_block(line).is_none() {
+            self.render_code_line(line)
+        } else {
+            self.render_line_inner(line, &self.md_syntax)
+        };
+        output.unwrap_or_else(|| line.to_string())
+    }
+
+    pub fn is_code_block(&self) -> bool {
+        matches!(
+            self.prev_line_type,
+            LineType::CodeBegin | LineType::CodeInner
+        )
+    }
+
+    fn render_line(&mut self, line: &str) -> Option<String> {
         if let Some(lang) = detect_code_block(line) {
-            match self.line_type {
+            match self.prev_line_type {
                 LineType::Normal | LineType::CodeEnd => {
-                    self.line_type = LineType::CodeBegin;
+                    self.prev_line_type = LineType::CodeBegin;
                     self.code_syntax = if lang.is_empty() {
                         None
                     } else {
@@ -55,20 +70,20 @@ impl MarkdownRender {
                     };
                 }
                 LineType::CodeBegin | LineType::CodeInner => {
-                    self.line_type = LineType::CodeEnd;
+                    self.prev_line_type = LineType::CodeEnd;
                     self.code_syntax = None;
                 }
             }
             self.render_line_inner(line, &self.md_syntax)
         } else {
-            match self.line_type {
+            match self.prev_line_type {
                 LineType::Normal => self.render_line_inner(line, &self.md_syntax),
                 LineType::CodeEnd => {
-                    self.line_type = LineType::Normal;
+                    self.prev_line_type = LineType::Normal;
                     self.render_line_inner(line, &self.md_syntax)
                 }
                 LineType::CodeBegin => {
-                    self.line_type = LineType::CodeInner;
+                    self.prev_line_type = LineType::CodeInner;
                     self.render_code_line(line)
                 }
                 LineType::CodeInner => self.render_code_line(line),
@@ -76,25 +91,14 @@ impl MarkdownRender {
         }
     }
 
-    pub fn render_line_stateless(&self, line: &str) -> String {
-        let output = if detect_code_block(line).is_some() {
-            self.render_line_inner(line, &self.md_syntax)
-        } else {
-            match self.line_type {
-                LineType::Normal | LineType::CodeEnd => {
-                    self.render_line_inner(line, &self.md_syntax)
-                }
-                _ => self.render_code_line(line),
-            }
-        };
-
-        output.unwrap_or_else(|| line.to_string())
-    }
-
     fn render_line_inner(&self, line: &str, syntax: &SyntaxReference) -> Option<String> {
+        let ws: String = line.chars().take_while(|c| c.is_whitespace()).collect();
+        let trimed_line = &line[ws.len()..];
         let mut highlighter = HighlightLines::new(syntax, &self.md_theme);
-        let ranges = highlighter.highlight_line(line, &self.syntax_set).ok()?;
-        Some(as_terminal_escaped(&ranges))
+        let ranges = highlighter
+            .highlight_line(trimed_line, &self.syntax_set)
+            .ok()?;
+        Some(format!("{ws}{}", as_terminal_escaped(&ranges)))
     }
 
     fn render_code_line(&self, line: &str) -> Option<String> {
@@ -112,7 +116,7 @@ impl MarkdownRender {
 }
 
 #[derive(Debug, Clone, Copy, PartialEq, Eq)]
-enum LineType {
+pub enum LineType {
     Normal,
     CodeBegin,
     CodeInner,
@@ -184,10 +188,16 @@ fn get_code_color(theme: &Theme) -> Color {
         .unwrap_or_else(|| Color::Yellow)
 }
 
-#[test]
-fn test_assets() {
-    let syntax_set: SyntaxSet = bincode::deserialize_from(SYNTAXES).expect("invalid syntaxes.bin");
-    assert!(syntax_set.find_syntax_by_extension("md").is_some());
-    let md_theme: Theme = bincode::deserialize_from(MD_THEME).expect("invalid md_theme binary");
-    assert_eq!(md_theme.name, Some("Monokai Extended".into()));
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn test_assets() {
+        let syntax_set: SyntaxSet =
+            bincode::deserialize_from(SYNTAXES).expect("invalid syntaxes.bin");
+        assert!(syntax_set.find_syntax_by_extension("md").is_some());
+        let md_theme: Theme = bincode::deserialize_from(MD_THEME).expect("invalid md_theme binary");
+        assert_eq!(md_theme.name, Some("Monokai Extended".into()));
+    }
 }