Skip to main content

fret_ui_headless/table/
grouping.rs

1use std::collections::{HashMap, HashSet};
2use std::sync::Arc;
3
4use super::{ColumnDef, ColumnId, RowId, RowKey, RowModel, TableOptions};
5
6/// TanStack-compatible grouping state: an ordered list of grouped column ids.
7pub type GroupingState = Vec<ColumnId>;
8
9#[derive(Debug, Clone, Copy, PartialEq, Eq)]
10pub enum GroupedColumnMode {
11    /// Do not reorder or remove grouped columns.
12    None,
13    /// Move grouped columns to the start of the leaf column list (TanStack default).
14    Reorder,
15    /// Remove grouped columns from the leaf column list.
16    Remove,
17}
18
19pub fn is_column_grouped(state: &GroupingState, column: &ColumnId) -> bool {
20    state.iter().any(|c| c.as_ref() == column.as_ref())
21}
22
23pub fn grouped_index(state: &GroupingState, column: &ColumnId) -> Option<usize> {
24    state.iter().position(|c| c.as_ref() == column.as_ref())
25}
26
27pub fn set_grouping(state: &mut GroupingState, ids: impl IntoIterator<Item = ColumnId>) {
28    state.clear();
29    for id in ids {
30        if state.iter().any(|c| c.as_ref() == id.as_ref()) {
31            continue;
32        }
33        state.push(id);
34    }
35}
36
37pub fn toggle_column_grouping(state: &mut GroupingState, column: &ColumnId) {
38    if let Some(i) = grouped_index(state, column) {
39        state.remove(i);
40    } else {
41        state.push(column.clone());
42    }
43}
44
45pub fn toggle_column_grouping_value(
46    state: &mut GroupingState,
47    column: &ColumnId,
48    grouped: Option<bool>,
49) {
50    let should_group = grouped.unwrap_or_else(|| !is_column_grouped(state, column));
51    let is_grouped = is_column_grouped(state, column);
52
53    match (is_grouped, should_group) {
54        (true, false) => {
55            if let Some(i) = grouped_index(state, column) {
56                state.remove(i);
57            }
58        }
59        (false, true) => state.push(column.clone()),
60        _ => {}
61    }
62}
63
64pub fn toggled_column_grouping(state: &GroupingState, column: &ColumnId) -> GroupingState {
65    let mut next = state.clone();
66    toggle_column_grouping(&mut next, column);
67    next
68}
69
70pub fn toggled_column_grouping_value(
71    state: &GroupingState,
72    column: &ColumnId,
73    grouped: Option<bool>,
74) -> GroupingState {
75    let mut next = state.clone();
76    toggle_column_grouping_value(&mut next, column, grouped);
77    next
78}
79
80pub fn order_columns_for_grouping<'c, TData>(
81    leaf_columns: &'c [ColumnDef<TData>],
82    grouping: &[ColumnId],
83    mode: GroupedColumnMode,
84) -> Vec<&'c ColumnDef<TData>> {
85    if leaf_columns.is_empty() {
86        return Vec::new();
87    }
88    if grouping.is_empty() || mode == GroupedColumnMode::None {
89        return leaf_columns.iter().collect();
90    }
91
92    let is_grouped = |c: &&ColumnDef<TData>| grouping.iter().any(|id| id.as_ref() == c.id.as_ref());
93
94    let non_grouped: Vec<&ColumnDef<TData>> =
95        leaf_columns.iter().filter(|c| !is_grouped(c)).collect();
96    if mode == GroupedColumnMode::Remove {
97        return non_grouped;
98    }
99
100    let by_id: HashMap<&str, &ColumnDef<TData>> =
101        leaf_columns.iter().map(|c| (c.id.as_ref(), c)).collect();
102    let mut grouped_cols: Vec<&ColumnDef<TData>> = Vec::new();
103    for id in grouping {
104        if let Some(col) = by_id.get(id.as_ref()).copied() {
105            grouped_cols.push(col);
106        }
107    }
108
109    grouped_cols.extend(non_grouped);
110    grouped_cols
111}
112
113pub fn order_column_refs_for_grouping<'c, TData>(
114    leaf_columns: &'c [&'c ColumnDef<TData>],
115    grouping: &[ColumnId],
116    mode: GroupedColumnMode,
117) -> Vec<&'c ColumnDef<TData>> {
118    if leaf_columns.is_empty() {
119        return Vec::new();
120    }
121    if grouping.is_empty() || mode == GroupedColumnMode::None {
122        return leaf_columns.to_vec();
123    }
124
125    let is_grouped = |c: &&ColumnDef<TData>| grouping.iter().any(|id| id.as_ref() == c.id.as_ref());
126
127    let non_grouped: Vec<&ColumnDef<TData>> = leaf_columns
128        .iter()
129        .copied()
130        .filter(|c| !is_grouped(c))
131        .collect();
132    if mode == GroupedColumnMode::Remove {
133        return non_grouped;
134    }
135
136    let by_id: HashMap<&str, &ColumnDef<TData>> = leaf_columns
137        .iter()
138        .copied()
139        .map(|c| (c.id.as_ref(), c))
140        .collect();
141
142    let mut grouped_cols: Vec<&ColumnDef<TData>> = Vec::new();
143    for id in grouping {
144        if let Some(col) = by_id.get(id.as_ref()).copied() {
145            grouped_cols.push(col);
146        }
147    }
148
149    grouped_cols.extend(non_grouped);
150    grouped_cols
151}
152
153pub fn column_can_group<TData>(options: TableOptions, column: &ColumnDef<TData>) -> bool {
154    if !(options.enable_grouping && column.enable_grouping) {
155        return false;
156    }
157    column.facet_key_fn.is_some() || column.facet_str_fn.is_some()
158}
159
160pub type GroupedRowIndex = usize;
161
162#[derive(Debug, Clone)]
163pub enum GroupedRowKind {
164    Group {
165        grouping_column: ColumnId,
166        grouping_value: u64,
167        first_leaf_row_key: RowKey,
168        leaf_row_count: usize,
169    },
170    Leaf {
171        row_key: RowKey,
172    },
173}
174
175#[derive(Debug, Clone)]
176pub struct GroupedRow {
177    pub id: RowId,
178    pub key: RowKey,
179    pub kind: GroupedRowKind,
180    pub depth: usize,
181    pub parent: Option<GroupedRowIndex>,
182    pub parent_key: Option<RowKey>,
183    pub sub_rows: Vec<GroupedRowIndex>,
184}
185
186#[derive(Debug, Clone, Default)]
187pub struct GroupedRowModel {
188    root_rows: Vec<GroupedRowIndex>,
189    flat_rows: Vec<GroupedRowIndex>,
190    rows_by_key: HashMap<RowKey, GroupedRowIndex>,
191    rows_by_id: HashMap<RowId, GroupedRowIndex>,
192    arena: Vec<GroupedRow>,
193}
194
195impl GroupedRowModel {
196    pub fn root_rows(&self) -> &[GroupedRowIndex] {
197        &self.root_rows
198    }
199
200    pub fn flat_rows(&self) -> &[GroupedRowIndex] {
201        &self.flat_rows
202    }
203
204    pub fn row(&self, index: GroupedRowIndex) -> Option<&GroupedRow> {
205        self.arena.get(index)
206    }
207
208    pub fn row_by_key(&self, key: RowKey) -> Option<GroupedRowIndex> {
209        self.rows_by_key.get(&key).copied()
210    }
211
212    pub fn row_by_id(&self, id: &str) -> Option<GroupedRowIndex> {
213        self.rows_by_id.get(id).copied()
214    }
215
216    pub fn is_leaf(&self, index: GroupedRowIndex) -> bool {
217        self.row(index)
218            .is_some_and(|r| matches!(r.kind, GroupedRowKind::Leaf { .. }))
219    }
220}
221
222fn fnv1a64_bytes(bytes: &[u8]) -> u64 {
223    const OFFSET: u64 = 0xcbf29ce484222325;
224    const PRIME: u64 = 0x00000100000001B3;
225
226    let mut h = OFFSET;
227    for &b in bytes {
228        h ^= b as u64;
229        h = h.wrapping_mul(PRIME);
230    }
231    h
232}
233
234fn mix64(mut h: u64, v: u64) -> u64 {
235    const PRIME: u64 = 0x00000100000001B3;
236    h ^= v;
237    h = h.wrapping_mul(PRIME);
238    h ^= h >> 32;
239    h
240}
241
242fn group_row_key_hash(parent: Option<RowKey>, column: &ColumnId, value: u64, attempt: u64) -> u64 {
243    let mut h = fnv1a64_bytes(column.as_ref().as_bytes());
244    h = mix64(h, parent.map(|k| k.0).unwrap_or(0));
245    h = mix64(h, value);
246    mix64(h, attempt)
247}
248
249fn alloc_group_row_key(
250    used: &mut HashSet<RowKey>,
251    parent: Option<RowKey>,
252    column: &ColumnId,
253    value: u64,
254) -> RowKey {
255    for attempt in 0u64.. {
256        let candidate = RowKey(group_row_key_hash(parent, column, value, attempt));
257        if used.insert(candidate) {
258            return candidate;
259        }
260    }
261    unreachable!("u64 key space exhausted")
262}
263
264#[derive(Debug, Clone)]
265struct GroupingValue {
266    key: u64,
267    label: Arc<str>,
268}
269
270fn grouping_value_for_row<TData>(column: &ColumnDef<TData>, row: &TData) -> GroupingValue {
271    if let Some(f) = column.facet_key_fn.as_ref() {
272        let key = f(row);
273        return GroupingValue {
274            key,
275            label: Arc::from(key.to_string()),
276        };
277    }
278    if let Some(f) = column.facet_str_fn.as_ref() {
279        let label = f(row);
280        return GroupingValue {
281            key: fnv1a64_bytes(label.as_bytes()),
282            label: Arc::from(label),
283        };
284    }
285    GroupingValue {
286        key: 0,
287        label: Arc::from(""),
288    }
289}
290
291pub fn grouped_row_model_from_leaf<'a, TData>(row_model: &RowModel<'a, TData>) -> GroupedRowModel {
292    let mut out = GroupedRowModel::default();
293    if row_model.root_rows().is_empty() {
294        return out;
295    }
296
297    out.arena.reserve(row_model.root_rows().len());
298    for &root in row_model.root_rows() {
299        let Some(row) = row_model.row(root) else {
300            continue;
301        };
302        let index = out.arena.len();
303        let id = row.id.clone();
304        out.arena.push(GroupedRow {
305            id: id.clone(),
306            key: row.key,
307            kind: GroupedRowKind::Leaf { row_key: row.key },
308            depth: 0,
309            parent: None,
310            parent_key: None,
311            sub_rows: Vec::new(),
312        });
313        out.root_rows.push(index);
314        out.flat_rows.push(index);
315        out.rows_by_key.insert(row.key, index);
316        out.rows_by_id.insert(id, index);
317    }
318
319    out
320}
321
322pub fn group_row_model<'a, TData>(
323    row_model: &RowModel<'a, TData>,
324    columns: &[ColumnDef<TData>],
325    grouping: &[ColumnId],
326) -> GroupedRowModel {
327    if grouping.is_empty() {
328        return grouped_row_model_from_leaf(row_model);
329    }
330
331    let mut columns_by_id: HashMap<&str, &ColumnDef<TData>> = HashMap::new();
332    for col in columns {
333        columns_by_id.insert(col.id.as_ref(), col);
334    }
335
336    let mut used_keys: HashSet<RowKey> = HashSet::new();
337    for &root in row_model.root_rows() {
338        if let Some(r) = row_model.row(root) {
339            used_keys.insert(r.key);
340        }
341    }
342
343    let mut out = GroupedRowModel::default();
344
345    fn build_groups<'a, TData>(
346        row_model: &RowModel<'a, TData>,
347        columns_by_id: &HashMap<&str, &ColumnDef<TData>>,
348        grouping: &[ColumnId],
349        used_keys: &mut HashSet<RowKey>,
350        parent: Option<GroupedRowIndex>,
351        parent_key: Option<RowKey>,
352        parent_id: Option<RowId>,
353        depth: usize,
354        row_keys: &[RowKey],
355        out: &mut GroupedRowModel,
356    ) -> Vec<GroupedRowIndex> {
357        let Some(grouping_column_id) = grouping.first() else {
358            return Vec::new();
359        };
360        let Some(column) = columns_by_id.get(grouping_column_id.as_ref()).copied() else {
361            return Vec::new();
362        };
363
364        let mut buckets: Vec<(GroupingValue, Vec<RowKey>)> = Vec::new();
365        let mut bucket_index_by_label: HashMap<Arc<str>, usize> = HashMap::new();
366
367        for &row_key in row_keys {
368            let Some(i) = row_model.row_by_key(row_key) else {
369                continue;
370            };
371            let Some(row) = row_model.row(i) else {
372                continue;
373            };
374            let value = grouping_value_for_row(column, row.original);
375
376            let bucket_idx = match bucket_index_by_label.get(&value.label).copied() {
377                Some(i) => i,
378                None => {
379                    let i = buckets.len();
380                    buckets.push((value.clone(), Vec::new()));
381                    bucket_index_by_label.insert(value.label.clone(), i);
382                    i
383                }
384            };
385            buckets[bucket_idx].1.push(row_key);
386        }
387
388        let mut out_children: Vec<GroupedRowIndex> = Vec::with_capacity(buckets.len());
389
390        for (value, rows) in buckets {
391            let first_leaf_row_key = rows.first().copied().unwrap_or(RowKey(0));
392            let leaf_row_count = rows.len();
393            let group_key = alloc_group_row_key(used_keys, parent_key, &column.id, value.key);
394            let segment = format!("{}:{}", column.id.as_ref(), value.label.as_ref());
395            let id = match parent_id.as_ref() {
396                None => RowId::new(segment),
397                Some(parent_id) => RowId::new(format!("{}>{segment}", parent_id.as_str())),
398            };
399
400            let index = out.arena.len();
401            out.arena.push(GroupedRow {
402                id: id.clone(),
403                key: group_key,
404                kind: GroupedRowKind::Group {
405                    grouping_column: column.id.clone(),
406                    grouping_value: value.key,
407                    first_leaf_row_key,
408                    leaf_row_count,
409                },
410                depth,
411                parent,
412                parent_key,
413                sub_rows: Vec::new(),
414            });
415            out.rows_by_key.insert(group_key, index);
416            out.rows_by_id.insert(id.clone(), index);
417
418            let sub_rows: Vec<GroupedRowIndex> = if grouping.len() > 1 {
419                build_groups(
420                    row_model,
421                    columns_by_id,
422                    &grouping[1..],
423                    used_keys,
424                    Some(index),
425                    Some(group_key),
426                    Some(id),
427                    depth + 1,
428                    &rows,
429                    out,
430                )
431            } else {
432                let mut leaf_nodes: Vec<GroupedRowIndex> = Vec::with_capacity(rows.len());
433                for &leaf_key in &rows {
434                    let leaf_id = row_model
435                        .row_by_key(leaf_key)
436                        .and_then(|i| row_model.row(i).map(|r| r.id.clone()))
437                        .unwrap_or_else(|| RowId::new(leaf_key.0.to_string()));
438                    let leaf_index = out.arena.len();
439                    out.arena.push(GroupedRow {
440                        id: leaf_id.clone(),
441                        key: leaf_key,
442                        kind: GroupedRowKind::Leaf { row_key: leaf_key },
443                        depth: depth + 1,
444                        parent: Some(index),
445                        parent_key: Some(group_key),
446                        sub_rows: Vec::new(),
447                    });
448                    out.rows_by_key.insert(leaf_key, leaf_index);
449                    out.rows_by_id.insert(leaf_id, leaf_index);
450                    leaf_nodes.push(leaf_index);
451                    // TanStack-compatible: leaf rows are included in flat rows during recursion,
452                    // and may appear again when their parent group appends its `subRows`.
453                    out.flat_rows.push(leaf_index);
454                }
455                leaf_nodes
456            };
457
458            // TanStack-compatible: group rows append their sub rows (not themselves) to the flat
459            // row list.
460            out.flat_rows.extend(sub_rows.iter().copied());
461            if let Some(node) = out.arena.get_mut(index) {
462                node.sub_rows = sub_rows;
463            }
464            out_children.push(index);
465        }
466
467        out_children
468    }
469
470    let mut root_keys: Vec<RowKey> = Vec::with_capacity(row_model.root_rows().len());
471    for &root in row_model.root_rows() {
472        if let Some(r) = row_model.row(root) {
473            root_keys.push(r.key);
474        }
475    }
476
477    out.root_rows = build_groups(
478        row_model,
479        &columns_by_id,
480        grouping,
481        &mut used_keys,
482        None,
483        None,
484        None,
485        0,
486        &root_keys,
487        &mut out,
488    );
489
490    // TanStack-compatible: root grouped rows are appended to the flat row list at the end.
491    out.flat_rows.extend(out.root_rows.iter().copied());
492
493    out
494}
495
496#[cfg(test)]
497mod tests {
498    use super::*;
499    use crate::table::{Table, create_column_helper};
500    use std::sync::Arc;
501
502    #[derive(Debug, Clone)]
503    struct Item {
504        role: Arc<str>,
505        score: u64,
506    }
507
508    #[test]
509    fn grouping_state_toggle_adds_and_removes() {
510        let mut g: GroupingState = Vec::new();
511        let a: ColumnId = "a".into();
512        let b: ColumnId = "b".into();
513
514        toggle_column_grouping(&mut g, &a);
515        assert_eq!(g.iter().map(|c| c.as_ref()).collect::<Vec<_>>(), vec!["a"]);
516        toggle_column_grouping(&mut g, &b);
517        assert_eq!(
518            g.iter().map(|c| c.as_ref()).collect::<Vec<_>>(),
519            vec!["a", "b"]
520        );
521        toggle_column_grouping(&mut g, &a);
522        assert_eq!(g.iter().map(|c| c.as_ref()).collect::<Vec<_>>(), vec!["b"]);
523    }
524
525    #[test]
526    fn group_row_model_groups_by_facet_key_and_preserves_first_seen_order() {
527        let data = vec![
528            Item {
529                role: "Admin".into(),
530                score: 1,
531            },
532            Item {
533                role: "Member".into(),
534                score: 2,
535            },
536            Item {
537                role: "Admin".into(),
538                score: 3,
539            },
540        ];
541
542        let helper = create_column_helper::<Item>();
543        let columns = vec![
544            helper
545                .clone()
546                .accessor("role", |it| it.role.clone())
547                .facet_key_by(|it| fnv1a64_bytes(it.role.as_bytes())),
548            helper.accessor("score", |it| it.score),
549        ];
550
551        let mut state = crate::table::TableState::default();
552        state.grouping = vec![ColumnId::from("role")];
553        let table = Table::builder(&data).columns(columns).state(state).build();
554
555        let grouped = table.grouped_row_model();
556        assert_eq!(grouped.root_rows().len(), 2);
557
558        let g0 = grouped.row(grouped.root_rows()[0]).unwrap();
559        let g1 = grouped.row(grouped.root_rows()[1]).unwrap();
560
561        assert!(matches!(g0.kind, GroupedRowKind::Group { .. }));
562        assert!(matches!(g1.kind, GroupedRowKind::Group { .. }));
563
564        // First group should be "Admin" (first seen), then "Member".
565        let GroupedRowKind::Group {
566            grouping_value: v0,
567            first_leaf_row_key: first0,
568            leaf_row_count: count0,
569            ..
570        } = &g0.kind
571        else {
572            panic!("expected group row");
573        };
574        let GroupedRowKind::Group {
575            grouping_value: v1,
576            first_leaf_row_key: first1,
577            leaf_row_count: count1,
578            ..
579        } = &g1.kind
580        else {
581            panic!("expected group row");
582        };
583
584        assert_eq!(*v0, fnv1a64_bytes("Admin".as_bytes()));
585        assert_eq!(*v1, fnv1a64_bytes("Member".as_bytes()));
586        assert_eq!((*first0, *count0), (RowKey(0), 2));
587        assert_eq!((*first1, *count1), (RowKey(1), 1));
588    }
589}