From 919b01f92e52f7a274fc17fab3b6a62610332e48 Mon Sep 17 00:00:00 2001
From: Maciej Jastrzebski <mdjastrzebski@gmail.com>
Date: Mon, 21 Oct 2024 10:29:48 +0200
Subject: [PATCH] refactor: improve

---
 packages/measure/src/measure-renders.tsx | 18 ++++++++++++++++--
 1 file changed, 16 insertions(+), 2 deletions(-)

diff --git a/packages/measure/src/measure-renders.tsx b/packages/measure/src/measure-renders.tsx
index c053630c..749f34d6 100644
--- a/packages/measure/src/measure-renders.tsx
+++ b/packages/measure/src/measure-renders.tsx
@@ -60,7 +60,6 @@ async function measureRendersInternal(
   const testingLibrary = getTestingLibrary();
 
   showFlagsOutputIfNeeded();
-  ensurePreciseMeasurements();
 
   const runResults: RunResult[] = [];
   let hasTooLateRender = false;
@@ -68,6 +67,8 @@ async function measureRendersInternal(
   const renderJsonTrees: ElementJsonTree[] = [];
   let initialRenderCount = 0;
 
+  installPerformanceNow();
+
   for (let i = 0; i < runs + warmupRuns; i += 1) {
     let duration = 0;
     let count = 0;
@@ -119,6 +120,8 @@ async function measureRendersInternal(
     runResults.push({ duration, count });
   }
 
+  restorePerformanceNow();
+
   if (hasTooLateRender) {
     const testName = expect.getState().currentTestName;
     logger.warn(
@@ -149,7 +152,18 @@ export function buildUiToRender(
   return Wrapper ? <Wrapper>{uiWithProfiler}</Wrapper> : uiWithProfiler;
 }
 
+let originalNow: () => number;
+
 //https://github.com/facebook/react/blob/65a56d0e99261481c721334a3ec4561d173594cd/packages/react-devtools-shared/src/backend/fiber/renderer.js#L294
-function ensurePreciseMeasurements() {
+function installPerformanceNow() {
+  originalNow = globalThis.performance.now;
   globalThis.performance.now = () => perf.now();
 }
+
+function restorePerformanceNow() {
+  if (originalNow == null) {
+    throw new Error('Called "restorePerformanceNow" without calling "installPerformanceNow" first');
+  }
+
+  globalThis.performance.now = originalNow;
+}