src/utils.mjs: Add take() and range()

2 files changed, 137 insertions, 0 deletions

diff --git a/src/utils.mjs b/src/utils.mjs
index 9a600e7..19c8e37 100644
--- a/src/utils.mjs
+++ b/src/utils.mjs
@@ -58,3 +58,43 @@ export const first   = a => a[0];
 export const rest    = a => a.slice(1);
 export const butlast = a => a.slice(a, a.length - 1);
 export const last    = a => a[a.length - 1];
+
+export const take = function*(n, gen) {
+	let i = 0n;
+	for (const x of gen) {
+		if (i >= n) {
+			break;
+		}
+		i++;
+		yield x;
+	}
+};
+
+export const range = function*(x, y, step = 1n) {
+	if (x === undefined) {
+		let i = 0n;
+		while (true) {
+			yield i++;
+		}
+	}
+	const [from, to] = y === undefined ?
+		[0n, x] :
+		[x, y];
+	const fromn = BigInt(from);
+	const ton   = BigInt(to);
+	const stepn = BigInt(step);
+	if (stepn === 0n) {
+		while (true) {
+			yield fromn;
+		}
+	}
+	if (step < 0n) {
+		for (let i = fromn; i > ton; i+= stepn) {
+			yield i;
+		}
+	} else {
+		for (let i = fromn; i < ton; i += stepn) {
+			yield i;
+		}
+	}
+};
diff --git a/tests/js/utils.mjs b/tests/js/utils.mjs
index 33bdac2..69fd54a 100644
--- a/tests/js/utils.mjs
+++ b/tests/js/utils.mjs
@@ -15,6 +15,8 @@ import {
 	rest,
 	butlast,
 	last,
+	take,
+	range,
 } from "../../src/utils.mjs";
 
 
@@ -359,6 +361,99 @@ const test_last = async t => {
 	});
 };
 
+const test_take = async t => {
+	t.start("take()");
+
+	await t.test("example usage", () => {
+		assert.deepEqual(Array.from(take(3, [1, 2, 3, 4, 5, 6])), [1, 2, 3]);
+		assert.deepEqual([...take(3, [1, 2, 3, 4, 5, 6])], [1, 2, 3]);
+
+		const gen = function*() {
+			yield* [1, 2, 3, 4, 5, 6];
+		}
+		assert.deepEqual([...take(3, gen())], [1, 2, 3]);
+
+		assert.deepEqual([...take(3, [1, 2])], [1, 2]);
+
+		assert.deepEqual([...take(1,  [])],  []);
+		assert.deepEqual([...take(0,  [1])], []);
+		assert.deepEqual([...take(-1, [1])], []);
+	});
+};
+
+const test_range = async t => {
+	t.start("range()");
+
+	await t.test("empty values", () => {
+		const [] = range();
+
+		const [ a ] = range();
+		assert.equal(a, 0n);
+
+		const [ b, c, d, e ] = range();
+		assert.deepEqual(
+			[ b,  c,  d,  e  ],
+			[ 0n, 1n, 2n, 3n ],
+		);
+
+		assert.deepEqual(
+			Array.from(take(5, range())),
+			[ 0n, 1n, 2n, 3n, 4n ],
+		);
+
+		assert.deepEqual(
+			Array.from(take(1, range())),
+			[ 0n ],
+		);
+
+		assert.deepEqual(
+			Array.from(take(0, range())),
+			[],
+		);
+	});
+
+	await t.test("example usage", () => {
+		assert.deepEqual(
+			[...range(-5, 5)],
+			[ -5n, -4n, -3n, -2n, -1n, 0n, 1n, 2n, 3n, 4n ],
+		);
+		assert.deepEqual(
+			[...range(-100, 100, 10)],
+			[
+				-100n, -90n, -80n, -70n, -60n, -50n, -40n, -30n,
+				-20n, -10n, 0n, 10n, 20n, 30n, 40n, 50n, 60n,
+				70n, 80n, 90n,
+			],
+		);
+
+		assert.deepEqual([...range(0, 4, 2)], [0n, 2n]);
+		assert.deepEqual([...range(0, 5, 2)], [0n, 2n, 4n]);
+		assert.deepEqual([...range(0, 6, 2)], [0n, 2n, 4n]);
+		assert.deepEqual([...range(0, 7, 2)], [0n, 2n, 4n, 6n]);
+
+		assert.deepEqual(
+			[...range(100, 0, -10)],
+			[ 100n, 90n, 80n, 70n, 60n, 50n, 40n, 30n, 20n, 10n ]
+		);
+		assert.deepEqual(
+			[...range(10, -10, -1)],
+			[
+				10n, 9n, 8n, 7n, 6n, 5n, 4n, 3n, 2n, 1n, 0n,
+				-1n, -2n, -3n, -4n, -5n, -6n, -7n, -8n, -9n,
+			],
+		);
+
+		assert.deepEqual(
+			[...take(3, range(1, 10, 0))],
+			[ 1n, 1n, 1n ],
+		);
+		assert.deepEqual(
+			[...take(3, range(10, 1, 0))],
+			[ 10n, 10n, 10n ],
+		);
+	});
+};
+
 
 await runner.runTests([
 	test_keys,
@@ -374,4 +469,6 @@ await runner.runTests([
 	test_rest,
 	test_butlast,
 	test_last,
+	test_take,
+	test_range,
 ]);