From 459433b582f4680c4670e6ebf21973dfda9c8085 Mon Sep 17 00:00:00 2001
From: zombiej <smith3816@gmail.com>
Date: Tue, 27 Apr 2021 17:37:47 +0800
Subject: [PATCH 1/2] feat: raf support times

---
 src/raf.ts        | 35 ++++++++++++++++++++++++++++++++---
 tests/raf.test.js | 38 ++++++++++++++++++++++++++++++++++++++
 2 files changed, 70 insertions(+), 3 deletions(-)

diff --git a/src/raf.ts b/src/raf.ts
index 25c2dd3c..96980a24 100644
--- a/src/raf.ts
+++ b/src/raf.ts
@@ -7,8 +7,37 @@ if (typeof window !== 'undefined' && 'requestAnimationFrame' in window) {
   caf = (handle: number) => window.cancelAnimationFrame(handle);
 }
 
-export default function wrapperRaf(callback: () => void): number {
-  return raf(callback);
+let rafUUID = 0;
+const rafIds = new Map<number, number>();
+
+export default function wrapperRaf(callback: () => void, times = 1): number {
+  rafUUID += 1;
+  const id = rafUUID;
+
+  function callRef(leftTimes: number) {
+    if (leftTimes === 0) {
+      // Clean up
+      rafIds.delete(id);
+
+      // Trigger
+      callback();
+    } else {
+      // Next raf
+      const realId = raf(() => {
+        callRef(leftTimes - 1);
+      });
+
+      // Bind real raf id
+      rafIds.set(id, realId);
+    }
+  }
+
+  callRef(times);
+
+  return id;
 }
 
-wrapperRaf.cancel = caf;
+wrapperRaf.cancel = (id: number) => {
+  const realId = rafIds.get(id);
+  return caf(realId);
+};
diff --git a/tests/raf.test.js b/tests/raf.test.js
index 374053e5..5a17076c 100644
--- a/tests/raf.test.js
+++ b/tests/raf.test.js
@@ -16,4 +16,42 @@ describe('raf', () => {
       done();
     });
   });
+
+  it('cancel', done => {
+    let bamboo = false;
+
+    const id = raf(() => {
+      bamboo = true;
+    }, 2);
+
+    raf.cancel(id);
+
+    requestAnimationFrame(() => {
+      requestAnimationFrame(() => {
+        requestAnimationFrame(() => {
+          expect(bamboo).toBeFalsy();
+          done();
+        });
+      });
+    });
+  });
+
+  it('multiple times', done => {
+    let bamboo = false;
+
+    raf(() => {
+      bamboo = true;
+    }, 2);
+
+    expect(bamboo).toBeFalsy();
+
+    requestAnimationFrame(() => {
+      expect(bamboo).toBeFalsy();
+
+      requestAnimationFrame(() => {
+        expect(bamboo).toBeTruthy();
+        done();
+      });
+    });
+  });
 });

From 1c4c83e032c37324c9c45a5b90d3e97105d01b9c Mon Sep 17 00:00:00 2001
From: zombiej <smith3816@gmail.com>
Date: Tue, 27 Apr 2021 17:39:58 +0800
Subject: [PATCH 2/2] chore: clean up

---
 src/raf.ts | 7 ++++++-
 1 file changed, 6 insertions(+), 1 deletion(-)

diff --git a/src/raf.ts b/src/raf.ts
index 96980a24..40e0ae6e 100644
--- a/src/raf.ts
+++ b/src/raf.ts
@@ -10,6 +10,10 @@ if (typeof window !== 'undefined' && 'requestAnimationFrame' in window) {
 let rafUUID = 0;
 const rafIds = new Map<number, number>();
 
+function cleanup(id: number) {
+  rafIds.delete(id);
+}
+
 export default function wrapperRaf(callback: () => void, times = 1): number {
   rafUUID += 1;
   const id = rafUUID;
@@ -17,7 +21,7 @@ export default function wrapperRaf(callback: () => void, times = 1): number {
   function callRef(leftTimes: number) {
     if (leftTimes === 0) {
       // Clean up
-      rafIds.delete(id);
+      cleanup(id);
 
       // Trigger
       callback();
@@ -39,5 +43,6 @@ export default function wrapperRaf(callback: () => void, times = 1): number {
 
 wrapperRaf.cancel = (id: number) => {
   const realId = rafIds.get(id);
+  cleanup(realId);
   return caf(realId);
 };