From 7ed79b13655e015442d5b2c69bbbd1522a014c9a Mon Sep 17 00:00:00 2001
From: Dastan-glitch <Dastan.ManDan@gmail.com>
Date: Thu, 31 Aug 2023 04:51:55 +0300
Subject: [PATCH] event_graph: [bug fix] rewrite get_offspring() method

---
 src/event_graph/model.rs          | 92 +++++++++++++++++++++++--------
 src/event_graph/protocol_event.rs |  3 +-
 2 files changed, 71 insertions(+), 24 deletions(-)

diff --git a/src/event_graph/model.rs b/src/event_graph/model.rs
index 753c000d1..e083ef108 100644
--- a/src/event_graph/model.rs
+++ b/src/event_graph/model.rs
@@ -239,26 +239,20 @@ where
     }
 
     /// Return all the offsprings (including branches if any) of a given EventID.
-    pub fn get_offspring(&self, event: &EventId) -> Result<Vec<Event<T>>> {
-        let mut offspring = vec![];
-        let mut event = *event;
-        let head = self.get_head_hash()?;
-        loop {
-            if event == head {
-                break
-            }
-            if let Some(ev) = self.event_map.get(&event) {
-                for child in ev.children.iter() {
-                    let child = self.event_map.get(child).unwrap();
-                    offspring.push(child.event.clone());
-                    event = child.event.hash();
-                }
-            } else {
-                break
-            }
+    pub fn get_offspring(&self, event: &EventId, offspring: &mut Vec<Event<T>>) -> Result<()> {
+        let node = self.event_map.get(event).ok_or(Error::EventNotFound("child node".into()))?;
+        offspring.push(node.event.clone());
+
+        // is a leaf
+        if node.children.is_empty() {
+            return Ok(())
         }
 
-        Ok(offspring)
+        for child in node.children.iter() {
+            self.get_offspring(child, offspring)?;
+        }
+
+        Ok(())
     }
 
     async fn reorganize(&mut self) -> Result<()> {
@@ -450,11 +444,10 @@ where
                 return Ok(self.current_root)
             }
             if self.event_map.get(&node_a).unwrap().children.len() > 1 {
-                let offsprings = self
-                    .get_offspring(&node_a)?
-                    .iter()
-                    .map(|event| event.hash())
-                    .collect::<Vec<EventId>>();
+                let mut offsprings = vec![];
+                self.get_offspring(&node_a, &mut offsprings)?;
+                let offsprings =
+                    offsprings.iter().map(|event| event.hash()).collect::<Vec<EventId>>();
                 if offsprings.contains(&node_b) {
                     return Ok(node_a)
                 }
@@ -678,6 +671,59 @@ mod tests {
         })
     }
 
+    #[test]
+    fn offspring_test() -> Result<()> {
+        smol::block_on(async {
+            let events_queue = EventsQueue::new();
+            let mut model = Model::new(events_queue);
+            let root_id = model.current_root;
+
+            // event_node 1
+            // Fill this node with 100 events
+            let mut event_node_1_ids = vec![];
+            let mut id1 = root_id;
+            let timestamp = Timestamp::current_time().0 - 500;
+            for i in 0..100 {
+                let node = create_message(id1, Timestamp(timestamp + i));
+                id1 = node.hash();
+                model.add(node).await?;
+                event_node_1_ids.push(id1);
+            }
+            sleep(1).await;
+
+            // event_node 2
+            // Fill this node with 150 events
+            let mut id2 = root_id;
+            let timestamp = Timestamp::current_time().0 - 100;
+            for i in 0..150 {
+                let node = create_message(id2, Timestamp(timestamp + i));
+                id2 = node.hash();
+                model.add(node).await?;
+            }
+            sleep(1).await;
+
+            // event_node 3
+            // Fill this node with 75 events
+            let mut id3 = root_id;
+            let timestamp = Timestamp::current_time().0 + 150;
+            for i in 0..75 {
+                let node = create_message(id3, Timestamp(timestamp + i));
+                id3 = node.hash();
+                model.add(node).await?;
+            }
+            sleep(1).await;
+
+            let mut offspring = vec![];
+
+            model.get_offspring(&model.current_root, &mut offspring)?;
+
+            // root event + 100 + 150 + 75 = 326
+            assert_eq!(offspring.len(), 326_usize);
+
+            Ok(())
+        })
+    }
+
     #[test]
     fn save_load_model() -> Result<()> {
         smol::block_on(async {
diff --git a/src/event_graph/protocol_event.rs b/src/event_graph/protocol_event.rs
index fea3c9855..84dbe2d3d 100644
--- a/src/event_graph/protocol_event.rs
+++ b/src/event_graph/protocol_event.rs
@@ -216,7 +216,8 @@ where
                     continue
                 }
 
-                let children = model.get_offspring(leaf)?;
+                let mut children = vec![];
+                model.get_offspring(leaf, &mut children)?;
 
                 for child in children {
                     self.channel.send(&child).await?;