Fix Wrong mantis for bitwise/flaggable enums in G (#4131)

* fix multivalue enums for go * Update changelog * Format code * Format code * minor: refactor if else
microsoft · Feb 5, 2024 · 6077e9f · 6077e9f
1 parent 86b3a62
commit 6077e9f
Show file tree

Hide file tree

Showing 4 changed files with 27 additions and 14 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -10,6 +10,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 ### Added
 
 ### Changed
+- Fixed mantis for bitwise enums in Go. [#3936](https://github.com/microsoft/kiota/issues/3936)
 
 ## [1.11.1] - 2024-02-05
 

diff --git a/src/Kiota.Builder/Refiners/GoRefiner.cs b/src/Kiota.Builder/Refiners/GoRefiner.cs
@@ -594,6 +594,7 @@ x.Type is CodeType pType &&
             "github.com/microsoft/kiota-abstractions-go/store", "BackingStoreFactory"),
         new (static x => x is CodeMethod method && method.IsOfKind(CodeMethodKind.RequestExecutor, CodeMethodKind.RequestGenerator) && method.Parameters.Any(static y => y.IsOfKind(CodeParameterKind.RequestBody) && y.Type.Name.Equals(MultipartBodyClassName, StringComparison.OrdinalIgnoreCase)),
             AbstractionsNamespaceName, MultipartBodyClassName),
+        new (static x => x is CodeEnum @enum && @enum.Flags,"", "math"),
     };
     private const string MultipartBodyClassName = "MultipartBody";
     private const string AbstractionsNamespaceName = "github.com/microsoft/kiota-abstractions-go";

diff --git a/src/Kiota.Builder/Writers/Go/CodeEnumWriter.cs b/src/Kiota.Builder/Writers/Go/CodeEnumWriter.cs
@@ -28,20 +28,25 @@ public override void WriteCodeElement(CodeEnum codeElement, LanguageWriter write
                         string.Empty,
                         "const (");
         writer.IncreaseIndent();
-        var iotaSuffix = $" {typeName} = iota";
+        var isMultiValue = codeElement.Flags;
+
         var enumOptions = codeElement.Options;
+        int power = 0;
         foreach (var item in enumOptions)
         {
             if (!string.IsNullOrEmpty(item.Documentation.Description))
                 writer.WriteLine($"// {item.Documentation.Description}");
-            writer.WriteLine($"{item.Name.ToUpperInvariant()}_{typeName.ToUpperInvariant()}{iotaSuffix}");
-            if (!string.IsNullOrEmpty(iotaSuffix))
-                iotaSuffix = string.Empty;
+
+            if (isMultiValue)
+                writer.WriteLine($"{item.Name.ToUpperInvariant()}_{typeName.ToUpperInvariant()} = {(int)Math.Pow(2, power)}");
+            else
+                writer.WriteLine($"{item.Name.ToUpperInvariant()}_{typeName.ToUpperInvariant()}{(power == 0 ? $" {typeName} = iota" : string.Empty)}");
+
+            power++;
         }
         writer.DecreaseIndent();
         writer.WriteLines(")", string.Empty);
 
-        var isMultiValue = codeElement.Flags;
         WriteStringFunction(codeElement, writer, isMultiValue);
         WriteParsableEnum(codeElement, writer, isMultiValue);
         WriteSerializeFunction(codeElement, writer, isMultiValue);
@@ -51,7 +56,7 @@ public override void WriteCodeElement(CodeEnum codeElement, LanguageWriter write
     private void WriteStringFunction(CodeEnum codeElement, LanguageWriter writer, Boolean isMultiValue)
     {
         var typeName = codeElement.Name.ToFirstCharacterUpperCase();
-        var enumOptions = codeElement.Options;
+        var enumOptions = codeElement.Options.ToList();
 
         if (isMultiValue)
         {
@@ -60,9 +65,11 @@ private void WriteStringFunction(CodeEnum codeElement, LanguageWriter writer, Bo
             var literalOptions = enumOptions
                 .Select(x => $"\"{x.WireName}\"")
                 .Aggregate((x, y) => x + ", " + y);
-            writer.StartBlock($"for p := {typeName}(1); p <= {enumOptions.Last().Name.ToUpperInvariant()}_{typeName.ToUpperInvariant()}; p <<= 1 {{");
-            writer.StartBlock($"if i&p == p {{");
-            writer.WriteLine($"values = append(values, []string{{{literalOptions}}}[p])");
+            writer.WriteLine($"options := []string{{{literalOptions}}}");
+            writer.StartBlock($"for p := 0; p < {enumOptions.Count}; p++ {{");
+            writer.WriteLine($"mantis := {typeName}(int(math.Pow(2, float64(p))))");
+            writer.StartBlock($"if i&mantis == mantis {{");
+            writer.WriteLine($"values = append(values, options[p])");
             writer.CloseBlock();
             writer.CloseBlock();
             writer.WriteLine("return strings.Join(values, \",\")");

diff --git a/tests/Kiota.Builder.Tests/Writers/Go/CodeEnumWriterTests.cs b/tests/Kiota.Builder.Tests/Writers/Go/CodeEnumWriterTests.cs
@@ -70,18 +70,22 @@ public void WritesMultiValueEnum()
             Flags = true
         }).First();
         const string optionName = "option1";
-        myEnum.AddOption(new CodeEnumOption { Name = optionName });
+        myEnum.AddOption(new CodeEnumOption { Name = optionName }, new CodeEnumOption { Name = "option2" }, new CodeEnumOption { Name = "option3" });
         writer.Write(myEnum);
 
         var result = tw.ToString();
         Assert.Contains($"type MultiValueEnum int", result);
         Assert.Contains("const (", result);
-        Assert.Contains("OPTION1_MULTIVALUEENUM MultiValueEnum = iota", result);
+        Assert.Contains("OPTION1_MULTIVALUEENUM = 1", result);
+        Assert.Contains("OPTION2_MULTIVALUEENUM = 2", result);
+        Assert.Contains("OPTION3_MULTIVALUEENUM = 4", result);
         Assert.Contains("func (i", result);
         Assert.Contains("String() string {", result);
-        Assert.Contains("for p := MultiValueEnum(1); p <= OPTION1_MULTIVALUEENUM; p <<= 1 {", result);
-        Assert.Contains("if i&p == p {", result);
-        Assert.Contains("values = append(values, []string{\"option1\"}[p])", result);
+        Assert.Contains("options := []string{\"option1\", \"option2\", \"option3\"}", result);
+        Assert.Contains("for p := 0; p < 3; p++ {", result);
+        Assert.Contains("mantis := MultiValueEnum(int(math.Pow(2, float64(p))))", result);
+        Assert.Contains("if i&mantis == mantis {", result);
+        Assert.Contains("values = append(values, options[p])", result);
         Assert.Contains("for _, str := range values {", result);
         Assert.Contains("strings.Join(values", result);
         Assert.Contains("result |= OPTION1_MULTIVALUEENUM", result);