Skip to content

Commit e12080e

Browse files
SamboyCodingSamboyCoding
committed
Support static fields as roots
1 parent d0b3d6c commit e12080e

File tree

5 files changed

+126
-65
lines changed

5 files changed

+126
-65
lines changed

‎UMS.Analysis/SnapshotFile.cs

Lines changed: 92 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -16,9 +16,12 @@ public class SnapshotFile : LowLevelSnapshotFile
1616

1717
public readonly Dictionary<int, int> TypeIndexIndices = new(); // TypeIndex -> Index in the TypeIndex array. What.
1818

19+
public readonly Dictionary<int, int> StaticFieldsToOwningTypes = new();
20+
1921
public readonly WellKnownTypeHelper WellKnownTypes;
2022

2123
private readonly Dictionary<int, BasicFieldInfoCache[]> _nonStaticFieldIndicesByTypeIndex = new();
24+
private readonly Dictionary<int, BasicFieldInfoCache[]> _staticFieldIndicesByTypeIndex = new();
2225

2326
private readonly Dictionary<int, BasicTypeInfoCache> _typeInfoCacheByTypeIndex = new();
2427

@@ -45,6 +48,74 @@ public SnapshotFile(string path) : base(path)
4548
TypeIndexIndices.Add(indices[i], i);
4649
}
4750

51+
public void LoadManagedObjectsFromGcRoots()
52+
{
53+
var gcRoots = GcHandles;
54+
var start = DateTime.Now;
55+
56+
Console.WriteLine($"Processing {gcRoots.Length} GC roots...");
57+
foreach (var gcHandle in gcRoots)
58+
GetOrCreateManagedClassInstance(gcHandle);
59+
60+
Console.WriteLine($"Found {_managedClassInstanceCache.Count} managed objects in {(DateTime.Now - start).TotalMilliseconds}ms");
61+
}
62+
63+
public void LoadManagedObjectsFromStaticFields()
64+
{
65+
var allStaticFields = ReadValueTypeArrayChapter<byte>(EntryType.TypeDescriptions_StaticFieldBytes, 0, -1);
66+
67+
var start = DateTime.Now;
68+
var initialCount = _managedClassInstanceCache.Count;
69+
70+
Console.WriteLine($"Processing static field info for {allStaticFields.Length} types...");
71+
for (var typeIndex = 0; typeIndex < allStaticFields.Length; typeIndex++)
72+
{
73+
var typeFieldBytes = allStaticFields[typeIndex].AsSpan();
74+
if (typeFieldBytes.Length == 0)
75+
continue;
76+
77+
var typeInfo = GetTypeInfo(typeIndex);
78+
var staticFields = GetStaticFieldInfoForTypeIndex(typeIndex);
79+
80+
foreach (var field in staticFields)
81+
{
82+
StaticFieldsToOwningTypes[field.FieldIndex] = typeIndex;
83+
if(field.IsValueType)
84+
continue; //TODO
85+
86+
if(field.IsArray)
87+
continue;
88+
89+
var fieldOffset = field.FieldOffset;
90+
91+
if(fieldOffset < 0)
92+
continue; //Generics, mainly
93+
94+
var fieldPointer = MemoryMarshal.Read<ulong>(typeFieldBytes[fieldOffset..]);
95+
if (fieldPointer == 0)
96+
continue;
97+
98+
GetOrCreateManagedClassInstance(fieldPointer, reason: LoadedReason.StaticField, fieldOrArrayIdx: field.FieldIndex);
99+
}
100+
}
101+
102+
Console.WriteLine($"Found {_managedClassInstanceCache.Count - initialCount} additional managed objects from static fields in {(DateTime.Now - start).TotalMilliseconds}ms");
103+
}
104+
105+
public ManagedClassInstance? GetOrCreateManagedClassInstance(ulong address, ManagedClassInstance? parent = null, int depth = 0, LoadedReason reason = LoadedReason.GcRoot, int fieldOrArrayIdx = int.MinValue)
106+
{
107+
if (_managedClassInstanceCache.TryGetValue(address, out var ret))
108+
return ret;
109+
110+
var info = ParseManagedObjectInfo(address);
111+
if (!info.IsKnownType)
112+
return null;
113+
114+
var instance = new ManagedClassInstance(this, info, parent, depth, reason, fieldOrArrayIdx);
115+
_managedClassInstanceCache[address] = instance;
116+
return instance;
117+
}
118+
48119
public RawManagedObjectInfo ParseManagedObjectInfo(ulong address)
49120
{
50121
if (_managedObjectInfoCache.TryGetValue(address, out var ret))
@@ -78,20 +149,6 @@ public RawManagedObjectInfo ParseManagedObjectInfo(ulong address)
78149

79150
return info;
80151
}
81-
82-
public ManagedClassInstance? GetManagedClassInstance(ulong address, ManagedClassInstance? parent = null, int depth = 0, LoadedReason reason = LoadedReason.GcRoot, int fieldOrArrayIdx = int.MinValue)
83-
{
84-
if (_managedClassInstanceCache.TryGetValue(address, out var ret))
85-
return ret;
86-
87-
var info = ParseManagedObjectInfo(address);
88-
if (!info.IsKnownType)
89-
return null;
90-
91-
var instance = new ManagedClassInstance(this, info, parent, depth, reason, fieldOrArrayIdx);
92-
_managedClassInstanceCache[address] = instance;
93-
return instance;
94-
}
95152

96153
public int SizeOfObjectInBytes(RawManagedObjectInfo info, Span<byte> heap)
97154
{
@@ -247,30 +304,43 @@ public BasicTypeInfoCache GetTypeInfo(int typeIndex)
247304
return info;
248305
}
249306

250-
public BasicFieldInfoCache[] GetFieldInfoForTypeIndex(int typeIndex)
307+
public BasicFieldInfoCache[] GetInstanceFieldInfoForTypeIndex(int typeIndex)
251308
{
252309
if(_nonStaticFieldIndicesByTypeIndex.TryGetValue(typeIndex, out var indices))
253310
return indices;
254311

255-
_nonStaticFieldIndicesByTypeIndex[typeIndex] = indices = WalkFieldInfoForTypeIndex(typeIndex).ToArray();
312+
_nonStaticFieldIndicesByTypeIndex[typeIndex] = indices = WalkFieldInfoForTypeIndex(typeIndex, false).ToArray();
313+
314+
return indices;
315+
}
316+
317+
public BasicFieldInfoCache[] GetStaticFieldInfoForTypeIndex(int typeIndex)
318+
{
319+
if(_staticFieldIndicesByTypeIndex.TryGetValue(typeIndex, out var indices))
320+
return indices;
321+
322+
_staticFieldIndicesByTypeIndex[typeIndex] = indices = WalkFieldInfoForTypeIndex(typeIndex, true).ToArray();
256323

257324
return indices;
258325
}
259326

260-
public IEnumerable<BasicFieldInfoCache> WalkFieldInfoForTypeIndex(int typeIndex)
327+
public IEnumerable<BasicFieldInfoCache> WalkFieldInfoForTypeIndex(int typeIndex, bool wantStatic)
261328
{
262-
var baseTypeIndex = ReadSingleValueType<int>(EntryType.TypeDescriptions_BaseOrElementTypeIndex, typeIndex);
263-
if (baseTypeIndex != -1)
329+
if (!wantStatic)
264330
{
265-
foreach (var i in GetFieldInfoForTypeIndex(baseTypeIndex))
266-
yield return i;
331+
var baseTypeIndex = ReadSingleValueType<int>(EntryType.TypeDescriptions_BaseOrElementTypeIndex, typeIndex);
332+
if (baseTypeIndex != -1)
333+
{
334+
foreach (var i in GetInstanceFieldInfoForTypeIndex(baseTypeIndex))
335+
yield return i;
336+
}
267337
}
268338

269339
var fieldIndices = ReadSingleValueTypeArrayChapterElement<int>(EntryType.TypeDescriptions_FieldIndices, typeIndex).ToArray();
270340
foreach (var fieldIndex in fieldIndices)
271341
{
272342
var isStatic = ReadSingleValueType<bool>(EntryType.FieldDescriptions_IsStatic, fieldIndex);
273-
if (isStatic)
343+
if (isStatic != wantStatic)
274344
continue;
275345

276346
var fieldOffset = ReadSingleValueType<int>(EntryType.FieldDescriptions_Offset, fieldIndex);

‎UMS.Analysis/Structures/BasicFieldInfoCache.cs

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -9,4 +9,8 @@ public struct BasicFieldInfoCache
99
public TypeFlags Flags;
1010
public int FieldOffset;
1111
public int FieldTypeSize;
12+
13+
public bool IsValueType => (Flags & TypeFlags.ValueType) == TypeFlags.ValueType;
14+
15+
public bool IsArray => (Flags & TypeFlags.Array) == TypeFlags.Array;
1216
}

‎UMS.Analysis/Structures/Objects/ComplexFieldValue.cs

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -38,7 +38,7 @@ public ComplexFieldValue(SnapshotFile file, BasicFieldInfoCache info, ManagedCla
3838
return;
3939
}
4040

41-
var mci = file.GetManagedClassInstance(ptr, parent, depth, LoadedReason.InstanceField, info.FieldIndex);
41+
var mci = file.GetOrCreateManagedClassInstance(ptr, parent, depth, LoadedReason.InstanceField, info.FieldIndex);
4242

4343
if (mci == null)
4444
{

‎UMS.Analysis/Structures/Objects/ManagedClassInstance.cs

Lines changed: 26 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -20,15 +20,20 @@ public readonly struct ManagedClassInstance
2020

2121
private ManagedClassInstance TypedParent => _parent == null ? default : Unsafe.Unbox<ManagedClassInstance>(_parent!);
2222

23+
public bool IsArray => (TypeDescriptionFlags & TypeFlags.Array) == TypeFlags.Array;
24+
25+
public bool IsValueType => (TypeDescriptionFlags & TypeFlags.ValueType) == TypeFlags.ValueType;
26+
2327
public ManagedClassInstance(SnapshotFile file, int typeDescriptionIndex, TypeFlags flags, int size, Span<byte> data, ManagedClassInstance parent, int depth, LoadedReason loadedReason, int fieldIndexOrArrayOffset = int.MinValue)
2428
{
25-
if ((flags & TypeFlags.ValueType) != TypeFlags.ValueType)
29+
TypeDescriptionFlags = flags;
30+
31+
if (!IsValueType)
2632
throw new("This constructor can only be used for value types");
2733

2834
_parent = parent;
2935
ObjectAddress = 0;
3036
TypeInfo = file.GetTypeInfo(typeDescriptionIndex);
31-
TypeDescriptionFlags = flags;
3237
IsInitialized = true;
3338
LoadedReason = loadedReason;
3439
FieldIndexOrArrayOffset = fieldIndexOrArrayOffset;
@@ -71,7 +76,7 @@ public ManagedClassInstance(SnapshotFile file, RawManagedObjectInfo info, Manage
7176

7277
var data = info.Data;
7378

74-
if ((TypeDescriptionFlags & TypeFlags.Array) == TypeFlags.Array)
79+
if (IsArray)
7580
{
7681
//TODO array items
7782
Fields = Array.Empty<IFieldValue>();
@@ -92,11 +97,11 @@ private IFieldValue[] ReadFields(SnapshotFile file, Span<byte> data, int depth)
9297
return Array.Empty<IFieldValue>();
9398
}
9499

95-
var fieldInfo = file.GetFieldInfoForTypeIndex(TypeInfo.TypeIndex);
100+
var fieldInfo = file.GetInstanceFieldInfoForTypeIndex(TypeInfo.TypeIndex);
96101

97102
var fields = new IFieldValue[fieldInfo.Length];
98103

99-
var isValueType = (TypeDescriptionFlags & TypeFlags.ValueType) == TypeFlags.ValueType;
104+
var isValueType = IsValueType;
100105

101106
for (var index = 0; index < fieldInfo.Length; index++)
102107
{
@@ -147,15 +152,18 @@ private bool CheckIfRecursiveReference()
147152
return false;
148153
}
149154

150-
public bool InheritsFromUnityEngineObject(SnapshotFile file)
155+
public bool InheritsFromUnityEngineObject(SnapshotFile file)
156+
=> InheritsFrom(file, file.WellKnownTypes.UnityEngineObject);
157+
158+
public bool InheritsFrom(SnapshotFile file, int baseTypeIndex)
151159
{
152160
if((TypeDescriptionFlags & TypeFlags.Array) == TypeFlags.Array)
153161
return false;
154162

155163
var baseClass = TypeInfo.BaseTypeIndex;
156164
while (baseClass != -1)
157165
{
158-
if (baseClass == file.WellKnownTypes.UnityEngineObject)
166+
if (baseClass == baseTypeIndex)
159167
return true;
160168

161169
baseClass = file.GetTypeInfo(baseClass).BaseTypeIndex;
@@ -193,16 +201,25 @@ private void AppendRetentionReason(StringBuilder sb, SnapshotFile file, ManagedC
193201
sb.Append("GC Root");
194202
return;
195203
case LoadedReason.StaticField:
196-
//TODO
204+
{
205+
var staticFieldDeclaringType = file.StaticFieldsToOwningTypes[child.FieldIndexOrArrayOffset];
206+
var fieldList = file.GetStaticFieldInfoForTypeIndex(staticFieldDeclaringType);
207+
var parentName = file.GetTypeName(staticFieldDeclaringType);
208+
var field = fieldList.First(f => f.FieldIndex == child.FieldIndexOrArrayOffset);
209+
sb.Append("Static Field ").Append(file.GetFieldName(field.FieldIndex)).Append(" of ");
210+
sb.Append(parentName);
197211
break;
212+
}
198213
case LoadedReason.InstanceField:
199-
var fieldList = file.GetFieldInfoForTypeIndex(parent.TypeInfo.TypeIndex);
214+
{
215+
var fieldList = file.GetInstanceFieldInfoForTypeIndex(parent.TypeInfo.TypeIndex);
200216
var parentName = file.GetTypeName(parent.TypeInfo.TypeIndex);
201217
var field = fieldList.First(f => f.FieldIndex == child.FieldIndexOrArrayOffset);
202218
sb.Append("Field ").Append(file.GetFieldName(field.FieldIndex)).Append(" of ");
203219
sb.Append(parentName).Append(" at 0x").Append(parent.ObjectAddress.ToString("X"));
204220
sb.Append(" <- ");
205221
break;
222+
}
206223
case LoadedReason.ArrayElement:
207224
//TODO
208225
break;

‎UnityMemorySnapshotThing/Program.cs

Lines changed: 3 additions & 33 deletions
Original file line numberDiff line numberDiff line change
@@ -57,41 +57,11 @@ public static void Main(string[] args)
5757
//
5858
// Console.WriteLine($"Querying large dynamic arrays took {(DateTime.Now - start).TotalMilliseconds} ms\n");
5959

60-
CrawlManagedObjects(file);
60+
file.LoadManagedObjectsFromGcRoots();
61+
file.LoadManagedObjectsFromStaticFields();
6162

6263
FindLeakedUnityObjects(file);
6364
}
64-
65-
private static void CrawlManagedObjects(SnapshotFile file)
66-
{
67-
//Start with GC Handles
68-
var gcHandles = file.GcHandles;
69-
70-
//Each of those is a pointer into a managed heap section which can be mapped to the data in that file
71-
//At that address there will be an object header which allows us to find the type index or type description
72-
//It also then contains the instance fields (which we can parse with type data), and we can potentially crawl to other objects using connection data
73-
//We want to find all the objects so we can then iterate on them easily
74-
75-
var validCount = 0;
76-
var start = DateTime.Now;
77-
var rootObjects = new List<ManagedClassInstance>(gcHandles.Length);
78-
79-
Console.WriteLine($"Processing {gcHandles.Length} GC roots...");
80-
// GCSettings.LatencyMode = GCLatencyMode.SustainedLowLatency;
81-
foreach (var gcHandle in gcHandles)
82-
{
83-
rootObjects.Add(file.GetManagedClassInstance(gcHandle)!.Value);
84-
85-
validCount++;
86-
87-
// if(validCount % 1000 == 0)
88-
// Console.WriteLine($"Processed {validCount} GC roots in {(DateTime.Now - start).TotalMilliseconds} ms");
89-
}
90-
91-
GCSettings.LatencyMode = GCLatencyMode.Interactive;
92-
93-
Console.WriteLine($"Found {validCount} valid GC roots out of {gcHandles.Length} total in {(DateTime.Now - start).TotalMilliseconds} ms");
94-
}
9565

9666
private static void FindLeakedUnityObjects(SnapshotFile file)
9767
{
@@ -114,7 +84,7 @@ private static void FindLeakedUnityObjects(SnapshotFile file)
11484
int numLeaked = 0;
11585
foreach (var managedClassInstance in unityEngineObjects)
11686
{
117-
var fields = file.GetFieldInfoForTypeIndex(managedClassInstance.TypeInfo.TypeIndex);
87+
var fields = file.GetInstanceFieldInfoForTypeIndex(managedClassInstance.TypeInfo.TypeIndex);
11888
for (var fieldNumber = 0; fieldNumber < fields.Length; fieldNumber++)
11989
{
12090
var basicFieldInfoCache = fields[fieldNumber];

0 commit comments

Comments
 (0)