raf support weakset

src/raf.ts

@@ -1,27 +1,30 @@
+const hasWin = typeof window !== 'undefined';
+const hasRaf = hasWin && typeof window.requestAnimationFrame !== 'undefined';
+const hasWeakSet = hasWin && typeof window.WeakSet !== 'undefined';
+
 let raf = (callback: FrameRequestCallback) => +setTimeout(callback, 16);
 let caf = (num: number) => clearTimeout(num);
 
-if (typeof window !== 'undefined' && 'requestAnimationFrame' in window) {
+if (hasRaf) {
   raf = (callback: FrameRequestCallback) =>
     window.requestAnimationFrame(callback);
   caf = (handle: number) => window.cancelAnimationFrame(handle);
 }
 
 let rafUUID = 0;
 const rafIds = new Map<number, number>();
+const rafKeys = new WeakSet<number[]>();
 
-function cleanup(id: number) {
-  rafIds.delete(id);
-}
+const cleanupByMap = (id: number) => rafIds.delete(id);
 
-export default function wrapperRaf(callback: () => void, times = 1): number {
+const useRafByMap = (callback: () => void, times = 1): number => {
   rafUUID += 1;
   const id = rafUUID;
 
   function callRef(leftTimes: number) {
     if (leftTimes === 0) {
       // Clean up
-      cleanup(id);
+      cleanupByMap(id);
 
       // Trigger
       callback();
@@ -39,10 +42,54 @@ export default function wrapperRaf(callback: () => void, times = 1): number {
   callRef(times);
 
   return id;
-}
+};
 
-wrapperRaf.cancel = (id: number) => {
-  const realId = rafIds.get(id);
-  cleanup(realId);
+useRafByMap.cancel = (key: number) => {
+  const realId = rafIds.get(key);
+  cleanupByMap(key);
   return caf(realId);
 };
+
+const cleanupByWeakSet = (key: number[] | number): number => {
+  let oldKey = typeof key === "number" ? [+key] : key;
+  const [timeId] = oldKey || [];
+  if (timeId) {
+    caf(timeId);
+    rafKeys.delete(oldKey);
+    oldKey = null;
+  }
+  return timeId;
+};
+
+const useRafByWeakSet = (callback: () => void, times = 1): number => {
+  let key: number[];
+
+  function callRef(leftTimes: number) {
+    if (leftTimes === 0) {
+      // Clean up
+      cleanupByWeakSet(key);
+
+      // Trigger
+      callback();
+    } else {
+      // Next raf
+      key = [
+        raf(() => {
+          callRef(leftTimes - 1);
+        }),
+      ];
+
+      // Bind real raf id
+      rafKeys.add(key);
+    }
+  }
+
+  callRef(times);
+
+  return +key;
+};
+
+useRafByWeakSet.cancel = (key: number[]) => cleanupByWeakSet(key);
+const useRaf = hasWeakSet ? useRafByWeakSet : useRafByMap;
+
+export default useRaf;