|
| 1 | +package tree |
| 2 | + |
| 3 | +import ( |
| 4 | + "math" |
| 5 | + "strings" |
| 6 | +) |
| 7 | + |
| 8 | +// BuildHeadingPaths reconciles the parser's section tree with the |
| 9 | +// LLM-built TOC tree into one canonical lookup: section ID → logical |
| 10 | +// heading path (the chain of TOC titles from the document root down to |
| 11 | +// the most specific TOC node that covers the section's pages). |
| 12 | +// |
| 13 | +// This is the bridge HAL-109 introduces. Ingestion builds two |
| 14 | +// independent structures: |
| 15 | +// |
| 16 | +// - the parser's Section tree, which carries content, summaries, and |
| 17 | +// candidate questions, and whose IDs every citation resolves to; its |
| 18 | +// Title fields are whatever the parser recovered (often empty or a |
| 19 | +// non-semantic chunk label for content leaves), and |
| 20 | +// - the LLM-built TOC tree ([]TOCNode, persisted on documents.toc_tree), |
| 21 | +// which carries the document's logical outline with the clean heading |
| 22 | +// vocabulary clients actually expect ("Item 8" → "Balance Sheet") and |
| 23 | +// page anchors, but no content. |
| 24 | +// |
| 25 | +// Because the two are never reconciled, the map a citation resolves |
| 26 | +// against (parser titles) and the map that holds the real headings (the |
| 27 | +// TOC) can — and do — diverge. BuildHeadingPaths closes that gap without |
| 28 | +// merging the two trees: it returns, for every section, the canonical |
| 29 | +// heading path it belongs under, so a citation can carry a real |
| 30 | +// structural path instead of a parser chunk label. |
| 31 | +// |
| 32 | +// # Matching: page-range containment |
| 33 | +// |
| 34 | +// A section belongs under the TOC node whose effective page span best |
| 35 | +// covers the section's own [PageStart, PageEnd]. Among the TOC nodes that |
| 36 | +// overlap a section, the winner is chosen by, in order: |
| 37 | +// |
| 38 | +// 1. containment — a node that fully contains the section beats one that |
| 39 | +// merely overlaps it (the section sits cleanly inside that heading); |
| 40 | +// 2. depth — the deeper (more specific) heading wins, so a section under |
| 41 | +// "Item 8 → Balance Sheet" maps to both, ending at "Balance Sheet" |
| 42 | +// rather than stopping at "Item 8"; |
| 43 | +// 3. overlap — more shared pages wins; |
| 44 | +// 4. span — the tighter node wins (more specific); |
| 45 | +// 5. start page — earlier wins, purely to make the result deterministic. |
| 46 | +// |
| 47 | +// # Degradation |
| 48 | +// |
| 49 | +// Sections with no page range (PageStart/PageEnd <= 0 — the normal state |
| 50 | +// for non-paginated formats), and every section when the TOC is empty or |
| 51 | +// nil, are simply absent from the returned map. Callers treat a missing |
| 52 | +// entry as "no canonical heading path known" and fall back to existing |
| 53 | +// behaviour, so wiring this in never makes a citation worse than today. |
| 54 | +// |
| 55 | +// The returned map is keyed by SectionID and never nil (an empty map is |
| 56 | +// returned when nothing could be mapped) so callers can index it without |
| 57 | +// a nil check. |
| 58 | +func BuildHeadingPaths(root *Section, toc []TOCNode) map[SectionID][]string { |
| 59 | + out := make(map[SectionID][]string) |
| 60 | + if root == nil || len(toc) == 0 { |
| 61 | + return out |
| 62 | + } |
| 63 | + |
| 64 | + maxPage := documentMaxPage(root, toc) |
| 65 | + entries := flattenTOC(toc, nil, maxPage) |
| 66 | + if len(entries) == 0 { |
| 67 | + return out |
| 68 | + } |
| 69 | + |
| 70 | + root.Walk(func(s *Section) bool { |
| 71 | + if s == nil || s.PageStart <= 0 || s.PageEnd <= 0 { |
| 72 | + return true |
| 73 | + } |
| 74 | + if path, ok := bestHeadingPath(entries, s.PageStart, s.PageEnd); ok { |
| 75 | + out[s.ID] = path |
| 76 | + } |
| 77 | + return true |
| 78 | + }) |
| 79 | + return out |
| 80 | +} |
| 81 | + |
| 82 | +// tocEntry is a flattened TOC node with its effective (resolved) page |
| 83 | +// span and the full heading path leading to it. depth is the node's |
| 84 | +// 0-indexed nesting level, used to prefer more specific headings. |
| 85 | +type tocEntry struct { |
| 86 | + start int |
| 87 | + end int |
| 88 | + depth int |
| 89 | + path []string |
| 90 | +} |
| 91 | + |
| 92 | +// span is the inclusive page count the entry covers. A malformed node |
| 93 | +// (end < start) that survived resolution reports the maximum span so it |
| 94 | +// sorts last in specificity comparisons, regardless of int width. |
| 95 | +func (e tocEntry) span() int { |
| 96 | + if e.end < e.start { |
| 97 | + return math.MaxInt |
| 98 | + } |
| 99 | + return e.end - e.start + 1 |
| 100 | +} |
| 101 | + |
| 102 | +// flattenTOC walks the TOC forest depth-first, resolving each node's |
| 103 | +// effective end page and accumulating the heading path. parentPath is |
| 104 | +// the chain of titles above this level; parentEnd bounds open-ended |
| 105 | +// nodes (a node whose EndPage is 0 runs until the next sibling's start, |
| 106 | +// or — for the last sibling — its parent's end, or the document end). |
| 107 | +// |
| 108 | +// Empty titles are skipped in the accumulated path so a structural |
| 109 | +// wrapper node with no heading doesn't inject a blank segment, but its |
| 110 | +// children still inherit the correct ancestry. |
| 111 | +func flattenTOC(nodes []TOCNode, parentPath []string, parentEnd int) []tocEntry { |
| 112 | + return flattenTOCAt(nodes, parentPath, parentEnd, 0) |
| 113 | +} |
| 114 | + |
| 115 | +func flattenTOCAt(nodes []TOCNode, parentPath []string, parentEnd, depth int) []tocEntry { |
| 116 | + var out []tocEntry |
| 117 | + for i, n := range nodes { |
| 118 | + start := n.StartPage |
| 119 | + end := resolveEndPage(nodes, i, parentEnd) |
| 120 | + |
| 121 | + path := parentPath |
| 122 | + if t := normaliseTitle(n.Title); t != "" { |
| 123 | + // Copy so sibling branches never share/alias the backing array. |
| 124 | + path = append(append([]string(nil), parentPath...), t) |
| 125 | + } |
| 126 | + |
| 127 | + if start > 0 && end >= start { |
| 128 | + out = append(out, tocEntry{start: start, end: end, depth: depth, path: path}) |
| 129 | + } |
| 130 | + if len(n.Nodes) > 0 { |
| 131 | + // A child can't extend past its parent's end; pass end down so |
| 132 | + // an open-ended deepest child is bounded by its ancestor. |
| 133 | + childBound := end |
| 134 | + if childBound <= 0 { |
| 135 | + childBound = parentEnd |
| 136 | + } |
| 137 | + out = append(out, flattenTOCAt(n.Nodes, path, childBound, depth+1)...) |
| 138 | + } |
| 139 | + } |
| 140 | + return out |
| 141 | +} |
| 142 | + |
| 143 | +// resolveEndPage computes the effective inclusive end page for nodes[i]. |
| 144 | +// An explicit EndPage wins. Otherwise the node runs until the page |
| 145 | +// before the next sibling that carries a StartPage; if there is no such |
| 146 | +// sibling it inherits parentEnd (the enclosing node's end, or the |
| 147 | +// document's last page at the top level). |
| 148 | +func resolveEndPage(nodes []TOCNode, i, parentEnd int) int { |
| 149 | + if nodes[i].EndPage > 0 { |
| 150 | + return nodes[i].EndPage |
| 151 | + } |
| 152 | + for j := i + 1; j < len(nodes); j++ { |
| 153 | + if nodes[j].StartPage > 0 { |
| 154 | + if nodes[j].StartPage-1 >= nodes[i].StartPage { |
| 155 | + return nodes[j].StartPage - 1 |
| 156 | + } |
| 157 | + return nodes[i].StartPage // degenerate ordering: single page |
| 158 | + } |
| 159 | + } |
| 160 | + if parentEnd > 0 { |
| 161 | + return parentEnd |
| 162 | + } |
| 163 | + return nodes[i].StartPage |
| 164 | +} |
| 165 | + |
| 166 | +// bestHeadingPath picks the heading path for a section spanning |
| 167 | +// [secStart, secEnd] using the precedence documented on |
| 168 | +// BuildHeadingPaths. Returns ok=false when no TOC entry overlaps the |
| 169 | +// section at all. |
| 170 | +func bestHeadingPath(entries []tocEntry, secStart, secEnd int) ([]string, bool) { |
| 171 | + bestIdx := -1 |
| 172 | + for i, e := range entries { |
| 173 | + ov := overlapPages(e.start, e.end, secStart, secEnd) |
| 174 | + if ov <= 0 { |
| 175 | + continue |
| 176 | + } |
| 177 | + if bestIdx < 0 || lessSpecific(entries[bestIdx], e, secStart, secEnd) { |
| 178 | + bestIdx = i |
| 179 | + } |
| 180 | + } |
| 181 | + if bestIdx < 0 || len(entries[bestIdx].path) == 0 { |
| 182 | + return nil, false |
| 183 | + } |
| 184 | + // Defensive copy so callers can't mutate our internal slices. |
| 185 | + return append([]string(nil), entries[bestIdx].path...), true |
| 186 | +} |
| 187 | + |
| 188 | +// lessSpecific reports whether the current best entry a is a WORSE match |
| 189 | +// for the section than candidate b — i.e. b should replace a. The |
| 190 | +// ordering mirrors the BuildHeadingPaths precedence list. |
| 191 | +func lessSpecific(a, b tocEntry, secStart, secEnd int) bool { |
| 192 | + aContains := contains(a.start, a.end, secStart, secEnd) |
| 193 | + bContains := contains(b.start, b.end, secStart, secEnd) |
| 194 | + if aContains != bContains { |
| 195 | + return bContains // prefer the container |
| 196 | + } |
| 197 | + if a.depth != b.depth { |
| 198 | + return b.depth > a.depth // prefer deeper / more specific |
| 199 | + } |
| 200 | + aOv := overlapPages(a.start, a.end, secStart, secEnd) |
| 201 | + bOv := overlapPages(b.start, b.end, secStart, secEnd) |
| 202 | + if aOv != bOv { |
| 203 | + return bOv > aOv // prefer more overlap |
| 204 | + } |
| 205 | + if a.span() != b.span() { |
| 206 | + return b.span() < a.span() // prefer the tighter node |
| 207 | + } |
| 208 | + return b.start < a.start // deterministic tie-break |
| 209 | +} |
| 210 | + |
| 211 | +// contains reports whether [oStart,oEnd] fully encloses [iStart,iEnd]. |
| 212 | +func contains(oStart, oEnd, iStart, iEnd int) bool { |
| 213 | + return oStart <= iStart && iEnd <= oEnd |
| 214 | +} |
| 215 | + |
| 216 | +// overlapPages returns the count of shared inclusive pages between two |
| 217 | +// ranges, or 0 when they don't intersect. |
| 218 | +func overlapPages(aStart, aEnd, bStart, bEnd int) int { |
| 219 | + if aStart <= 0 || aEnd <= 0 || bStart <= 0 || bEnd <= 0 { |
| 220 | + return 0 |
| 221 | + } |
| 222 | + lo := max(aStart, bStart) |
| 223 | + hi := min(aEnd, bEnd) |
| 224 | + if hi < lo { |
| 225 | + return 0 |
| 226 | + } |
| 227 | + return hi - lo + 1 |
| 228 | +} |
| 229 | + |
| 230 | +// documentMaxPage is the highest page the document is known to reach, |
| 231 | +// used to bound open-ended top-level TOC nodes. It takes the max across |
| 232 | +// section PageEnds and TOC StartPages so a TOC whose last node has no |
| 233 | +// EndPage still resolves to something sane. |
| 234 | +func documentMaxPage(root *Section, toc []TOCNode) int { |
| 235 | + hi := 0 |
| 236 | + if root != nil { |
| 237 | + root.Walk(func(s *Section) bool { |
| 238 | + if s != nil { |
| 239 | + hi = max(hi, s.PageEnd) |
| 240 | + } |
| 241 | + return true |
| 242 | + }) |
| 243 | + } |
| 244 | + var scan func(nodes []TOCNode) |
| 245 | + scan = func(nodes []TOCNode) { |
| 246 | + for _, n := range nodes { |
| 247 | + hi = max(hi, n.EndPage, n.StartPage) |
| 248 | + scan(n.Nodes) |
| 249 | + } |
| 250 | + } |
| 251 | + scan(toc) |
| 252 | + return hi |
| 253 | +} |
| 254 | + |
| 255 | +// normaliseTitle trims a TOC title for use as a path segment. It only |
| 256 | +// strips surrounding whitespace — the bench's anchor matcher already |
| 257 | +// handles case/punctuation/ordinal normalisation, so we keep the |
| 258 | +// heading verbatim here and let the consumer normalise for comparison. |
| 259 | +func normaliseTitle(s string) string { |
| 260 | + return strings.TrimSpace(s) |
| 261 | +} |
0 commit comments