Skip to content

Commit 92615db

Browse files
K0lb3K0lb3
committed
refractor(TypeTreeGenerator): create abstract TypeTreeGenerator interface, move AssetStudio specialisation into TypeTreeGenerator_AssetStudio
1 parent b27ee58 commit 92615db

4 files changed

Lines changed: 134 additions & 89 deletions

File tree

‎TypeTreeGeneratorAPI/NativeAPI.cs‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,7 @@ public static IntPtr TypeTreeGenerator_create(IntPtr unityVersionPtr)
1515
}
1616
try
1717
{
18-
var typeTreeGenerator = new TypeTreeGenerator(unityVersion);
18+
var typeTreeGenerator = new TypeTreeGenerator_AssetStudio(unityVersion);
1919
return GCHandle.ToIntPtr(GCHandle.Alloc(typeTreeGenerator));
2020
}
2121
catch
@@ -33,7 +33,7 @@ unsafe public static int TypeTreeGenerator_loadDLL(IntPtr typeTreeGeneratorPtr,
3333
}
3434
try
3535
{
36-
var typeTreeGenerator = (TypeTreeGenerator)GCHandle.FromIntPtr(typeTreeGeneratorPtr).Target!;
36+
var typeTreeGenerator = (TypeTreeGenerator_AssetStudio)GCHandle.FromIntPtr(typeTreeGeneratorPtr).Target!;
3737
using (var stream = new UnmanagedMemoryStream(dllPtr, dllLength, dllLength, FileAccess.ReadWrite))
3838
{
3939
if (stream == null)
@@ -61,7 +61,7 @@ unsafe public static int TypeTreeGenerator_loadIL2CPP(IntPtr typeTreeGeneratorPt
6161

6262
try
6363
{
64-
var typeTreeGenerator = (TypeTreeGenerator)GCHandle.FromIntPtr(typeTreeGeneratorPtr).Target!;
64+
var typeTreeGenerator = (TypeTreeGenerator_AssetStudio)GCHandle.FromIntPtr(typeTreeGeneratorPtr).Target!;
6565
byte[] assemblyData = new byte[assemblyDataLength];
6666
byte[] metadataData = new byte[metadataDataLength];
6767

@@ -115,7 +115,7 @@ public static int TypeTreeGenerator_generateTypeTreeNodesJson(IntPtr typeTreeGen
115115
}
116116
try
117117
{
118-
var typeTreeGenerator = (TypeTreeGenerator)GCHandle.FromIntPtr(typeTreeGeneratorPtr).Target!;
118+
var typeTreeGenerator = (TypeTreeGenerator_AssetStudio)GCHandle.FromIntPtr(typeTreeGeneratorPtr).Target!;
119119

120120
var typeTreeNodes = typeTreeGenerator.GenerateTreeNodes(assemblyName, fullName);
121121
if (typeTreeNodes == null)
@@ -144,7 +144,7 @@ public static int TypeTreeGenerator_generateTypeTreeNodesRaw(IntPtr typeTreeGene
144144
}
145145
try
146146
{
147-
var typeTreeGenerator = (TypeTreeGenerator)GCHandle.FromIntPtr(typeTreeGeneratorPtr).Target!;
147+
var typeTreeGenerator = (TypeTreeGenerator_AssetStudio)GCHandle.FromIntPtr(typeTreeGeneratorPtr).Target!;
148148
var typeTreeNodes = typeTreeGenerator.GenerateTreeNodes(assemblyName, fullName);
149149
if (typeTreeNodes == null)
150150
{
@@ -187,16 +187,16 @@ public static int TypeTreeGenerator_getMonoBehaviorDefinitions(IntPtr typeTreeGe
187187
}
188188
try
189189
{
190-
var typeTreeGenerator = (TypeTreeGenerator)GCHandle.FromIntPtr(typeTreeGeneratorPtr).Target!;
190+
var typeTreeGenerator = (TypeTreeGenerator_AssetStudio)GCHandle.FromIntPtr(typeTreeGeneratorPtr).Target!;
191191

192192
var typeNames = typeTreeGenerator.GetMonoBehaviourDefinitions();
193193

194194
var arrayLength = typeNames.Count;
195195
var arrayPtr = Marshal.AllocCoTaskMem(Marshal.SizeOf<IntPtr>() * arrayLength * 2);
196196
for (int i = 0; i < arrayLength; i++)
197197
{
198-
string module = typeNames[i].Module.Name;
199-
string fullName = typeNames[i].FullName;
198+
string module = typeNames[i].Item1;
199+
string fullName = typeNames[i].Item2;
200200
Marshal.WriteIntPtr(arrayPtr, (i * 2) * Marshal.SizeOf<IntPtr>(), Marshal.StringToCoTaskMemUTF8(module));
201201
Marshal.WriteIntPtr(arrayPtr, (i * 2 + 1) * Marshal.SizeOf<IntPtr>(), Marshal.StringToCoTaskMemUTF8(fullName));
202202
}

‎TypeTreeGeneratorAPI/TypeTreeGenerator.cs‎

Lines changed: 33 additions & 80 deletions
Original file line numberDiff line numberDiff line change
@@ -2,44 +2,39 @@
22
using AsmResolver.PE.Builder;
33
using AssetRipper.Primitives;
44
using AssetStudio;
5-
using Cpp2IL.Core.OutputFormats;
65
using Cpp2IL.Core;
6+
using Cpp2IL.Core.OutputFormats;
77
using Mono.Cecil;
88
using System.Diagnostics.CodeAnalysis;
99

1010
namespace TypeTreeGeneratorAPI
1111
{
12-
public class TypeTreeGenerator
12+
abstract public class TypeTreeGenerator
1313
{
14-
private Dictionary<string, ModuleDefinition> moduleDic;
15-
private MyAssemblyResolver resolver;
16-
private ReaderParameters readerParameters;
17-
private string unityVersion;
18-
private SerializedTypeHelper serializedTypeHelper;
14+
protected readonly string unityVersionString;
15+
protected readonly UnityVersion unityVersion;
16+
protected readonly MyAssemblyResolver resolver;
17+
protected readonly ReaderParameters readerParameters;
1918

2019
public TypeTreeGenerator(string unityVersion)
2120
{
22-
this.unityVersion = unityVersion;
23-
moduleDic = new Dictionary<string, ModuleDefinition>();
24-
resolver = new MyAssemblyResolver();
25-
readerParameters = new ReaderParameters();
26-
readerParameters.InMemory = true;
27-
readerParameters.ReadWrite = false;
28-
readerParameters.AssemblyResolver = resolver;
21+
this.unityVersionString = unityVersion;
22+
this.unityVersion = UnityVersion.Parse(unityVersionString);
2923

30-
var arVersion = UnityVersion.Parse(unityVersion);
31-
serializedTypeHelper = new SerializedTypeHelper([arVersion.Major, arVersion.Minor, arVersion.Build, arVersion.TypeNumber]);
32-
}
33-
34-
~TypeTreeGenerator()
35-
{
36-
foreach (var pair in moduleDic)
24+
resolver = new MyAssemblyResolver();
25+
readerParameters = new ReaderParameters
3726
{
38-
pair.Value.Dispose();
39-
}
40-
moduleDic.Clear();
27+
InMemory = true,
28+
ReadWrite = false,
29+
AssemblyResolver = resolver
30+
};
4131
}
4232

33+
34+
abstract public List<Tuple<string, string>> GetMonoBehaviourDefinitions();
35+
abstract public List<TypeTreeNode>? GenerateTreeNodes(string assemblyName, string fullName);
36+
abstract public void LoadDLL(Stream dllStream);
37+
4338
public void LoadDLL(byte[] dll)
4439
{
4540
using (var dllStream = new MemoryStream(dll))
@@ -48,19 +43,12 @@ public void LoadDLL(byte[] dll)
4843
}
4944
}
5045

51-
public void LoadDLL(Stream dllStream)
52-
{
53-
var assembly = AssemblyDefinition.ReadAssembly(dllStream, readerParameters);
54-
resolver.Register(assembly);
55-
moduleDic.Add(assembly.MainModule.Name, assembly.MainModule);
56-
}
57-
5846
[UnconditionalSuppressMessage("Trimming", "IL2026:Members annotated with 'RequiresUnreferencedCodeAttribute' require dynamic access otherwise can break functionality when trimming application code", Justification = "<Pending>")]
5947
public void LoadIL2CPP(byte[] assemblyData, byte[] metadataData)
6048
{
6149
Cpp2IlApi.Init();
6250
Cpp2IlApi.ConfigureLib(false);
63-
Cpp2IlApi.InitializeLibCpp2Il(assemblyData, metadataData, UnityVersion.Parse(unityVersion), false);
51+
Cpp2IlApi.InitializeLibCpp2Il(assemblyData, metadataData, this.unityVersion, false);
6452

6553
var dllGen = new AsmResolverDllOutputFormatEmpty();
6654

@@ -80,61 +68,26 @@ public void LoadIL2CPP(byte[] assemblyData, byte[] metadataData)
8068
Cpp2IlPluginManager.CallOnFinish();
8169
}
8270

83-
public List<TypeDefinition> GetMonoBehaviourDefinitions()
71+
protected static bool IsMonoBehaviour(TypeDefinition type)
8472
{
85-
var monoBehaviourDefs = new List<TypeDefinition>();
86-
foreach (var (name, module) in moduleDic)
73+
while (type != null)
8774
{
88-
foreach (var type in module.Types)
75+
if (type.BaseType == null)
76+
return false;
77+
if (type.BaseType.FullName == "UnityEngine.MonoBehaviour")
78+
return true;
79+
try
8980
{
90-
if (type.BaseType?.FullName == "UnityEngine.MonoBehaviour")
91-
{
92-
monoBehaviourDefs.Add(type);
93-
}
81+
// Resolve the base type to continue up the hierarchy
82+
type = type.BaseType.Resolve();
9483
}
95-
}
96-
return monoBehaviourDefs;
97-
}
98-
99-
public TypeDefinition? GetTypeDefinition(string assemblyName, string fullName)
100-
{
101-
if (moduleDic.TryGetValue(assemblyName, out var module))
102-
{
103-
var typeDef = module.GetType(fullName);
104-
if (typeDef == null && assemblyName == "UnityEngine.dll")
84+
catch
10585
{
106-
foreach (var pair in moduleDic)
107-
{
108-
typeDef = pair.Value.GetType(fullName);
109-
if (typeDef != null)
110-
{
111-
break;
112-
}
113-
}
86+
// If we can't resolve, break out
87+
break;
11488
}
115-
return typeDef;
11689
}
117-
return null;
118-
}
119-
120-
public List<TypeTreeNode>? GenerateTreeNodes(string assemblyName, string fullName)
121-
{
122-
var typeDef = GetTypeDefinition(assemblyName, fullName);
123-
if (typeDef != null)
124-
{
125-
return GenerateTreeNodes(typeDef);
126-
}
127-
return null;
128-
}
129-
130-
public List<TypeTreeNode> GenerateTreeNodes(TypeDefinition typeDef)
131-
{
132-
// from AssetStudioUtility.MonoBehaviourConverter
133-
var m_Nodes = new List<TypeTreeNode>();
134-
serializedTypeHelper.AddMonoBehaviour(m_Nodes, 0);
135-
var converter = new TypeDefinitionConverter(typeDef, serializedTypeHelper, 1);
136-
m_Nodes.AddRange(converter.ConvertToTypeTreeNodes());
137-
return m_Nodes;
90+
return false;
13891
}
13992

14093
}

‎TypeTreeGeneratorAPI/TypeTreeGenerator_AssetStudio.cs‎

Lines changed: 92 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,92 @@
1+
using AssetRipper.Primitives;
2+
using AssetStudio;
3+
using Mono.Cecil;
4+
5+
namespace TypeTreeGeneratorAPI
6+
{
7+
public class TypeTreeGenerator_AssetStudio : TypeTreeGenerator
8+
{
9+
private Dictionary<string, ModuleDefinition> moduleDic;
10+
private SerializedTypeHelper serializedTypeHelper;
11+
12+
public TypeTreeGenerator_AssetStudio(string unityVersion) : base(unityVersion)
13+
{
14+
moduleDic = new Dictionary<string, ModuleDefinition>();
15+
var arVersion = UnityVersion.Parse(unityVersion);
16+
serializedTypeHelper = new SerializedTypeHelper([arVersion.Major, arVersion.Minor, arVersion.Build, arVersion.TypeNumber]);
17+
}
18+
19+
~TypeTreeGenerator_AssetStudio()
20+
{
21+
foreach (var pair in moduleDic)
22+
{
23+
pair.Value.Dispose();
24+
}
25+
moduleDic.Clear();
26+
}
27+
28+
public override void LoadDLL(Stream dllStream)
29+
{
30+
var assembly = AssemblyDefinition.ReadAssembly(dllStream, readerParameters);
31+
resolver.Register(assembly);
32+
moduleDic.Add(assembly.MainModule.Name, assembly.MainModule);
33+
}
34+
35+
public override List<Tuple<string, string>> GetMonoBehaviourDefinitions()
36+
{
37+
var monoBehaviourDefs = new List<Tuple<string, string>>();
38+
foreach (var (moduleName, module) in moduleDic)
39+
{
40+
foreach (var type in module.Types)
41+
{
42+
if (IsMonoBehaviour(type))
43+
{
44+
monoBehaviourDefs.Add(new Tuple<string, string>(moduleName, type.FullName));
45+
}
46+
}
47+
}
48+
return monoBehaviourDefs;
49+
}
50+
51+
public override List<TypeTreeNode>? GenerateTreeNodes(string assemblyName, string fullName)
52+
{
53+
var typeDef = GetTypeDefinition(assemblyName, fullName);
54+
if (typeDef != null)
55+
{
56+
return GenerateTreeNodes(typeDef);
57+
}
58+
return null;
59+
}
60+
61+
public List<TypeTreeNode> GenerateTreeNodes(TypeDefinition typeDef)
62+
{
63+
// from AssetStudioUtility.MonoBehaviourConverter
64+
var m_Nodes = new List<TypeTreeNode>();
65+
serializedTypeHelper.AddMonoBehaviour(m_Nodes, 0);
66+
var converter = new TypeDefinitionConverter(typeDef, serializedTypeHelper, 1);
67+
m_Nodes.AddRange(converter.ConvertToTypeTreeNodes());
68+
return m_Nodes;
69+
}
70+
71+
public TypeDefinition? GetTypeDefinition(string assemblyName, string fullName)
72+
{
73+
if (moduleDic.TryGetValue(assemblyName, out var module))
74+
{
75+
var typeDef = module.GetType(fullName);
76+
if (typeDef == null && assemblyName == "UnityEngine.dll")
77+
{
78+
foreach (var pair in moduleDic)
79+
{
80+
typeDef = pair.Value.GetType(fullName);
81+
if (typeDef != null)
82+
{
83+
break;
84+
}
85+
}
86+
}
87+
return typeDef;
88+
}
89+
return null;
90+
}
91+
}
92+
}

‎TypeTreeGeneratorCLI/Program.cs‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -7,7 +7,7 @@
77
return;
88
}
99

10-
var generator = new TypeTreeGenerator(args[0]);
10+
var generator = new TypeTreeGenerator_AssetStudio(args[0]);
1111

1212
if (args.Length == 2)
1313
{

0 commit comments

Comments
 (0)