Fix some things related to AllowedOn attribute

xADDBx · Nov 2, 2024 · f0d846c · f0d846c
1 parent dc0e2a5
commit f0d846c
Show file tree

Hide file tree

Showing 2 changed files with 21 additions and 12 deletions.
diff --git a/ToyBox/Classes/MainUI/PatchTool/PatchToolUI.cs b/ToyBox/Classes/MainUI/PatchTool/PatchToolUI.cs
@@ -22,6 +22,7 @@ public static class PatchToolUI {
     private static Dictionary<((object, FieldInfo), int), bool> _listToggleStates = new();
     private static Dictionary<(object, FieldInfo), AddItemState> _addItemStates = new();
     private static Dictionary<Type, List<Type>> _compatibleTypes = new();
+    private static Dictionary<Type, List<Type>> _allowedTypes = new();
     private static HashSet<object> _visited = new();
     // private static string _target = "649ae43543fd4b47ae09a6547e67bcfc";
     private static string _target = "";
@@ -60,14 +61,19 @@ public static AddItemState Create(object parent, FieldInfo info, object @object,
             _addItemStates[(parent, info)] = state;
 
             if (!_compatibleTypes.ContainsKey(elementType)) {
-                _compatibleTypes[elementType] = PatchToolUtils.GetInstantiableTypes(elementType).ToList();
+                (var all, var allowed) = PatchToolUtils.GetInstantiableTypes(elementType, parent);
+                if (allowed != null) {
+                    state.ToAddBrowser.DisplayShowAllGUI = true;
+                }
+                _allowedTypes[elementType] = allowed?.ToList();
+                _compatibleTypes[elementType] = all.ToList();
             }
 
             return state;
         }
         public void AddItemGUI() {
             using (VerticalScope()) {
-                ToAddBrowser.OnGUI(_compatibleTypes[ElementType], () => _compatibleTypes[ElementType], d => d, t => $"{t.Name}", t => [$"{t.Name}"], null,
+                ToAddBrowser.OnGUI(_allowedTypes[ElementType] ?? _compatibleTypes[ElementType], () => _compatibleTypes[ElementType], d => d, t => $"{t.Name}", t => [$"{t.Name}"], null,
                     (type, maybeType) => {
                         Label(type.Name, Width(400));
                         Space(200);

diff --git a/ToyBox/Classes/MainUI/PatchTool/PatchToolUtils.cs b/ToyBox/Classes/MainUI/PatchTool/PatchToolUtils.cs
@@ -73,8 +73,10 @@ public static PatchOperation AddOperation(this PatchOperation head, PatchOperati
             return copy;
         }
     }
-    public static HashSet<Type> GetInstantiableTypes(Type elementType) {
-        HashSet<Type> instantiableTypes = new();
+    public static (HashSet<Type>, HashSet<Type>) GetInstantiableTypes(Type elementType, object maybeParent) {
+        HashSet<Type> allowedinstantiableTypes = typeof(BlueprintComponent).IsAssignableFrom(elementType) ? new() : null;
+        HashSet<Type> allinstantiableTypes = new();
+        Type parentType = maybeParent?.GetType();
         foreach (var assembly in AppDomain.CurrentDomain.GetAssemblies()) {
             Type[] types;
             try {
@@ -87,21 +89,22 @@ public static HashSet<Type> GetInstantiableTypes(Type elementType) {
                 if (type == null) continue;
 
                 if (elementType.IsAssignableFrom(type) && !type.IsAbstract && !type.IsInterface) {
-                    var attributes = type.GetCustomAttributes(typeof(AllowedOnAttribute), inherit: false);
-                    if (attributes.Length > 0) {
-                        foreach (AllowedOnAttribute attr in attributes) {
-                            if (attr.Type == elementType) {
-                                instantiableTypes.Add(elementType);
+                    if (parentType != null && allowedinstantiableTypes != null) {
+                        var attributes = type.GetCustomAttributes(typeof(AllowedOnAttribute), inherit: false);
+                        if (attributes.Length > 0) {
+                            foreach (AllowedOnAttribute attr in attributes) {
+                                if (attr.Type.IsAssignableFrom(parentType)) {
+                                    allowedinstantiableTypes.Add(type);
+                                }
                             }
                         }
-                    } else {
-                        instantiableTypes.Add(elementType);
                     }
+                    allinstantiableTypes.Add(type);
                 }
             }
         }
 
-        return instantiableTypes;
+        return (allinstantiableTypes, allowedinstantiableTypes);
     }
 
 }