From 1c1845fc762511471087b3f9300d73c3a33a966c Mon Sep 17 00:00:00 2001
From: Mika Vilpas <mika.vilpas@gmail.com>
Date: Tue, 14 May 2024 22:23:33 +0300
Subject: [PATCH] feat: allow themes to add a custom style for file extensions

---
 .luarc.json                                   |   7 +
 yazi-config/src/theme/theme.rs                |   7 +-
 yazi-fm/src/lives/file.rs                     |  36 ++-
 .../file/highlightable_filename_builder.rs    | 261 ++++++++++++++++++
 yazi-plugin/preset/components/file.lua        |  22 +-
 yazi-plugin/src/elements/span.rs              |   4 +-
 yazi-plugin/src/elements/style.rs             |   4 +-
 yazi-plugin/src/url/url.rs                    |   3 +
 yazi-shared/src/theme/color.rs                |   8 +-
 yazi-shared/src/theme/style.rs                |  12 +-
 10 files changed, 325 insertions(+), 39 deletions(-)
 create mode 100644 .luarc.json
 create mode 100644 yazi-fm/src/lives/file/highlightable_filename_builder.rs

diff --git a/.luarc.json b/.luarc.json
new file mode 100644
index 00000000..771b5265
--- /dev/null
+++ b/.luarc.json
@@ -0,0 +1,7 @@
+{
+    "diagnostics.globals": [
+        "ui",
+        "THEME",
+        "MANAGER"
+    ]
+}
\ No newline at end of file
diff --git a/yazi-config/src/theme/theme.rs b/yazi-config/src/theme/theme.rs
index 00508ded..9d170e4a 100644
--- a/yazi-config/src/theme/theme.rs
+++ b/yazi-config/src/theme/theme.rs
@@ -49,13 +49,16 @@ impl Default for Theme {
 pub struct Manager {
 	cwd: Style,
 
+	/// The style for files that have an extension
+	pub extension: Option<Style>,
+
 	// Hovered
 	hovered:         Style,
 	preview_hovered: Style,
 
 	// Find
-	find_keyword:  Style,
-	find_position: Style,
+	find_keyword:      Style,
+	pub find_position: Style,
 
 	// Marker
 	marker_copied:   Style,
diff --git a/yazi-fm/src/lives/file.rs b/yazi-fm/src/lives/file.rs
index 8037fd22..6778df5d 100644
--- a/yazi-fm/src/lives/file.rs
+++ b/yazi-fm/src/lives/file.rs
@@ -2,11 +2,14 @@ use std::ops::Deref;
 
 use mlua::{AnyUserData, IntoLua, Lua, UserDataFields, UserDataMethods};
 use yazi_config::THEME;
-use yazi_plugin::{bindings::{Cast, Cha, Icon, Range}, elements::Style, url::Url};
+use yazi_plugin::{bindings::{Cast, Cha, Icon}, elements::{Span, Style}, url::Url};
 use yazi_shared::MIME_DIR;
 
+use self::highlightable_filename_builder::HighlightedFilenameBuilder;
 use super::{CtxRef, SCOPE};
 
+mod highlightable_filename_builder;
+
 pub(super) struct File {
 	idx:    usize,
 	folder: *const yazi_core::folder::Folder,
@@ -123,17 +126,30 @@ impl File {
 			});
 			reg.add_method("highlights", |lua, me, ()| {
 				let cx = lua.named_registry_value::<CtxRef>("cx")?;
-				let Some(finder) = &cx.manager.active().finder else {
-					return Ok(None);
-				};
-				if me.folder().cwd != me.tab().current.cwd {
-					return Ok(None);
-				}
-				let Some(h) = me.name().and_then(|n| finder.filter.highlighted(n)) else {
-					return Ok(None);
+
+				let stem = me.url.file_stem().unwrap_or_default().to_string_lossy().to_string();
+				let extension_with_leading_dot =
+					me.url.extension().map(|ext| ".".to_string() + ext.to_string_lossy().as_ref());
+
+				let mut builder = HighlightedFilenameBuilder::new(stem, extension_with_leading_dot);
+				if let Some(style) = THEME.manager.extension {
+					if !me.is_dir() {
+						builder.add_extension_highlight(style);
+					}
 				};
 
-				Ok(Some(h.into_iter().map(Range::from).collect::<Vec<_>>()))
+				if let Some(finder) = &cx.manager.active().finder {
+					if me.folder().cwd == me.tab().current.cwd {
+						if let Some(highlighted) = me.name().and_then(|n| finder.filter.highlighted(n)) {
+							highlighted
+								.into_iter()
+								.for_each(|range| builder.add_highlight(range, THEME.manager.find_position));
+						};
+					}
+				};
+
+				let spans: Vec<Span> = builder.build_spans().into_iter().map(|span| span.into()).collect();
+				Ok(Some(spans))
 			});
 		})?;
 
diff --git a/yazi-fm/src/lives/file/highlightable_filename_builder.rs b/yazi-fm/src/lives/file/highlightable_filename_builder.rs
new file mode 100644
index 00000000..34d73071
--- /dev/null
+++ b/yazi-fm/src/lives/file/highlightable_filename_builder.rs
@@ -0,0 +1,261 @@
+#![allow(unused_variables)]
+#![allow(dead_code)]
+
+use std::{collections::HashMap, ops::Range};
+
+use yazi_shared::theme::Style;
+
+/// Represents the name and extension of a file or directory, that can
+/// optionally be highlighted (e.g. when searching, to highlight the matched
+/// characters), and converted to `ui.Span` elements for rendering. Consecutive
+/// highlights with the same style are merged into a single span.
+pub struct HighlightedFilenameBuilder {
+	pub(crate) stem:      String,
+	/// The extension of the file, if any. Must have a leading dot.
+	/// Example: `".txt"`
+	pub(crate) extension: Option<String>,
+
+	/// Maps the index of a character in the full text to its highlight style.
+	pub(crate) highlights: HashMap<usize, Style>,
+}
+
+// NOTE: Since rust strings are utf-8 strings, the characters are not
+// necessarily one byte long. "".chars().count() will return the number of
+// characters in the string, not the number of bytes.
+impl HighlightedFilenameBuilder {
+	pub fn new(stem: String, extension_with_leading_dot: Option<String>) -> Self {
+		assert!(
+			extension_with_leading_dot.is_none()
+				|| extension_with_leading_dot.as_ref().unwrap().starts_with('.')
+		);
+		Self { stem, extension: extension_with_leading_dot, highlights: HashMap::new() }
+	}
+
+	/// Adds a default color to the extension. It might be partially or fully
+	/// overwritten by highlights that are added later.
+	pub fn add_extension_highlight(&mut self, style: Style) {
+		if let Some(extension) = &self.extension {
+			let start = self.stem.len();
+			let end = self.stem.len() + extension.len();
+
+			self.add_highlight(start..end, style);
+		}
+	}
+
+	/// Makes the indices specified by the given Range as having the given
+	/// highlight style. Old highlights are overwritten.
+	pub fn add_highlight(&mut self, range: Range<usize>, style: Style) {
+		for i in range {
+			self.highlights.insert(i, style);
+		}
+	}
+
+	/// Consumes the builder and produces styled spans of text for rendering. The
+	/// stem and extension are separated into different spans.
+	pub fn build_spans(&mut self) -> Vec<HighlightedSpan> {
+		let (stem_highlights, extension_highlights) = {
+			let mut highlights = Vec::new();
+
+			self.stem.chars().enumerate().for_each(|(index, character)| {
+				let style = self.highlights.remove(&index);
+				highlights.insert(index, HighlightedCharacter { character, style });
+			});
+
+			if let Some(extension) = &self.extension {
+				extension.chars().enumerate().for_each(|(index, character)| {
+					let next = self.stem.chars().count() + index;
+					let style = self.highlights.remove(&next);
+					highlights
+						.insert(self.stem.chars().count() + index, HighlightedCharacter { character, style });
+				});
+			}
+
+			let extension = highlights.split_off(self.stem.chars().count());
+			let stem = highlights;
+			(stem, extension)
+		};
+
+		let mut spans = spanify(stem_highlights);
+		spans.extend(spanify(extension_highlights));
+
+		spans
+	}
+}
+
+fn spanify(highlights: Vec<HighlightedCharacter>) -> Vec<HighlightedSpan> {
+	let mut spans = Vec::new();
+	let mut current_index = 0;
+
+	while let Some(this) = highlights.get(current_index) {
+		let next_with_same_style =
+			highlights[current_index + 1..].iter().take_while(|that| this.style == that.style);
+
+		let content =
+			this.character.to_string() + &next_with_same_style.map(|c| c.character).collect::<String>();
+
+		current_index += content.len();
+		let span = HighlightedSpan { content, style: this.style };
+		spans.push(span);
+	}
+
+	spans
+}
+
+#[derive(Debug)]
+pub struct HighlightedSpan {
+	content: String,
+	style:   Option<Style>,
+}
+
+impl From<HighlightedSpan> for yazi_plugin::elements::Span {
+	fn from(value: HighlightedSpan) -> Self {
+		yazi_plugin::elements::Span(ratatui::text::Span {
+			content: value.content.into(),
+			style:   value.style.map(Into::into).unwrap_or_default(),
+		})
+	}
+}
+
+#[derive(Debug)]
+struct HighlightedCharacter {
+	character: char,
+	style:     Option<Style>,
+}
+
+#[cfg(test)]
+mod tests {
+	use yazi_shared::theme::StyleShadow;
+
+	use super::*;
+
+	fn create_style(fgcolor: ratatui::style::Color) -> yazi_shared::theme::Style {
+		Style::from(StyleShadow { fg: Some(yazi_shared::theme::Color(fgcolor)), ..Default::default() })
+	}
+
+	#[test]
+	fn test_highlight_stem() {
+		let mut hl = HighlightedFilenameBuilder::new("filename".to_string(), None);
+		// when there is no extension, the whole name is the stem
+
+		let style = create_style(ratatui::style::Color::Blue);
+		hl.add_highlight(0.."file".len(), style);
+		let spans = hl.build_spans();
+		assert_eq!(spans.len(), 2);
+
+		assert_eq!(spans[0].content, "file");
+		assert_eq!(spans[0].style, Some(style));
+
+		assert_eq!(spans[1].content, "name");
+		assert_eq!(spans[1].style, None);
+	}
+
+	#[test]
+	fn test_highlight_extension() {
+		let mut hl = HighlightedFilenameBuilder::new("filename".to_string(), Some(".txt".to_string()));
+
+		let style = create_style(ratatui::style::Color::Blue);
+		{
+			// simulate the user highlighting the extension without the dot
+			let stem_length = "filename.".chars().count();
+			let full_length = stem_length + "txt".chars().count();
+			hl.add_highlight(stem_length..full_length, style);
+		}
+
+		let spans = hl.build_spans();
+		assert_eq!(spans.len(), 3);
+
+		assert_eq!(spans[0].content, "filename");
+		assert_eq!(spans[0].style, None);
+
+		assert_eq!(spans[1].content, ".");
+		assert_eq!(spans[1].style, None);
+
+		assert_eq!(spans[2].content, "txt");
+		assert_eq!(spans[2].style, Some(style));
+	}
+
+	#[test]
+	fn test_highlight_nothing() {
+		// verify that the stem and extension are separate spans even if they are
+		// not highlighted
+		let mut hl = HighlightedFilenameBuilder::new("notes".to_string(), Some(".txt".to_string()));
+		let spans = hl.build_spans();
+
+		assert_eq!(spans.len(), 2);
+
+		assert_eq!(spans[0].content, "notes");
+		assert_eq!(spans[0].style, None);
+
+		assert_eq!(spans[1].content, ".txt");
+		assert_eq!(spans[1].style, None);
+	}
+
+	#[test]
+	fn test_empty_name() {
+		// just to make sure it doesn't crash
+		let mut hl = HighlightedFilenameBuilder::new("".to_string(), None);
+		let spans = hl.build_spans();
+		assert_eq!(spans.len(), 0);
+	}
+
+	#[test]
+	fn test_highlight_across_extension_boundary() {
+		// for "filename.txt", if the user highlights "name.t"
+		let mut hl = HighlightedFilenameBuilder::new("filename".to_string(), Some(".txt".to_string()));
+		let style = create_style(ratatui::style::Color::Blue);
+
+		{
+			let start = "file".chars().count();
+			let end = "filename.t".chars().count();
+			hl.add_highlight(start..end, style);
+		}
+
+		let spans = hl.build_spans();
+		assert_eq!(spans.len(), 4);
+
+		assert_eq!(spans[0].content, "file");
+		assert_eq!(spans[0].style, None);
+
+		assert_eq!(spans[1].content, "name");
+		assert_eq!(spans[1].style, Some(style));
+
+		assert_eq!(spans[2].content, ".t");
+		assert_eq!(spans[2].style, Some(style));
+
+		assert_eq!(spans[3].content, "xt");
+		assert_eq!(spans[3].style, None);
+	}
+
+	#[test]
+	fn test_add_extension_highlight() {
+		let mut hl = HighlightedFilenameBuilder::new("filename".to_string(), Some(".txt".to_string()));
+		let style = create_style(ratatui::style::Color::Blue);
+		hl.add_extension_highlight(style);
+
+		let spans = hl.build_spans();
+		assert_eq!(spans.len(), 2);
+
+		assert_eq!(spans[0].content, "filename");
+		assert_eq!(spans[0].style, None);
+
+		assert_eq!(spans[1].content, ".txt");
+		assert_eq!(spans[1].style, Some(style));
+	}
+
+	#[test]
+	fn test_multibyte_characters() {
+		// multi-byte characters must not cause panics even if the file name is very
+		// complex. Test both the stem and the extension.
+		let mut hl = HighlightedFilenameBuilder::new(
+			"pokémon-listing-🤔-スゴイ-H̐͌̃ē̐̅l̐͑̚l͋͋́o̐̒̓,̍̑̀ ̓̈́̚w͊̀̕o̓͆͝r̎͝͝l͑̑̇d̎̐̕! ".to_string(),
+			Some(".🤔-スH̐͌̃ē̐̅l̐͑̚txt".to_string()),
+		);
+		let style = create_style(ratatui::style::Color::Blue);
+		hl.add_highlight(0.."poké".chars().count(), style);
+		hl.add_extension_highlight(style);
+
+		let spans = hl.build_spans();
+
+		assert_eq!(spans.len(), 3);
+	}
+}
diff --git a/yazi-plugin/preset/components/file.lua b/yazi-plugin/preset/components/file.lua
index 1e7cabd2..782e63c9 100644
--- a/yazi-plugin/preset/components/file.lua
+++ b/yazi-plugin/preset/components/file.lua
@@ -16,26 +16,8 @@ function File:prefix(file)
 	return prefix == "" and {} or { ui.Span(prefix .. "/") }
 end
 
-function File:highlights(file)
-	local name = file.name:gsub("\r", "?", 1)
-	local highlights = file:highlights()
-	if not highlights or #highlights == 0 then
-		return { ui.Span(name) }
-	end
-
-	local spans, last = {}, 0
-	for _, h in ipairs(highlights) do
-		if h[1] > last then
-			spans[#spans + 1] = ui.Span(name:sub(last + 1, h[1]))
-		end
-		spans[#spans + 1] = ui.Span(name:sub(h[1] + 1, h[2])):style(THEME.manager.find_keyword)
-		last = h[2]
-	end
-	if last < #name then
-		spans[#spans + 1] = ui.Span(name:sub(last + 1))
-	end
-	return spans
-end
+--- Returns the name for the file, optionally highlighted
+function File:highlights(file) return file:highlights() end
 
 function File:found(file)
 	if not file:is_hovered() then
diff --git a/yazi-plugin/src/elements/span.rs b/yazi-plugin/src/elements/span.rs
index 96549f3c..241ea6d3 100644
--- a/yazi-plugin/src/elements/span.rs
+++ b/yazi-plugin/src/elements/span.rs
@@ -3,8 +3,8 @@ use yazi_shared::theme::Color;
 
 use super::Style;
 
-#[derive(Clone, FromLua)]
-pub struct Span(pub(super) ratatui::text::Span<'static>);
+#[derive(Clone, FromLua, Debug)]
+pub struct Span(pub ratatui::text::Span<'static>);
 
 impl Span {
 	pub fn install(lua: &Lua, ui: &Table) -> mlua::Result<()> {
diff --git a/yazi-plugin/src/elements/style.rs b/yazi-plugin/src/elements/style.rs
index 67070b99..3f2080a2 100644
--- a/yazi-plugin/src/elements/style.rs
+++ b/yazi-plugin/src/elements/style.rs
@@ -3,8 +3,8 @@ use std::str::FromStr;
 use mlua::{AnyUserData, ExternalError, ExternalResult, Lua, Table, UserData, UserDataMethods, Value};
 use yazi_shared::theme::Color;
 
-#[derive(Clone, Copy, Default)]
-pub struct Style(pub(super) ratatui::style::Style);
+#[derive(Clone, Copy, Default, PartialEq, Debug)]
+pub struct Style(pub ratatui::style::Style);
 
 impl Style {
 	pub fn install(lua: &Lua, ui: &Table) -> mlua::Result<()> {
diff --git a/yazi-plugin/src/url/url.rs b/yazi-plugin/src/url/url.rs
index 30c5b7c6..bfb1bbca 100644
--- a/yazi-plugin/src/url/url.rs
+++ b/yazi-plugin/src/url/url.rs
@@ -20,6 +20,9 @@ impl Url {
 			reg.add_method("stem", |lua, me, ()| {
 				me.file_stem().map(|s| lua.create_string(s.as_encoded_bytes())).transpose()
 			});
+			reg.add_method("extension", |lua, me, ()| {
+				me.extension().map(|s| lua.create_string(s.as_encoded_bytes())).transpose()
+			});
 			reg.add_method("join", |lua, me, other: UrlRef| Self::cast(lua, me.join(&*other)));
 			reg.add_method("parent", |lua, me, ()| {
 				me.parent_url().map(|u| Self::cast(lua, u)).transpose()
diff --git a/yazi-shared/src/theme/color.rs b/yazi-shared/src/theme/color.rs
index 5cde2172..c561a667 100644
--- a/yazi-shared/src/theme/color.rs
+++ b/yazi-shared/src/theme/color.rs
@@ -3,9 +3,9 @@ use std::str::FromStr;
 use anyhow::Result;
 use serde::{Deserialize, Serialize};
 
-#[derive(Clone, Copy, Debug, Deserialize)]
+#[derive(Clone, Copy, Debug, Deserialize, PartialEq)]
 #[serde(try_from = "String")]
-pub struct Color(ratatui::style::Color);
+pub struct Color(pub ratatui::style::Color);
 
 impl FromStr for Color {
 	type Err = anyhow::Error;
@@ -25,6 +25,10 @@ impl From<Color> for ratatui::style::Color {
 	fn from(value: Color) -> Self { value.0 }
 }
 
+impl From<ratatui::style::Color> for Color {
+	fn from(value: ratatui::style::Color) -> Self { Self(value) }
+}
+
 impl Serialize for Color {
 	fn serialize<S>(&self, serializer: S) -> Result<S::Ok, S::Error>
 	where
diff --git a/yazi-shared/src/theme/style.rs b/yazi-shared/src/theme/style.rs
index 8054d4a4..aff83733 100644
--- a/yazi-shared/src/theme/style.rs
+++ b/yazi-shared/src/theme/style.rs
@@ -3,7 +3,7 @@ use serde::{ser::SerializeMap, Deserialize, Serialize, Serializer};
 
 use super::Color;
 
-#[derive(Clone, Copy, Debug, Deserialize)]
+#[derive(Clone, Copy, Debug, Deserialize, PartialEq)]
 #[serde(from = "StyleShadow")]
 pub struct Style {
 	pub fg:       Option<Color>,
@@ -36,6 +36,16 @@ impl From<Style> for ratatui::style::Style {
 	}
 }
 
+impl From<ratatui::style::Style> for Style {
+	fn from(value: ratatui::style::Style) -> Self {
+		Self {
+			fg:       value.fg.map(Into::into),
+			bg:       value.bg.map(Into::into),
+			modifier: value.add_modifier,
+		}
+	}
+}
+
 #[derive(Default, Deserialize)]
 pub struct StyleShadow {
 	#[serde(default)]