Fix InlineUIContainer focus (#14590)

* Fix TextBlock MeasureOverride visual child handling

* Make sure InlineUIContainer's child retains focus on measure

* Resolve merge error

Author: Gillibald

src/Avalonia.Controls/Documents/IInlineHost.cs

@@ -1,9 +1,12 @@
-using Avalonia.LogicalTree;
+using Avalonia.Collections;
+using Avalonia.LogicalTree;
 
 namespace Avalonia.Controls.Documents
 {
     internal interface IInlineHost : ILogical
     {
         void Invalidate();
+
+        IAvaloniaList<Visual> VisualChildren { get; }
     }
 }

src/Avalonia.Controls/Documents/InlineCollection.cs

@@ -26,13 +26,27 @@ public InlineCollection()
                 x =>
                 {
                     x.InlineHost = InlineHost;
+
                     LogicalChildren?.Add(x);
+
+                    if (x is InlineUIContainer container)
+                    {
+                        InlineHost?.VisualChildren.Add(container.Child);
+                    }
+
                     Invalidate();
                 },
                 x =>
                 {
                     LogicalChildren?.Remove(x);
-                    x.InlineHost = InlineHost;
+
+                    if(x is InlineUIContainer container)
+                    {
+                        InlineHost?.VisualChildren.Remove(container.Child);
+                    }
+
+                    x.InlineHost = null;
+
                     Invalidate();
                 },
                 () => throw new NotSupportedException());

src/Avalonia.Controls/TextBlock.cs

@@ -2,6 +2,7 @@
 using System.Collections.Generic;
 using System.Diagnostics;
 using Avalonia.Automation.Peers;
+using Avalonia.Collections;
 using Avalonia.Controls.Documents;
 using Avalonia.Layout;
 using Avalonia.Media;
@@ -685,16 +686,18 @@ protected virtual TextLayout CreateTextLayout(string? text)
         /// Invalidates <see cref="TextLayout"/>.
         /// </summary>
         protected void InvalidateTextLayout()
+        {
+            InvalidateMeasure();
+        }
+
+        protected override void OnMeasureInvalidated()
         {
             _textLayout?.Dispose();
             _textLayout = null;
-            
-            VisualChildren.Clear();
 
             _textRuns = null;
 
-            InvalidateVisual();
-            InvalidateMeasure();
+            base.OnMeasureInvalidated();
         }
 
         protected override Size MeasureOverride(Size availableSize)
@@ -703,15 +706,11 @@ protected override Size MeasureOverride(Size availableSize)
             var padding = LayoutHelper.RoundLayoutThickness(Padding, scale, scale);
 
             _constraint = availableSize.Deflate(padding);
-            _textLayout?.Dispose();
-            _textLayout = null;
 
             var inlines = Inlines;
 
             if (HasComplexContent)
             {
-                VisualChildren.Clear();
-
                 var textRuns = new List<TextRun>();
 
                 foreach (var inline in inlines!)
@@ -720,21 +719,6 @@ protected override Size MeasureOverride(Size availableSize)
                 }
 
                 _textRuns = textRuns;
-
-                foreach (var textLine in TextLayout.TextLines)
-                {
-                    foreach (var run in textLine.TextRuns)
-                    {
-                        if (run is DrawableTextRun drawable)
-                        {
-                            if (drawable is EmbeddedControlRun controlRun
-                                && controlRun.Control is Control control)
-                            {
-                                VisualChildren.Add(control);
-                            }
-                        }
-                    }
-                }
             }
 
             var width = TextLayout.OverhangLeading + TextLayout.WidthIncludingTrailingWhitespace + TextLayout.OverhangTrailing;
@@ -847,26 +831,30 @@ protected override void OnPropertyChanged(AvaloniaPropertyChangedEventArgs chang
 
         private void OnInlinesChanged(InlineCollection? oldValue, InlineCollection? newValue)
         {
+            VisualChildren.Clear();
+
             if (oldValue is not null)
             {
                 oldValue.LogicalChildren = null;
                 oldValue.InlineHost = null;
-                oldValue.Invalidated -= (s, e) => InvalidateTextLayout();
+                oldValue.Invalidated -= (s, e) => InvalidateMeasure();
             }
 
             if (newValue is not null)
             {
                 newValue.LogicalChildren = LogicalChildren;
                 newValue.InlineHost = this;
-                newValue.Invalidated += (s, e) => InvalidateTextLayout();
+                newValue.Invalidated += (s, e) => InvalidateMeasure();
             }
         }
 
         void IInlineHost.Invalidate()
         {
-            InvalidateTextLayout();
+            InvalidateMeasure();
         }
 
+        IAvaloniaList<Visual> IInlineHost.VisualChildren => VisualChildren;
+
         protected readonly record struct SimpleTextSource : ITextSource
         {
             private readonly string _text;

tests/Avalonia.Controls.UnitTests/TextBlockTests.cs

@@ -2,6 +2,7 @@
 using Avalonia.Controls.Documents;
 using Avalonia.Controls.Templates;
 using Avalonia.Data;
+using Avalonia.Input;
 using Avalonia.Media;
 using Avalonia.Metadata;
 using Avalonia.Rendering;
@@ -50,6 +51,55 @@ public void Changing_InlinesCollection_Should_Invalidate_Measure()
             }
         }
 
+        [Fact]
+        public void Can_Call_Measure_Without_InvalidateTextLayout()
+        {
+            using (UnitTestApplication.Start(TestServices.MockPlatformRenderInterface))
+            {
+                var target = new TextBlock();
+
+                target.Inlines.Add(new TextBox { Text = "Hello"});
+
+                target.Measure(Size.Infinity);
+
+                target.InvalidateMeasure();
+
+                target.Measure(Size.Infinity);
+            }
+        }
+
+        [Fact]
+        public void Embedded_Control_Should_Keep_Focus()
+        {
+            using (UnitTestApplication.Start(TestServices.RealFocus))
+            {
+                var target = new TextBlock();
+
+                var root = new TestRoot
+                {
+                    Child = target
+                };
+
+                var textBox = new TextBox { Text = "Hello", Template = TextBoxTests.CreateTemplate() };
+
+                target.Inlines.Add(textBox);
+
+                target.Measure(Size.Infinity);
+
+                textBox.Focus();
+
+                Assert.Same(textBox, root.FocusManager.GetFocusedElement());
+
+                target.InvalidateMeasure();
+
+                Assert.Same(textBox, root.FocusManager.GetFocusedElement());
+
+                target.Measure(Size.Infinity);
+
+                Assert.Same(textBox, root.FocusManager.GetFocusedElement());
+            }
+        }
+
         [Fact]
         public void Changing_Inlines_Properties_Should_Invalidate_Measure()
         {
@@ -115,6 +165,25 @@ public void Changing_Inlines_Should_Reset_Inlines_Parent()
             }
         }
 
+        [Fact]
+        public void Changing_Inlines_Should_Reset_VisualChildren()
+        {
+            using (UnitTestApplication.Start(TestServices.MockPlatformRenderInterface))
+            {
+                var target = new TextBlock();
+
+                target.Inlines.Add(new Border());
+
+                target.Measure(Size.Infinity);
+
+                Assert.NotEmpty(target.VisualChildren);
+
+                target.Inlines = null;
+
+                Assert.Empty(target.VisualChildren);
+            }
+        }
+
         [Fact]
         public void Changing_Inlines_Should_Reset_InlineUIContainer_VisualParent_On_Measure()
         {

tests/Avalonia.Controls.UnitTests/TextBoxTests.cs

@@ -1484,7 +1484,7 @@ public void Should_Throw_ArgumentOutOfRange()
             textShaperImpl: new HeadlessTextShaperStub(), 
             fontManagerImpl: new HeadlessFontManagerStub());
 
-        private IControlTemplate CreateTemplate()
+        internal static IControlTemplate CreateTemplate()
         {
             return new FuncControlTemplate<TextBox>((control, scope) =>
             new ScrollViewer