From cbf1810b72a9aec68b7f3a0a097476b4a4884d04 Mon Sep 17 00:00:00 2001
From: jquense <monastic.panic@gmail.com>
Date: Wed, 6 Dec 2023 12:57:21 -0500
Subject: [PATCH] feat: add leading, and maxWait options to debounce hooks

---
 src/useDebouncedCallback.ts        | 144 +++++++++++++++++++--
 src/useDebouncedState.ts           |  10 +-
 src/useDebouncedValue.ts           |  36 +++++-
 src/useTimeout.ts                  |   1 +
 test/useDebouncedCallback.test.tsx | 198 ++++++++++++++++++++++++++++-
 test/useDebouncedValue.test.tsx    |  59 ++++++++-
 6 files changed, 417 insertions(+), 31 deletions(-)

diff --git a/src/useDebouncedCallback.ts b/src/useDebouncedCallback.ts
index 2e09e59..18f6138 100644
--- a/src/useDebouncedCallback.ts
+++ b/src/useDebouncedCallback.ts
@@ -1,23 +1,141 @@
-import { useCallback } from 'react'
+import { useCallback, useMemo, useRef } from 'react'
 import useTimeout from './useTimeout'
+import useMounted from './useMounted'
+
+export interface UseDebouncedCallbackOptions {
+  wait: number
+  leading?: boolean
+  trailing?: boolean
+  maxWait?: number
+}
 
 /**
  * Creates a debounced function that will invoke the input function after the
- * specified delay.
+ * specified wait.
  *
  * @param fn a function that will be debounced
- * @param delay The milliseconds delay before invoking the function
+ * @param waitOrOptions a wait in milliseconds or a debounce configuration
  */
 export default function useDebouncedCallback<
-  TCallback extends (...args: any[]) => any
->(fn: TCallback, delay: number): (...args: Parameters<TCallback>) => void {
+  TCallback extends (...args: any[]) => any,
+>(
+  fn: TCallback,
+  waitOrOptions: number | UseDebouncedCallbackOptions,
+): (...args: Parameters<TCallback>) => void {
+  const lastCallTimeRef = useRef<number | null>(null)
+  const lastInvokeTimeRef = useRef(0)
+
+  const isTimerSetRef = useRef(false)
+  const lastArgsRef = useRef<unknown[] | null>(null)
+
+  const {
+    wait,
+    maxWait,
+    leading = false,
+    trailing = true,
+  } = typeof waitOrOptions === 'number'
+    ? ({ wait: waitOrOptions } as UseDebouncedCallbackOptions)
+    : waitOrOptions
+
   const timeout = useTimeout()
-  return useCallback(
-    (...args: any[]) => {
-      timeout.set(() => {
-        fn(...args)
-      }, delay)
-    },
-    [fn, delay],
-  )
+
+  return useMemo(() => {
+    const hasMaxWait = !!maxWait
+
+    function leadingEdge(time: number) {
+      // Reset any `maxWait` timer.
+      lastInvokeTimeRef.current = time
+
+      // Start the timer for the trailing edge.
+      isTimerSetRef.current = true
+      timeout.set(timerExpired, wait)
+
+      // Invoke the leading edge.
+      if (leading) {
+        invokeFunc(time)
+      }
+    }
+
+    function trailingEdge(time: number) {
+      isTimerSetRef.current = false
+
+      // Only invoke if we have `lastArgs` which means `func` has been
+      // debounced at least once.
+      if (trailing && lastArgsRef.current) {
+        return invokeFunc(time)
+      }
+
+      lastArgsRef.current = null
+    }
+
+    function timerExpired() {
+      var time = Date.now()
+
+      if (shouldInvoke(time)) {
+        return trailingEdge(time)
+      }
+
+      const timeSinceLastCall = time - (lastCallTimeRef.current ?? 0)
+      const timeSinceLastInvoke = time - lastInvokeTimeRef.current
+      const timeWaiting = wait - timeSinceLastCall
+
+      // Restart the timer.
+      timeout.set(
+        timerExpired,
+        hasMaxWait
+          ? Math.min(timeWaiting, maxWait - timeSinceLastInvoke)
+          : timeWaiting,
+      )
+    }
+
+    function invokeFunc(time: number) {
+      const args = lastArgsRef.current ?? []
+
+      lastArgsRef.current = null
+      lastInvokeTimeRef.current = time
+
+      return fn(...args)
+    }
+
+    function shouldInvoke(time: number) {
+      const timeSinceLastCall = time - (lastCallTimeRef.current ?? 0)
+      const timeSinceLastInvoke = time - lastInvokeTimeRef.current
+
+      // Either this is the first call, activity has stopped and we're at the
+      // trailing edge, the system time has gone backwards and we're treating
+      // it as the trailing edge, or we've hit the `maxWait` limit.
+      return (
+        lastCallTimeRef.current === null ||
+        timeSinceLastCall >= wait ||
+        timeSinceLastCall < 0 ||
+        (hasMaxWait && timeSinceLastInvoke >= maxWait)
+      )
+    }
+
+    return (...args: any[]) => {
+      const time = Date.now()
+      const isInvoking = shouldInvoke(time)
+
+      lastArgsRef.current = args
+      lastCallTimeRef.current = time
+
+      if (isInvoking) {
+        if (!isTimerSetRef.current) {
+          return leadingEdge(lastCallTimeRef.current)
+        }
+
+        if (hasMaxWait) {
+          // Handle invocations in a tight loop.
+          isTimerSetRef.current = true
+          setTimeout(timerExpired, wait)
+          return invokeFunc(lastCallTimeRef.current)
+        }
+      }
+
+      if (!isTimerSetRef.current) {
+        isTimerSetRef.current = true
+        setTimeout(timerExpired, wait)
+      }
+    }
+  }, [fn, wait, maxWait, leading, trailing])
 }
diff --git a/src/useDebouncedState.ts b/src/useDebouncedState.ts
index 18c463a..27be8ad 100644
--- a/src/useDebouncedState.ts
+++ b/src/useDebouncedState.ts
@@ -1,5 +1,7 @@
 import { useState, Dispatch, SetStateAction } from 'react'
-import useDebouncedCallback from './useDebouncedCallback'
+import useDebouncedCallback, {
+  UseDebouncedCallbackOptions,
+} from './useDebouncedCallback'
 
 /**
  * Similar to `useState`, except the setter function is debounced by
@@ -12,16 +14,16 @@ import useDebouncedCallback from './useDebouncedCallback'
  * ```
  *
  * @param initialState initial state value
- * @param delay The milliseconds delay before a new value is set
+ * @param delayOrOptions The milliseconds delay before a new value is set, or options object
  */
 export default function useDebouncedState<T>(
   initialState: T,
-  delay: number,
+  delayOrOptions: number | UseDebouncedCallbackOptions,
 ): [T, Dispatch<SetStateAction<T>>] {
   const [state, setState] = useState(initialState)
   const debouncedSetState = useDebouncedCallback<Dispatch<SetStateAction<T>>>(
     setState,
-    delay,
+    delayOrOptions,
   )
   return [state, debouncedSetState]
 }
diff --git a/src/useDebouncedValue.ts b/src/useDebouncedValue.ts
index 6e9193b..37e4565 100644
--- a/src/useDebouncedValue.ts
+++ b/src/useDebouncedValue.ts
@@ -1,6 +1,12 @@
-import { delay } from 'lodash'
-import { useEffect, useDebugValue } from 'react'
+import { useEffect, useDebugValue, useRef } from 'react'
 import useDebouncedState from './useDebouncedState'
+import { UseDebouncedCallbackOptions } from './useDebouncedCallback'
+
+const defaultIsEqual = (a: any, b: any) => a === b
+
+export type UseDebouncedValueOptions = UseDebouncedCallbackOptions & {
+  isEqual?: (a: any, b: any) => boolean
+}
 
 /**
  * Debounce a value change by a specified number of milliseconds. Useful
@@ -8,17 +14,33 @@ import useDebouncedState from './useDebouncedState'
  * to defer changes until the changes reach some level of infrequency.
  *
  * @param value
- * @param delayMs
+ * @param waitOrOptions
  * @returns
  */
-function useDebouncedValue<TValue>(value: TValue, delayMs = 500): TValue {
-  const [debouncedValue, setDebouncedValue] = useDebouncedState(value, delayMs)
+function useDebouncedValue<TValue>(
+  value: TValue,
+  waitOrOptions: number | UseDebouncedValueOptions = 500,
+): TValue {
+  const previousValueRef = useRef<TValue | null>(value)
+
+  const isEqual =
+    typeof waitOrOptions === 'object'
+      ? waitOrOptions.isEqual || defaultIsEqual
+      : defaultIsEqual
+
+  const [debouncedValue, setDebouncedValue] = useDebouncedState(
+    value,
+    waitOrOptions,
+  )
 
   useDebugValue(debouncedValue)
 
   useEffect(() => {
-    setDebouncedValue(value)
-  }, [value, delayMs])
+    if (!isEqual || !isEqual(previousValueRef.current, value)) {
+      previousValueRef.current = value
+      setDebouncedValue(value)
+    }
+  })
 
   return debouncedValue
 }
diff --git a/src/useTimeout.ts b/src/useTimeout.ts
index 279d05b..1c61d72 100644
--- a/src/useTimeout.ts
+++ b/src/useTimeout.ts
@@ -73,6 +73,7 @@ export default function useTimeout() {
     return {
       set,
       clear,
+      handleRef,
     }
   }, [])
 }
diff --git a/test/useDebouncedCallback.test.tsx b/test/useDebouncedCallback.test.tsx
index f7fa4e3..4bd9ae9 100644
--- a/test/useDebouncedCallback.test.tsx
+++ b/test/useDebouncedCallback.test.tsx
@@ -1,12 +1,17 @@
+/*! tests an impl adapted from https://github.com/xnimorz/use-debounce/blob/master/test/useDebouncedCallback.test.tsx itself adapted from lodash*/
+
 import useDebouncedCallback from '../src/useDebouncedCallback'
 import { renderHook, act } from '@testing-library/react-hooks'
 
 describe('useDebouncedCallback', () => {
-  it('should return a function that debounces input callback', () => {
+  beforeEach(() => {
     jest.useFakeTimers()
-    const spy = jest.fn()
+  })
 
-    const { result } = renderHook(() => useDebouncedCallback(spy, 500))
+  it('should return a function that debounces input callback', () => {
+    const callback = jest.fn()
+
+    const { result } = renderHook(() => useDebouncedCallback(callback, 500))
 
     act(() => {
       result.current(1)
@@ -14,11 +19,192 @@ describe('useDebouncedCallback', () => {
       result.current(3)
     })
 
-    expect(spy).not.toHaveBeenCalled()
+    expect(callback).not.toHaveBeenCalled()
 
     jest.runOnlyPendingTimers()
 
-    expect(spy).toHaveBeenCalledTimes(1)
-    expect(spy).toHaveBeenCalledWith(3)
+    expect(callback).toHaveBeenCalledTimes(1)
+    expect(callback).toHaveBeenCalledWith(3)
+  })
+
+  it('will call leading callback immediately (but only once, as trailing is set to false)', () => {
+    const callback = jest.fn()
+
+    const { result } = renderHook(() =>
+      useDebouncedCallback(callback, {
+        wait: 1000,
+        leading: true,
+        trailing: false,
+      }),
+    )
+
+    act(() => {
+      result.current(1)
+    })
+
+    expect(callback).toHaveBeenCalledTimes(1)
+
+    act(() => {
+      jest.runOnlyPendingTimers()
+    })
+
+    expect(callback).toHaveBeenCalledTimes(1)
+  })
+
+  it('will call leading callback as well as next debounced call', () => {
+    const callback = jest.fn()
+
+    const { result } = renderHook(() =>
+      useDebouncedCallback(callback, {
+        wait: 1000,
+        leading: true,
+      }),
+    )
+
+    act(() => {
+      result.current()
+      result.current()
+      result.current()
+    })
+
+    expect(callback).toHaveBeenCalledTimes(1)
+
+    act(() => {
+      jest.runOnlyPendingTimers()
+    })
+
+    expect(callback).toHaveBeenCalledTimes(2)
+  })
+
+  it('will call three callbacks if no debounced callbacks are pending', () => {
+    const callback = jest.fn()
+
+    const { result } = renderHook(() =>
+      useDebouncedCallback(callback, {
+        wait: 1000,
+        leading: true,
+      }),
+    )
+
+    act(() => {
+      result.current()
+      result.current()
+      result.current()
+
+      setTimeout(() => {
+        result.current()
+      }, 1001)
+    })
+
+    expect(callback).toHaveBeenCalledTimes(1)
+
+    act(() => {
+      jest.advanceTimersByTime(1001)
+    })
+
+    expect(callback).toHaveBeenCalledTimes(3)
+  })
+
+  it('will call a second leading callback if no debounced callbacks are pending with trailing false', () => {
+    const callback = jest.fn()
+
+    const { result } = renderHook(() =>
+      useDebouncedCallback(callback, {
+        wait: 1000,
+        leading: true,
+        trailing: false,
+      }),
+    )
+
+    act(() => {
+      result.current()
+
+      setTimeout(() => {
+        result.current()
+      }, 1001)
+    })
+
+    expect(callback).toHaveBeenCalledTimes(1)
+
+    act(() => {
+      jest.advanceTimersByTime(1001)
+    })
+
+    expect(callback).toHaveBeenCalledTimes(2)
+  })
+
+  it("won't call both on the leading edge and on the trailing edge if leading and trailing are set up to true and function call is only once", () => {
+    const callback = jest.fn()
+
+    const { result } = renderHook(() =>
+      useDebouncedCallback(callback, {
+        wait: 1000,
+        leading: true,
+      }),
+    )
+
+    act(() => {
+      result.current()
+    })
+
+    expect(callback).toHaveBeenCalledTimes(1)
+
+    act(() => {
+      jest.runAllTimers()
+    })
+
+    expect(callback).toHaveBeenCalledTimes(1)
+  })
+
+  it('will call both on the leading edge and on the trailing edge if leading and trailing are set up to true and there are more than 1 function call', () => {
+    const callback = jest.fn()
+
+    const { result } = renderHook(() =>
+      useDebouncedCallback(callback, {
+        wait: 1000,
+        leading: true,
+      }),
+    )
+
+    act(() => {
+      result.current()
+      result.current()
+    })
+
+    expect(callback).toHaveBeenCalledTimes(1)
+
+    act(() => {
+      jest.runAllTimers()
+    })
+
+    expect(callback).toHaveBeenCalledTimes(2)
+  })
+
+  it('will call callback if maxWait time exceed', () => {
+    const callback = jest.fn()
+
+    const { result } = renderHook(() =>
+      useDebouncedCallback(callback, {
+        wait: 500,
+        maxWait: 600,
+      }),
+    )
+
+    expect(callback).toHaveBeenCalledTimes(0)
+
+    act(() => {
+      result.current()
+      jest.advanceTimersByTime(400)
+    })
+
+    expect(callback).toHaveBeenCalledTimes(0)
+
+    act(() => {
+      result.current()
+
+      jest.advanceTimersByTime(400)
+    })
+
+    expect(callback).toHaveBeenCalledTimes(1)
   })
 })
diff --git a/test/useDebouncedValue.test.tsx b/test/useDebouncedValue.test.tsx
index d434a9b..7620523 100644
--- a/test/useDebouncedValue.test.tsx
+++ b/test/useDebouncedValue.test.tsx
@@ -4,9 +4,11 @@ import useDebouncedValue from '../src/useDebouncedValue'
 import { act, render } from '@testing-library/react'
 
 describe('useDebouncedValue', () => {
-  it('should return a function that debounces input callback', () => {
+  beforeEach(() => {
     jest.useFakeTimers()
+  })
 
+  it('should return a function that debounces input callback', () => {
     let count = 0
     function Wrapper({ value }) {
       const debouncedValue = useDebouncedValue(value, 500)
@@ -37,4 +39,59 @@ describe('useDebouncedValue', () => {
       expect(count).toBe(2)
     })
   })
+
+  it('will update value immediately if leading is set to true', () => {
+    function Wrapper({ text }) {
+      const value = useDebouncedValue(text, { wait: 1000, leading: true })
+      return <div role="test">{value}</div>
+    }
+    const { rerender, getByText } = render(<Wrapper text={'Hello'} />)
+
+    expect(getByText('Hello')).toBeTruthy()
+
+    act(() => {
+      rerender(<Wrapper text={'Hello world'} />)
+    })
+
+    // value should be set immediately by first leading call
+
+    expect(getByText('Hello world')).toBeTruthy()
+
+    act(() => {
+      rerender(<Wrapper text={'Hello again'} />)
+    })
+
+    // timeout shouldn't have been called yet after leading call was executed
+    expect(getByText('Hello world')).toBeTruthy()
+
+    act(() => {
+      jest.runAllTimers()
+    })
+
+    expect(getByText('Hello again')).toBeTruthy()
+  })
+
+  it('should use isEqual function if supplied', () => {
+    const isEqual = jest.fn((_left: string, _right: string): boolean => true)
+
+    function Wrapper({ text }) {
+      const value = useDebouncedValue(text, { wait: 1000, isEqual })
+      return <div role="test">{value}</div>
+    }
+
+    const { rerender, getByRole } = render(<Wrapper text={'Hello'} />)
+
+    expect(isEqual).toHaveBeenCalledTimes(1)
+
+    act(() => {
+      rerender(<Wrapper text="Test" />)
+
+      jest.runAllTimers()
+    })
+
+    expect(isEqual).toHaveBeenCalledTimes(2)
+    expect(isEqual).toHaveBeenCalledWith('Hello', 'Test')
+
+    expect(getByRole('test').textContent).toEqual('Hello')
+  })
 })