Fixed text layout buffer reset for Stack and Dictionary (+ tests)

src/Avalonia.Base/Media/TextFormatting/FormattingBufferHelper.cs

@@ -1,4 +1,5 @@
 using System.Collections.Generic;
+using System.Numerics;
 using System.Runtime.CompilerServices;
 using Avalonia.Utilities;
 
@@ -13,7 +14,7 @@ namespace Avalonia.Media.TextFormatting
         {
             arrayBuilder.Clear();
 
-            if (IsBufferTooLarge<T>(arrayBuilder.Capacity))
+            if (IsBufferTooLarge<T>((uint) arrayBuilder.Capacity))
             {
                 arrayBuilder = default;
             }
@@ -23,7 +24,7 @@ namespace Avalonia.Media.TextFormatting
         {
             list.Clear();
 
-            if (IsBufferTooLarge<T>(list.Capacity))
+            if (IsBufferTooLarge<T>((uint) list.Capacity))
             {
                 list.TrimExcess();
             }
@@ -31,9 +32,11 @@ namespace Avalonia.Media.TextFormatting
 
         public static void ClearThenResetIfTooLarge<T>(Stack<T> stack)
         {
+            var approximateCapacity = RoundUpToPowerOf2((uint)stack.Count);
+
             stack.Clear();
 
-            if (IsBufferTooLarge<T>(stack.Count))
+            if (IsBufferTooLarge<T>(approximateCapacity))
             {
                 stack.TrimExcess();
             }
@@ -42,10 +45,12 @@ namespace Avalonia.Media.TextFormatting
         public static void ClearThenResetIfTooLarge<TKey, TValue>(ref Dictionary<TKey, TValue> dictionary)
             where TKey : notnull
         {
+            var approximateCapacity = RoundUpToPowerOf2((uint)dictionary.Count);
+
             dictionary.Clear();
 
             // dictionary is in fact larger than that: it has entries and buckets, but let's only count our data here
-            if (IsBufferTooLarge<KeyValuePair<TKey, TValue>>(dictionary.Count))
+            if (IsBufferTooLarge<KeyValuePair<TKey, TValue>>(approximateCapacity))
             {
 #if NET6_0_OR_GREATER
                 dictionary.TrimExcess();
@@ -56,7 +61,24 @@ namespace Avalonia.Media.TextFormatting
         }
 
         [MethodImpl(MethodImplOptions.AggressiveInlining)]
-        private static bool IsBufferTooLarge<T>(int length)
-            => (long)Unsafe.SizeOf<T>() * length > MaxKeptBufferSizeInBytes;
+        private static bool IsBufferTooLarge<T>(uint capacity)
+            => (long) (uint) Unsafe.SizeOf<T>() * capacity > MaxKeptBufferSizeInBytes;
+
+        [MethodImpl(MethodImplOptions.AggressiveInlining)]
+        private static uint RoundUpToPowerOf2(uint value)
+        {
+#if NET6_0_OR_GREATER
+            return BitOperations.RoundUpToPowerOf2(value);
+#else
+            // Based on https://graphics.stanford.edu/~seander/bithacks.html#RoundUpPowerOf2
+            --value;
+            value |= value >> 1;
+            value |= value >> 2;
+            value |= value >> 4;
+            value |= value >> 8;
+            value |= value >> 16;
+            return value + 1;
+#endif
+        }
     }
 }

src/Avalonia.Base/Media/TextFormatting/Unicode/BiDiData.cs

@@ -150,6 +150,8 @@ namespace Avalonia.Media.TextFormatting.Unicode
         /// </remarks>
         public void SaveTypes()
         {
+            _hasCleanState = false;
+
             // Capture the types data
             _savedClasses.Clear();
             _savedClasses.Add(_classes.AsSlice());
@@ -162,6 +164,8 @@ namespace Avalonia.Media.TextFormatting.Unicode
         /// </summary>
         public void RestoreTypes()
         {
+            _hasCleanState = false;
+
             _classes.Clear();
             _classes.Add(_savedClasses.AsSlice());
             _pairedBracketTypes.Clear();

tests/Avalonia.Base.UnitTests/Media/TextFormatting/FormattingBufferHelperTests.cs

@@ -0,0 +1,151 @@
+using System;
+using System.Collections.Generic;
+using System.Reflection;
+using Avalonia.Media.TextFormatting;
+using Avalonia.Utilities;
+using Xunit;
+
+namespace Avalonia.Base.UnitTests.Media.TextFormatting
+{
+    public class FormattingBufferHelperTests
+    {
+        public static TheoryData<int> SmallSizes => new() { 1, 500, 10_000, 125_000 };
+        public static TheoryData<int> LargeSizes => new() { 500_000, 1_000_000 };
+
+        [Theory]
+        [MemberData(nameof(SmallSizes))]
+        public void Should_Keep_Small_Buffer_List(int itemCount)
+        {
+            var capacity = FillAndClearList(itemCount);
+
+            Assert.True(capacity >= itemCount);
+        }
+
+        [Theory]
+        [MemberData(nameof(LargeSizes))]
+        public void Should_Reset_Large_Buffer_List(int itemCount)
+        {
+            var capacity = FillAndClearList(itemCount);
+
+            Assert.Equal(0, capacity);
+        }
+
+        private static int FillAndClearList(int itemCount)
+        {
+            var list = new List<int>();
+
+            for (var i = 0; i < itemCount; ++i)
+            {
+                list.Add(i);
+            }
+
+            FormattingBufferHelper.ClearThenResetIfTooLarge(list);
+
+            return list.Capacity;
+        }
+
+        [Theory]
+        [MemberData(nameof(SmallSizes))]
+        public void Should_Keep_Small_Buffer_ArrayBuilder(int itemCount)
+        {
+            var capacity = FillAndClearArrayBuilder(itemCount);
+
+            Assert.True(capacity >= itemCount);
+        }
+
+        [Theory]
+        [MemberData(nameof(LargeSizes))]
+        public void Should_Reset_Large_Buffer_ArrayBuilder(int itemCount)
+        {
+            var capacity = FillAndClearArrayBuilder(itemCount);
+
+            Assert.Equal(0, capacity);
+        }
+
+        private static int FillAndClearArrayBuilder(int itemCount)
+        {
+            var arrayBuilder = new ArrayBuilder<int>();
+
+            for (var i = 0; i < itemCount; ++i)
+            {
+                arrayBuilder.AddItem(i);
+            }
+
+            FormattingBufferHelper.ClearThenResetIfTooLarge(ref arrayBuilder);
+
+            return arrayBuilder.Capacity;
+        }
+
+        [Theory]
+        [MemberData(nameof(SmallSizes))]
+        public void Should_Keep_Small_Buffer_Stack(int itemCount)
+        {
+            var capacity = FillAndClearStack(itemCount);
+
+            Assert.True(capacity >= itemCount);
+        }
+
+        [Theory]
+        [MemberData(nameof(LargeSizes))]
+        public void Should_Reset_Large_Buffer_Stack(int itemCount)
+        {
+            var capacity = FillAndClearStack(itemCount);
+
+            Assert.Equal(0, capacity);
+        }
+
+        private static int FillAndClearStack(int itemCount)
+        {
+            var stack = new Stack<int>();
+
+            for (var i = 0; i < itemCount; ++i)
+            {
+                stack.Push(i);
+            }
+
+            FormattingBufferHelper.ClearThenResetIfTooLarge(stack);
+
+            var array = (Array) stack.GetType()
+                .GetField("_array", BindingFlags.NonPublic | BindingFlags.Instance)!
+                .GetValue(stack)!;
+
+            return array.Length;
+        }
+
+        [Theory]
+        [MemberData(nameof(SmallSizes))]
+        public void Should_Keep_Small_Buffer_Dictionary(int itemCount)
+        {
+            var capacity = FillAndClearDictionary(itemCount);
+
+            Assert.True(capacity >= itemCount);
+        }
+
+        [Theory]
+        [MemberData(nameof(LargeSizes))]
+        public void Should_Reset_Large_Buffer_Dictionary(int itemCount)
+        {
+            var capacity = FillAndClearDictionary(itemCount);
+
+            Assert.True(capacity <= 3); // dictionary trims to the nearest prime starting with 3
+        }
+
+        private static int FillAndClearDictionary(int itemCount)
+        {
+            var dictionary = new Dictionary<int, int>();
+
+            for (var i = 0; i < itemCount; ++i)
+            {
+                dictionary.Add(i, i);
+            }
+
+            FormattingBufferHelper.ClearThenResetIfTooLarge(ref dictionary);
+
+            var array = (Array) dictionary.GetType()
+                .GetField("_entries", BindingFlags.NonPublic | BindingFlags.Instance)!
+                .GetValue(dictionary)!;
+
+            return array.Length;
+        }
+    }
+}