Move rooms with no explicit parent after any space with contains them

src/buffers/mod.rs

@@ -20,14 +20,14 @@ use futures::StreamExt;
 use matrix_sdk::async_trait;
 use nonempty::NonEmpty;
 use ratatui::text::Text;
-use std::collections::HashSet;
+use std::collections::{HashMap, HashSet};
 
 mod log;
 pub use log::LogBuffer;
 mod room;
 pub use room::RoomBuffer;
 
-#[derive(Clone, Debug, PartialEq, Eq, Hash)]
+#[derive(Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
 pub enum BufferId {
   /// The main/home buffer
   Log,
@@ -57,6 +57,9 @@ pub trait Buffer: Send + Sync + memuse::DynamicUsage {
   fn parent(&self) -> Option<BufferId> {
     None
   }
+  fn children(&self) -> Option<Vec<BufferId>> {
+    None
+  }
   /// Returns if there are any updates to apply.
   async fn poll_updates(&mut self);
   fn content<'a>(&'a self) -> Box<dyn Iterator<Item = BufferItem<'a>> + 'a>;
@@ -75,6 +78,12 @@ pub trait Buffer: Send + Sync + memuse::DynamicUsage {
 
 pub struct Buffers {
   buffers: Vec<Box<dyn Buffer>>,
+  parents: HashMap<BufferId, HashSet<BufferId>>,
+  children: HashMap<BufferId, HashSet<BufferId>>,
+  /// Set of buffers already placed after a parent space, so other spaces should not
+  /// steal them, even if they are also their child (or it would cause buffers to move
+  /// every time we re-sort the buffer list)
+  attached_to_parent: HashSet<BufferId>,
   active_index: usize,
 }
 
@@ -82,6 +91,9 @@ impl Buffers {
   pub fn new(initial_buffer: Box<dyn Buffer>) -> Self {
     Self {
       buffers: vec![initial_buffer],
+      children: HashMap::new(),
+      parents: HashMap::new(),
+      attached_to_parent: HashSet::new(),
       active_index: 0,
     }
   }
@@ -98,6 +110,9 @@ impl Buffers {
     // Reorder buffers in case we just got an update on space relationships
     // FIXME: do this only when needed
     let mut buffers = Vec::new();
+    //self.children.clear();
+    //self.parents.clear();
+    self.attached_to_parent.clear();
     std::mem::swap(&mut self.buffers, &mut buffers);
     for buf in buffers.into_iter() {
       self.push(buf);
@@ -108,6 +123,14 @@ impl Buffers {
     self.buffers.len()
   }
 
+  pub fn parents(&self) -> &HashMap<BufferId, HashSet<BufferId>> {
+    &self.parents
+  }
+
+  pub fn children(&self) -> &HashMap<BufferId, HashSet<BufferId>> {
+    &self.children
+  }
+
   pub fn iter(&self) -> impl Iterator<Item = &dyn Buffer> {
     self.buffers.iter().map(|buffer_box| &**buffer_box)
   }
@@ -120,31 +143,72 @@ impl Buffers {
   pub fn push(&mut self, new_buffer: Box<dyn Buffer>) {
     let new_buffer_id = new_buffer.id();
 
-    match new_buffer
+    match new_buffer.children() {
-      .parent()
+      Some(children) => {
-      .and_then(|parent_id| self.buffers.iter().position(|buf| buf.id() == parent_id))
+        for child in children.iter() {
-    {
+          self
+            .parents
+            .entry(child.clone())
+            .or_insert_with(HashSet::new)
+            .insert(new_buffer_id.clone());
+        }
+        self
+          .children
+          .entry(new_buffer_id.clone())
+          .or_insert_with(HashSet::new)
+          .extend(children.into_iter());
+      },
+      None => {},
+    };
+    let parent = match new_buffer.parent() {
+      Some(parent) => {
+        self
+          .parents
+          .entry(new_buffer_id.clone())
+          .or_insert_with(HashSet::new)
+          .insert(parent.clone());
+        self
+          .children
+          .entry(parent.clone())
+          .or_insert_with(HashSet::new)
+          .insert(new_buffer_id.clone());
+        Some(parent)
+      },
+      None => self
+        .parents
+        .get(&new_buffer_id)
+        .and_then(|parents| parents.iter().min().cloned()), // .min() for determinism
+    };
+    let children = self
+      .children
+      .entry(new_buffer_id.clone())
+      .or_insert_with(HashSet::new);
+
+    match parent.and_then(|parent_id| self.buffers.iter().position(|buf| buf.id() == parent_id)) {
       None => self.buffers.push(new_buffer),
       Some(parent_position) => {
-        // iterator through buffers after the parent, and record on the stack the chain
+        // iterate through buffers after the parent, and record on the stack the chain
         // of successor from the parent to the current buffer. As soon as we see a buffer
         // with a parent not on the stack (or with no parent), it means we left the subtree,
         // and should insert the new buffer there.
         let mut stack = vec![self.buffers[parent_position].id()];
         let mut new_buffer = Some(new_buffer);
         for (i, buf) in self.buffers.iter().enumerate().skip(parent_position + 1) {
-          match buf.parent() {
+          match self
-            None => {}, // root buffer
+            .parents
-            Some(parent_id) => {
+            .entry(buf.id())
-              match stack.iter().position(|id| *id == parent_id) {
+            .or_insert_with(HashSet::new)
-                None => {}, // parent is not in the subtree
+            .iter()
-                Some(parent_position) => {
+            .flat_map(|parent_id| stack.iter().position(|id| *id == *parent_id))
-                  // parent is in the subtree
+            .min() // for determinism
-                  stack.truncate(parent_position + 1);
+          {
-                  stack.push(buf.id());
+            None => {}, // root buffer, or parent is not in the subtree
-                  continue;
+            Some(parent_position) => {
-                },
+              // parent is in the subtree
-              }
+              stack.truncate(parent_position + 1);
+              stack.push(buf.id());
+              self.attached_to_parent.insert(new_buffer_id.clone());
+              continue;
             },
           }
           // if we are here, it means buf.parent() is not in the subtree (or buf is orphan)
@@ -165,39 +229,21 @@ impl Buffers {
     // Find existing children and move them after this buffer (while preserving their
     // relative order.
 
-    // 1. collect the set of children
+    // 1. pop children from the list of buffers
-    let mut children_ids = HashSet::new();
+    // FIXME: should be recursive
-    children_ids.insert(new_buffer_id.clone());
-    let mut added_children_ids = true;
-    while added_children_ids {
-      added_children_ids = false;
-      for buf in self.buffers.iter() {
-        if children_ids.contains(&buf.id()) {
-          // Already collected
-          continue;
-        }
-        if let Some(parent) = buf.parent() {
-          if children_ids.contains(&parent) {
-            children_ids.insert(buf.id());
-          }
-        }
-      }
-    }
-    children_ids.remove(&new_buffer_id);
-
-    // 2. pop them from the list of buffers
     let mut all_buffers = Vec::new();
     let mut child_buffers = Vec::new();
     std::mem::swap(&mut all_buffers, &mut self.buffers);
     for buf in all_buffers.into_iter() {
-      if children_ids.contains(&buf.id()) {
+      if children.contains(&buf.id()) && !self.attached_to_parent.contains(&buf.id()) {
-        child_buffers.push(buf)
+        self.attached_to_parent.insert(buf.id());
+        child_buffers.push(buf);
       } else {
         self.buffers.push(buf);
       }
     }
 
-    // 3. split the set of buffers, and insert the children after the new buffer
+    // 2. split the set of buffers, and insert the children after the new buffer
     let after_children = self.buffers.split_off(
       self
         .buffers

src/buffers/room.rs

@@ -14,6 +14,7 @@
  * along with this program.  If not, see <https://www.gnu.org/licenses/>.
  */
 
+use std::collections::HashSet;
 use std::sync::atomic::{AtomicU16, Ordering};
 use std::sync::{Arc, OnceLock};
 
@@ -25,7 +26,10 @@ use futures::stream::FuturesUnordered;
 use futures::{FutureExt, Stream, StreamExt};
 use itertools::Itertools;
 use matrix_sdk::async_trait;
+use matrix_sdk::deserialized_responses::SyncOrStrippedState;
 use matrix_sdk::room::ParentSpace;
+use matrix_sdk::ruma::events::space::child::SpaceChildEventContent;
+use matrix_sdk::ruma::events::SyncStateEvent;
 use matrix_sdk::ruma::{OwnedRoomId, RoomId};
 use matrix_sdk::{Client, DisplayName, Room, RoomInfo};
 use matrix_sdk_ui::timeline::{
@@ -331,6 +335,7 @@ pub struct RoomBuffer {
 
   initialized_roominfo: bool,
   parent: Option<OwnedRoomId>,
+  children: Option<Vec<OwnedRoomId>>,
   display_name: Option<DisplayName>,
 
   // It's unlikely users will join the same room with more than one account;
@@ -362,6 +367,7 @@ impl RoomBuffer {
       room_id,
       initialized_roominfo: false,
       parent: None,
+      children: None,
       display_name: None,
       buffers: SmallVec::new(),
     };
@@ -402,8 +408,9 @@ impl RoomBuffer {
   }
 
   async fn update_room_info(&mut self, room: &Room) {
-    (self.parent, self.display_name) = tokio::join!(
+    (self.parent, self.children, self.display_name) = tokio::join!(
       async {
+        // Get parent space
         match room.parent_spaces().await {
           Ok(parents) => {
             parents
@@ -424,6 +431,37 @@ impl RoomBuffer {
           },
         }
       },
+      async {
+        // Get child rooms
+        match room
+          .get_state_events_static::<SpaceChildEventContent>()
+          .await
+        {
+          Ok(child_events) => {
+            Some(
+              child_events
+                .into_iter()
+                // Extract state key (ie. the child's id) and sender
+                .flat_map(|parent_event| {
+                  match parent_event.deserialize() {
+                    Ok(SyncOrStrippedState::Sync(SyncStateEvent::Original(e))) => {
+                      Some(e.state_key.to_owned())
+                    },
+                    Ok(SyncOrStrippedState::Sync(SyncStateEvent::Redacted(_))) => None,
+                    Ok(SyncOrStrippedState::Stripped(e)) => Some(e.state_key.to_owned()),
+
+                    Err(_) => None, // Ignore deserialization errors
+                  }
+                })
+                .collect(),
+            )
+          },
+          Err(e) => {
+            tracing::error!("Failed to get child rooms of {}: {:?}", self.room_id, e);
+            None
+          },
+        }
+      },
       async {
         match room.display_name().await {
           Ok(dn) => {
@@ -480,6 +518,14 @@ impl Buffer for RoomBuffer {
       .as_ref()
       .map(|parent| BufferId::Room(parent.to_owned()))
   }
+  fn children(&self) -> Option<Vec<BufferId>> {
+    self.children.as_ref().map(|children| {
+      children
+        .iter()
+        .map(|child| BufferId::Room(child.to_owned()))
+        .collect()
+    })
+  }
 
   async fn poll_updates(&mut self) {
     let room = if self.initialized_roominfo {

src/components/buflist.rs

@@ -72,10 +72,15 @@ impl Component for Buflist {
           .iter()
           .enumerate()
           .map(|(i, buf)| {
-            match buf.parent() {
+            match buffers
-              Some(parent) => {
+              .parents()
-                stack.truncate(stack.iter().position(|id| *id == parent).unwrap_or(0) + 1)
+              .get(&buf.id())
-              },
+              .into_iter()
+              .flatten()
+              .flat_map(|parent| stack.iter().position(|id| *id == *parent))
+              .max() // if in multiple spaces, the last one is most likely to matter
+            {
+              Some(parent_position) => stack.truncate(parent_position + 1),
               None => stack.clear(),
             }
             stack.push(buf.id());