fresh/model/

piece_tree_diff.rs

use std::ops::Range;
use std::sync::Arc;

use crate::model::piece_tree::{LeafData, PieceTreeNode};

/// Summary of differences between two piece tree roots.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct PieceTreeDiff {
    /// Whether the two trees represent identical piece sequences.
    pub equal: bool,
    /// Changed byte ranges in the "after" tree (exclusive end). Empty when `equal` is true.
    pub byte_ranges: Vec<Range<usize>>,
    /// Number of tree nodes visited during the diff walk.
    /// When `Arc::ptr_eq` short-circuits effectively, this should be
    /// proportional to the edited region, not the entire tree.
    pub nodes_visited: usize,
}

/// Compute a diff between two piece tree roots.
///
/// Uses structural sharing (Arc::ptr_eq) to skip identical subtrees in O(1),
/// falling back to leaf-level comparison only for subtrees that actually differ.
/// After path-copying edits, this is O(changed_path) instead of O(all_leaves).
pub fn diff_piece_trees(before: &Arc<PieceTreeNode>, after: &Arc<PieceTreeNode>) -> PieceTreeDiff {
    // Fast path: identical subtree (same Arc pointer)
    if Arc::ptr_eq(before, after) {
        return PieceTreeDiff {
            equal: true,
            byte_ranges: Vec::new(),
            nodes_visited: 1,
        };
    }

    // Collect leaves only from differing subtrees using structural walk.
    // Spans include document-absolute byte offsets.
    let mut before_spans = Vec::new();
    let mut after_spans = Vec::new();
    let mut nodes_visited: usize = 0;
    let mut before_doc_offset: usize = 0;
    let mut after_doc_offset: usize = 0;
    diff_collect_leaves(
        before,
        after,
        &mut before_spans,
        &mut after_spans,
        &mut nodes_visited,
        &mut before_doc_offset,
        &mut after_doc_offset,
    );

    let before_spans = normalize_spans(before_spans);
    let after_spans = normalize_spans(after_spans);

    // Fast-path: identical leaf sequences (same content, different tree structure).
    if span_slices_equal(&before_spans, &after_spans) {
        return PieceTreeDiff {
            equal: true,
            byte_ranges: Vec::new(),
            nodes_visited,
        };
    }

    // Longest common prefix at byte granularity.
    let prefix = common_prefix_bytes(&before_spans, &after_spans);
    // Longest common suffix without overlapping prefix.
    let suffix = common_suffix_bytes(&before_spans, &after_spans, prefix);

    let ranges = collect_diff_ranges(&before_spans, &after_spans, prefix, suffix);

    PieceTreeDiff {
        equal: false,
        byte_ranges: ranges,
        nodes_visited,
    }
}

/// Total bytes stored under a node (without calling private methods).
fn node_bytes(node: &PieceTreeNode) -> usize {
    match node {
        PieceTreeNode::Internal {
            left_bytes, right, ..
        } => left_bytes + node_bytes(right),
        PieceTreeNode::Leaf { bytes, .. } => *bytes,
    }
}

/// Parallel tree walk that uses Arc::ptr_eq to skip identical subtrees.
/// Only collects leaves from subtrees that differ between before and after.
/// Tracks running document byte offsets so collected spans have absolute positions.
fn diff_collect_leaves(
    before: &Arc<PieceTreeNode>,
    after: &Arc<PieceTreeNode>,
    before_out: &mut Vec<Span>,
    after_out: &mut Vec<Span>,
    nodes_visited: &mut usize,
    before_doc_offset: &mut usize,
    after_doc_offset: &mut usize,
) {
    *nodes_visited += 2; // counting both before and after nodes

    // Identical subtree - skip entirely, advancing document offsets
    if Arc::ptr_eq(before, after) {
        let bytes = node_bytes(before);
        *before_doc_offset += bytes;
        *after_doc_offset += bytes;
        return;
    }

    match (before.as_ref(), after.as_ref()) {
        // Both internal: recurse into children
        (
            PieceTreeNode::Internal {
                left: b_left,
                right: b_right,
                ..
            },
            PieceTreeNode::Internal {
                left: a_left,
                right: a_right,
                ..
            },
        ) => {
            diff_collect_leaves(
                b_left,
                a_left,
                before_out,
                after_out,
                nodes_visited,
                before_doc_offset,
                after_doc_offset,
            );
            diff_collect_leaves(
                b_right,
                a_right,
                before_out,
                after_out,
                nodes_visited,
                before_doc_offset,
                after_doc_offset,
            );
        }
        // Structure mismatch - fall back to full leaf collection for both subtrees
        _ => {
            collect_leaves_with_offsets(before, before_out, nodes_visited, before_doc_offset);
            collect_leaves_with_offsets(after, after_out, nodes_visited, after_doc_offset);
        }
    }
}

fn collect_leaves_with_offsets(
    node: &Arc<PieceTreeNode>,
    out: &mut Vec<Span>,
    nodes_visited: &mut usize,
    doc_offset: &mut usize,
) {
    *nodes_visited += 1;
    match node.as_ref() {
        PieceTreeNode::Internal { left, right, .. } => {
            collect_leaves_with_offsets(left, out, nodes_visited, doc_offset);
            collect_leaves_with_offsets(right, out, nodes_visited, doc_offset);
        }
        PieceTreeNode::Leaf {
            location,
            offset,
            bytes,
            line_feed_cnt,
        } => {
            let leaf = LeafData::new(*location, *offset, *bytes, *line_feed_cnt);
            out.push(Span {
                leaf,
                doc_offset: *doc_offset,
            });
            *doc_offset += bytes;
        }
    }
}

#[derive(Clone)]
struct Span {
    leaf: LeafData,
    doc_offset: usize,
}

fn spans_equal(a: &Span, b: &Span) -> bool {
    a.leaf.location == b.leaf.location
        && a.leaf.offset == b.leaf.offset
        && a.leaf.bytes == b.leaf.bytes
}

fn span_slices_equal(a: &[Span], b: &[Span]) -> bool {
    a.len() == b.len() && a.iter().zip(b.iter()).all(|(x, y)| spans_equal(x, y))
}

fn normalize_spans(spans: Vec<Span>) -> Vec<Span> {
    if spans.is_empty() {
        return spans;
    }

    let mut it = spans.into_iter();
    let mut current = it.next().unwrap();
    let mut normalized = Vec::new();

    for span in it {
        let contiguous = current.leaf.location == span.leaf.location
            && current.leaf.offset + current.leaf.bytes == span.leaf.offset
            && current.doc_offset + current.leaf.bytes == span.doc_offset;
        if contiguous {
            current.leaf.bytes += span.leaf.bytes;
            current.leaf.line_feed_cnt = match (current.leaf.line_feed_cnt, span.leaf.line_feed_cnt)
            {
                (Some(a), Some(b)) => Some(a + b),
                _ => None,
            };
        } else {
            normalized.push(current);
            current = span;
        }
    }

    normalized.push(current);
    normalized
}

fn common_prefix_bytes(before: &[Span], after: &[Span]) -> usize {
    let mut b_idx = 0;
    let mut a_idx = 0;
    let mut b_off = 0;
    let mut a_off = 0;
    let mut consumed = 0;

    while b_idx < before.len() && a_idx < after.len() {
        let b_span = &before[b_idx];
        let a_span = &after[a_idx];
        let b = &b_span.leaf;
        let a = &a_span.leaf;

        let b_pos = b.offset + b_off;
        let a_pos = a.offset + a_off;

        // Must also ensure they are at the same document relative position
        // if they were separated by gaps.
        if b.location == a.location
            && b_pos == a_pos
            && (b_span.doc_offset + b_off) == (a_span.doc_offset + a_off)
        {
            let b_rem = b.bytes - b_off;
            let a_rem = a.bytes - a_off;
            let take = b_rem.min(a_rem);

            consumed += take;
            b_off += take;
            a_off += take;

            if b_off == b.bytes {
                b_idx += 1;
                b_off = 0;
            }
            if a_off == a.bytes {
                a_idx += 1;
                a_off = 0;
            }
        } else {
            break;
        }
    }

    consumed
}

fn common_suffix_bytes(before: &[Span], after: &[Span], prefix_bytes: usize) -> usize {
    let total_before = before
        .last()
        .map(|s| s.doc_offset + s.leaf.bytes)
        .unwrap_or(0);
    let total_after = after
        .last()
        .map(|s| s.doc_offset + s.leaf.bytes)
        .unwrap_or(0);

    let mut b_idx: isize = before.len() as isize - 1;
    let mut a_idx: isize = after.len() as isize - 1;
    let mut b_off = 0;
    let mut a_off = 0;
    let mut consumed = 0;

    while b_idx >= 0
        && a_idx >= 0
        && (total_before - consumed) > prefix_bytes
        && (total_after - consumed) > prefix_bytes
    {
        let b_span = &before[b_idx as usize];
        let a_span = &after[a_idx as usize];
        let b_leaf = &b_span.leaf;
        let a_leaf = &a_span.leaf;

        let b_pos = b_leaf.offset + b_leaf.bytes - b_off;
        let a_pos = a_leaf.offset + a_leaf.bytes - a_off;

        // Compare by buffer identity only (location + offset). Suffix bytes are
        // at different doc_offsets in before vs after when insertions/deletions
        // change the total size, but they still contain the same data.
        if b_leaf.location == a_leaf.location && b_pos == a_pos {
            let b_rem = b_leaf.bytes - b_off;
            let a_rem = a_leaf.bytes - a_off;
            let take = b_rem.min(a_rem);

            consumed += take;
            b_off += take;
            a_off += take;

            if b_off == b_leaf.bytes {
                b_idx -= 1;
                b_off = 0;
            }
            if a_off == a_leaf.bytes {
                a_idx -= 1;
                a_off = 0;
            }
        } else {
            break;
        }
    }

    consumed.min(total_after.saturating_sub(prefix_bytes))
}

fn collect_diff_ranges(
    before: &[Span],
    after: &[Span],
    prefix: usize,
    suffix: usize,
) -> Vec<Range<usize>> {
    let mut ranges = Vec::new();
    let mut b_idx = 0;
    let mut a_idx = 0;
    let mut b_off = 0;
    let mut a_off = 0;
    let mut matched_prefix = 0;

    // Skip matching prefix
    while matched_prefix < prefix && b_idx < before.len() && a_idx < after.len() {
        let b = &before[b_idx].leaf;
        let a = &after[a_idx].leaf;
        let b_rem = b.bytes - b_off;
        let a_rem = a.bytes - a_off;
        let take = b_rem.min(a_rem).min(prefix - matched_prefix);
        matched_prefix += take;
        b_off += take;
        a_off += take;
        if b_off == b.bytes {
            b_idx += 1;
            b_off = 0;
        }
        if a_off == a.bytes {
            a_idx += 1;
            a_off = 0;
        }
    }

    // compare_limit in document-absolute space: end of collected spans minus suffix
    let doc_end = after
        .last()
        .map(|s| s.doc_offset + s.leaf.bytes)
        .unwrap_or(0);
    let compare_limit = doc_end.saturating_sub(suffix);

    // prefix_start in document-absolute space
    let doc_start = after.first().map(|s| s.doc_offset).unwrap_or(0);

    let mut current_start: Option<usize> = None;
    let mut current_end: usize = 0;

    while a_idx < after.len() {
        let a = &after[a_idx];
        let pos = a.doc_offset + a_off;
        if pos >= compare_limit {
            break;
        }

        // If we have a current range but there's a gap in document offset,
        // it means there was an identical subtree that was skipped.
        // We must break the range here.
        if let Some(start) = current_start {
            if pos > current_end {
                ranges.push(start..current_end);
                current_start = None;
            }
        }

        let matches = if b_idx < before.len() {
            let b = &before[b_idx].leaf;
            let b_pos = b.offset + b_off;
            let a_pos = a.leaf.offset + a_off;
            // Compare by buffer identity only. Insertions/deletions shift
            // doc_offsets, but matching buffer location + offset means same data.
            b.location == a.leaf.location && b_pos == a_pos
        } else {
            false
        };

        if matches {
            if let Some(start) = current_start.take() {
                ranges.push(start..current_end);
            }

            let b = &before[b_idx].leaf;
            let b_rem = b.bytes - b_off;
            let a_rem = a.leaf.bytes - a_off;
            let take = b_rem.min(a_rem).min(compare_limit.saturating_sub(pos));

            b_off += take;
            a_off += take;

            if b_off == b.bytes {
                b_idx += 1;
                b_off = 0;
            }
            if a_off == a.leaf.bytes {
                a_idx += 1;
                a_off = 0;
            }
        } else {
            if current_start.is_none() {
                current_start = Some(pos);
            }
            let take = (a.leaf.bytes - a_off).min(compare_limit.saturating_sub(pos));
            current_end = pos + take;
            a_off += take;
            if a_off == a.leaf.bytes {
                a_idx += 1;
                a_off = 0;
            }
        }
    }

    if let Some(start) = current_start {
        ranges.push(start..current_end);
    }

    // Any trailing unmatched "after" spans up to suffix boundary
    while a_idx < after.len() {
        let start = after[a_idx].doc_offset + a_off;
        if start >= compare_limit {
            break;
        }
        let end = (after[a_idx].doc_offset + after[a_idx].leaf.bytes).min(compare_limit);
        ranges.push(start..end);
        a_idx += 1;
        a_off = 0;
    }

    if ranges.is_empty() {
        // Anchor range: either there's content between prefix and suffix, or
        // the trees differ only in size (deletion) — report the anchor point.
        ranges.push((doc_start + prefix)..compare_limit);
    }

    ranges
}

#[cfg(test)]
#[allow(clippy::single_range_in_vec_init)]
mod tests {
    use super::*;
    use crate::model::piece_tree::BufferLocation;

    fn sum_bytes(leaves: &[LeafData]) -> usize {
        leaves.iter().map(|leaf| leaf.bytes).sum()
    }

    fn leaf(loc: BufferLocation, offset: usize, bytes: usize, lfs: Option<usize>) -> LeafData {
        LeafData::new(loc, offset, bytes, lfs)
    }

    // Minimal balanced builder for tests.
    fn build(leaves: &[LeafData]) -> Arc<PieceTreeNode> {
        if leaves.is_empty() {
            return Arc::new(PieceTreeNode::Leaf {
                location: BufferLocation::Stored(0),
                offset: 0,
                bytes: 0,
                line_feed_cnt: Some(0),
            });
        }
        if leaves.len() == 1 {
            let l = leaves[0];
            return Arc::new(PieceTreeNode::Leaf {
                location: l.location,
                offset: l.offset,
                bytes: l.bytes,
                line_feed_cnt: l.line_feed_cnt,
            });
        }

        let mid = leaves.len() / 2;
        let left = build(&leaves[..mid]);
        let right = build(&leaves[mid..]);

        Arc::new(PieceTreeNode::Internal {
            left_bytes: sum_bytes(&leaves[..mid]),
            lf_left: leaves[..mid]
                .iter()
                .map(|l| l.line_feed_cnt)
                .try_fold(0usize, |acc, v| v.map(|b| acc + b)),
            left,
            right,
        })
    }

    #[test]
    fn detects_identical_trees() {
        let leaves = vec![leaf(BufferLocation::Stored(0), 0, 10, Some(0))];
        let before = build(&leaves);
        let after = build(&leaves);

        let diff = diff_piece_trees(&before, &after);
        assert!(diff.equal);
        assert!(diff.byte_ranges.is_empty());
    }

    #[test]
    fn detects_single_line_change() {
        let before = build(&[leaf(BufferLocation::Stored(0), 0, 5, Some(0))]);
        let after = build(&[leaf(BufferLocation::Added(1), 0, 5, Some(0))]);

        let diff = diff_piece_trees(&before, &after);
        assert!(!diff.equal);
        assert_eq!(diff.byte_ranges, vec![0..5]);
    }

    #[test]
    fn tracks_newlines_in_changed_span() {
        let before = build(&[leaf(BufferLocation::Stored(0), 0, 6, Some(0))]);
        let after = build(&[leaf(BufferLocation::Added(1), 0, 6, Some(1))]);

        let diff = diff_piece_trees(&before, &after);
        assert!(!diff.equal);
        assert_eq!(diff.byte_ranges, vec![0..6]);
    }

    #[test]
    fn handles_deletion_by_marking_anchor() {
        let before = build(&[
            leaf(BufferLocation::Stored(0), 0, 6, Some(1)),
            leaf(BufferLocation::Stored(0), 6, 4, Some(0)),
        ]);
        let after = build(&[leaf(BufferLocation::Stored(0), 0, 6, Some(1))]);

        let diff = diff_piece_trees(&before, &after);
        assert!(!diff.equal);
        assert_eq!(diff.byte_ranges, vec![6..6]);
    }

    /// Uses real PieceTree::insert (path-copy) near EOF.
    /// The diff must produce document-absolute offsets.
    #[test]
    fn diff_after_path_copy_insert_at_eof() {
        use crate::model::piece_tree::{PieceTree, StringBuffer};

        let chunk_size = 1000;
        let total = 10_000;
        let content: Vec<u8> = (0..total)
            .map(|i| if i % 100 == 99 { b'\n' } else { b'A' })
            .collect();
        let buf = StringBuffer::new_loaded(0, content, false);

        let mut saved_tree = PieceTree::new(BufferLocation::Stored(0), 0, total, None);
        saved_tree.split_leaves_to_chunk_size(chunk_size);
        let lf_updates: Vec<(usize, usize)> = (0..10).map(|i| (i, 10)).collect();
        saved_tree.update_leaf_line_feeds(&lf_updates);
        let saved_root = saved_tree.root();

        let mut after_tree = saved_tree;
        let insert_offset = total - 100;
        let insert_buf = StringBuffer::new_loaded(1, b"HELLO".to_vec(), false);
        after_tree.insert(
            insert_offset,
            BufferLocation::Added(1),
            0,
            5,
            Some(0),
            &[buf, insert_buf],
        );
        let after_root = after_tree.root();

        let diff = diff_piece_trees(&saved_root, &after_root);
        assert!(!diff.equal);

        assert!(
            diff.byte_ranges[0].start >= total - 200,
            "byte_ranges should be document-absolute (near EOF): got {:?}, expected near {}",
            diff.byte_ranges,
            insert_offset,
        );
    }

    /// After rebalance, Arc sharing is destroyed. The diff must still produce
    /// the same byte_ranges as the path-copy version.
    #[test]
    fn diff_after_rebalance_matches_path_copy_diff() {
        use crate::model::piece_tree::{PieceTree, StringBuffer};

        let chunk_size = 1000;
        let total = 10_000;
        let content: Vec<u8> = (0..total)
            .map(|i| if i % 100 == 99 { b'\n' } else { b'A' })
            .collect();
        let buf = StringBuffer::new_loaded(0, content, false);

        let mut saved_tree = PieceTree::new(BufferLocation::Stored(0), 0, total, None);
        saved_tree.split_leaves_to_chunk_size(chunk_size);
        let lf_updates: Vec<(usize, usize)> = (0..10).map(|i| (i, 10)).collect();
        saved_tree.update_leaf_line_feeds(&lf_updates);
        let saved_root = saved_tree.root();

        let mut after_tree = saved_tree;
        let insert_buf = StringBuffer::new_loaded(1, b"HELLO".to_vec(), false);
        after_tree.insert(
            total - 100,
            BufferLocation::Added(1),
            0,
            5,
            Some(0),
            &[buf.clone(), insert_buf.clone()],
        );

        let diff_shared = diff_piece_trees(&saved_root, &after_tree.root());

        after_tree.rebalance();
        let diff_rebalanced = diff_piece_trees(&saved_root, &after_tree.root());

        assert!(!diff_shared.equal);
        assert!(!diff_rebalanced.equal);
        assert_eq!(
            diff_shared.byte_ranges, diff_rebalanced.byte_ranges,
            "byte_ranges should be identical whether or not Arc sharing exists"
        );
    }

    #[test]
    fn tolerates_split_leaves_with_same_content_prefix() {
        let before = build(&[leaf(BufferLocation::Stored(0), 0, 100, Some(1))]);
        let after = build(&[
            leaf(BufferLocation::Stored(0), 0, 50, Some(0)),
            leaf(BufferLocation::Added(1), 0, 10, Some(0)),
            leaf(BufferLocation::Stored(0), 50, 50, Some(1)),
        ]);

        let diff = diff_piece_trees(&before, &after);
        assert!(!diff.equal);
        assert_eq!(diff.byte_ranges, vec![50..60]);
    }

    #[test]
    fn diff_with_disjoint_changes() {
        let leaf1_before = leaf(BufferLocation::Stored(0), 0, 10, Some(0));
        let leaf2 = leaf(BufferLocation::Stored(0), 10, 10, Some(1));
        let leaf3_before = leaf(BufferLocation::Stored(0), 20, 10, Some(0));

        let leaf1_after = leaf(BufferLocation::Added(1), 0, 10, Some(0));
        let leaf3_after = leaf(BufferLocation::Added(1), 10, 10, Some(0));

        let leaf2_arc = Arc::new(PieceTreeNode::Leaf {
            location: leaf2.location,
            offset: leaf2.offset,
            bytes: leaf2.bytes,
            line_feed_cnt: leaf2.line_feed_cnt,
        });

        let before = Arc::new(PieceTreeNode::Internal {
            left_bytes: 10,
            lf_left: Some(0),
            left: Arc::new(PieceTreeNode::Leaf {
                location: leaf1_before.location,
                offset: leaf1_before.offset,
                bytes: leaf1_before.bytes,
                line_feed_cnt: leaf1_before.line_feed_cnt,
            }),
            right: Arc::new(PieceTreeNode::Internal {
                left_bytes: 10,
                lf_left: Some(1),
                left: Arc::clone(&leaf2_arc),
                right: Arc::new(PieceTreeNode::Leaf {
                    location: leaf3_before.location,
                    offset: leaf3_before.offset,
                    bytes: leaf3_before.bytes,
                    line_feed_cnt: leaf3_before.line_feed_cnt,
                }),
            }),
        });

        let after = Arc::new(PieceTreeNode::Internal {
            left_bytes: 10,
            lf_left: Some(0),
            left: Arc::new(PieceTreeNode::Leaf {
                location: leaf1_after.location,
                offset: leaf1_after.offset,
                bytes: leaf1_after.bytes,
                line_feed_cnt: leaf1_after.line_feed_cnt,
            }),
            right: Arc::new(PieceTreeNode::Internal {
                left_bytes: 10,
                lf_left: Some(1),
                left: Arc::clone(&leaf2_arc), // Shared!
                right: Arc::new(PieceTreeNode::Leaf {
                    location: leaf3_after.location,
                    offset: leaf3_after.offset,
                    bytes: leaf3_after.bytes,
                    line_feed_cnt: leaf3_after.line_feed_cnt,
                }),
            }),
        });

        let diff = diff_piece_trees(&before, &after);

        assert_eq!(diff.byte_ranges, vec![0..10, 20..30]);
    }
}