From 823f527e3d4c78de565fa397f6d72a38cf260eb2 Mon Sep 17 00:00:00 2001
From: Roger Qiu <roger.qiu@polyhack.io>
Date: Fri, 22 Apr 2022 18:18:37 +1000
Subject: [PATCH] Parameterising resource acquisitions

---
 src/types.ts        |  8 ++++----
 src/utils.ts        |  4 ++--
 tests/index.test.ts | 35 ++++++++++++++++++++++++++++++++++-
 3 files changed, 40 insertions(+), 7 deletions(-)

diff --git a/src/types.ts b/src/types.ts
index 072181d..0ae8482 100644
--- a/src/types.ts
+++ b/src/types.ts
@@ -1,10 +1,10 @@
-type ResourceAcquire<Resource = void> = () => Promise<
-  readonly [ResourceRelease, Resource?]
->;
+type ResourceAcquire<Resource> = (
+  resources: readonly any[],
+) => Promise<readonly [ResourceRelease, Resource?]>;
 
 type ResourceRelease = (e?: Error) => Promise<void>;
 
-type Resources<T extends readonly ResourceAcquire<any>[]> = {
+type Resources<T extends readonly ResourceAcquire<unknown>[]> = {
   [K in keyof T]: T[K] extends ResourceAcquire<infer R> ? R : never;
 };
 
diff --git a/src/utils.ts b/src/utils.ts
index 738854e..7472419 100644
--- a/src/utils.ts
+++ b/src/utils.ts
@@ -17,7 +17,7 @@ async function withF<
   let e_: Error | undefined;
   try {
     for (const acquire of acquires) {
-      const [release, resource] = await acquire();
+      const [release, resource] = await acquire(resources);
       releases.push(release);
       resources.push(resource);
     }
@@ -54,7 +54,7 @@ async function* withG<
   let e_: Error | undefined;
   try {
     for (const acquire of acquires) {
-      const [release, resource] = await acquire();
+      const [release, resource] = await acquire(resources);
       releases.push(release);
       resources.push(resource);
     }
diff --git a/tests/index.test.ts b/tests/index.test.ts
index e30bf8e..123672f 100644
--- a/tests/index.test.ts
+++ b/tests/index.test.ts
@@ -82,7 +82,7 @@ describe('index', () => {
     // Multiple resources outside requires type declaration
     const resourceAcquires6: [
       ResourceAcquire<number>,
-      ResourceAcquire,
+      ResourceAcquire<void>,
       ResourceAcquire<string>,
     ] = [
       async () => {
@@ -246,4 +246,37 @@ describe('index', () => {
     expect(acquireOrder).toStrictEqual([lock1, lock2]);
     expect(releaseOrder).toStrictEqual([lock2, lock1]);
   });
+  test('withF parameterised resources', async () => {
+    await withF(
+      [
+        async () => [async () => {}, 1],
+        async ([a]) => [async () => {}, a + 1],
+        async ([, b]) => [async () => {}, b + 1],
+      ],
+      async ([a, b, c]) => {
+        expect(a).toBe(1);
+        expect(b).toBe(2);
+        expect(c).toBe(3);
+      },
+    );
+  });
+  test('withG parameterised resources', async () => {
+    const g = withG(
+      [
+        async () => [async () => {}, 1],
+        async ([a]) => [async () => {}, a + 1],
+        async ([, b]) => [async () => {}, b + 1],
+      ],
+      async function* ([a, b, c]): AsyncGenerator<number> {
+        yield a;
+        yield b;
+        yield c;
+      },
+    );
+    let counter = 1;
+    for await (const r of g) {
+      expect(r).toBe(counter);
+      counter++;
+    }
+  });
 });