From 855485dcc8f5d1dac4b481f79fa4757cdcd43a97 Mon Sep 17 00:00:00 2001
From: Jonathan Eskew <jeskew@microsoft.com>
Date: Wed, 27 Mar 2024 18:09:39 -0400
Subject: [PATCH] Fix loading tagged unions from compiled JSON

---
 ...serDefinedDiscriminatedObjectUnionTests.cs | 88 +++++++++++++++++++
 .../ArmDeclarationToExpressionConverter.cs    | 16 ++++
 .../ArmReferenceCollector.cs                  |  8 ++
 3 files changed, 112 insertions(+)

diff --git a/src/Bicep.Core.IntegrationTests/UserDefinedDiscriminatedObjectUnionTests.cs b/src/Bicep.Core.IntegrationTests/UserDefinedDiscriminatedObjectUnionTests.cs
index c21b23d8e81..fcd06ec0e70 100644
--- a/src/Bicep.Core.IntegrationTests/UserDefinedDiscriminatedObjectUnionTests.cs
+++ b/src/Bicep.Core.IntegrationTests/UserDefinedDiscriminatedObjectUnionTests.cs
@@ -773,5 +773,93 @@ public void User_defined_discriminated_objects_can_amend_resource_derived_discri
                 }
                 """);
         }
+
+        [TestMethod]
+        public void Tagged_unions_can_be_imported_from_json_templates()
+        {
+            var test1Bicep = """
+                @export()
+                type testType = {
+                  subType: subType[]
+                }
+
+                @discriminator('type')
+                type subType = testSub1 | testSub2 | testSub3
+
+                type testSub1 = {
+                  type: '1'
+                  subOption1: string
+                }
+
+                type testSub2 = {
+                  type: '2'
+                  subOption2: int
+                }
+
+                type testSub3 = {
+                  type: '3'
+                  subOption3: bool
+                }
+                """;
+
+            static string expectedSubTypeSchema(string extension) => $$"""
+                {
+                  "type": "object",
+                  "discriminator": {
+                    "propertyName": "type",
+                    "mapping": {
+                      "1": {
+                        "$ref": "#/definitions/_1.testSub1"
+                      },
+                      "2": {
+                        "$ref": "#/definitions/_1.testSub2"
+                      },
+                      "3": {
+                        "$ref": "#/definitions/_1.testSub3"
+                      }
+                    }
+                  },
+                  "metadata": {
+                    "__bicep_imported_from!": {
+                      "sourceTemplate": "test1.{{extension}}"
+                    }
+                  }
+                }
+                """;
+
+            static string mainTypesBicep(string extension) => $$"""
+                import { testType } from 'test1.{{extension}}'
+
+                @export()
+                type mainType = {
+                  name: string
+                  test: testType[]?
+                }
+                """;
+
+            var mainBicep = """
+                import { mainType } from 'main.types.bicep'
+
+                param main mainType
+
+                output mainOut object = main
+                """;
+
+            var resultFromBicep = CompilationHelper.Compile(
+                ("test1.bicep", test1Bicep),
+                ("main.types.bicep", mainTypesBicep("bicep")),
+                ("main.bicep", mainBicep));
+
+            resultFromBicep.Template.Should().NotBeNull();
+            resultFromBicep.Template.Should().HaveJsonAtPath("$.definitions['_1.subType']", expectedSubTypeSchema("bicep"));
+
+            var resultFromJson = CompilationHelper.Compile(
+                ("test1.json", CompilationHelper.Compile(test1Bicep).Template!.ToString()),
+                ("main.types.bicep", mainTypesBicep("json")),
+                ("main.bicep", mainBicep));
+
+            resultFromJson.Template.Should().NotBeNull();
+            resultFromJson.Template.Should().HaveJsonAtPath("$.definitions['_1.subType']", expectedSubTypeSchema("json"));
+        }
     }
 }
diff --git a/src/Bicep.Core/Emit/CompileTimeImports/ArmDeclarationToExpressionConverter.cs b/src/Bicep.Core/Emit/CompileTimeImports/ArmDeclarationToExpressionConverter.cs
index 18ed5c4c0af..2d351bfc741 100644
--- a/src/Bicep.Core/Emit/CompileTimeImports/ArmDeclarationToExpressionConverter.cs
+++ b/src/Bicep.Core/Emit/CompileTimeImports/ArmDeclarationToExpressionConverter.cs
@@ -1,6 +1,8 @@
 // Copyright (c) Microsoft Corporation.
 // Licensed under the MIT License.
+
 using System.Collections.Immutable;
+using System.Linq;
 using Azure.Deployments.Core.Definitions.Schema;
 using Azure.Deployments.Core.Entities;
 using Azure.Deployments.Expression.Engines;
@@ -361,6 +363,20 @@ private TypeExpression ConvertObjectNodeToTypeExpression(ITemplateSchemaNode sch
                 modifiers.Sealed);
         }
 
+        if (schemaNode.Discriminator is { } discriminatorConstraint)
+        {
+            var unionMembers = discriminatorConstraint.Mapping.OrderByAscendingOrdinalInsensitively(kvp => kvp.Key)
+                .Select(kvp => ConvertToTypeExpression(kvp.Value))
+                .ToImmutableArray();
+
+            return new DiscriminatedObjectTypeExpression(sourceSyntax,
+                new(string.Empty,
+                    TypeSymbolValidationFlags.Default,
+                    discriminatorConstraint.PropertyName.Value,
+                    unionMembers.Select(expression => expression.ExpressedType)),
+                unionMembers);
+        }
+
         return new ObjectTypeExpression(sourceSyntax,
             new(string.Empty,
                 TypeSymbolValidationFlags.Default,
diff --git a/src/Bicep.Core/Emit/CompileTimeImports/ArmReferenceCollector.cs b/src/Bicep.Core/Emit/CompileTimeImports/ArmReferenceCollector.cs
index 84dd7a9cf28..25488e3d7f3 100644
--- a/src/Bicep.Core/Emit/CompileTimeImports/ArmReferenceCollector.cs
+++ b/src/Bicep.Core/Emit/CompileTimeImports/ArmReferenceCollector.cs
@@ -132,5 +132,13 @@ private static IEnumerable<ArmIdentifier> EnumerateReferencesUsedIn(ITemplateSch
                 yield return nested;
             }
         }
+
+        if (schemaNode.Discriminator is { } discriminatorConstraint)
+        {
+            foreach (var nested in discriminatorConstraint.Mapping.Values.SelectMany(EnumerateReferencesUsedIn))
+            {
+                yield return nested;
+            }
+        }
     }
 }