more unit tests, better highlighting

src/prompt/highlight.rs

@@ -1,4 +1,4 @@
-use std::{env, os::unix::fs::PermissionsExt, path::{Path, PathBuf}, sync::Arc};
+use std::{env, mem, os::unix::fs::PermissionsExt, path::{Path, PathBuf}, sync::Arc};
 use crate::builtin::BUILTINS;
 
 use rustyline::highlight::Highlighter;
@@ -26,24 +26,48 @@ impl FernHighlighter {
 	pub fn highlight_subsh(&self, token: Tk) -> String {
 		if token.flags.contains(TkFlags::IS_SUBSH) {
 			let raw = token.as_str();
-			let body = &raw[1..raw.len() - 1];
-			let sub_hl = FernHighlighter::new(body.to_string());
-			let body_highlighted = sub_hl.hl_input();
-			let open_paren = "(".styled(Style::BrightBlue);
-			let close_paren = ")".styled(Style::BrightBlue);
-			format!("{open_paren}{body_highlighted}{close_paren}")
+			Self::hl_subsh_raw(raw)
 		} else if token.flags.contains(TkFlags::IS_CMDSUB) {
 			let raw = token.as_str();
-			let body = &raw[2..raw.len() - 1];
-			let sub_hl = FernHighlighter::new(body.to_string());
-			let body_highlighted = sub_hl.hl_input();
-			let dollar_paren = "$(".styled(Style::BrightBlue);
-			let close_paren = ")".styled(Style::BrightBlue);
-			format!("{dollar_paren}{body_highlighted}{close_paren}")
+			Self::hl_cmdsub_raw(raw)
 		} else {
 			unreachable!()
 		}
 	}
+	pub fn hl_subsh_raw(raw: &str) -> String {
+		let mut body = &raw[1..];
+		let mut closed = false;
+		if body.ends_with(')') {
+			body = &body[..body.len() - 1];
+			closed = true;
+		}
+		let sub_hl = FernHighlighter::new(body.to_string());
+		let body_highlighted = sub_hl.hl_input();
+		let open_paren = "(".styled(Style::BrightBlue);
+		let close_paren = ")".styled(Style::BrightBlue);
+		let mut result = format!("{open_paren}{body_highlighted}");
+		if closed {
+			result.push_str(&close_paren);
+		}
+		result
+	}
+	pub fn hl_cmdsub_raw(raw: &str) -> String {
+		let mut body = &raw[2..];
+		let mut closed = false;
+		if body.ends_with(')') {
+			body = &body[..body.len() - 1];
+			closed = true;
+		}
+		let sub_hl = FernHighlighter::new(body.to_string());
+		let body_highlighted = sub_hl.hl_input();
+		let dollar_paren = "$(".styled(Style::BrightBlue);
+		let close_paren = ")".styled(Style::BrightBlue);
+		let mut result = format!("{dollar_paren}{body_highlighted}");
+		if closed {
+			result.push_str(&close_paren);
+		}
+		result
+	}
 	pub fn hl_command(&self, token: Tk) -> String {
 		let raw = token.as_str();
 		let paths = env::var("PATH")
@@ -78,11 +102,82 @@ impl FernHighlighter {
 			raw.styled(Style::Bold | Style::Red)
 		}
 	}
+	pub fn hl_dquote(&self, token: Tk) -> String {
+		let raw = token.as_str();
+		let mut chars = raw.chars().peekable();
+		const YELLOW: &str = "\x1b[33m";
+		const RESET: &str = "\x1b[0m";
+		let mut result = String::new();
+		let mut dquote_count = 0;
+
+		result.push_str(YELLOW);
+
+		while let Some(ch) = chars.next() {
+			match ch {
+				'\\' => {
+					result.push(ch);
+					if let Some(ch) = chars.next() {
+						result.push(ch);
+					}
+				}
+				'"' => {
+					dquote_count += 1;
+					result.push(ch);
+					if dquote_count >= 2 {
+						break
+					}
+				}
+				'$' if chars.peek() == Some(&'(')  => {
+					let mut raw_cmd_sub = String::new();
+					raw_cmd_sub.push(ch);
+					raw_cmd_sub.push(chars.next().unwrap());
+					let mut cmdsub_count = 1;
+
+					while let Some(cmdsub_ch) = chars.next() {
+						match cmdsub_ch {
+							'\\' => {
+								raw_cmd_sub.push(cmdsub_ch);
+								if let Some(ch) = chars.next() {
+									raw_cmd_sub.push(ch);
+								}
+							}
+							'$' if chars.peek() == Some(&'(')  => {
+								cmdsub_count += 1;
+								raw_cmd_sub.push(cmdsub_ch);
+								raw_cmd_sub.push(chars.next().unwrap());
+							}
+							')' => {
+								cmdsub_count -= 1;
+								raw_cmd_sub.push(cmdsub_ch);
+								if cmdsub_count <= 0 {
+									let styled = Self::hl_cmdsub_raw(&mem::take(&mut raw_cmd_sub));
+									result.push_str(&styled);
+									result.push_str(YELLOW);
+									break
+								}
+							}
+							_ => raw_cmd_sub.push(cmdsub_ch)
+						}
+					}
+					if !raw_cmd_sub.is_empty() {
+						let styled = Self::hl_cmdsub_raw(&mem::take(&mut raw_cmd_sub));
+						result.push_str(&styled);
+						result.push_str(YELLOW);
+					}
+				}
+				_ => result.push(ch)
+			}
+		}
+
+		result.push_str(RESET);
+
+		result
+	}
 	pub fn hl_input(&self) -> String {
 		let mut output = self.input.clone();
 
 		// TODO: properly implement highlighting for unfinished input
-		let lex_results = LexStream::new(Arc::new(output.clone()), LexFlags::empty());
+		let lex_results = LexStream::new(Arc::new(output.clone()), LexFlags::LEX_UNFINISHED);
 		let mut tokens = vec![];
 
 		for result in lex_results {
@@ -107,6 +202,9 @@ impl FernHighlighter {
 					if token.flags.contains(TkFlags::IS_CMD) {
 						let styled = self.hl_command(token.clone());
 						output.replace_range(token.span.start..token.span.end, &styled);
+					} else if is_dquote(&token) {
+						let styled = self.hl_dquote(token.clone());
+						output.replace_range(token.span.start..token.span.end, &styled);
 					} else {
 						output.replace_range(token.span.start..token.span.end, &token.to_string());
 					}
@@ -143,3 +241,8 @@ impl Highlighter for FernReadline {
 		true
 	}
 }
+
+fn is_dquote(token: &Tk) -> bool {
+	let raw = token.as_str();
+	raw.starts_with('"') 
+}