use std::collections::HashSet;
use ratatui::{
buffer::Buffer,
layout::Rect,
style::{Color, Modifier, Style},
text::{Line, Span},
};
use crate::vcs::{DiffKind, DiffLine};
use crate::widgets::tree::{FlatNode, TreeNode, TreeWidget};
#[derive(Debug, Clone)]
pub(crate) enum DiffEntry {
Hunk { content: String },
ContextBlock { line_count: usize },
Line { raw_idx: usize, line: DiffLine },
}
#[derive(Debug)]
pub(crate) struct DiffWidget {
inner: TreeWidget<DiffEntry>,
staged: HashSet<usize>,
}
impl DiffWidget {
pub(crate) fn new(raw: &[DiffLine]) -> Self {
Self {
inner: TreeWidget::new(build_diff_tree(raw)),
staged: HashSet::new(),
}
}
pub(crate) fn move_up(&mut self) {
self.inner.move_up();
}
pub(crate) fn move_down(&mut self) {
self.inner.move_down();
}
pub(crate) fn page_up(&mut self, n: usize) {
self.inner.move_up_n(n);
}
pub(crate) fn page_down(&mut self, n: usize) {
self.inner.move_down_n(n);
}
pub(crate) fn toggle_stage_or_fold(&mut self) {
let entry = match self.inner.selected() {
Some(f) => f.node.data.clone(),
None => return,
};
match entry {
DiffEntry::Hunk { .. } | DiffEntry::ContextBlock { .. } => {
self.inner.toggle_selected();
}
DiffEntry::Line { raw_idx, .. } => {
if !self.staged.remove(&raw_idx) {
self.staged.insert(raw_idx);
}
}
}
}
pub(crate) fn has_staged(&self) -> bool {
!self.staged.is_empty()
}
pub(crate) fn staged_lines<'a>(&'a self, raw: &'a [DiffLine]) -> Vec<(usize, &'a DiffLine)> {
let mut pairs: Vec<(usize, &DiffLine)> = self
.staged
.iter()
.filter_map(|&i| raw.get(i).map(|l| (i, l)))
.collect();
pairs.sort_by_key(|(i, _)| *i);
pairs
}
pub(crate) fn cursor_flat_idx(&self) -> usize {
self.inner.cursor_flat_idx()
}
pub(crate) fn visible_count(&self) -> usize {
self.inner.flatten().len()
}
pub(crate) fn selected_line_no(&self) -> Option<usize> {
let flat = self.inner.selected()?;
match &flat.node.data {
DiffEntry::Line { line, .. } => line.line_no,
DiffEntry::Hunk { .. } | DiffEntry::ContextBlock { .. } => None,
}
}
pub(crate) fn render_into(&self, area: Rect, buf: &mut Buffer, active: bool) {
let staged = &self.staged;
self.inner.render(area, buf, |flat, selected| {
build_line(flat, selected, active, staged)
});
}
}
const CONTEXT_LINES: usize = 3;
fn build_diff_tree(raw: &[DiffLine]) -> Vec<TreeNode<DiffEntry>> {
let hunk_header = raw
.iter()
.find(|dl| dl.kind == DiffKind::HunkHeader)
.map(|dl| dl.content.clone())
.unwrap_or_default();
let content: Vec<(usize, &DiffLine)> = raw
.iter()
.enumerate()
.filter(|(_, dl)| dl.kind != DiffKind::HunkHeader)
.collect();
let n = content.len();
if n == 0 {
return vec![];
}
let changed: Vec<usize> = content
.iter()
.enumerate()
.filter(|(_, (_, dl))| matches!(dl.kind, DiffKind::Added | DiffKind::Removed))
.map(|(pos, _)| pos)
.collect();
if changed.is_empty() {
return vec![];
}
let mut visible = vec![false; n];
for &cp in &changed {
let lo = cp.saturating_sub(CONTEXT_LINES);
let hi = (cp + CONTEXT_LINES + 1).min(n);
for visible_item in visible.iter_mut().take(hi).skip(lo) {
*visible_item = true;
}
}
let mut children: Vec<TreeNode<DiffEntry>> = Vec::new();
let mut hidden_buf: Vec<TreeNode<DiffEntry>> = Vec::new();
for (pos, (raw_idx, dl)) in content.iter().enumerate() {
if visible[pos] {
if !hidden_buf.is_empty() {
let count = hidden_buf.len();
let mut block = TreeNode::inner(
DiffEntry::ContextBlock { line_count: count },
std::mem::take(&mut hidden_buf),
);
block.expanded = false;
children.push(block);
}
children.push(TreeNode::leaf(DiffEntry::Line {
raw_idx: *raw_idx,
line: (*dl).clone(),
}));
} else {
hidden_buf.push(TreeNode::leaf(DiffEntry::Line {
raw_idx: *raw_idx,
line: (*dl).clone(),
}));
}
}
if !hidden_buf.is_empty() {
let count = hidden_buf.len();
let mut block = TreeNode::inner(
DiffEntry::ContextBlock { line_count: count },
hidden_buf,
);
block.expanded = false;
children.push(block);
}
let mut root = TreeNode::inner(DiffEntry::Hunk { content: hunk_header }, children);
root.expanded = true;
vec![root]
}
fn build_line(
flat: &FlatNode<'_, DiffEntry>,
selected: bool,
active: bool,
staged: &HashSet<usize>,
) -> Line<'static> {
match &flat.node.data {
DiffEntry::Hunk { content } => {
let icon = if flat.node.expanded { "▾ " } else { "▸ " };
let text = format!("{}{}", icon, content);
let style = if selected && active {
Style::default()
.fg(Color::Yellow)
.add_modifier(Modifier::BOLD | Modifier::REVERSED)
} else {
Style::default()
.fg(Color::Yellow)
.add_modifier(Modifier::BOLD)
};
Line::from(Span::styled(text, style))
}
DiffEntry::ContextBlock { line_count } => {
let icon = if flat.node.expanded { "▾" } else { "▸" };
let label = if flat.node.expanded {
format!("{} context", icon)
} else {
let s = if *line_count == 1 { "" } else { "s" };
format!("{} {} unchanged line{}", icon, line_count, s)
};
let style = if selected && active {
Style::default().fg(Color::DarkGray).add_modifier(Modifier::REVERSED)
} else {
Style::default().fg(Color::DarkGray)
};
Line::from(Span::styled(label, style))
}
DiffEntry::Line { raw_idx, line } => {
let is_staged = staged.contains(raw_idx);
let sign = match line.kind {
DiffKind::Added => "+",
DiffKind::Removed => "-",
DiffKind::Context => " ",
DiffKind::HunkHeader => " ", };
let lineno = line
.line_no
.map(|n| format!("{:>4}", n))
.unwrap_or_else(|| " ".into());
let stage_mark = if is_staged { "●" } else { " " };
let content = format!(
"{} {} {} {}",
stage_mark,
lineno,
sign,
truncate(&line.content, 80)
);
let base = match line.kind {
DiffKind::Added => Style::default().fg(Color::Green),
DiffKind::Removed => Style::default().fg(Color::Red),
DiffKind::Context | DiffKind::HunkHeader => Style::default().fg(Color::DarkGray),
};
let style = match (is_staged, selected && active) {
(true, true) => {
base.add_modifier(Modifier::BOLD | Modifier::UNDERLINED | Modifier::REVERSED)
}
(true, false) => base.add_modifier(Modifier::BOLD | Modifier::UNDERLINED),
(false, true) => base.add_modifier(Modifier::REVERSED),
(false, false) => base,
};
Line::from(Span::styled(content, style))
}
}
}
fn truncate(s: &str, max_chars: usize) -> &str {
match s.char_indices().nth(max_chars) {
None => s,
Some((byte_idx, _)) => &s[..byte_idx],
}
}