Merged PR 48250: Fix #130768

Caching dimensions from USS

Author: anton-trackman

Runtime/Extensions/NodeExtensions.cs

@@ -1,4 +1,5 @@
 using System;
+using System.Collections;
 using System.Collections.Generic;
 using System.IO;
 using System.Linq;
@@ -25,6 +26,8 @@ record Metadata(UIDocument document, UxmlAttribute uxml, string path);
         static Dictionary<VisualElement, Metadata> rootMetadata = new();
         static List<VisualElement> search = new(256);
         static Dictionary<VisualElement, string> cloneMap = new(256);
+
+        static Dictionary<string, (float width, float height)> cachedDimensions = new();
         #endregion
 
         #region Properties
@@ -43,6 +46,8 @@ public static void Initialize(UIDocument document, IEnumerable<IRootElement> tar
         {
             foreach (IRootElement target in targets.Where(x => x is not null))
                 Initialize(document, target);
+
+            CacheDimensions(document);
         }
         public static void Initialize(UIDocument document, IRootElement target)
         {
@@ -505,8 +510,23 @@ public static async void MarginMe(this VisualElement value)
         {
             static int GetLines(VisualElement value, VisualElement parent, float spacing, bool horizontalDirection)
             {
-                float valueSize = horizontalDirection ? value.resolvedStyle.width : value.resolvedStyle.height;
-                float parentSize = horizontalDirection ? parent.resolvedStyle.width : parent.resolvedStyle.height;
+                string className = value.GetClasses().Last();
+                float valueSize;
+                if (cachedDimensions.ContainsKey(className))
+                    valueSize = horizontalDirection
+                        ? cachedDimensions[className].width
+                        : cachedDimensions[className].height;
+                else
+                    valueSize = horizontalDirection ? value.resolvedStyle.width : value.resolvedStyle.height;
+
+                string parentClassName = parent.GetClasses().Last();
+                float parentSize;
+                if (cachedDimensions.ContainsKey(parentClassName))
+                    parentSize = horizontalDirection
+                        ? cachedDimensions[parentClassName].width
+                        : cachedDimensions[parentClassName].height;
+                else
+                    parentSize = horizontalDirection ? parent.resolvedStyle.width : parent.resolvedStyle.height;
 
                 return (valueSize.Invalid() || valueSize == 0) ? parent.childCount : (int)(parentSize / ((2 * valueSize + (spacing.Invalid() ? 0 : spacing)) / 2));
             }
@@ -764,6 +784,54 @@ VisualElement FindIn(VisualElement root)
                 }
             }
         }
+
+        static void CacheDimensions(UIDocument document)
+        {
+            static object GetValue(object target, string name) => target.GetType().GetField(name, BindingFlags.Instance | BindingFlags.Public | BindingFlags.NonPublic).GetValue(target);
+
+            StyleSheet uss = document.visualTreeAsset.stylesheets.First();
+
+            object[] complexSelectors = (object[])GetValue(uss, "m_ComplexSelectors");
+            object[] rules = (object[])GetValue(uss, "m_Rules");
+            object[] dimensions = ((IEnumerable)GetValue(uss, "dimensions")).Cast<object>().ToArray();
+
+            float width = 0, height = 0;
+            for (int i = 0; i < complexSelectors.Length; ++i)
+            {
+                object[] properties = (object[])GetValue(rules[i], "m_Properties");
+                foreach (object property in properties)
+                {
+                    string propertyName = GetValue(property, "m_Name").ToString();
+                    if (propertyName is not "width" and not "height") continue;
+
+                    object propertyValue = ((IEnumerable)GetValue(property, "m_Values")).Cast<object>().First();
+                    string valueType = GetValue(propertyValue, "m_ValueType").ToString();
+                    if (valueType is not "Dimension") continue;
+
+                    int valueIndex = (int)GetValue(propertyValue, "valueIndex");
+
+                    string unit = GetValue(dimensions[(int)valueIndex], "unit").ToString();
+                    if (unit is not "Pixel") continue;
+
+                    float value = (float)GetValue(dimensions[valueIndex], "value");
+
+                    width = propertyName is "width" ? value : width;
+                    height = propertyName is "height"? value : height;
+                }
+
+                if (width != 0 || height != 0)
+                {
+                    object selector = ((object[])GetValue(complexSelectors[i], "m_Selectors")).First();
+                    object part = ((IEnumerable)GetValue(selector, "m_Parts")).Cast<object>().First();
+                    string className = GetValue(part, "m_Value").ToString();
+
+                    cachedDimensions.Add(className, new(width, height));
+
+                    width = 0;
+                    height = 0;
+                }
+            }
+        }
         #endregion
     }
 }