1 files changed, 32 insertions, 32 deletions

diff --git a/src/display.rs b/src/display.rs
index 45bbca9..d60c325 100644
--- a/src/display.rs
+++ b/src/display.rs
@@ -1,5 +1,3 @@
-#![allow(unused_variables)]
-
 use crate::state::*;
 use parse_wiki_text::*;
 
@@ -14,16 +12,13 @@ const skippable_headers: &[&str; 15] =
 // but it's fine because we're working with MUCH smaller strings lol
 pub fn display(definition: String, state: &State) {
     let definition = Configuration::default().parse(&definition);
-
     display_language(&definition, &state.lang);
 }
 
-// no overloading?? O_O
-// matching on an enum of structs SUCKS
-// functions as parameters is too hard
+/// Prints only the provided language if present, otherwise prints the first language
 fn display_language(definition: &Output, lang: &str) {
     let mut has_lang = false;
-    for (i, node) in definition.nodes.iter().enumerate() {
+    for node in definition.nodes.iter() {
         if let Node::Heading { nodes, level, .. } = node
         && let Some(Node::Text { value, .. }) = nodes.get(0) {
             if *level == 2 && *value == lang {
@@ -33,35 +28,40 @@ fn display_language(definition: &Output, lang: &str) {
         }
     }
 
-    let mut inside_heading = false;
     let mut skipping_heading = false;
-    for (i, node) in definition.nodes.iter().enumerate() {
-
-        if let Node::Heading { nodes, level, .. } = node
-        && let Some(Node::Text { value, .. }) = nodes.get(0) {
-            if inside_heading {
-                if *level == 2 {
-                    inside_heading = false;
-                } else if skippable_headers.contains(value) {
-                    skipping_heading = true;
-                } else {
-                    if skipping_heading && !skippable_headers.contains(value) {
-                        skipping_heading = false;
-                    }
-                    print!("\n{}\n", node);
+    let mut inside_main_heading = false;
+    let mut nodes = definition.nodes.iter().peekable();
+    while let Some(node) = nodes.next() {
+        if let Node::Heading { nodes, level, .. } = node && let Some(Node::Text { value, .. }) = nodes.get(0) {
+            // if at a language header
+            if *level == 2 {
+                // if we're done with the main heading, break
+                if inside_main_heading {
+                    break;
+                }
+                // otherwise, mark the start of the main heading
+                if *value == lang || !has_lang {
+                    inside_main_heading = true;
                 }
-            } else if *level == 2 && (*value == lang || has_lang == false) {
-                inside_heading = true;
-                print!("{}", node);
             }
-        } else if inside_heading && !skipping_heading {
-            if let Node::OrderedList { .. } | Node::UnorderedList { .. } | Node::DefinitionList { .. } = node {
-                print!("{}", format!("{}", node).trim());
-            } else {
-                print!("{}", node);
+            // if the header is in our skippable headers list: skip until the next header
+            if inside_main_heading {
+                skipping_heading = skippable_headers.contains(value);
+            }
+        }
+        if inside_main_heading && !skipping_heading {
+            print!("{}", node);
+            match node {
+                Node::Heading { .. } => println!(),
+                _ => match nodes.peek() {
+                    Some(Node::Heading { .. }) |
+                    Some(Node::OrderedList { .. }) |
+                    Some(Node::UnorderedList { .. }) |
+                    Some(Node::DefinitionList { .. }) |
+                    None => println!(),
+                    _ => (),
+                }
             }
         }
     }
-    println!();
 }
-