use super::pruning_context::HeadsonPruningContext;
use crate::grep::{
GrepShow, GrepState, compute_grep_state, reorder_priority_for_grep,
};
use crate::order::{NodeId, ObjectType};
use crate::utils::measure::{OutputStats, count_output_stats};
use crate::{GrepConfig, MatchSummary, PriorityOrder, RenderConfig};
use prunist::{
Budget, BudgetKind, Budgets, MustKeep, MustKeepStats, PruningConfig,
PruningResult, select_best_k,
};
use std::collections::VecDeque;
fn is_fileset_root(order_build: &PriorityOrder) -> bool {
order_build
.object_type
.get(crate::order::ROOT_PQ_ID)
.is_some_and(|t| *t == ObjectType::Fileset)
}
fn compute_match_summary(
grep_state: &Option<GrepState>,
inclusion_flags: &[u32],
render_set_id: u32,
) -> Option<MatchSummary> {
let state = grep_state.as_ref()?;
let total: usize = state.direct_matches.iter().filter(|b| **b).count();
let shown: usize = state
.direct_matches
.iter()
.enumerate()
.filter(|(i, b)| {
**b && inclusion_flags.get(*i).copied().unwrap_or(0)
== render_set_id
})
.count();
Some(MatchSummary {
shown,
hidden: total.saturating_sub(shown),
})
}
fn match_summary_zero_shown(
grep_state: &Option<GrepState>,
) -> Option<MatchSummary> {
let state = grep_state.as_ref()?;
let total: usize = state.direct_matches.iter().filter(|b| **b).count();
Some(MatchSummary {
shown: 0,
hidden: total,
})
}
pub fn find_largest_render_under_budgets(
order_build: &mut PriorityOrder,
config: &RenderConfig,
grep: &GrepConfig,
budgets: Budgets,
) -> (String, Option<MatchSummary>) {
let total = order_build.total_nodes;
if total == 0 {
let summary = grep.patterns.is_active().then_some(MatchSummary {
shown: 0,
hidden: 0,
});
return (String::new(), summary);
}
let root_is_fileset = is_fileset_root(order_build);
let mut grep_state = compute_grep_state(order_build, grep);
if strong_fileset_grep_without_matches(grep, &grep_state, root_is_fileset)
{
return (String::new(), match_summary_zero_shown(&grep_state));
}
filter_fileset_without_matches(
order_build,
&mut grep_state,
grep,
config.fileset_tree,
);
reorder_if_grep(order_build, &grep_state);
let fileset_slots = FilesetSlots::new(order_build);
let header_budgeting = header_budgeting_policy(order_build, config);
let measure_cfg = measure_config(order_build, config, header_budgeting);
let min_k = min_k_for(&grep_state, grep);
let must_keep_slice = guaranteed_nodes_slice(&grep_state, grep);
let must_keep = must_keep_slice.map(|flags| {
let free_allowance = effective_budgets_with_grep(
order_build,
&measure_cfg,
grep,
&grep_state,
fileset_slots.as_ref(),
budgets.measure_chars(),
);
let (mk_stats, mk_slots) = if let Some((mk, mk_slots)) = free_allowance
{
(mk, mk_slots)
} else {
measure_must_keep_with_slots(
order_build,
&measure_cfg,
flags,
budgets.measure_chars(),
fileset_slots.as_ref(),
)
};
let per_slot = if budgets.per_slot_active() && mk_slots.is_none() {
Some(vec![mk_stats])
} else {
mk_slots
};
MustKeep {
flags,
stats: MustKeepStats {
total: mk_stats,
per_slot,
},
}
});
let context = HeadsonPruningContext {
order_build,
measure_cfg: &measure_cfg,
fileset_slots: fileset_slots.as_ref(),
};
let selection =
select_best_k(PruningConfig::new(&context, budgets, min_k, must_keep));
let finalize_ctx = FinalizeContext {
budgets,
fileset_slots: fileset_slots.as_ref(),
measure_cfg: &measure_cfg,
grep,
grep_state: &grep_state,
must_keep: must_keep_slice,
root_is_fileset,
header_budgeting,
};
finalize_render_from_selection(
order_build,
config,
selection,
&finalize_ctx,
)
.unwrap_or_else(|| (String::new(), match_summary_zero_shown(&grep_state)))
}
struct FinalizeContext<'a> {
budgets: Budgets,
fileset_slots: Option<&'a FilesetSlots>,
measure_cfg: &'a RenderConfig,
grep: &'a GrepConfig,
grep_state: &'a Option<GrepState>,
must_keep: Option<&'a [bool]>,
root_is_fileset: bool,
header_budgeting: HeadersBudgeting,
}
fn finalize_render_from_selection(
order_build: &mut PriorityOrder,
config: &RenderConfig,
selection: PruningResult<NodeId>,
ctx: &FinalizeContext<'_>,
) -> Option<(String, Option<MatchSummary>)> {
let PruningResult {
top_k: k_opt,
mut inclusion_flags,
render_set_id,
selection_order,
} = selection;
let found_k = k_opt.is_some();
let k = k_opt.unwrap_or(0);
if should_short_circuit_after_selection(
&ctx.budgets,
ctx.must_keep,
ctx.root_is_fileset,
found_k,
k,
) {
return None;
}
inclusion_flags.fill(0);
let per_slot_caps_active = ctx.budgets.per_slot_active();
apply_selection(
order_build,
selection_order.as_deref(),
k,
&mut inclusion_flags,
render_set_id,
);
include_strong_grep_must_keep(
order_build,
ctx.grep,
ctx.grep_state,
&mut inclusion_flags,
render_set_id,
);
if per_slot_caps_active && !config.count_fileset_headers_in_budgets {
ensure_fileset_headers_for_empty_slots(
order_build,
render_set_id,
&mut inclusion_flags,
&ctx.budgets,
ctx.measure_cfg,
ctx.fileset_slots,
ctx.header_budgeting,
);
}
if should_short_circuit_zero_line_slots(
&ctx.budgets,
ctx.fileset_slots,
&inclusion_flags,
render_set_id,
ctx.root_is_fileset,
) {
return None;
}
if config.debug {
crate::debug::emit_render_debug(
order_build,
&inclusion_flags,
render_set_id,
config,
ctx.budgets,
k,
);
}
let text = crate::serialization::render_from_render_set(
order_build,
&inclusion_flags,
render_set_id,
&crate::RenderConfig {
grep_highlight: ctx.grep.patterns.highlight().cloned(),
..config.clone()
},
);
let summary =
compute_match_summary(ctx.grep_state, &inclusion_flags, render_set_id);
Some((text, summary))
}
fn strong_fileset_grep_without_matches(
grep: &GrepConfig,
state: &Option<GrepState>,
root_is_fileset: bool,
) -> bool {
if !grep.has_strong()
|| !matches!(grep.show, GrepShow::Matching)
|| !root_is_fileset
{
return false;
}
let Some(s) = state else {
return true;
};
s.guaranteed_count == 0
}
fn is_strong_grep(grep: &GrepConfig, state: &Option<GrepState>) -> bool {
let count = state.as_ref().map(|s| s.guaranteed_count).unwrap_or(0);
grep.has_strong() && grep.force_strong_inclusion && count > 0
}
fn apply_selection(
order_build: &PriorityOrder,
selection_order: Option<&[NodeId]>,
k: usize,
inclusion_flags: &mut Vec<u32>,
render_set_id: u32,
) {
if let Some(order) = selection_order {
mark_custom_top_k_and_ancestors(
order_build,
order,
k,
inclusion_flags.as_mut_slice(),
render_set_id,
);
} else {
crate::serialization::prepare_render_set_top_k_and_ancestors(
order_build,
k,
inclusion_flags,
render_set_id,
);
}
}
fn include_strong_grep_must_keep(
order_build: &PriorityOrder,
grep: &GrepConfig,
grep_state: &Option<GrepState>,
inclusion_flags: &mut [u32],
render_set_id: u32,
) {
if !is_strong_grep(grep, grep_state) {
return;
}
if let Some(state) = grep_state {
include_must_keep(
order_build,
inclusion_flags,
render_set_id,
&state.guaranteed_nodes,
);
}
}
fn should_short_circuit_after_selection(
budgets: &Budgets,
must_keep_slice: Option<&[bool]>,
root_is_fileset: bool,
found_k: bool,
k: usize,
) -> bool {
if budgets.per_slot_zero_cap() {
return true;
}
if k == 0
&& must_keep_slice.is_none()
&& !budgets.per_slot_active()
&& !root_is_fileset
{
return true;
}
if !found_k && must_keep_slice.is_none() && !root_is_fileset {
return true;
}
false
}
fn should_short_circuit_zero_line_slots(
budgets: &Budgets,
fileset_slots: Option<&FilesetSlots>,
inclusion_flags: &[u32],
render_set_id: u32,
root_is_fileset: bool,
) -> bool {
let Some(Budget {
kind: BudgetKind::Lines,
cap: 0,
}) = budgets.per_slot
else {
return false;
};
if let Some(slots) = fileset_slots {
let has_included_slot =
inclusion_flags.iter().enumerate().any(|(idx, flag)| {
*flag == render_set_id
&& slots.map.get(idx).and_then(|s| *s).is_some()
});
if !has_included_slot {
return true;
}
}
!root_is_fileset
}
fn reorder_if_grep(
order_build: &mut PriorityOrder,
state: &Option<GrepState>,
) {
if let Some(s) = state {
reorder_priority_for_grep(order_build, &s.matched_nodes);
}
}
#[allow(
clippy::cognitive_complexity,
reason = "Fileset filtering logic is easier to follow inline"
)]
fn filter_fileset_without_matches(
order_build: &mut PriorityOrder,
state: &mut Option<GrepState>,
grep: &GrepConfig,
keep_fileset_children_for_tree: bool,
) {
if !grep.has_strong() {
return;
}
let Some(s) = state.as_mut() else {
return;
};
if matches!(grep.show, crate::grep::GrepShow::All) {
return;
}
if order_build
.object_type
.get(crate::order::ROOT_PQ_ID)
.is_none_or(|t| *t != ObjectType::Fileset)
{
return;
}
let Some(fileset_slots) = order_build.fileset_render_slots() else {
return;
};
if fileset_slots.is_empty() {
return;
}
let Some(slot_map) = compute_fileset_slot_map(order_build) else {
return;
};
let mut keep_slots = vec![false; fileset_slots.len()];
for (idx, keep) in s.guaranteed_nodes.iter().enumerate() {
if !*keep {
continue;
}
if let Some(slot) = slot_map.get(idx).copied().flatten() {
keep_slots[slot] = true;
}
}
for (slot, child) in fileset_slots.iter().enumerate() {
if s.guaranteed_nodes.get(child.id.0).copied().unwrap_or(false) {
keep_slots[slot] = true;
}
}
let filtered_slots = if keep_fileset_children_for_tree {
None
} else {
let mut filtered_slots: Vec<crate::order::FilesetRenderSlot> =
Vec::new();
for (slot, child) in fileset_slots.iter().enumerate() {
if keep_slots.get(slot).copied().unwrap_or(false) {
filtered_slots.push(*child);
}
}
Some(filtered_slots)
};
order_build.by_priority.retain(|node| {
match slot_map.get(node.0).copied().flatten() {
Some(slot) => keep_slots.get(slot).copied().unwrap_or(false),
None => true,
}
});
if let Some(filtered_slots) = filtered_slots {
let filtered_len = filtered_slots.len();
order_build.fileset_render_slots = Some(filtered_slots);
if let Some(metrics) =
order_build.metrics.get_mut(crate::order::ROOT_PQ_ID)
{
metrics.object_len = Some(filtered_len);
}
}
for (idx, keep) in s.guaranteed_nodes.iter_mut().enumerate() {
if let Some(slot) = slot_map.get(idx).copied().flatten() {
if !keep_slots.get(slot).copied().unwrap_or(false) {
*keep = false;
}
}
}
s.guaranteed_count = s.guaranteed_nodes.iter().filter(|b| **b).count();
}
#[allow(
clippy::cognitive_complexity,
reason = "single DFS that is clearer in one routine than split helpers"
)]
pub(crate) fn compute_fileset_slot_map(
order_build: &PriorityOrder,
) -> Option<Vec<Option<usize>>> {
if order_build
.object_type
.get(crate::order::ROOT_PQ_ID)
.is_none_or(|t| *t != ObjectType::Fileset)
{
return None;
}
let children = order_build.fileset_render_slots()?;
if children.is_empty() {
return None;
}
let mut slots: Vec<Option<usize>> = vec![None; order_build.total_nodes];
for (slot, child) in children.iter().enumerate() {
let mut stack = vec![child.id.0];
while let Some(node_idx) = stack.pop() {
if slots[node_idx].is_some() {
continue;
}
slots[node_idx] = Some(slot);
if let Some(kids) = order_build.children.get(node_idx) {
stack.extend(kids.iter().map(|k| k.0));
}
}
}
Some(slots)
}
#[derive(Clone, Debug)]
pub(crate) struct FilesetSlots {
pub map: Vec<Option<usize>>,
pub count: usize,
pub names: Option<Vec<String>>,
}
#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum HeadersBudgeting {
Free,
Charged,
}
impl HeadersBudgeting {
pub fn is_charged(self) -> bool {
matches!(self, HeadersBudgeting::Charged)
}
}
fn header_budgeting_policy(
order_build: &PriorityOrder,
config: &RenderConfig,
) -> HeadersBudgeting {
if !is_fileset_root(order_build) || !config.show_fileset_headers {
return HeadersBudgeting::Free;
}
if config.count_fileset_headers_in_budgets {
HeadersBudgeting::Charged
} else {
HeadersBudgeting::Free
}
}
impl FilesetSlots {
pub(crate) fn new(order_build: &PriorityOrder) -> Option<Self> {
let map = compute_fileset_slot_map(order_build)?;
let count = map.iter().flatten().max().map(|s| *s + 1)?;
let names = fileset_slot_names(order_build);
Some(Self { map, count, names })
}
}
fn fileset_slot_names(order_build: &PriorityOrder) -> Option<Vec<String>> {
let children = order_build.fileset_render_slots()?;
if children.is_empty() {
return None;
}
let mut names = Vec::with_capacity(children.len());
for child in children {
let name = order_build
.nodes
.get(child.id.0)
.and_then(|n| n.key_in_object())
.unwrap_or_default()
.to_string();
names.push(name);
}
Some(names)
}
pub(super) fn round_robin_slot_priority(
order_build: &PriorityOrder,
slots: &FilesetSlots,
) -> Option<Vec<NodeId>> {
let slot_count = slots.count;
if slot_count == 0 {
return None;
}
let (mut buckets, unslotted) =
bucket_nodes_by_slot(order_build, &slots.map, slot_count);
let mut out = drain_round_robin(&mut buckets);
out.extend(unslotted);
Some(out)
}
fn bucket_nodes_by_slot(
order_build: &PriorityOrder,
slot_map: &[Option<usize>],
slot_count: usize,
) -> (Vec<VecDeque<NodeId>>, Vec<NodeId>) {
let mut buckets: Vec<VecDeque<NodeId>> = vec![VecDeque::new(); slot_count];
let mut unslotted: Vec<NodeId> = Vec::new();
for node in order_build.by_priority.iter().copied() {
if let Some(slot) = slot_map.get(node.0).and_then(|s| *s) {
if let Some(bucket) = buckets.get_mut(slot) {
bucket.push_back(node);
continue;
}
}
unslotted.push(node);
}
(buckets, unslotted)
}
fn drain_round_robin(buckets: &mut [VecDeque<NodeId>]) -> Vec<NodeId> {
let slot_count = buckets.len();
let total: usize = buckets.iter().map(VecDeque::len).sum();
let mut out: Vec<NodeId> = Vec::with_capacity(total);
let mut remaining = total;
let mut cursor = 0usize;
while remaining > 0 {
let slot = cursor % slot_count;
if let Some(node) = buckets.get_mut(slot).and_then(VecDeque::pop_front)
{
out.push(node);
remaining = remaining.saturating_sub(1);
}
cursor = cursor.saturating_add(1);
}
out
}
fn effective_budgets_with_grep(
order_build: &PriorityOrder,
measure_cfg: &RenderConfig,
grep: &GrepConfig,
state: &Option<GrepState>,
fileset_slots: Option<&FilesetSlots>,
measure_chars: bool,
) -> Option<(OutputStats, Option<Vec<OutputStats>>)> {
if !is_strong_grep(grep, state) {
return None;
}
let Some(s) = state else {
return None;
};
Some(measure_must_keep_with_slots(
order_build,
measure_cfg,
&s.guaranteed_nodes,
measure_chars,
fileset_slots,
))
}
fn min_k_for(state: &Option<GrepState>, grep: &GrepConfig) -> usize {
if is_strong_grep(grep, state) {
state
.as_ref()
.map(|s| s.guaranteed_count.max(1))
.unwrap_or(1)
} else {
1
}
}
fn guaranteed_nodes_slice<'a>(
state: &'a Option<GrepState>,
grep: &GrepConfig,
) -> Option<&'a [bool]> {
state
.as_ref()
.filter(|_| grep.has_strong() && grep.force_strong_inclusion)
.and_then(|s| {
(s.guaranteed_count > 0).then_some(s.guaranteed_nodes.as_slice())
})
}
#[allow(
clippy::cognitive_complexity,
reason = "Tiny budget summary checks are clearer inline than split helpers."
)]
pub(crate) fn constrained_dimensions(
budgets: Budgets,
stats: &crate::utils::measure::OutputStats,
slot_stats: Option<&[crate::utils::measure::OutputStats]>,
) -> Vec<&'static str> {
let mut dims: Vec<&'static str> = Vec::new();
if let Some(b) = budgets.global {
if b.exceeds(stats) {
dims.push(kind_str(b.kind, false));
}
}
if let Some(b) = budgets.per_slot {
if let Some(slot_vec) = slot_stats {
if slot_vec.iter().any(|st| b.exceeds(st)) {
dims.push(kind_str(b.kind, true));
}
} else if b.exceeds(stats) {
dims.push(kind_str(b.kind, true));
}
}
dims
}
fn kind_str(kind: BudgetKind, per_slot: bool) -> &'static str {
match (kind, per_slot) {
(BudgetKind::Bytes, false) => "bytes",
(BudgetKind::Chars, false) => "chars",
(BudgetKind::Lines, false) => "lines",
(BudgetKind::Bytes, true) => "per-file bytes",
(BudgetKind::Chars, true) => "per-file chars",
(BudgetKind::Lines, true) => "per-file lines",
}
}
fn measure_config(
order_build: &PriorityOrder,
config: &RenderConfig,
header_budgeting: HeadersBudgeting,
) -> RenderConfig {
let root_is_fileset = order_build
.object_type
.get(crate::order::ROOT_PQ_ID)
.is_some_and(|t| *t == crate::order::ObjectType::Fileset);
let mut measure_cfg = config.clone();
measure_cfg.color_enabled = false;
measure_cfg.count_fileset_headers_in_budgets =
header_budgeting.is_charged();
if config.fileset_tree {
measure_cfg.show_fileset_headers = header_budgeting.is_charged();
} else if config.show_fileset_headers
&& root_is_fileset
&& header_budgeting == HeadersBudgeting::Free
{
measure_cfg.show_fileset_headers = false;
}
measure_cfg
}
fn measure_must_keep_with_slots(
order_build: &PriorityOrder,
measure_cfg: &RenderConfig,
must_keep: &[bool],
measure_chars: bool,
fileset_slots: Option<&FilesetSlots>,
) -> (OutputStats, Option<Vec<OutputStats>>) {
let mut measure_cfg = measure_cfg.clone();
if matches!(
measure_cfg.template,
crate::OutputTemplate::Text | crate::OutputTemplate::Auto
) {
measure_cfg.style = crate::serialization::types::Style::Strict;
}
let mut inclusion_flags: Vec<u32> = vec![0; order_build.total_nodes];
let render_set_id: u32 = 1;
include_must_keep(
order_build,
&mut inclusion_flags,
render_set_id,
must_keep,
);
let mut recorder = fileset_slots.map(|slots| {
crate::serialization::output::SlotStatsRecorder::new(
slots.count,
measure_chars,
)
});
let (rendered, slot_stats) =
crate::serialization::render_from_render_set_with_slots(
order_build,
&inclusion_flags,
render_set_id,
&measure_cfg,
fileset_slots.map(|slots| slots.map.as_slice()),
recorder.take(),
);
(
crate::utils::measure::count_output_stats(&rendered, measure_chars),
slot_stats,
)
}
fn include_string_descendants(
order: &PriorityOrder,
id: usize,
flags: &mut [u32],
render_id: u32,
) {
if let Some(children) = order.children.get(id) {
for child in children {
let idx = child.0;
if flags[idx] != render_id {
flags[idx] = render_id;
include_string_descendants(order, idx, flags, render_id);
}
}
}
}
pub(super) fn include_must_keep(
order_build: &PriorityOrder,
inclusion_flags: &mut [u32],
render_set_id: u32,
must_keep: &[bool],
) {
for (idx, keep) in must_keep.iter().enumerate() {
if !*keep {
continue;
}
crate::utils::graph::mark_node_and_ancestors(
order_build,
crate::NodeId(idx),
inclusion_flags,
render_set_id,
);
if matches!(
order_build.nodes.get(idx),
Some(crate::RankedNode::SplittableLeaf { .. })
) {
include_string_descendants(
order_build,
idx,
inclusion_flags,
render_set_id,
);
}
}
}
pub(super) fn mark_custom_top_k_and_ancestors(
order_build: &PriorityOrder,
selection_order: &[NodeId],
top_k: usize,
inclusion_flags: &mut [u32],
render_id: u32,
) {
for node in selection_order.iter().take(top_k) {
crate::utils::graph::mark_node_and_ancestors(
order_build,
*node,
inclusion_flags,
render_id,
);
}
}
#[allow(
clippy::cognitive_complexity,
clippy::too_many_arguments,
reason = "Header insertion must juggle per-slot state, budgeting policy, and tree marking in one pass; splitting would hurt readability."
)]
fn ensure_fileset_headers_for_empty_slots(
order_build: &PriorityOrder,
render_id: u32,
inclusion_flags: &mut Vec<u32>,
budgets: &Budgets,
measure_cfg: &RenderConfig,
fileset_slots: Option<&FilesetSlots>,
header_budgeting: HeadersBudgeting,
) {
let Some(slots) = fileset_slots else {
return;
};
if slots.count == 0 {
return;
}
let Some(fileset_children) = order_build.fileset_render_slots() else {
return;
};
if inclusion_flags.len() < order_build.total_nodes {
inclusion_flags.resize(order_build.total_nodes, 0);
}
let measure_chars = budgets.measure_chars();
let newline_len = measure_cfg.newline.len();
let zero_per_slot =
matches!(budgets.per_slot, Some(Budget { cap: 0, .. }));
for slot_idx in 0..slots.count {
let slot_has_nodes =
inclusion_flags.iter().enumerate().any(|(idx, flag)| {
*flag == render_id
&& slots
.map
.get(idx)
.and_then(|s| *s)
.is_some_and(|s| s == slot_idx)
});
if slot_has_nodes || zero_per_slot {
continue;
}
let header_stats = header_stats_for_slot(
slot_idx,
slots.names.as_ref(),
measure_chars,
newline_len,
budgets,
);
if header_budgeting.is_charged() && header_stats.is_none() {
continue;
}
if let Some(file_node) = fileset_children.get(slot_idx) {
crate::utils::graph::mark_node_and_ancestors(
order_build,
file_node.id,
inclusion_flags,
render_id,
);
}
}
}
#[allow(
clippy::cognitive_complexity,
reason = "Header measurement branches on name presence and budget kinds; keeping it in one routine makes the cap checks traceable."
)]
fn header_stats_for_slot(
slot_idx: usize,
header_names: Option<&Vec<String>>,
measure_chars: bool,
newline_len: usize,
budgets: &Budgets,
) -> Option<OutputStats> {
let stats = match header_names.and_then(|n| n.get(slot_idx)) {
Some(name) => {
let mut s =
count_output_stats(&format!("==> {name} <=="), measure_chars);
s.lines = s.lines.max(1);
s.bytes = s.bytes.saturating_add(newline_len);
if measure_chars {
s.chars = s.chars.saturating_add(newline_len);
}
s
}
None => OutputStats {
bytes: newline_len,
chars: if measure_chars { newline_len } else { 0 },
lines: 1,
},
};
if let Some(cap) = budgets.per_slot {
let value = match cap.kind {
BudgetKind::Bytes => stats.bytes,
BudgetKind::Chars => stats.chars,
BudgetKind::Lines => stats.lines,
};
if value > cap.cap {
return None;
}
}
Some(stats)
}
#[cfg(test)]
mod tests {
}