From 99f10fda4c4c8b6952522a45936e9df3c36d0a69 Mon Sep 17 00:00:00 2001
From: Noah Hellman <noah@hllmn.net>
Date: Thu, 26 Jan 2023 20:16:05 +0100
Subject: [PATCH] tree: allow exiting empty containers

---
 src/tree.rs | 6 ++++--
 1 file changed, 4 insertions(+), 2 deletions(-)

diff --git a/src/tree.rs b/src/tree.rs
index de049cc..40632a2 100644
--- a/src/tree.rs
+++ b/src/tree.rs
@@ -209,8 +209,10 @@ impl<C: std::fmt::Debug, A: std::fmt::Debug> Builder<C, A> {
 
     pub(super) fn exit(&mut self) {
         self.depth -= 1;
-        if self.head.is_some() {
-            self.head = None;
+        if let Some(head) = self.head.take() {
+            if matches!(self.nodes[head.index()].kind, NodeKind::Container(..)) {
+                self.branch.push(head);
+            }
         } else {
             let last = self.branch.pop();
             assert_ne!(last, None);