From 445fac4d84afa1e7a5772bd66ed49533b731734e Mon Sep 17 00:00:00 2001
From: thecraftianman <64441307+thecraftianman@users.noreply.github.com>
Date: Tue, 27 Aug 2024 19:42:47 -0400
Subject: [PATCH] Use CFW for CalcMassRatio

Notable performance improvement for both CalcMassRatio functions. Should hopefully also fix an occasional error thrown by the Armor Properties tool on reload
---
 lua/acf/contraption/contraption_sv.lua | 38 ++++++++++++++------------
 lua/entities/acf_engine/init.lua       | 23 +++++-----------
 2 files changed, 27 insertions(+), 34 deletions(-)

diff --git a/lua/acf/contraption/contraption_sv.lua b/lua/acf/contraption/contraption_sv.lua
index 8798dc09..bf2fcfd9 100644
--- a/lua/acf/contraption/contraption_sv.lua
+++ b/lua/acf/contraption/contraption_sv.lua
@@ -87,17 +87,23 @@ function Contraption.GetAllChildren(Ent, Tab)
 end
 
 function Contraption.GetEnts(Ent)
-	local Ancestor 	= Ent:GetAncestor()
-	local Phys 		= Contraption.GetAllPhysicalEntities(Ancestor)
-	local Pare 		= {}
-
-	for K in pairs(Phys) do
-		Contraption.GetAllChildren(K, Pare)
-	end
+	local Con      = Ent:GetContraption()
+	local ConEnts  = Con and Con.ents or {[Ent] = true}
+	local Children = Ent:GetFamilyChildren()
+	local Phys     = {}
+	local Pare     = {}
+
+	for K in pairs(ConEnts) do
+		if Children[K] then
+			Pare[K] = true
+		else
+			local CurFamily = K:GetFamily()
 
-	for K in pairs(Phys) do -- Go through the all physical ents (There's probably less of those than the parented ones)
-		if Pare[K] then -- Remove them from parented table
-			Pare[K] = nil
+			if CurFamily and CurFamily:GetRoot() ~= K then
+				Pare[K] = true
+			else
+				Phys[K] = true
+			end
 		end
 	end
 
@@ -117,7 +123,7 @@ function Contraption.HasConstraints(Ent)
 end
 
 function Contraption.CalcMassRatio(Ent, Tally)
-	local TotMass  = 0
+	local Con      = Ent:GetContraption()
 	local PhysMass = 0
 	local Time     = CurTime()
 
@@ -161,10 +167,7 @@ function Contraption.CalcMassRatio(Ent, Tally)
 				PhysN = PhysN + 1
 			end
 
-
 			local Mass = Phys:GetMass()
-
-			TotMass  = TotMass + Mass
 			PhysMass = PhysMass + Mass
 		end
 	end
@@ -188,12 +191,11 @@ function Contraption.CalcMassRatio(Ent, Tally)
 
 				ParN = ParN + 1
 			end
-
-			TotMass = TotMass + Phys:GetMass()
-
 		end
 	end
 
+	local TotMass = Con and Con.totalMass or PhysMass
+
 	for K in pairs(Physical) do
 		K.acfphystotal      = PhysMass
 		K.acftotal          = TotMass
@@ -209,7 +211,7 @@ function Contraption.CalcMassRatio(Ent, Tally)
 	if Tally then
 		local Owner = Ent:CPPIGetOwner()
 
-		return Power, Fuel, PhysN, ParN, ConN, IsValid(Owner) and Owner:Name() or "Unknown", OthN
+		return Power, Fuel, PhysN, ParN, ConN, IsValid(Owner) and Owner:Name() or "Unknown", OthN, TotMass, PhysMass
 	end
 end
 
diff --git a/lua/entities/acf_engine/init.lua b/lua/entities/acf_engine/init.lua
index d3375561..20ded745 100644
--- a/lua/entities/acf_engine/init.lua
+++ b/lua/entities/acf_engine/init.lua
@@ -653,31 +653,22 @@ end
 
 -- specialized calcmassratio for engines
 function ENT:CalcMassRatio(SelfTbl)
-	SelfTbl = SelfTbl or self:GetTable()
-	local PhysMass 	= 0
-	local TotalMass = 0
-	local Physical, Parented = Contraption.GetEnts(self)
+	SelfTbl        = SelfTbl or self:GetTable()
+	local Con      = self:GetContraption()
+	local PhysMass = 0
+
+	local Physical = Contraption.GetEnts(self)
 
 	for K in pairs(Physical) do
 		local Phys = K:GetPhysicsObject() -- Should always exist, but just in case
 
 		if IsValid(Phys) then
 			local Mass = Phys:GetMass()
-
-			TotalMass = TotalMass + Mass
-			PhysMass  = PhysMass + Mass
+			PhysMass   = PhysMass + Mass
 		end
 	end
 
-	for K in pairs(Parented) do
-		if not Physical[K] then
-			local Phys = K:GetPhysicsObject()
-
-			if IsValid(Phys) then
-				TotalMass = TotalMass + Phys:GetMass()
-			end
-		end
-	end
+	local TotalMass = Con and Con.totalMass or PhysMass
 
 	SelfTbl.MassRatio = PhysMass / TotalMass
 	TotalMass = Round(TotalMass, 2)