From 216589f00c01c30301fc9e7fcc0ced0065233d1a Mon Sep 17 00:00:00 2001
From: Sterling Camden <rusty.camden@gmail.com>
Date: Mon, 16 Dec 2024 13:14:27 -1000
Subject: [PATCH] Update use-debounced-callback.ts to add a flush method to the
 returned callback as well as give an option to simply flush on unmount

---
 .../use-debounced-callback.test.ts            | 19 ++++++++++++
 .../use-debounced-callback.ts                 | 30 ++++++++++++++++---
 2 files changed, 45 insertions(+), 4 deletions(-)

diff --git a/packages/@mantine/hooks/src/use-debounced-callback/use-debounced-callback.test.ts b/packages/@mantine/hooks/src/use-debounced-callback/use-debounced-callback.test.ts
index e177e9bb602..cf6d85ce1d3 100644
--- a/packages/@mantine/hooks/src/use-debounced-callback/use-debounced-callback.test.ts
+++ b/packages/@mantine/hooks/src/use-debounced-callback/use-debounced-callback.test.ts
@@ -26,4 +26,23 @@ describe('@mantine/hooks/use-debounced-callback', () => {
     jest.advanceTimersByTime(100);
     expect(callback).toHaveBeenCalledWith(3);
   });
+
+  it('can be flushed immediately', () => {
+    const callback = jest.fn();
+    const { result } = renderHook(() => useDebouncedCallback(callback, 100));
+    result.current(1);
+    result.current(2);
+    result.current(3);
+    result.current.flush();
+    expect(callback).toHaveBeenCalledWith(3);
+  });
+  it('can flush on unmount', () => {
+    const callback = jest.fn();
+    const { result, unmount } = renderHook(() => useDebouncedCallback(callback, {delay : 100, flushOnUnmount : true}));
+    result.current(1);
+    result.current(2);
+    result.current(3);
+    unmount();
+    expect(callback).toHaveBeenCalledWith(3);
+  });
 });
diff --git a/packages/@mantine/hooks/src/use-debounced-callback/use-debounced-callback.ts b/packages/@mantine/hooks/src/use-debounced-callback/use-debounced-callback.ts
index a0089afc6ee..fa950db17f3 100644
--- a/packages/@mantine/hooks/src/use-debounced-callback/use-debounced-callback.ts
+++ b/packages/@mantine/hooks/src/use-debounced-callback/use-debounced-callback.ts
@@ -1,19 +1,41 @@
 import { useCallback, useEffect, useRef } from 'react';
 import { useCallbackRef } from '../use-callback-ref/use-callback-ref';
 
+const noop = () => {};
+
 export function useDebouncedCallback<T extends (...args: any[]) => any>(
   callback: T,
-  delay: number
+  options: number | { delay: number; flushOnUnmount?: boolean }
 ) {
+  const delay = typeof options === 'number' ? options : options.delay;
+  const flushOnUnmount = typeof options === 'number' ? false : options.flushOnUnmount;
   const handleCallback = useCallbackRef(callback);
   const debounceTimerRef = useRef(0);
-  useEffect(() => () => window.clearTimeout(debounceTimerRef.current), []);
 
-  return useCallback(
+  const lastCallback = Object.assign(useCallback(
     (...args: Parameters<T>) => {
       window.clearTimeout(debounceTimerRef.current);
-      debounceTimerRef.current = window.setTimeout(() => handleCallback(...args), delay);
+      const flush = () => {
+        if (debounceTimerRef.current !== 0) {
+          debounceTimerRef.current = 0;
+          handleCallback(...args);
+        }
+      };
+      lastCallback.flush = flush;
+      debounceTimerRef.current = window.setTimeout(flush, delay);
     },
     [handleCallback, delay]
+  ), {flush : noop});
+
+  useEffect(
+    () => () => {
+      window.clearTimeout(debounceTimerRef.current);
+      if (flushOnUnmount) {
+        lastCallback.flush();
+      }
+    },
+    [lastCallback, flushOnUnmount]
   );
+
+  return lastCallback;
 }