oxipdf-ir 0.1.0

//! IR tree partitioning for chapter-level streaming (§12.1).
//!
//! # Data Flow for Streaming Layout
//!
//! The streaming architecture processes a document in independent chunks.
//! The `render_chunked()` function in the facade crate uses this module
//! to partition the tree, process each chunk independently, and combine
//! the results into a single PDF.
//!
//! ```text
//! StyledTree
//!   │
//!   ├─ partition_tree() ──► Vec<ChunkBoundary>
//!   │     Identifies top-level section boundaries where the tree can be
//!   │     split without cross-chunk layout dependencies.
//!   │
//!   ├─ For each chunk:
//!   │     1. Extract sub-tree (root + section children)
//!   │     2. Shape text (oxipdf-shaping)
//!   │     3. Compute layout (oxipdf-layout)
//!   │     4. Fragment into pages (oxipdf-fragment)
//!   │     5. Record cross-references into DocumentSpine
//!   │     6. Drop layout tree — memory freed
//!   │
//!   └─ DocumentSpine
//!         Accumulates: page count, cross-ref map, PDF object index.
//!         After all chunks: resolve cross-references globally,
//!         then emit final PDF in a single sequential pass.
//! ```
//!
//! # Integration
//!
//! The chunked pipeline is wired into the render pipeline via
//! `oxipdf::render::render_chunked()` and `render_chunked_shaped()`.
//! These functions call `partition_tree()`, iterate over chunk boundaries,
//! extract subtrees, layout/paginate each independently, accumulate
//! cross-references in the `DocumentSpine`, then combine all chunk pages
//! into the final PDF output.

use std::collections::BTreeMap;

use crate::node::NodeId;
use crate::tree::StyledTree;

/// A boundary where the IR tree can be split for independent layout.
///
/// Each boundary identifies a contiguous range of top-level children
/// that can be laid out, fragmented, and emitted independently. Cross-chunk
/// references (e.g., TOC entries pointing to later sections) are resolved
/// after all chunks are processed.
#[derive(Debug, Clone, PartialEq, Eq)]
pub enum ChunkBoundary {
    /// A single top-level section that can be processed independently.
    /// Contains the `NodeId` of the section root (a direct child of the
    /// tree root) and the range of its subtree.
    Section {
        /// The root node of this chunk (a direct child of the tree root).
        section_root: NodeId,
    },

    /// A group of consecutive top-level children that must be processed
    /// together because they share layout dependencies (e.g., a heading
    /// with `keep_with_next` referencing the next sibling).
    Group {
        /// The first child in this group.
        first: NodeId,
        /// The last child in this group (inclusive).
        last: NodeId,
    },
}

impl ChunkBoundary {
    /// Return the top-level child `NodeId`s that belong to this chunk.
    ///
    /// The caller passes the full list of root's children; this method
    /// extracts the slice covered by this boundary.
    #[must_use]
    pub fn child_ids<'a>(&self, root_children: &'a [NodeId]) -> &'a [NodeId] {
        match self {
            ChunkBoundary::Section { section_root } => {
                let idx = root_children
                    .iter()
                    .position(|c| c == section_root)
                    .expect("section_root must be a child of the tree root");
                &root_children[idx..=idx]
            }
            ChunkBoundary::Group { first, last } => {
                let first_idx = root_children
                    .iter()
                    .position(|c| c == first)
                    .expect("group first must be a child of the tree root");
                let last_idx = root_children
                    .iter()
                    .position(|c| c == last)
                    .expect("group last must be a child of the tree root");
                &root_children[first_idx..=last_idx]
            }
        }
    }
}

/// The document spine: cross-chunk metadata accumulated during streaming.
///
/// Holds only the minimum data needed for global reference resolution and
/// final PDF assembly. Layout trees are released after each chunk is
/// processed.
#[derive(Debug, Clone, Default)]
pub struct DocumentSpine {
    /// Cross-reference map: element ID → (chunk index, page number within chunk).
    pub cross_refs: BTreeMap<String, CrossRefEntry>,
    /// Cumulative page count after each chunk.
    pub chunk_page_counts: Vec<u32>,
    /// PDF object byte offsets for each chunk's emitted objects, used to
    /// assemble the final cross-reference table.
    pub object_index: Vec<ObjectIndexEntry>,
}

impl DocumentSpine {
    /// Record the number of pages produced by the next chunk.
    ///
    /// Called once per chunk, in order. The chunk index is inferred
    /// from the current length of `chunk_page_counts`.
    pub fn add_chunk_pages(&mut self, count: u32) {
        self.chunk_page_counts.push(count);
    }

    /// Total pages across all chunks processed so far.
    #[must_use]
    pub fn total_pages(&self) -> u32 {
        self.chunk_page_counts.iter().sum()
    }

    /// Record a cross-reference from a chunk: element ID → (chunk, local page).
    ///
    /// `chunk_index` is the 0-based index of the chunk that contains the
    /// element. `local_page` is the 0-based page within that chunk.
    pub fn add_cross_ref(&mut self, element_id: String, chunk_index: usize, local_page: u32) {
        self.cross_refs.insert(
            element_id,
            CrossRefEntry {
                chunk_index,
                local_page,
            },
        );
    }

    /// Resolve an element ID to its absolute 1-indexed page number.
    ///
    /// Returns `None` if the element ID was not recorded by any chunk.
    #[must_use]
    pub fn resolve_element(&self, element_id: &str) -> Option<u32> {
        self.cross_refs
            .get(element_id)
            .map(|entry| entry.absolute_page(&self.chunk_page_counts) + 1)
    }
}

/// A cross-reference entry mapping an element ID to its location.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct CrossRefEntry {
    /// Which chunk (0-indexed) contains this element.
    pub chunk_index: usize,
    /// Page number within the chunk (0-indexed).
    pub local_page: u32,
}

impl CrossRefEntry {
    /// Compute the absolute page number given cumulative chunk page counts.
    #[must_use]
    pub fn absolute_page(&self, chunk_page_counts: &[u32]) -> u32 {
        let pages_before: u32 = chunk_page_counts.iter().take(self.chunk_index).sum();
        pages_before + self.local_page
    }
}

/// An entry in the PDF object index for final assembly.
#[derive(Debug, Clone)]
pub struct ObjectIndexEntry {
    /// Chunk that produced this object.
    pub chunk_index: usize,
    /// Byte offset of the object in the chunk's output buffer.
    pub byte_offset: u64,
    /// Length of the object in bytes.
    pub byte_length: u64,
}

/// Partition a `StyledTree` into independently layoutable chunks.
///
/// Identifies top-level section boundaries by examining direct children of
/// the tree root. Each child that is a `Container` node becomes its own
/// chunk boundary, unless it has a `keep_with_next` constraint that ties
/// it to the following sibling.
///
/// Returns an empty `Vec` if the tree has a single leaf root (no children
/// to partition).
#[must_use]
pub fn partition_tree(tree: &StyledTree) -> Vec<ChunkBoundary> {
    let root = tree.root();
    let children = tree.children(root);

    if children.is_empty() {
        return Vec::new();
    }

    let mut boundaries = Vec::new();
    let mut group_start: Option<NodeId> = None;

    for (i, &child_id) in children.iter().enumerate() {
        let node = tree.node(child_id);
        let keep_with_next = node.style.fragmentation.keep_with_next;

        if keep_with_next && i + 1 < children.len() {
            // This node must stay with the next — start or extend a group.
            if group_start.is_none() {
                group_start = Some(child_id);
            }
        } else if let Some(start) = group_start.take() {
            // End of a keep-with-next chain — emit as group.
            boundaries.push(ChunkBoundary::Group {
                first: start,
                last: child_id,
            });
        } else {
            // Standalone section.
            boundaries.push(ChunkBoundary::Section {
                section_root: child_id,
            });
        }
    }

    // If a group was open at the end (shouldn't happen with valid trees
    // since keep_with_next on the last child has no effect), close it.
    if let Some(start) = group_start {
        let last = *children.last().expect("children is non-empty");
        boundaries.push(ChunkBoundary::Group { first: start, last });
    }

    boundaries
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::node::{ContentVariant, TextContent};
    use crate::style::ResolvedStyle;
    use crate::tree::StyledTreeBuilder;
    use crate::version::IrVersion;

    #[test]
    fn single_leaf_produces_no_chunks() {
        let mut b = StyledTreeBuilder::new(IrVersion::new(1, 0));
        b.add_node(
            ContentVariant::Container,
            ResolvedStyle::default(),
            None,
            None,
        );
        let tree = b.build().unwrap();
        assert!(partition_tree(&tree).is_empty());
    }

    #[test]
    fn three_sections_produce_three_chunks() {
        let mut b = StyledTreeBuilder::new(IrVersion::new(1, 0));
        let root = b.add_node(
            ContentVariant::Container,
            ResolvedStyle::default(),
            None,
            None,
        );
        for i in 0..3 {
            b.add_child(
                root,
                ContentVariant::Text(TextContent::new(format!("Section {i}"))),
                ResolvedStyle::default(),
                None,
                None,
            );
        }
        let tree = b.build().unwrap();
        let chunks = partition_tree(&tree);
        assert_eq!(chunks.len(), 3);
        for chunk in &chunks {
            assert!(matches!(chunk, ChunkBoundary::Section { .. }));
        }
    }

    #[test]
    fn keep_with_next_creates_group() {
        let mut b = StyledTreeBuilder::new(IrVersion::new(1, 0));
        let root = b.add_node(
            ContentVariant::Container,
            ResolvedStyle::default(),
            None,
            None,
        );
        // First child: keep_with_next = true
        let mut style1 = ResolvedStyle::default();
        style1.fragmentation.keep_with_next = true;
        b.add_child(
            root,
            ContentVariant::Text(TextContent::new("Heading")),
            style1,
            None,
            None,
        );
        // Second child: standalone
        b.add_child(
            root,
            ContentVariant::Text(TextContent::new("Body")),
            ResolvedStyle::default(),
            None,
            None,
        );
        // Third child: standalone
        b.add_child(
            root,
            ContentVariant::Text(TextContent::new("Footer")),
            ResolvedStyle::default(),
            None,
            None,
        );
        let tree = b.build().unwrap();
        let chunks = partition_tree(&tree);
        assert_eq!(chunks.len(), 2);
        assert!(matches!(chunks[0], ChunkBoundary::Group { .. }));
        assert!(matches!(chunks[1], ChunkBoundary::Section { .. }));
    }

    #[test]
    fn cross_ref_absolute_page() {
        let entry = CrossRefEntry {
            chunk_index: 2,
            local_page: 3,
        };
        let counts = vec![10, 5, 8];
        // pages before chunk 2 = 10 + 5 = 15; absolute = 15 + 3 = 18
        assert_eq!(entry.absolute_page(&counts), 18);
    }

    #[test]
    fn document_spine_default() {
        let spine = DocumentSpine::default();
        assert!(spine.cross_refs.is_empty());
        assert!(spine.chunk_page_counts.is_empty());
        assert!(spine.object_index.is_empty());
    }

    #[test]
    fn partition_produces_valid_boundaries_for_sections() {
        // Build a tree with 5 top-level sections, the second with keep_with_next.
        let mut b = StyledTreeBuilder::new(IrVersion::new(1, 0));
        let root = b.add_node(
            ContentVariant::Container,
            ResolvedStyle::default(),
            None,
            None,
        );
        for i in 0..5 {
            let mut style = ResolvedStyle::default();
            if i == 1 {
                style.fragmentation.keep_with_next = true;
            }
            b.add_child(
                root,
                ContentVariant::Text(TextContent::new(format!("Section {i}"))),
                style,
                None,
                None,
            );
        }
        let tree = b.build().unwrap();
        let boundaries = partition_tree(&tree);

        // Should produce: Section(0), Group(1,2), Section(3), Section(4)
        assert_eq!(boundaries.len(), 4);
        assert!(matches!(boundaries[0], ChunkBoundary::Section { .. }));
        assert!(matches!(boundaries[1], ChunkBoundary::Group { .. }));
        assert!(matches!(boundaries[2], ChunkBoundary::Section { .. }));
        assert!(matches!(boundaries[3], ChunkBoundary::Section { .. }));

        // Verify all child NodeIds are covered by exactly one boundary.
        let children = tree.children(tree.root());
        let mut covered = vec![false; children.len()];
        for boundary in &boundaries {
            match boundary {
                ChunkBoundary::Section { section_root } => {
                    let idx = children.iter().position(|c| c == section_root).unwrap();
                    assert!(!covered[idx], "node covered twice");
                    covered[idx] = true;
                }
                ChunkBoundary::Group { first, last } => {
                    let first_idx = children.iter().position(|c| c == first).unwrap();
                    let last_idx = children.iter().position(|c| c == last).unwrap();
                    for item in &mut covered[first_idx..=last_idx] {
                        assert!(!*item, "node covered twice");
                        *item = true;
                    }
                }
            }
        }
        assert!(covered.iter().all(|&c| c), "all children must be covered");
    }

    #[test]
    fn document_spine_constructed_without_full_tree() {
        // Demonstrates that DocumentSpine can be built incrementally
        // without holding the full tree — each chunk contributes its
        // page count and cross-refs, then the chunk's layout data can
        // be dropped.
        let mut spine = DocumentSpine::default();

        // Simulate chunk 0: 3 pages, one cross-ref.
        spine.chunk_page_counts.push(3);
        spine.cross_refs.insert(
            "intro".into(),
            CrossRefEntry {
                chunk_index: 0,
                local_page: 0,
            },
        );
        spine.object_index.push(ObjectIndexEntry {
            chunk_index: 0,
            byte_offset: 0,
            byte_length: 1024,
        });
        // At this point, chunk 0's layout tree could be dropped.

        // Simulate chunk 1: 5 pages, one cross-ref.
        spine.chunk_page_counts.push(5);
        spine.cross_refs.insert(
            "details".into(),
            CrossRefEntry {
                chunk_index: 1,
                local_page: 2,
            },
        );
        spine.object_index.push(ObjectIndexEntry {
            chunk_index: 1,
            byte_offset: 1024,
            byte_length: 2048,
        });
        // Chunk 1's layout tree could also be dropped now.

        // Verify global resolution works with only the spine.
        let intro_page = spine.cross_refs["intro"].absolute_page(&spine.chunk_page_counts);
        assert_eq!(intro_page, 0); // chunk 0, page 0 → absolute page 0

        let details_page = spine.cross_refs["details"].absolute_page(&spine.chunk_page_counts);
        assert_eq!(details_page, 5); // 3 pages before chunk 1, local page 2 → absolute 5

        assert_eq!(spine.object_index.len(), 2);
        assert_eq!(
            spine.chunk_page_counts.iter().sum::<u32>(),
            8,
            "total pages across all chunks"
        );
    }
}