From 31e1d8493b77ec2c8c6179e6fb450b96d0f753af Mon Sep 17 00:00:00 2001
From: Sam Byass <sam@samboycoding.me>
Date: Thu, 17 Aug 2023 15:32:19 +0900
Subject: [PATCH] Log if an object is itself a LMS in the retention path

---
 .../Objects/ManagedClassInstance.cs           | 41 +++++++++++++++++++
 UnityMemorySnapshotThing/Program.cs           | 37 ++++++-----------
 2 files changed, 53 insertions(+), 25 deletions(-)

diff --git a/UMS.Analysis/Structures/Objects/ManagedClassInstance.cs b/UMS.Analysis/Structures/Objects/ManagedClassInstance.cs
index b2496dc..44c03d3 100644
--- a/UMS.Analysis/Structures/Objects/ManagedClassInstance.cs
+++ b/UMS.Analysis/Structures/Objects/ManagedClassInstance.cs
@@ -262,6 +262,17 @@ private void AppendRetentionReason(StringBuilder sb, SnapshotFile file, ManagedC
                 var field = fieldList.First(f => f.FieldIndex == child.FieldIndexOrArrayOffset);
                 sb.Append("Field ").Append(file.GetFieldName(field.FieldIndex)).Append(" of ");
                 sb.Append(parentName).Append(" at 0x").Append(parent.ObjectAddress.ToString("X"));
+
+                if (parent.InheritsFromUnityEngineObject(file))
+                {
+                    var parentInst = file.GetOrCreateManagedClassInstance(parent.ObjectAddress);
+                    
+                    if (parentInst.HasValue && parentInst.Value.IsLeakedManagedShell(file))
+                        sb.Append(" (leaked managed shell)");
+                    else
+                        sb.Append(" (unity object, non-leaked)");
+                }
+
                 sb.Append(" <- ");
                 break;
             }
@@ -277,4 +288,34 @@ private void AppendRetentionReason(StringBuilder sb, SnapshotFile file, ManagedC
                 throw new ArgumentOutOfRangeException(nameof(child), "Invalid LoadedReason");
         }
     }
+    
+    public bool IsLeakedManagedShell(SnapshotFile file)
+    {
+        if (!InheritsFromUnityEngineObject(file))
+            //Can't be a leaked managed shell if it's not a managed shell at all
+            return false;
+
+        // if (Fields == null)
+        //     return false; //Can't check
+
+        var fields = file.GetInstanceFieldInfoForTypeIndex(TypeInfo.TypeIndex);
+        for (var fieldNumber = 0; fieldNumber < fields.Length; fieldNumber++)
+        {
+            var basicFieldInfoCache = fields[fieldNumber];
+            var name = file.GetFieldName(basicFieldInfoCache.FieldIndex);
+
+            if (name == "m_CachedPtr")
+            {
+                var value = Fields[fieldNumber];
+
+                if (value is not IntegerFieldValue integerFieldValue)
+                    throw new Exception("Expected integer field value");
+
+                return integerFieldValue.Value == 0;
+            }
+        }
+
+        //Couldn't find the m_CachedPtr field. Weird, but return false.
+        return false;
+    }
 }
\ No newline at end of file
diff --git a/UnityMemorySnapshotThing/Program.cs b/UnityMemorySnapshotThing/Program.cs
index ddf4bb2..baddef0 100644
--- a/UnityMemorySnapshotThing/Program.cs
+++ b/UnityMemorySnapshotThing/Program.cs
@@ -87,36 +87,21 @@ private static void FindLeakedUnityObjects(SnapshotFile file)
         var leakedTypes = new Dictionary<string, int>();
         foreach (var managedClassInstance in unityEngineObjects)
         {
-            var fields = file.GetInstanceFieldInfoForTypeIndex(managedClassInstance.TypeInfo.TypeIndex);
-            for (var fieldNumber = 0; fieldNumber < fields.Length; fieldNumber++)
+            if (managedClassInstance.IsLeakedManagedShell(file))
             {
-                var basicFieldInfoCache = fields[fieldNumber];
-                var name = file.GetFieldName(basicFieldInfoCache.FieldIndex);
+                var typeName = file.GetTypeName(managedClassInstance.TypeInfo.TypeIndex);
 
-                if (name == "m_CachedPtr")
-                {
-                    var value = managedClassInstance.Fields[fieldNumber];
-                    
-                    if(value is not IntegerFieldValue integerFieldValue)
-                        throw new Exception("Expected integer field value");
-                    
-                    if (integerFieldValue.Value == 0)
-                    {
-                        var typeName = file.GetTypeName(managedClassInstance.TypeInfo.TypeIndex);
+                str = $"Found leaked managed object of type: {typeName} at memory address 0x{managedClassInstance.ObjectAddress:X}";
+                Console.WriteLine(str);
+                ret.AppendLine(str);
 
-                        str = $"Found leaked managed object of type: {typeName} at memory address 0x{managedClassInstance.ObjectAddress:X}";
-                        Console.WriteLine(str);
-                        ret.AppendLine(str);
-
-                        str = $"    Retention Path: {managedClassInstance.GetFirstObservedRetentionPath(file)}";
-                        Console.WriteLine(str);
-                        ret.AppendLine(str);
+                str = $"    Retention Path: {managedClassInstance.GetFirstObservedRetentionPath(file)}";
+                Console.WriteLine(str);
+                ret.AppendLine(str);
                         
-                        leakedTypes[typeName] = leakedTypes.GetValueOrDefault(typeName) + 1;
+                leakedTypes[typeName] = leakedTypes.GetValueOrDefault(typeName) + 1;
                         
-                        numLeaked++;
-                    }
-                }
+                numLeaked++;
             }
         }
 
@@ -131,4 +116,6 @@ private static void FindLeakedUnityObjects(SnapshotFile file)
         
         File.WriteAllText("leaked_objects.txt", ret.ToString());
     }
+
+   
 }
\ No newline at end of file