optimize immutable container command generation

lunditoph-siago · Oct 27, 2023 · 98b9313 · 98b9313
1 parent 98e3767
commit 98b9313
Show file tree

Hide file tree

Showing 2 changed files with 83 additions and 80 deletions.
diff --git a/Sia.CodeGenerators/Common.cs b/Sia.CodeGenerators/Common.cs
@@ -17,18 +17,13 @@ public record PropertyInfo(
         string Name, ITypeSymbol Type, string DisplayType,
         ImmutableArray<string> TypeArguments, ImmutableDictionary<string, TypedConstant> Arguments)
     {
-        public bool IsImmutableDictionary =>
-            DisplayType.StartsWith("global::System.Collections.Immutable.ImmutableDictionary<");
-        public bool IsImmutableList =>
-            DisplayType.StartsWith("global::System.Collections.Immutable.ImmutableList<");
-        public bool IsImmutableArray =>
-            DisplayType.StartsWith("global::System.Collections.Immutable.ImmutableArray<");
-        public bool IsImmutableHashSet =>
-            DisplayType.StartsWith("global::System.Collections.Immutable.ImmutableHashSet<");
-        public bool IsImmutableQueue =>
-            DisplayType.StartsWith("global::System.Collections.Immutable.ImmutableQueue<");
-        public bool IsImmutableStack =>
-            DisplayType.StartsWith("global::System.Collections.Immutable.ImmutableStack<");
+        public string? ImmutableContainerType =>
+            DisplayType.StartsWith(ImmutableContainerHead)
+                ? DisplayType[ImmutableContainerHead.Length
+                    ..DisplayType.IndexOf('<', ImmutableContainerHead.Length)]
+                : null;
+
+        private const string ImmutableContainerHead = "global::System.Collections.Immutable.Immutable";
 
         public PropertyInfo(string name, ITypeSymbol symbol, IEnumerable<AttributeData> attributes)
             : this(name, symbol, symbol.ToDisplayString(SymbolDisplayFormat.FullyQualifiedFormat),
@@ -49,7 +44,6 @@ public T GetArgument<T>(string name, in T defaultValue)
                 ? (T)value.Value! : defaultValue;
     }
 
-
     public static readonly AssemblyName AssemblyName = typeof(Common).Assembly.GetName();
     public static readonly string GeneratedCodeAttribute =
         $@"global::System.CodeDom.Compiler.GeneratedCodeAttribute(""{AssemblyName.Name}"", ""{AssemblyName.Version}"")";

diff --git a/Sia.CodeGenerators/SiaPropertyGenerator.cs b/Sia.CodeGenerators/SiaPropertyGenerator.cs
@@ -122,77 +122,86 @@ public static void GeneratePropertyCommands(
         if (property.GetArgument("GenerateSetCommand", true)) {
             GenerateSetCommand(source, property, componentType, componentTypeParams);
         }
-
-        if (property.IsImmutableDictionary) {
-            var itemName = property.GetArgument("Item", "");
-            if (itemName != "") {
-                var keyType = property.TypeArguments[0];
-                var valueType = property.TypeArguments[1];
-
-                if (property.GetArgument("GenerateAddItemCommand", true)) {
-                    source.WriteLine();
-                    GenerateImmutableContainerCommand(
-                        source, property, componentType, componentTypeParams,
-                        "Add" + itemName, $"{keyType} Key, {valueType} Value", ".Add(Key, Value);");
-                }
-                if (property.GetArgument("GenerateSetItemCommand", true)) {
-                    source.WriteLine();
-                    GenerateImmutableContainerCommand(
-                        source, property, componentType, componentTypeParams,
-                        "Set" + itemName, $"{keyType} Key, {valueType} Value", ".SetItem(Key, Value);");
-                }
-                if (property.GetArgument("GenerateRemoveItemCommand", true)) {
-                    source.WriteLine();
-                    GenerateImmutableContainerCommand(
-                        source, property, componentType, componentTypeParams,
-                        "Remove" + itemName, $"{keyType} Key", ".Remove(Key);");
-                }
+
+        var immutableType = property.ImmutableContainerType;
+        if (immutableType == null) { return; }
+
+        var itemName = property.GetArgument("Item", "");
+        if (itemName == "") { return; }
+
+        string? keyType, valueType;
+
+        switch (immutableType) {
+        case "Dictionary":
+            keyType = property.TypeArguments[0];
+            valueType = property.TypeArguments[1];
+
+            if (property.GetArgument("GenerateAddItemCommand", true)) {
+                source.WriteLine();
+                GenerateImmutableContainerCommand(
+                    source, property, componentType, componentTypeParams,
+                    "Add" + itemName, $"{keyType} Key, {valueType} Value", ".Add(Key, Value);");
             }
-        }
-        else if (property.IsImmutableHashSet || property.IsImmutableList || property.IsImmutableArray) {
-            var itemName = property.GetArgument("Item", "");
-            if (itemName != "") {
-                var valueType = property.TypeArguments[0];
-
-                if (property.GetArgument("GenerateAddItemCommand", true)) {
-                    source.WriteLine();
-                    GenerateImmutableContainerCommand(
-                        source, property, componentType, componentTypeParams,
-                        "Add" + itemName, $"{valueType} Value", ".Add(Value);");
-                }
-                if (property.GetArgument("GenerateRemoveItemCommand", true)) {
-                    source.WriteLine();
-                    GenerateImmutableContainerCommand(
-                        source, property, componentType, componentTypeParams,
-                        "Remove" + itemName, $"{valueType} Value", ".Remove(Value);");
-                }
+            if (property.GetArgument("GenerateSetItemCommand", true)) {
+                source.WriteLine();
+                GenerateImmutableContainerCommand(
+                    source, property, componentType, componentTypeParams,
+                    "Set" + itemName, $"{keyType} Key, {valueType} Value", ".SetItem(Key, Value);");
             }
-        }
-        else if (property.IsImmutableQueue) {
-            var itemName = property.GetArgument("Item", "");
-            if (itemName != "") {
-                var valueType = property.TypeArguments[0];
-
-                if (property.GetArgument("GenerateAddItemCommand", true)) {
-                    source.WriteLine();
-                    GenerateImmutableContainerCommand(
-                        source, property, componentType, componentTypeParams,
-                        "Enqueue" + itemName, $"{valueType} Value", ".Enqueue(Value);");
-                }
+            if (property.GetArgument("GenerateRemoveItemCommand", true)) {
+                source.WriteLine();
+                GenerateImmutableContainerCommand(
+                    source, property, componentType, componentTypeParams,
+                    "Remove" + itemName, $"{keyType} Key", ".Remove(Key);");
             }
-        }
-        else if (property.IsImmutableStack) {
-            var itemName = property.GetArgument("Item", "");
-            if (itemName != "") {
-                var valueType = property.TypeArguments[0];
-
-                if (property.GetArgument("GenerateAddItemCommand", true)) {
-                    source.WriteLine();
-                    GenerateImmutableContainerCommand(
-                        source, property, componentType, componentTypeParams,
-                        "Push" + itemName, $"{valueType} Value", ".Push(Value);");
-                }
+            break;
+
+        case "Queue":
+            valueType = property.TypeArguments[0];
+
+            if (property.GetArgument("GenerateAddItemCommand", true)) {
+                source.WriteLine();
+                GenerateImmutableContainerCommand(
+                    source, property, componentType, componentTypeParams,
+                    "Enqueue" + itemName, $"{valueType} Value", ".Enqueue(Value);");
+            }
+            break;
+
+        case "Stack":
+            valueType = property.TypeArguments[0];
+
+            if (property.GetArgument("GenerateAddItemCommand", true)) {
+                source.WriteLine();
+                GenerateImmutableContainerCommand(
+                    source, property, componentType, componentTypeParams,
+                    "Push" + itemName, $"{valueType} Value", ".Push(Value);");
+            }
+            break;
+
+        case "HashSet":
+        case "List":
+        case "Array":
+            valueType = property.TypeArguments[0];
+
+            if (property.GetArgument("GenerateAddItemCommand", true)) {
+                source.WriteLine();
+                GenerateImmutableContainerCommand(
+                    source, property, componentType, componentTypeParams,
+                    "Add" + itemName, $"{valueType} Value", ".Add(Value);");
+            }
+            if (immutableType != "HashSet" && property.GetArgument("GenerateSetItemCommand", true)) {
+                source.WriteLine();
+                GenerateImmutableContainerCommand(
+                    source, property, componentType, componentTypeParams,
+                    "Set" + itemName, $"int Index, {valueType} Value", ".SetItem(Index, Value);");
+            }
+            if (property.GetArgument("GenerateRemoveItemCommand", true)) {
+                source.WriteLine();
+                GenerateImmutableContainerCommand(
+                    source, property, componentType, componentTypeParams,
+                    "Remove" + itemName, $"{valueType} Value", ".Remove(Value);");
             }
+            break;
         }
     }