From 7c51eff47dd17a522a83d34673ce2cd43dca25b3 Mon Sep 17 00:00:00 2001
From: Jonathan Eskew <jeskew@microsoft.com>
Date: Tue, 5 Mar 2024 14:53:24 -0500
Subject: [PATCH] Preserve nullability when loading resource-derived type from
 ARM JSON template

---
 .../UserDefinedTypeTests.cs                   | 30 +++++++++++++
 .../TypeSystem/ArmTemplateTypeLoader.cs       | 42 +++++++++----------
 2 files changed, 51 insertions(+), 21 deletions(-)

diff --git a/src/Bicep.Core.IntegrationTests/UserDefinedTypeTests.cs b/src/Bicep.Core.IntegrationTests/UserDefinedTypeTests.cs
index 60d7806318a..8ce07bbcc15 100644
--- a/src/Bicep.Core.IntegrationTests/UserDefinedTypeTests.cs
+++ b/src/Bicep.Core.IntegrationTests/UserDefinedTypeTests.cs
@@ -1550,4 +1550,34 @@ public void Using_a_complete_resource_body_as_a_type_should_not_throw_exception(
             ("BCP394", DiagnosticLevel.Error, "Resource-derived type expressions must derefence a property within the resource body. Using the entire resource body type is not permitted."),
         });
     }
+
+    [TestMethod]
+    public void Resource_derived_type_nullability_should_be_preserved_when_loading_from_ARM_JSON()
+    {
+        var result = CompilationHelper.Compile(new ServiceBuilder().WithFeatureOverrides(new(TestContext, ResourceDerivedTypesEnabled: true)),
+            ("main.bicep", """
+                module mod 'mod.json' = {
+                    name: 'mod'
+                }
+                """),
+            ("mod.json", $$"""
+                {
+                    "$schema": "https://schema.management.azure.com/schemas/2019-04-01/deploymentTemplate.json#",
+                    "languageVersion": "2.0",
+                    "contentVersion": "1.0.0.0",
+                    "parameters": {
+                        "foo": {
+                            "type": "string",
+                            "metadata": {
+                                "{{LanguageConstants.MetadataResourceDerivedTypePropertyName}}": "Microsoft.Storage/storageAccounts@2022-09-01#properties/sku/properties/name"
+                            },
+                            "nullable": true
+                        }
+                    },
+                    "resources": []
+                }
+                """));
+
+        result.Should().NotHaveAnyDiagnostics();
+    }
 }
diff --git a/src/Bicep.Core/TypeSystem/ArmTemplateTypeLoader.cs b/src/Bicep.Core/TypeSystem/ArmTemplateTypeLoader.cs
index 3bb3fe4d011..d203b2ef15c 100644
--- a/src/Bicep.Core/TypeSystem/ArmTemplateTypeLoader.cs
+++ b/src/Bicep.Core/TypeSystem/ArmTemplateTypeLoader.cs
@@ -24,27 +24,7 @@ public static ITypeReference ToTypeReference(SchemaValidationContext context, IT
         {
             var resolved = TemplateEngine.ResolveSchemaReferences(context, armTemplateSchemaNode);
 
-            if (TryGetResourceDerivedType(context, resolved, flags) is ITypeReference resourceDerivedType)
-            {
-                return resourceDerivedType;
-            }
-
-            if (resolved.Type.Value == TemplateParameterType.SecureString || resolved.Type.Value == TemplateParameterType.SecureObject)
-            {
-                flags = TypeSymbolValidationFlags.IsSecure | flags;
-            }
-
-            var bicepType = resolved.Type.Value switch
-            {
-                TemplateParameterType.String or
-                TemplateParameterType.SecureString => GetStringType(resolved, flags),
-                TemplateParameterType.Int => GetIntegerType(resolved, flags),
-                TemplateParameterType.Bool => GetBooleanType(resolved, flags),
-                TemplateParameterType.Array => GetArrayType(context, resolved),
-                TemplateParameterType.Object or
-                TemplateParameterType.SecureObject => GetObjectType(context, resolved, flags),
-                _ => ErrorType.Empty(),
-            };
+            var bicepType = ToTypeReferenceIgnoringNullability(context, resolved, flags);
 
             if (resolved.Nullable?.Value == true)
             {
@@ -59,6 +39,26 @@ TemplateParameterType.Object or
         }
     }
 
+    private static ITypeReference ToTypeReferenceIgnoringNullability(SchemaValidationContext context, ITemplateSchemaNode withResolvedRefs, TypeSymbolValidationFlags flags)
+    {
+        if (TryGetResourceDerivedType(context, withResolvedRefs, flags) is ITypeReference resourceDerivedType)
+        {
+            return resourceDerivedType;
+        }
+
+        return withResolvedRefs.Type.Value switch
+        {
+            TemplateParameterType.String => GetStringType(withResolvedRefs, flags),
+            TemplateParameterType.SecureString => GetStringType(withResolvedRefs, flags | TypeSymbolValidationFlags.IsSecure),
+            TemplateParameterType.Int => GetIntegerType(withResolvedRefs, flags),
+            TemplateParameterType.Bool => GetBooleanType(withResolvedRefs, flags),
+            TemplateParameterType.Array => GetArrayType(context, withResolvedRefs),
+            TemplateParameterType.Object => GetObjectType(context, withResolvedRefs, flags),
+            TemplateParameterType.SecureObject => GetObjectType(context, withResolvedRefs, flags | TypeSymbolValidationFlags.IsSecure),
+            _ => ErrorType.Empty(),
+        };
+    }
+
     private static ITypeReference? TryGetResourceDerivedType(SchemaValidationContext context, ITemplateSchemaNode schemaNode, TypeSymbolValidationFlags flags)
     {
         if (schemaNode.Metadata?.Value is JObject metadataObject &&