Fix DType.Add overloads (#2856)

`DType.Add` overload was broken and it did not work correctly for LazyTypes leading to errors and NullReferenceExceptions in some cases. this bug came to light when investigating a customer reported incident for new analysis. This PR fixes the `DType.Add` overload and adds tests to ensure the fix works as expected.
microsoft · Feb 20, 2025 · 2403e39 · 2403e39
1 parent 5132996
commit 2403e39
Show file tree

Hide file tree

Showing 2 changed files with 35 additions and 3 deletions.
diff --git a/src/libraries/Microsoft.PowerFx.Core/Types/DType.cs b/src/libraries/Microsoft.PowerFx.Core/Types/DType.cs
@@ -1341,6 +1341,11 @@ public DType Add(ref bool fError, DPath path, DName name, DType type)
 
             Contracts.Assert(typeOuter.IsRecord || typeOuter.IsTable);
 
+            if (typeOuter.IsLazyType)
+            {
+                typeOuter = typeOuter.LazyTypeProvider.GetExpandedType(typeOuter.IsTable);
+            }
+
             if (typeOuter.TypeTree.TryGetValue(name, out var typeCur))
             {
                 fError = true;
@@ -1371,9 +1376,15 @@ public DType Add(DName name, DType type)
                 fullType = LazyTypeProvider.GetExpandedType(IsTable);
             }
 
-            Contracts.Assert(!TypeTree.Contains(name));
-            var tree = TypeTree.SetItem(name, type);
-            var newType = new DType(Kind, tree, AssociatedDataSources, DisplayNameProvider);
+            Contracts.Assert(!fullType.TypeTree.Contains(name));
+
+            var tree = fullType.TypeTree.SetItem(name, type);
+            var newType = new DType(fullType.Kind, tree, AssociatedDataSources, fullType.DisplayNameProvider);
+
+            if (fullType.HasExpandInfo)
+            {
+                newType = CopyExpandInfo(newType, fullType);
+            }
 
             return newType;
         }

diff --git a/src/tests/Microsoft.PowerFx.Core.Tests.Shared/TypeSystemTests/LazyTypeTests.cs b/src/tests/Microsoft.PowerFx.Core.Tests.Shared/TypeSystemTests/LazyTypeTests.cs
@@ -268,6 +268,27 @@ public void AddField()
             Assert.Equal(3, _getter1CalledCount);
 
             Assert.Equal("![Bar:s, Baz:b, Foo:n, Test:O]", type1.ToString());
+
+            // test Add overload without path 
+            // using lazy record
+            var type2 = _lazyRecord1._type.Add(new DName("ThisRecord"), TestUtils.DT("![a: n]"));
+            Assert.Equal("r!", _lazyRecord1._type.ToString());
+            Assert.Equal("![Bar:s, Baz:b, Foo:n, ThisRecord:![a:n]]", type2.ToString());
+
+            // using lazy table
+            var type3 = _lazyTable2._type.Add(new DName("ThisRecord"), TestUtils.DT("![Value:O]"));
+            Assert.Equal("r*", _lazyTable2._type.ToString());
+            Assert.Equal(2, _getter2CalledCount);
+            Assert.Equal("*[Nested:r!, Qux:n, ThisRecord:![Value:O]]", type3.ToString());
+
+            // test Add overload with path and LazyType as field
+            // using lazy record
+            var type4 = _lazyRecord2._type.Add(ref fError, DPath.Root.Append(new DName("Nested")), new DName("New"), DType.Number);
+            Assert.Equal("![Nested:![Bar:s, Baz:b, Foo:n, New:n], Qux:n]", type4.ToString());
+
+            // using lazy table
+            var type5 = _lazyTable2._type.Add(ref fError, DPath.Root.Append(new DName("Nested")), new DName("NewCol"), TestUtils.DT("![a:n]"));
+            Assert.Equal("*[Nested:![Bar:s, Baz:b, Foo:n, NewCol:![a:n]], Qux:n]", type5.ToString());
         }
 
         [Fact]