tree reparent

2 files changed, 28 insertions, 6 deletions

diff --git a/shared/src/packets.rs b/shared/src/packets.rs
index 1f0dab1..1c457d3 100644
--- a/shared/src/packets.rs
+++ b/shared/src/packets.rs
@@ -42,7 +42,7 @@ impl<T> Hash for Resource<T> {
     }
 }
 
-#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
+#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, PartialOrd, Ord)]
 pub struct Object(pub u128);
 
 #[derive(Clone)]
diff --git a/shared/src/tree.rs b/shared/src/tree.rs
index 18d55b3..2de02d0 100644
--- a/shared/src/tree.rs
+++ b/shared/src/tree.rs
@@ -19,7 +19,7 @@ use crate::{
     resources::Prefab,
 };
 use glam::Vec3A;
-use std::collections::HashMap;
+use std::collections::{BTreeSet, HashMap};
 
 pub struct SceneTree {
     pub objects: HashMap<Object, ObjectData>,
@@ -28,6 +28,7 @@ pub struct ObjectData {
     pub pos: Vec3A,
     pub rot: Vec3A,
     pub parent: Object,
+    pub children: BTreeSet<Object>,
     pub pose: Vec<f32>,
     pub res: Resource<Prefab>,
 }
@@ -48,10 +49,15 @@ impl SceneTree {
                     rot: Vec3A::ZERO,
                     pose: Vec::new(),
                     res: res.clone(),
+                    children: BTreeSet::new(),
                 });
             }
             Packet::Remove(object) => {
-                self.objects.remove(&object);
+                if let Some(o) = self.objects.remove(&object) {
+                    for c in o.children {
+                        self.reparent(o.parent, c);
+                    }
+                }
             }
             Packet::Position(object, pos, rot) => {
                 if let Some(o) = self.objects.get_mut(&object) {
@@ -65,14 +71,30 @@ impl SceneTree {
                 }
             }
             Packet::Parent(parent, child) => {
-                if let Some(o) = self.objects.get_mut(&parent) {
-                    o.parent = *child
-                }
+                self.reparent(*parent, *child);
             }
             _ => (),
         }
     }
 
+    pub fn reparent(&mut self, parent: Object, child: Object) {
+        if !self.objects.contains_key(&parent) || !self.objects.contains_key(&child) {
+            return;
+        }
+        if let Some(co) = self.objects.get(&child) {
+            let old_parent = co.parent;
+            if let Some(po) = self.objects.get_mut(&old_parent) {
+                po.children.remove(&child);
+            }
+        }
+        if let Some(po) = self.objects.get_mut(&parent) {
+            po.children.insert(child);
+        }
+        if let Some(co) = self.objects.get_mut(&child) {
+            co.parent = parent;
+        }
+    }
+
     pub fn prime_client(&self) -> impl Iterator<Item = Packet> {
         self.objects
             .iter()