From d4f6f86cfd75db56da8b3fc387cffa354dc17366 Mon Sep 17 00:00:00 2001
From: Hung Thai <hungth.it@gmail.com>
Date: Mon, 21 Aug 2023 11:32:30 +0200
Subject: [PATCH] Record computation time

Record computation time along with re-computation count
---
 src/index.ts          | 18 +++++++++++++++--
 src/types.ts          |  4 +++-
 test/reselect.spec.ts | 45 ++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 63 insertions(+), 4 deletions(-)

diff --git a/src/index.ts b/src/index.ts
index 66067705c..6b0df584c 100644
--- a/src/index.ts
+++ b/src/index.ts
@@ -46,6 +46,14 @@ export function setInputStabilityCheckEnabled(enabled: StabilityCheck) {
   globalStabilityCheck = enabled
 }
 
+function getTime() {
+  if (performance && typeof performance.now === 'function') {
+    return performance.now()
+  }
+
+  return Date.now()
+}
+
 function getDependencies(funcs: unknown[]) {
   const dependencies = Array.isArray(funcs[0]) ? funcs[0] : funcs
 
@@ -79,6 +87,7 @@ export function createSelectorCreator<
 ) {
   const createSelector = (...funcs: Function[]) => {
     let recomputations = 0
+    let computationTime = 0
     let lastResult: unknown
 
     // Due to the intricacies of rest params, we can't do an optional arg after `...funcs`.
@@ -121,8 +130,12 @@ export function createSelectorCreator<
     const memoizedResultFunc = memoize(
       function recomputationWrapper() {
         recomputations++
+        const computationStartTime = getTime()
         // apply arguments instead of spreading for performance.
-        return resultFunc!.apply(null, arguments)
+        const resultFuncOutput = resultFunc!.apply(null, arguments)
+        computationTime += getTime() - computationStartTime
+
+        return resultFuncOutput
       } as F,
       ...finalMemoizeOptions
     )
@@ -204,7 +217,8 @@ export function createSelectorCreator<
       dependencies,
       lastResult: () => lastResult,
       recomputations: () => recomputations,
-      resetRecomputations: () => (recomputations = 0)
+      resetRecomputations: () => (recomputations = 0, computationTime = 0),
+      computationTime: () => computationTime
     })
 
     return selector
diff --git a/src/types.ts b/src/types.ts
index c1ca68641..b4228a87b 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -37,7 +37,9 @@ export interface OutputSelectorFields<Combiner extends UnknownFunction, Keys> {
   dependencies: SelectorArray
   /** Counts the number of times the output has been recalculated */
   recomputations: () => number
-  /** Resets the count of recomputations count to 0 */
+  /** Total computation time of everytime the output has been recalculated */
+  computationTime: () => number
+  /** Resets the count of recomputations count and computation time to 0 */
   resetRecomputations: () => number
 }
 
diff --git a/test/reselect.spec.ts b/test/reselect.spec.ts
index 47c4c00ce..c69df41c4 100644
--- a/test/reselect.spec.ts
+++ b/test/reselect.spec.ts
@@ -9,7 +9,7 @@ import {
   weakMapMemoize
 } from 'reselect'
 import lodashMemoize from 'lodash/memoize'
-import { vi } from 'vitest'
+import { afterAll, afterEach, beforeEach, describe, vi } from 'vitest'
 import { configureStore, createSlice, PayloadAction } from '@reduxjs/toolkit'
 
 // Construct 1E6 states for perf test outside of the perf test so as to not change the execute time of the test function
@@ -101,6 +101,49 @@ describe('Basic selector behavior', () => {
     )
   })
 
+  describe('computation time', () => {
+    const performanceSpy = vi.spyOn(performance, 'now')
+
+    beforeEach(() => {
+      performanceSpy
+        .mockReturnValueOnce(1500)
+        .mockReturnValueOnce(3000)
+        .mockReturnValueOnce(4500)
+        .mockReturnValueOnce(6000)
+    })
+
+    it('should records computation time for every re-computation', () => {
+      const selector = createSelector(
+        (state: StateAB) => state.a,
+        (state: StateAB) => state.b,
+        (a, b) => a + b
+      )
+      const state1 = { a: 1, b: 2 }
+      selector(state1)
+      expect(selector.computationTime()).toBe(1500)
+      const state2 = { a: 3, b: 2 }
+      selector(state2)
+      expect(selector.computationTime()).toBe(3000)
+    })
+
+    it('should reset computation time when calling resetRecomputations', () => {
+      const selector = createSelector(
+        (state: StateAB) => state.a,
+        (state: StateAB) => state.b,
+        (a, b) => a + b
+      )
+      const state1 = { a: 1, b: 2 }
+      selector(state1)
+      expect(selector.computationTime()).toBe(1500)
+      selector.resetRecomputations()
+      expect(selector.computationTime()).toBe(0)
+    })
+
+    afterEach(() => {
+      performanceSpy.mockClear()
+    })
+  })
+
   describe('performance checks', () => {
     const originalEnv = process.env.NODE_ENV