From 4695b029a0b578238696bcfdd3af68eda8f8147f Mon Sep 17 00:00:00 2001
From: southerntofu <52931252+southerntofu@users.noreply.github.com>
Date: Thu, 15 Aug 2019 08:19:32 +0000
Subject: [PATCH] Fix ToC generation for heading levels > 3 (bugfix) (#774)

* Fix ToC generation for heading levels > 3

* typo

* Add tests for deep ToCs

* Code style change
---
 components/rendering/src/table_of_contents.rs | 100 +++++++++++++++---
 1 file changed, 84 insertions(+), 16 deletions(-)

diff --git a/components/rendering/src/table_of_contents.rs b/components/rendering/src/table_of_contents.rs
index 377b94d5..4b60b973 100644
--- a/components/rendering/src/table_of_contents.rs
+++ b/components/rendering/src/table_of_contents.rs
@@ -27,33 +27,58 @@ impl Default for Header {
     }
 }
 
+// Takes a potential (mutable) parent and a header to try and insert into
+// Returns true when it performed the insertion, false otherwise
+fn insert_into_parent(potential_parent: Option<&mut Header>, header: &Header) -> bool {
+    match potential_parent {
+        None => {
+            // No potential parent to insert into so it needs to be insert higher
+            return false;
+        },
+        Some(parent) => {
+            let diff = header.level - parent.level;
+            if diff <= 0 {
+                // Heading is same level or higher so we don't insert here
+                return false;
+            } 
+            if diff == 1 {
+                // We have a direct child of the parent
+                parent.children.push(header.clone());
+                return true;
+            }
+            // We need to go deeper
+            if !insert_into_parent(parent.children.iter_mut().last(), header) {
+                // No, we need to insert it here
+                parent.children.push(header.clone());
+            }
+            return true;
+        }
+    }
+}
+
 /// Converts the flat temp headers into a nested set of headers
 /// representing the hierarchy
 pub fn make_table_of_contents(headers: Vec<Header>) -> Vec<Header> {
     let mut toc = vec![];
-    'parent: for header in headers {
+    for header in headers {
         if toc.is_empty() {
+            // First header, nothing to compare it with
             toc.push(header);
             continue;
         }
 
-        // See if we have to insert as a child of a previous header
-        for h in toc.iter_mut().rev() {
-            // Look in its children first
-            for child in h.children.iter_mut().rev() {
-                if header.level > child.level {
-                    child.children.push(header);
-                    continue 'parent;
-                }
-            }
-            if header.level > h.level {
-                h.children.push(header);
-                continue 'parent;
+        // We try to insert the current header in a previous one
+        match insert_into_parent(toc.iter_mut().last(), &header) {
+            true => {
+                // Header was successfully inserted as a child of a previous element
+                continue;
+            },
+            false => {
+                // Couldn't insert in a previous header, so it's a top-level header
+                toc.push(header);
+                continue;
             }
         }
-
-        // Nop, just insert it
-        toc.push(header)
     }
 
     toc
@@ -91,6 +116,49 @@ mod tests {
         assert_eq!(toc[1].children[0].children.len(), 2);
     }
 
+    #[test]
+    fn can_make_deep_toc() {
+        let input = vec![
+            Header::new(1),
+            Header::new(2),
+            Header::new(3),
+            Header::new(4),
+            Header::new(5),
+            Header::new(4),
+        ];
+        let toc = make_table_of_contents(input);
+        assert_eq!(toc.len(), 1);
+        assert_eq!(toc[0].children.len(), 1);
+        assert_eq!(toc[0].children[0].children.len(), 1);
+        assert_eq!(toc[0].children[0].children[0].children.len(), 2);
+        assert_eq!(toc[0].children[0].children[0].children[0].children.len(), 1);
+    }
+
+    #[test]
+    fn can_make_deep_messy_toc() {
+        let input = vec![
+            Header::new(2), // toc[0]
+            Header::new(3),
+            Header::new(4),
+            Header::new(5),
+            Header::new(4),
+            Header::new(2), // toc[1]
+            Header::new(1), // toc[2]
+            Header::new(2),
+            Header::new(3),
+            Header::new(4),
+        ];
+        let toc = make_table_of_contents(input);
+        assert_eq!(toc.len(), 3);
+        assert_eq!(toc[0].children.len(), 1);
+        assert_eq!(toc[0].children[0].children.len(), 2);
+        assert_eq!(toc[0].children[0].children[0].children.len(), 1);
+        assert_eq!(toc[1].children.len(), 0);
+        assert_eq!(toc[2].children.len(), 1);
+        assert_eq!(toc[2].children[0].children.len(), 1);
+        assert_eq!(toc[2].children[0].children[0].children.len(), 1);
+    }
+
     #[test]
     fn can_make_messy_toc() {
         let input = vec![