escape angle brackets <> for matrix

2 files changed, 24 insertions, 9 deletions

diff --git a/src/matrix.rs b/src/matrix.rs
index 0c359d1..981d906 100644
--- a/src/matrix.rs
+++ b/src/matrix.rs
@@ -2,7 +2,7 @@ use pyo3::prelude::*;
 
 use crate::parser::parse_with_limits;
 
-const MATRIX_FORMATS: &[(&'static str, (&'static str, &'static str))] = &[
+const DUAL_TAGS: &[(&'static str, (&'static str, &'static str))] = &[
     ("_", ("<em>", "</em>")),
     ("*", ("<strong>", "</strong>")),
     ("~", ("<strike>", "</strike>")),
@@ -13,6 +13,14 @@ const MATRIX_FORMATS: &[(&'static str, (&'static str, &'static str))] = &[
     ("||", ("<span data-mx-spoiler>", "</span>")),
 ];
 
+const SINGLE_TAGS: &[(&'static str, &'static str)] = &[
+    (">>", ""),
+    ("```>", ""),
+    ("\\", ""),
+    ("&lt;", "&lt;"),
+    ("&gt;", "&gt;"),
+];
+
 #[pyfunction]
 pub fn format_for_matrix(body: String, mentions: Option<Vec<(String, usize, usize)>>) -> PyResult<String> {
     let mut chars: Vec<char> = body.chars().collect();
@@ -24,19 +32,19 @@ pub fn format_for_matrix(body: String, mentions: Option<Vec<(String, usize, usiz
     let styles: Vec<(String, usize, usize, usize, usize)> = parse_with_limits(&chars, 0, chars.len() - 1, 0);
     let mut tags: Vec<(usize, String, usize)> = Vec::with_capacity(styles.len() * 2);
     for (keyword, start, remove_start, end, remove_end) in styles {
-        if MATRIX_FORMATS.iter().any(|&(k, _)| k == keyword) {
+        if DUAL_TAGS.iter().any(|&(k, _)| k == keyword) {
             let opening_tag = if keyword == "```language" {
-                MATRIX_FORMATS.iter().find(|&&(k, _)| k == keyword).unwrap().1.0.clone()
+                DUAL_TAGS.iter().find(|&&(k, _)| k == keyword).unwrap().1.0.clone()
                 .replace("{}", &chars[start+3..remove_start-1]
                 .into_iter()
                 .collect::<String>())
             } else {
-                MATRIX_FORMATS.iter().find(|&&(k, _)| k == keyword).unwrap().1.0.clone().to_owned()
+                DUAL_TAGS.iter().find(|&&(k, _)| k == keyword).unwrap().1.0.clone().to_owned()
             };
             tags.push((start, opening_tag, remove_start));
-            tags.push((end, MATRIX_FORMATS.iter().find(|&&(k, _)| k == keyword).unwrap().1.1.clone().to_owned(), remove_end));
-        } else if keyword == ">>" || keyword == "```>" || keyword == "\\" {
-            tags.push((start, String::new(), start+1));
+            tags.push((end, DUAL_TAGS.iter().find(|&&(k, _)| k == keyword).unwrap().1.1.clone().to_owned(), remove_end));
+        } else if SINGLE_TAGS.iter().any(|&(k, _)| k == keyword) {
+            tags.push((start, SINGLE_TAGS.iter().find(|&&(k, _)| k == keyword).unwrap().1.to_owned(), start+1));
         }
     }
     for (mxid, start, end) in mentions {
diff --git a/src/parser.rs b/src/parser.rs
index 404bb15..b44d2d0 100644
--- a/src/parser.rs
+++ b/src/parser.rs
@@ -22,14 +22,21 @@ pub fn parse_with_limits(chars: &Vec<char>, start: usize, end: usize, depth: usi
                 styles.append(&mut parse_with_limits(chars, index + 1, to, depth + 1));
                 index = to;
                 continue;
-            }
-            if is_nested_quote(chars, index, depth) {
+            } else if is_nested_quote(chars, index, depth) {
                 styles.push((">>".to_owned(), index, index + 1, index + 1, index + 1));
+            } else {
+                styles.push(("&gt;".to_owned(), index, index + 1, index + 1, index + 1));
             }
             index += 1;
             continue;
         }
 
+        if c == '<' {
+            styles.push(("&lt;".to_owned(), index, index + 1, index + 1, index + 1));
+            index += 1;
+            continue;
+        }
+
         if c == '`' && is_char_repeating(chars, c, 2, index + 1, end) {
             let end_of_line = seek_end_of_line(chars, index + 1, end);
             if end_of_line == end {