[charts] Batch string size measurement

Author: bernardobelchior

packages/x-charts/src/ChartsXAxis/ChartsXAxis.tsx

@@ -8,7 +8,7 @@ import { useRtl } from '@mui/system/RtlProvider';
 import { clampAngle } from '../internals/clampAngle';
 import { useIsHydrated } from '../hooks/useIsHydrated';
 import { doesTextFitInRect, ellipsize } from '../internals/ellipsize';
-import { getStringSize } from '../internals/domUtils';
+import { getStringSize, warmUpStringCache } from '../internals/domUtils';
 import { useTicks, TickItemType } from '../hooks/useTicks';
 import { AxisConfig, ChartsXAxisProps, ComputedXAxis } from '../models/axis';
 import { getAxisUtilityClass } from '../ChartsAxis/axisClasses';
@@ -82,6 +82,14 @@ function getVisibleLabels(
   let previousTextLimit = 0;
   const direction = reverse ? -1 : 1;
 
+  /* Avoid warming up the cache for too many values because we know not all of them will fit, so we'd be doing useless work */
+  if (isMounted && xTicks.length < 100) {
+    warmUpStringCache(
+      xTicks.flatMap((t) => t.formattedValue?.split('\n')).filter((t) => t != null),
+      style,
+    );
+  }
+
   return new Set(
     xTicks.filter((item, labelIndex) => {
       const { offset, labelOffset } = item;
@@ -153,6 +161,14 @@ function shortenLabels(
     [leftBoundFactor, rightBoundFactor] = [rightBoundFactor, leftBoundFactor];
   }
 
+  // Measure strings so it's cached
+  warmUpStringCache(
+    Array.from(visibleLabels)
+      .map((item) => item.formattedValue)
+      .filter((item) => item != null),
+    tickLabelStyle,
+  );
+
   for (const item of visibleLabels) {
     if (item.formattedValue) {
       // That maximum width of the tick depends on its proximity to the axis bounds.

packages/x-charts/src/ChartsYAxis/ChartsYAxis.tsx

@@ -8,7 +8,7 @@ import { useRtl } from '@mui/system/RtlProvider';
 import { useIsHydrated } from '../hooks/useIsHydrated';
 import { getDefaultBaseline, getDefaultTextAnchor } from '../ChartsText/defaultTextPlacement';
 import { doesTextFitInRect, ellipsize } from '../internals/ellipsize';
-import { getStringSize } from '../internals/domUtils';
+import { getStringSize, warmUpStringCache } from '../internals/domUtils';
 import { TickItemType, useTicks } from '../hooks/useTicks';
 import { ChartDrawingArea, useDrawingArea } from '../hooks/useDrawingArea';
 import { AxisConfig, ChartsYAxisProps } from '../models/axis';
@@ -73,6 +73,14 @@ function shortenLabels(
     [topBoundFactor, bottomBoundFactor] = [bottomBoundFactor, topBoundFactor];
   }
 
+  /* Avoid warming up the cache for too many values because we know not all of them will fit, so we'd be doing useless work */
+  if (visibleLabels.length < 100) {
+    warmUpStringCache(
+      visibleLabels.flatMap((t) => t.formattedValue).filter((t) => t != null),
+      tickLabelStyle,
+    );
+  }
+
   for (const item of visibleLabels) {
     if (item.formattedValue) {
       // That maximum height of the tick depends on its proximity to the axis bounds.

packages/x-charts/src/internals/domUtils.ts

@@ -1,19 +1,19 @@
 // DOM utils taken from
 // https://github.com/recharts/recharts/blob/master/src/util/DOMUtils.ts
 
+import * as React from 'react';
+
 function isSsr(): boolean {
   return typeof window === 'undefined';
 }
 
-interface StringCache {
-  widthCache: Record<string, { width: number; height: number }>;
-  cacheCount: number;
-}
+type Size = { width: number; height: number };
+
+const cache = new Map<string, Size>();
+
+// @ts-ignore
+window.xChartsClearStringCache = () => cache.clear();
 
-const stringCache: StringCache = {
-  widthCache: {},
-  cacheCount: 0,
-};
 const MAX_CACHE_NUM = 2000;
 const SPAN_STYLE = {
   position: 'absolute',
@@ -45,7 +45,8 @@ const STYLE_LIST = [
   'marginTop',
   'marginBottom',
 ];
-export const MEASUREMENT_SPAN_ID = 'mui_measurement_span';
+
+const MEASUREMENT_DIV_ID = 'mui_measurement_div';
 
 /**
  *
@@ -114,58 +115,101 @@ export const getStringSize = (text: string | number, style: React.CSSProperties
   const styleString = getStyleString(style);
   const cacheKey = `${str}-${styleString}`;
 
-  if (stringCache.widthCache[cacheKey]) {
-    return stringCache.widthCache[cacheKey];
+  const result = cache.get(cacheKey);
+
+  if (result) {
+    return result;
   }
 
+  warmUpStringCache([text], style);
+
+  return cache.get(cacheKey) ?? { width: 0, height: 0 };
+};
+
+export function warmUpStringCache(
+  texts: Iterable<string | number>,
+  style: React.CSSProperties = {},
+) {
+  if (isSsr()) {
+    return;
+  }
+
+  const styleString = getStyleString(style);
+
   try {
-    let measurementSpan = document.getElementById(MEASUREMENT_SPAN_ID);
-    if (measurementSpan === null) {
-      measurementSpan = document.createElement('span');
-      measurementSpan.setAttribute('id', MEASUREMENT_SPAN_ID);
-      measurementSpan.setAttribute('aria-hidden', 'true');
-      document.body.appendChild(measurementSpan);
+    let measurementDiv = document.getElementById(MEASUREMENT_DIV_ID);
+    if (measurementDiv === null) {
+      measurementDiv = document.createElement('div');
+      measurementDiv.setAttribute('id', MEASUREMENT_DIV_ID);
+      measurementDiv.setAttribute('aria-hidden', 'true');
+      document.body.appendChild(measurementDiv);
     }
+
     // Need to use CSS Object Model (CSSOM) to be able to comply with Content Security Policy (CSP)
     // https://en.wikipedia.org/wiki/Content_Security_Policy
-    const measurementSpanStyle: Record<string, any> = { ...SPAN_STYLE, ...style };
+    const measurementDivStyle: Record<string, any> = { ...SPAN_STYLE, ...style };
 
-    Object.keys(measurementSpanStyle).map((styleKey) => {
-      (measurementSpan!.style as Record<string, any>)[camelToMiddleLine(styleKey)] =
-        autoCompleteStyle(styleKey, measurementSpanStyle[styleKey]);
+    Object.keys(style).map((styleKey) => {
+      (measurementDiv!.style as Record<string, any>)[camelToMiddleLine(styleKey)] =
+        autoCompleteStyle(styleKey, measurementDivStyle[styleKey]);
       return styleKey;
     });
-    measurementSpan.textContent = str;
-    const rect = measurementSpan.getBoundingClientRect();
-    const result = { width: rect.width, height: rect.height };
 
-    stringCache.widthCache[cacheKey] = result;
+    const spans = [];
+    for (const text of texts) {
+      const str = `${text}`;
+      const cacheKey = `${str}-${styleString}`;
 
-    if (stringCache.cacheCount + 1 > MAX_CACHE_NUM) {
-      stringCache.cacheCount = 0;
-      stringCache.widthCache = {};
-    } else {
-      stringCache.cacheCount += 1;
+      if (cache.has(cacheKey)) {
+        // If already cached, skip the measurement
+        continue;
+      }
+
+      const span = document.createElement('span');
+      span.textContent = str;
+      spans.push(span);
+    }
+
+    measurementDiv.replaceChildren(...spans);
+
+    let i = 0;
+    for (const text of texts) {
+      const str = `${text}`;
+      const cacheKey = `${str}-${styleString}`;
+
+      if (cache.has(cacheKey)) {
+        // The string is already cached. Do not increment the index because a span wasn't created for cached strings
+        continue;
+      }
+
+      const span = spans[i];
+      const rect = span.getBoundingClientRect();
+
+      cache.set(cacheKey, { width: rect.width, height: rect.height });
+
+      if (cache.size + 1 > MAX_CACHE_NUM) {
+        // This is very inefficient, we should be smarter about which entries to purge
+        cache.clear();
+      }
+      i += 1;
     }
 
     if (process.env.NODE_ENV === 'test') {
-      // In test environment, we clean the measurement span immediately
-      measurementSpan.textContent = '';
+      // In test environment, we clean the measurement div immediately
+      measurementDiv.replaceChildren();
     } else {
       if (domCleanTimeout) {
         clearTimeout(domCleanTimeout);
       }
       domCleanTimeout = setTimeout(() => {
         // Limit node cleaning to once per render cycle
-        measurementSpan.textContent = '';
+        measurementDiv.replaceChildren();
       }, 0);
     }
-
-    return result;
   } catch {
-    return { width: 0, height: 0 };
+    /* Intentionally do nothing */
   }
-};
+}
 
 // eslint-disable-next-line @typescript-eslint/naming-convention
 export function unstable_cleanupDOM() {