1 files changed, 60 insertions, 13 deletions

diff --git a/shared/src/tree.rs b/shared/src/tree.rs
index 2de02d0..a0bbf43 100644
--- a/shared/src/tree.rs
+++ b/shared/src/tree.rs
@@ -19,6 +19,7 @@ use crate::{
     resources::Prefab,
 };
 use glam::Vec3A;
+use log::warn;
 use std::collections::{BTreeSet, HashMap};
 
 pub struct SceneTree {
@@ -43,21 +44,10 @@ impl SceneTree {
     pub fn packet(&mut self, p: &Packet) {
         match p {
             Packet::Add(object, res) => {
-                self.objects.insert(*object, ObjectData {
-                    parent: Object(0),
-                    pos: Vec3A::ZERO,
-                    rot: Vec3A::ZERO,
-                    pose: Vec::new(),
-                    res: res.clone(),
-                    children: BTreeSet::new(),
-                });
+                self.add(*object, res.to_owned());
             }
             Packet::Remove(object) => {
-                if let Some(o) = self.objects.remove(&object) {
-                    for c in o.children {
-                        self.reparent(o.parent, c);
-                    }
-                }
+                self.remove_reparent(*object, &mut ());
             }
             Packet::Position(object, pos, rot) => {
                 if let Some(o) = self.objects.get_mut(&object) {
@@ -77,8 +67,36 @@ impl SceneTree {
         }
     }
 
+    pub fn add(&mut self, object: Object, res: Resource<Prefab>) {
+        self.objects.insert(object, ObjectData {
+            parent: Object(0),
+            pos: Vec3A::ZERO,
+            rot: Vec3A::ZERO,
+            pose: Vec::new(),
+            res: res.clone(),
+            children: BTreeSet::new(),
+        });
+    }
+
     pub fn reparent(&mut self, parent: Object, child: Object) {
+        fn check_parent_loop(tree: &SceneTree, o: Object, test: Object) -> bool {
+            if o == test {
+                true
+            } else {
+                let parent = tree.objects[&o].parent;
+                if parent == o {
+                    false
+                } else {
+                    check_parent_loop(tree, parent, test)
+                }
+            }
+        }
         if !self.objects.contains_key(&parent) || !self.objects.contains_key(&child) {
+            warn!("reparent of missing objects");
+            return;
+        }
+        if check_parent_loop(&self, parent, child) {
+            warn!("cyclic parenting prevented");
             return;
         }
         if let Some(co) = self.objects.get(&child) {
@@ -94,6 +112,23 @@ impl SceneTree {
             co.parent = parent;
         }
     }
+    pub fn remove_reparent(&mut self, object: Object, ps: &mut dyn PacketSink) {
+        if let Some(o) = self.objects.remove(&object) {
+            for c in o.children {
+                self.reparent(o.parent, c);
+                ps.push(Packet::Parent(o.parent, c));
+            }
+            ps.push(Packet::Remove(object));
+        }
+    }
+    pub fn remove_recursive(&mut self, object: Object, ps: &mut dyn PacketSink) {
+        if let Some(o) = self.objects.remove(&object) {
+            for c in o.children {
+                self.remove_recursive(c, ps);
+            }
+            ps.push(Packet::Remove(object));
+        }
+    }
 
     pub fn prime_client(&self) -> impl Iterator<Item = Packet> {
         self.objects
@@ -114,3 +149,15 @@ impl SceneTree {
             .flatten()
     }
 }
+
+pub trait PacketSink {
+    fn push(&mut self, p: Packet);
+}
+impl PacketSink for Vec<Packet> {
+    fn push(&mut self, p: Packet) {
+        Vec::push(self, p);
+    }
+}
+impl PacketSink for () {
+    fn push(&mut self, _: Packet) {}
+}