Implement ENT:GetFamilyChildren + add mass for families

lua/cfw/classes/family_sv.lua

@@ -8,6 +8,7 @@ function CFW.Classes.Family.create(ancestor)
         count    = 0,
         ents     = {},
         ancestor = ancestor,
+        children = {},
         color    = ColorRand()
     }
 
@@ -54,6 +55,7 @@ do -- Class def
             if child.CFW_NO_FAMILY_TRAVERSAL then continue end
 
             self:Add(child)
+            self.children[v] = true
         end
     end
 
@@ -73,6 +75,7 @@ do -- Class def
             if child.CFW_NO_FAMILY_TRAVERSAL then continue end
 
             self:Sub(child)
+            self.children[v] = nil
         end
     end
 end
@@ -85,7 +88,8 @@ do
     end
 
     function ENT:GetAncestor()
-        return self._family and self._family.ancestor or self
+        local Family = self._family
+        return Family and Family.ancestor or self
     end
 
     function ENT:SetFamily(newFamily)
@@ -107,6 +111,7 @@ do
     end
 
     function ENT:GetFamilyChildren()
-        return self._family.lookup[self].children
+        local Family = self._family
+        return Family and Family.children or self:GetChildren()
     end
 end

lua/cfw/extensions/mass_sv.lua

@@ -1,4 +1,4 @@
--- Tracks the total mass of a contraption
+-- Tracks the total mass of a contraption or family
 
 local PHYS    = FindMetaTable("PhysObj")
 local setMass = setMass or PHYS.SetMass
@@ -11,36 +11,45 @@ function PHYS:SetMass(newMass)
 
     setMass(self, newMass)
 
-    local con = self:GetEntity():GetContraption()
+    local con = ent:GetContraption()
 
     if con then
         con.totalMass = con.totalMass + (newMass - oldMass)
     end
 end
 
-hook.Add("cfw.contraption.created", "CFW_Mass", function(con)
-    con.totalMass = 0
-end)
+local function InitMass(Class)
+    Class.totalMass = 0
+end
+
+hook.Add("cfw.contraption.created", "CFW_Mass", InitMass)
+hook.Add("cfw.family.created", "CFW_Mass", InitMass)
 
-hook.Add("cfw.contraption.entityAdded", "CFW_Mass", function(con, ent)
-    if not IsValid(ent) then return end
+local function AddMass(Class, Ent)
+    if not IsValid(Ent) then return end
 
-    local obj = ent:GetPhysicsObject()
+    local PhysObj = Ent:GetPhysicsObject()
 
-    if IsValid(obj) then
-        local mass = obj:GetMass()
+    if IsValid(PhysObj) then
+        local Mass = PhysObj:GetMass()
 
-        ent._mass     = mass
-        con.totalMass = con.totalMass + mass
+        Ent._mass     = Mass
+        Class.totalMass = Class.totalMass + Mass
     end
-end)
+end
 
-hook.Add("cfw.contraption.entityRemoved", "CFW_Mass", function(con, ent)
-    if not IsValid(ent) then return end
+hook.Add("cfw.contraption.entityAdded", "CFW_Mass", AddMass)
+hook.Add("cfw.family.added", "CFW_Mass", AddMass)
 
-    local obj = ent:GetPhysicsObject()
+local function SubMass(Class, Ent)
+    if not IsValid(Ent) then return end
 
-    if IsValid(obj) then
-        con.totalMass = con.totalMass - obj:GetMass()
+    local PhysObj = Ent:GetPhysicsObject()
+
+    if IsValid(PhysObj) then
+        Class.totalMass = Class.totalMass - PhysObj:GetMass()
     end
-end)
+end
+
+hook.Add("cfw.contraption.entityRemoved", "CFW_Mass", SubMass)
+hook.Add("cfw.family.subbed", "CFW_Mass", SubMass)