perf: avoid unnecessary allocations in code highlighting (#3804)

yazi-shim/Cargo.toml

@@ -19,6 +19,10 @@ yazi-macro = { path = "../yazi-macro", version = "26.2.2" }
 crossterm = { workspace = true }
 ratatui   = { workspace = true }
 twox-hash = { workspace = true }
+unicode-width = { workspace = true }
+
+[dependencies.unicode-segmentation]
+version = "1"
 
 [target.'cfg(target_os = "macos")'.dependencies]
 crossterm = { workspace = true, features = [ "use-dev-tty", "libc" ] }

yazi-shim/src/ratatui/grapheme.rs

@@ -0,0 +1,47 @@
+// Copied from https://github.com/ratatui/ratatui/blob/main/ratatui-core/src/text/grapheme.rs
+use std::borrow::Cow;
+
+use ratatui::style::{Style, Styled};
+
+const NBSP: &str = "\u{00a0}";
+const ZWSP: &str = "\u{200b}";
+
+/// A grapheme associated to a style.
+/// Note that, although `StyledGrapheme` is the smallest divisible unit of text,
+/// it actually is not a member of the text type hierarchy (`Text` -> `Line` ->
+/// `Span`). It is a separate type used mostly for rendering purposes. A `Span`
+/// consists of components that can be split into `StyledGrapheme`s, but it does
+/// not contain a collection of `StyledGrapheme`s.
+#[derive(Debug, Default, Clone, Eq, PartialEq, Hash)]
+pub struct StyledGrapheme<'a> {
+	pub symbol: Cow<'a, str>,
+	pub style:  Style,
+}
+
+impl<'a> StyledGrapheme<'a> {
+	/// Creates a new `StyledGrapheme` with the given symbol and style.
+	///
+	/// `style` accepts any type that is convertible to [`Style`] (e.g. [`Style`],
+	/// [`Color`], or your own type that implements [`Into<Style>`]).
+	///
+	/// [`Color`]: crate::style::Color
+	pub fn new<S: Into<Cow<'a, str>>>(symbol: S, style: Style) -> Self {
+		Self { symbol: symbol.into(), style }
+	}
+
+	pub fn is_whitespace(&self) -> bool {
+		let symbol = &*self.symbol;
+		symbol == ZWSP || symbol.chars().all(char::is_whitespace) && symbol != NBSP
+	}
+}
+
+impl Styled for StyledGrapheme<'_> {
+	type Item = Self;
+
+	fn style(&self) -> Style { self.style }
+
+	fn set_style<S: Into<Style>>(mut self, style: S) -> Self::Item {
+		self.style = style.into();
+		self
+	}
+}

yazi-shim/src/ratatui/line.rs

@@ -0,0 +1,83 @@
+use std::{boxed::Box, iter, mem, str::{self, Lines}, vec::IntoIter};
+
+use ratatui::{layout::Alignment, text::Line, widgets::Wrap};
+
+use super::wrapper::{LineComposer, WordWrapper};
+use crate::ratatui::SpanIter;
+
+type WrappedLines<'text> = Box<dyn Iterator<Item = (SpanIter<'text, 'text>, Alignment)> + 'text>;
+
+pub struct LineIter<'text> {
+	inner:    LineIterInner<'text>,
+	tab_size: u8,
+}
+
+enum LineIterInner<'text> {
+	Source { empty: bool, lines: Lines<'text> },
+	Parsed(IntoIter<Line<'text>>),
+	Wrapped(WordWrapper<'text, WrappedLines<'text>, SpanIter<'text, 'text>>),
+	Unlimited(WrappedLines<'text>),
+}
+
+impl<'text> LineIter<'text> {
+	pub fn source(source: &'text str, tab_size: u8) -> Self {
+		Self {
+			inner: LineIterInner::Source { empty: source.is_empty(), lines: source.lines() },
+			tab_size,
+		}
+	}
+
+	pub fn parsed(mut text: Vec<Line<'text>>, tab_size: u8) -> Self {
+		if text.is_empty() {
+			text.push(Line::from(""));
+		}
+		Self { inner: LineIterInner::Parsed(text.into_iter()), tab_size }
+	}
+
+	pub fn wrapped(mut self, wrap: Wrap, width: u16) -> Self {
+		let lines = Box::new(iter::from_fn(move || self.next_owned()));
+		Self {
+			inner:    if width == 0 {
+				LineIterInner::Unlimited(lines)
+			} else {
+				LineIterInner::Wrapped(WordWrapper::new(lines, width, wrap.trim))
+			},
+			tab_size: 0,
+		}
+	}
+
+	pub fn next<'lend>(&'lend mut self) -> Option<(SpanIter<'lend, 'text>, Alignment)> {
+		match self.inner {
+			LineIterInner::Source { .. } | LineIterInner::Parsed(_) => self.next_owned(),
+			LineIterInner::Wrapped(ref mut wrapper) => {
+				let wrapped = wrapper.next_line()?;
+				Some((SpanIter::Wrapped(wrapped.graphemes.iter()), wrapped.alignment))
+			}
+			LineIterInner::Unlimited(ref mut lines) => lines.next(),
+		}
+	}
+
+	fn next_owned(&mut self) -> Option<(SpanIter<'text, 'text>, Alignment)> {
+		match &mut self.inner {
+			LineIterInner::Source { empty, lines } => {
+				let line = if mem::replace(empty, false) {
+					Some(Line::from(""))
+				} else {
+					lines.next().map(Line::from)
+				}?;
+				let alignment = line.alignment.unwrap_or(Alignment::Left);
+
+				Some((SpanIter::new(line, self.tab_size), alignment))
+			}
+			LineIterInner::Parsed(lines) => {
+				let line = lines.next()?;
+				let alignment = line.alignment.unwrap_or(Alignment::Left);
+
+				Some((SpanIter::new(line, self.tab_size), alignment))
+			}
+			LineIterInner::Wrapped(_) | LineIterInner::Unlimited(_) => {
+				unreachable!(); // This branch is handled by next() and should never call next_owned()
+			}
+		}
+	}
+}

yazi-shim/src/ratatui/mod.rs

@@ -1 +1 @@
-yazi_macro::mod_flat!(paragraph);
+yazi_macro::mod_flat!(grapheme line span wrapper);

yazi-shim/src/ratatui/paragraph.rs

@@ -1,52 +0,0 @@
-use ratatui::{text::Text, widgets::{Paragraph, Wrap}};
-
-pub fn line_count<'a, T, W, I>(text: T, width: u16, indent: I, wrap: W) -> usize
-where
-	T: Into<Text<'a>>,
-	I: AsRef<str>,
-	W: Into<Option<Wrap>>,
-{
-	line_count_impl(text.into(), width, indent.as_ref(), wrap.into())
-}
-
-fn line_count_impl(mut text: Text<'_>, mut width: u16, indent: &str, wrap: Option<Wrap>) -> usize {
-	width = width.max(1);
-
-	let Some(wrap) = wrap else {
-		return Paragraph::new(text).line_count(width);
-	};
-
-	if indent.len() == 1 {
-		return Paragraph::new(text).wrap(wrap).line_count(width);
-	}
-
-	let extra = indent.len().saturating_sub(1);
-	for line in &mut text.lines {
-		for span in &mut line.spans {
-			let mut out = None::<String>;
-			let mut start = 0;
-			for (idx, b) in span.content.bytes().enumerate() {
-				if b != b'\t' {
-					continue;
-				}
-
-				let out = out.get_or_insert_with(|| String::with_capacity(span.content.len() + extra));
-				if start < idx {
-					out.push_str(unsafe { span.content.get_unchecked(start..idx) });
-				}
-
-				out.push_str(indent);
-				start = idx + 1;
-			}
-
-			if let Some(mut out) = out {
-				if start < span.content.len() {
-					out.push_str(unsafe { span.content.get_unchecked(start..) });
-				}
-				span.content = out.into();
-			}
-		}
-	}
-
-	Paragraph::new(text).wrap(wrap).line_count(width)
-}

yazi-shim/src/ratatui/span.rs

@@ -0,0 +1,101 @@
+use std::{borrow::Cow, slice::Iter, vec::IntoIter};
+
+use ratatui::{style::Style, text::{Line, Span}};
+use unicode_segmentation::{Graphemes, UnicodeSegmentation};
+
+use crate::ratatui::StyledGrapheme;
+
+#[allow(private_interfaces)]
+pub enum SpanIter<'lend, 'text> {
+	Line {
+		spans:         IntoIter<Span<'text>>,
+		line_style:    Style,
+		current:       Option<CurrentSpan<'text>>,
+		tab_size:      u8,
+		pending_tabs:  u8,
+		pending_style: Style,
+	},
+	Wrapped(Iter<'lend, StyledGrapheme<'text>>),
+}
+
+impl<'lend, 'text> SpanIter<'lend, 'text> {
+	pub(super) fn new(line: Line<'text>, tab_size: u8) -> Self {
+		Self::Line {
+			spans: line.spans.into_iter(),
+			line_style: line.style,
+			current: None,
+			tab_size,
+			pending_tabs: 0,
+			pending_style: Style::default(),
+		}
+	}
+
+	pub fn into_static_line(self) -> Line<'static> {
+		Line::from_iter(self.map(|g| Span { style: g.style, content: g.symbol.into_owned().into() }))
+	}
+}
+
+impl<'lend, 'text> Iterator for SpanIter<'lend, 'text> {
+	type Item = StyledGrapheme<'text>;
+
+	fn next(&mut self) -> Option<Self::Item> {
+		match self {
+			Self::Wrapped(inner) => inner.next().cloned(),
+			Self::Line { spans, line_style, current, tab_size, pending_tabs, pending_style } => loop {
+				if *pending_tabs > 0 {
+					*pending_tabs -= 1;
+					return Some(StyledGrapheme::new(" ", *pending_style));
+				}
+
+				if let Some(span) = current
+					&& let Some(symbol) = span.next_symbol()
+				{
+					if symbol == "\t" {
+						if *tab_size == 0 {
+							continue;
+						}
+
+						*pending_tabs = tab_size.saturating_sub(1);
+						*pending_style = span.style();
+						return Some(StyledGrapheme::new(" ", span.style()));
+					}
+
+					return Some(StyledGrapheme::new(symbol, span.style()));
+				}
+
+				let span = spans.next()?;
+				*current = Some(match span.content {
+					Cow::Borrowed(content) => CurrentSpan::Borrowed {
+						style:     line_style.patch(span.style),
+						graphemes: content.graphemes(true),
+					},
+					Cow::Owned(content) => CurrentSpan::Owned {
+						style:     line_style.patch(span.style),
+						graphemes: content.graphemes(true).map(str::to_owned).collect::<Vec<_>>().into_iter(),
+					},
+				});
+			},
+		}
+	}
+}
+
+// --- CurrentSpan
+enum CurrentSpan<'text> {
+	Borrowed { style: Style, graphemes: Graphemes<'text> },
+	Owned { style: Style, graphemes: IntoIter<String> },
+}
+
+impl<'text> CurrentSpan<'text> {
+	fn style(&self) -> Style {
+		match self {
+			Self::Borrowed { style, .. } | Self::Owned { style, .. } => *style,
+		}
+	}
+
+	fn next_symbol(&mut self) -> Option<Cow<'text, str>> {
+		match self {
+			Self::Borrowed { graphemes, .. } => graphemes.next().map(Cow::Borrowed),
+			Self::Owned { graphemes, .. } => graphemes.next().map(Cow::Owned),
+		}
+	}
+}

yazi-shim/src/ratatui/wrapper.rs

@@ -0,0 +1,196 @@
+// Copied from https://github.com/ratatui/ratatui/blob/main/ratatui-widgets/src/reflow.rs
+use std::{collections::VecDeque, mem};
+
+use ratatui::layout::Alignment;
+use unicode_width::UnicodeWidthStr;
+
+use crate::ratatui::StyledGrapheme;
+
+pub trait LineComposer<'a> {
+	fn next_line<'lend>(&'lend mut self) -> Option<WrappedLine<'lend, 'a>>;
+}
+
+pub struct WrappedLine<'lend, 'text> {
+	pub graphemes: &'lend [StyledGrapheme<'text>],
+	pub width:     u16,
+	pub alignment: Alignment,
+}
+
+#[derive(Debug, Default, Clone)]
+pub struct WordWrapper<'a, O, I>
+where
+	O: Iterator<Item = (I, Alignment)>,
+	I: Iterator<Item = StyledGrapheme<'a>>,
+{
+	input_lines:       O,
+	max_line_width:    u16,
+	wrapped_lines:     VecDeque<Vec<StyledGrapheme<'a>>>,
+	current_alignment: Alignment,
+	current_line:      Vec<StyledGrapheme<'a>>,
+	trim:              bool,
+
+	pending_word:       Vec<StyledGrapheme<'a>>,
+	pending_whitespace: VecDeque<StyledGrapheme<'a>>,
+	pending_line_pool:  Vec<Vec<StyledGrapheme<'a>>>,
+}
+
+impl<'a, O, I> WordWrapper<'a, O, I>
+where
+	O: Iterator<Item = (I, Alignment)>,
+	I: Iterator<Item = StyledGrapheme<'a>>,
+{
+	pub const fn new(lines: O, max_line_width: u16, trim: bool) -> Self {
+		Self {
+			input_lines: lines,
+			max_line_width,
+			wrapped_lines: VecDeque::new(),
+			current_alignment: Alignment::Left,
+			current_line: vec![],
+			trim,
+
+			pending_word: Vec::new(),
+			pending_whitespace: VecDeque::new(),
+			pending_line_pool: Vec::new(),
+		}
+	}
+
+	fn process_input(&mut self, line_symbols: impl IntoIterator<Item = StyledGrapheme<'a>>) {
+		let mut pending_line = self.pending_line_pool.pop().unwrap_or_default();
+		let mut line_width = 0;
+		let mut word_width = 0;
+		let mut whitespace_width = 0;
+		let mut non_whitespace_previous = false;
+
+		self.pending_word.clear();
+		self.pending_whitespace.clear();
+		pending_line.clear();
+
+		for grapheme in line_symbols {
+			let is_whitespace = grapheme.is_whitespace();
+			let symbol_width = grapheme.symbol.width() as u16;
+
+			if symbol_width > self.max_line_width {
+				continue;
+			}
+
+			let word_found = non_whitespace_previous && is_whitespace;
+			let trimmed_overflow =
+				pending_line.is_empty() && self.trim && word_width + symbol_width > self.max_line_width;
+			let whitespace_overflow = pending_line.is_empty()
+				&& self.trim
+				&& whitespace_width + symbol_width > self.max_line_width;
+			let untrimmed_overflow = pending_line.is_empty()
+				&& !self.trim
+				&& word_width + whitespace_width + symbol_width > self.max_line_width;
+
+			if word_found || trimmed_overflow || whitespace_overflow || untrimmed_overflow {
+				if !pending_line.is_empty() || !self.trim {
+					pending_line.extend(self.pending_whitespace.drain(..));
+					line_width += whitespace_width;
+				}
+
+				pending_line.append(&mut self.pending_word);
+				line_width += word_width;
+
+				self.pending_whitespace.clear();
+				whitespace_width = 0;
+				word_width = 0;
+			}
+
+			let line_full = line_width >= self.max_line_width;
+			let pending_word_overflow =
+				symbol_width > 0 && line_width + whitespace_width + word_width >= self.max_line_width;
+
+			if line_full || pending_word_overflow {
+				let mut remaining_width = u16::saturating_sub(self.max_line_width, line_width);
+
+				self.wrapped_lines.push_back(mem::take(&mut pending_line));
+				line_width = 0;
+
+				while let Some(grapheme) = self.pending_whitespace.front() {
+					let width = grapheme.symbol.width() as u16;
+
+					if width > remaining_width {
+						break;
+					}
+
+					whitespace_width -= width;
+					remaining_width -= width;
+					self.pending_whitespace.pop_front();
+				}
+
+				if is_whitespace && self.pending_whitespace.is_empty() {
+					continue;
+				}
+			}
+
+			if is_whitespace {
+				whitespace_width += symbol_width;
+				self.pending_whitespace.push_back(grapheme);
+			} else {
+				word_width += symbol_width;
+				self.pending_word.push(grapheme);
+			}
+
+			non_whitespace_previous = !is_whitespace;
+		}
+
+		if pending_line.is_empty()
+			&& self.pending_word.is_empty()
+			&& !self.pending_whitespace.is_empty()
+			&& self.trim
+		{
+			self.wrapped_lines.push_back(vec![]);
+		}
+		if !pending_line.is_empty() || !self.trim {
+			pending_line.extend(self.pending_whitespace.drain(..));
+		}
+		pending_line.append(&mut self.pending_word);
+
+		#[expect(clippy::else_if_without_else)]
+		if !pending_line.is_empty() {
+			self.wrapped_lines.push_back(pending_line);
+		} else if pending_line.capacity() > 0 {
+			self.pending_line_pool.push(pending_line);
+		}
+		if self.wrapped_lines.is_empty() {
+			self.wrapped_lines.push_back(vec![]);
+		}
+	}
+
+	fn replace_current_line(&mut self, line: Vec<StyledGrapheme<'a>>) {
+		let cache = mem::replace(&mut self.current_line, line);
+		if cache.capacity() > 0 {
+			self.pending_line_pool.push(cache);
+		}
+	}
+}
+
+impl<'a, O, I> LineComposer<'a> for WordWrapper<'a, O, I>
+where
+	O: Iterator<Item = (I, Alignment)>,
+	I: Iterator<Item = StyledGrapheme<'a>>,
+{
+	fn next_line<'lend>(&'lend mut self) -> Option<WrappedLine<'lend, 'a>> {
+		if self.max_line_width == 0 {
+			return None;
+		}
+
+		loop {
+			if let Some(line) = self.wrapped_lines.pop_front() {
+				let line_width = line.iter().map(|grapheme| grapheme.symbol.width() as u16).sum();
+
+				self.replace_current_line(line);
+				return Some(WrappedLine {
+					graphemes: &self.current_line,
+					width:     line_width,
+					alignment: self.current_alignment,
+				});
+			}
+
+			let (line_symbols, line_alignment) = self.input_lines.next()?;
+			self.current_alignment = line_alignment;
+			self.process_input(line_symbols);
+		}
+	}
+}