CellsPresenter and use the RowsPresenter viewport instead of estimating

src/Avalonia.Controls.TreeDataGrid/Primitives/TreeDataGridCellsPresenter.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Linq;
 using System.Xml.Linq;
 using Avalonia.Controls.Models.TreeDataGrid;
 using Avalonia.Controls.Selection;
@@ -63,6 +64,13 @@ public void UpdateRowIndex(int index)
             }
         }
 
+        protected override Rect? GetParentPresenterViewPort()
+        {
+            var parentRowPresenter = this.GetVisualAncestors().OfType<TreeDataGridRowsPresenter>().FirstOrDefault();
+
+            return parentRowPresenter?.Viewport;
+        }
+
         protected override Size MeasureOverride(Size availableSize)
         {
             return RowIndex == -1 ? default : base.MeasureOverride(availableSize);

src/Avalonia.Controls.TreeDataGrid/Primitives/TreeDataGridPresenterBase.cs

@@ -88,7 +88,8 @@ public IReadOnlyList<TItem>? Items
         internal IReadOnlyList<Control?> RealizedElements => _realizedElements?.Elements ?? Array.Empty<Control>();
 
         protected abstract Orientation Orientation { get; }
-        protected Rect Viewport { get; private set; } = s_invalidViewport;
+
+        internal Rect Viewport { get; private set; } = s_invalidViewport;
 
         public Control? BringIntoView(int index, Rect? rect = null)
         {
@@ -652,8 +653,18 @@ private double EstimateElementSizeU()
             return _lastEstimatedElementSizeU;
         }
 
+        protected virtual Rect? GetParentPresenterViewPort()
+        {
+            return null;
+        }
+
         private Rect EstimateViewport(Size availableSize)
         {
+            if (GetParentPresenterViewPort() is { } parentViewport && parentViewport != s_invalidViewport)
+            {
+                return parentViewport;
+            }
+
             var c = this.GetVisualParent();
 
             if (c is null)

tests/Avalonia.Controls.TreeDataGrid.Tests/TestTemplates.cs

@@ -54,7 +54,7 @@ public static IControlTemplate ScrollViewerTemplate()
             Setters = { new Setter(TemplatedControl.TemplateProperty, ScrollViewerTemplate()) }
         };
 
-        public static IControlTemplate TreeDataGridTemplate()
+        public static IControlTemplate TreeDataGridTemplate(double headerPresenterSize = 0)
         {
             return new FuncControlTemplate<TreeDataGrid>((x, ns) =>
                 new DockPanel
@@ -65,7 +65,7 @@ public static IControlTemplate TreeDataGridTemplate()
                         {
                             Name = "PART_HeaderScrollViewer",
                             Template = ScrollViewerTemplate(),
-                            Height = 0,
+                            Height = headerPresenterSize,
                             HorizontalScrollBarVisibility = ScrollBarVisibility.Hidden,
                             VerticalScrollBarVisibility = ScrollBarVisibility.Disabled,
                             [DockPanel.DockProperty] = Dock.Top,

tests/Avalonia.Controls.TreeDataGrid.Tests/TreeDataGridTests_Flat.cs

@@ -494,6 +494,107 @@ static void AssertRealizedCells(TreeDataGrid target)
             target.UpdateLayout();
             AssertRealizedCells(target);
         }
+
+        [AvaloniaFact(Timeout = 10000)]
+        public void Ensure_That_Falling_Back_To_The_RowPresenters_ViewPort_Instead_Of_Estimating_Prevents_A_Full_Realization_of_Columns()
+        {
+            var (target, items) = CreateTarget(headerPresenterSize: 10, columns:
+            [
+                new TextColumn<Model, int>("ID", x => x.Id),
+                new TextColumn<Model, string?>("Title1", x => x.Title),
+                new TextColumn<Model, string?>("Title2", x => x.Title),
+                new TextColumn<Model, string?>("Title3", x => x.Title),
+                new TextColumn<Model, string?>("Title4", x => x.Title),
+                new TextColumn<Model, string?>("Title5", x => x.Title),
+                new TextColumn<Model, string?>("Title6", x => x.Title),
+                new TextColumn<Model, string?>("Title7", x => x.Title),
+                new TextColumn<Model, string?>("Title8", x => x.Title),
+                new TextColumn<Model, string?>("Title9", x => x.Title),
+                new TextColumn<Model, string?>("Title10", x => x.Title),
+                new TextColumn<Model, string?>("Title11", x => x.Title),
+                new TextColumn<Model, string?>("Title12", x => x.Title),
+                new TextColumn<Model, string?>("Title13", x => x.Title),
+                new TextColumn<Model, string?>("Title14", x => x.Title),
+                new TextColumn<Model, string?>("Title15", x => x.Title),
+            ]);
+
+            // Scroll all the way to the right.
+            target.Scroll!.Offset = target.Scroll.Offset.WithX(target.Scroll.Extent.Width);
+            target.UpdateLayout();
+            Dispatcher.UIThread.RunJobs();
+
+            // Replace the source, so that non  of the columns have an actual width.
+            var newSource = new FlatTreeDataGridSource<Model>(items)
+            {
+                Columns =
+                {
+                    new TextColumn<Model, int>("ID", x => x.Id),
+                    new TextColumn<Model, string?>("Title1", x => x.Title),
+                    new TextColumn<Model, string?>("Title2", x => x.Title),
+                    new TextColumn<Model, string?>("Title3", x => x.Title),
+                    new TextColumn<Model, string?>("Title4", x => x.Title),
+                    new TextColumn<Model, string?>("Title5", x => x.Title),
+                    new TextColumn<Model, string?>("Title6", x => x.Title),
+                    new TextColumn<Model, string?>("Title7", x => x.Title),
+                    new TextColumn<Model, string?>("Title8", x => x.Title),
+                    new TextColumn<Model, string?>("Title9", x => x.Title),
+                    new TextColumn<Model, string?>("Title10", x => x.Title),
+                    new TextColumn<Model, string?>("Title11", x => x.Title),
+                    new TextColumn<Model, string?>("Title12", x => x.Title),
+                    new TextColumn<Model, string?>("Title13", x => x.Title),
+                    new TextColumn<Model, string?>("Title14", x => x.Title),
+                    new TextColumn<Model, string?>("Title15", x => x.Title),
+                }
+            };
+
+            target.Source = newSource;
+            Dispatcher.UIThread.RunJobs();
+
+            Assert.True(double.IsNaN(newSource.Columns[1].ActualWidth));
+            Assert.True(double.IsNaN(newSource.Columns[2].ActualWidth));
+            Assert.True(double.IsNaN(newSource.Columns[3].ActualWidth));
+            Assert.True(double.IsNaN(newSource.Columns[4].ActualWidth));
+
+            AssertRealizedCells(target, (columnHeaders, cells) =>
+            {
+                Assert.Equal(4, cells.Count);
+                Assert.Equal(12, cells[0].ColumnIndex);
+                Assert.Equal(13, cells[1].ColumnIndex);
+                Assert.Equal(14, cells[2].ColumnIndex);
+                Assert.Equal(15, cells[3].ColumnIndex);
+
+                Assert.Equal(12, columnHeaders[0].ColumnIndex);
+                Assert.Equal(13, columnHeaders[1].ColumnIndex);
+                Assert.Equal(14, columnHeaders[2].ColumnIndex);
+                Assert.Equal(15, columnHeaders[3].ColumnIndex);
+
+                Assert.Equal(columnHeaders[0].Bounds.Left, cells[0].Bounds.Left);
+                Assert.Equal(columnHeaders[1].Bounds.Left, cells[1].Bounds.Left);
+                Assert.Equal(columnHeaders[2].Bounds.Left, cells[2].Bounds.Left);
+                Assert.Equal(columnHeaders[3].Bounds.Left, cells[3].Bounds.Left);
+            });
+            return;
+
+            static void AssertRealizedCells(TreeDataGrid target, Action<List<TreeDataGridColumnHeader>, List<TreeDataGridCell>> assert)
+            {
+                var columnHeaders = target.ColumnHeadersPresenter!.GetRealizedElements()
+                    .Cast<TreeDataGridColumnHeader>()
+                    .OrderBy(x=>x.ColumnIndex)
+                    .ToList();
+
+                var rows = target.RowsPresenter!.GetRealizedElements().Cast<TreeDataGridRow>();
+
+                foreach (var row in rows)
+                {
+                    var cells = row.CellsPresenter!.GetRealizedElements()
+                        .Cast<TreeDataGridCell>()
+                        .OrderBy(x => x.ColumnIndex)
+                        .ToList();
+
+                    assert(columnHeaders, cells);
+                }
+            }
+        }
 
         [AvaloniaFact(Timeout = 10000)]
         public void Should_Use_TextCell_StringFormat()
@@ -814,7 +915,8 @@ private static void AssertInteractionSelection(TreeDataGrid target, params int[]
         private static (TreeDataGrid, AvaloniaList<Model>) CreateTarget(IEnumerable<Model>? models = null,
             IEnumerable<IColumn<Model>>? columns = null,
             int itemCount = 100,
-            bool runLayout = true)
+            bool runLayout = true,
+            double headerPresenterSize = 0)
         {
             AvaloniaList<Model>? items = null;
             if (models == null)
@@ -848,7 +950,7 @@ private static (TreeDataGrid, AvaloniaList<Model>) CreateTarget(IEnumerable<Mode
 
             var target = new TreeDataGrid
             {
-                Template = TestTemplates.TreeDataGridTemplate(),
+                Template = TestTemplates.TreeDataGridTemplate(headerPresenterSize),
                 Source = source,
             };