From c41ab742d37051b9eef31f42b2f80e8a4b77c72c Mon Sep 17 00:00:00 2001
From: Joscha <joscha@plugh.de>
Date: Tue, 9 Aug 2022 15:12:49 +0200
Subject: [PATCH] Fix message count in folded info

---
 src/store.rs               |  9 +++++++++
 src/ui/chat/tree/layout.rs | 21 ++++++++++++---------
 2 files changed, 21 insertions(+), 9 deletions(-)

diff --git a/src/store.rs b/src/store.rs
index 89bb530..24b95d0 100644
--- a/src/store.rs
+++ b/src/store.rs
@@ -94,6 +94,15 @@ impl<M: Msg> Tree<M> {
         self.children.get(id).map(|c| c as &[M::Id])
     }
 
+    pub fn subtree_size(&self, id: &M::Id) -> usize {
+        let children = self.children(id).unwrap_or_default();
+        let mut result = children.len();
+        for child in children {
+            result += self.subtree_size(child);
+        }
+        result
+    }
+
     pub fn siblings(&self, id: &M::Id) -> Option<&[M::Id]> {
         if let Some(parent) = self.parent(id) {
             self.children(&parent)
diff --git a/src/ui/chat/tree/layout.rs b/src/ui/chat/tree/layout.rs
index 32f8ecf..0122dcf 100644
--- a/src/ui/chat/tree/layout.rs
+++ b/src/ui/chat/tree/layout.rs
@@ -90,12 +90,13 @@ impl<M: Msg + ChatMsg, S: MsgStore<M>> InnerTreeViewState<M, S> {
             blocks.blocks_mut().push_back(block);
         }
 
+        // Last part of message body if message is folded
         let folded = self.folded.contains(id);
-        let children = tree.children(id);
-        let folded_info = children
-            .filter(|_| folded)
-            .map(|c| c.len())
-            .filter(|c| *c > 0);
+        let folded_info = if folded {
+            Some(tree.subtree_size(id)).filter(|s| *s > 0)
+        } else {
+            None
+        };
 
         // Main message body
         let highlighted = self.cursor.refers_to(id);
@@ -107,10 +108,12 @@ impl<M: Msg + ChatMsg, S: MsgStore<M>> InnerTreeViewState<M, S> {
         let block = Block::new(frame, BlockId::Msg(id.clone()), widget);
         blocks.blocks_mut().push_back(block);
 
-        // Children recursively (if not folded)
-        if let Some(children) = children.filter(|_| !folded) {
-            for child in children {
-                self.layout_subtree(nick, frame, tree, indent + 1, child, blocks);
+        // Children, recursively
+        if !folded {
+            if let Some(children) = tree.children(id) {
+                for child in children {
+                    self.layout_subtree(nick, frame, tree, indent + 1, child, blocks);
+                }
             }
         }