make extendMainline more composable

niklasf · Feb 5, 2024 · bc6ba6f · bc6ba6f
1 parent 763e825
commit bc6ba6f
Show file tree

Hide file tree

Showing 2 changed files with 31 additions and 26 deletions.
diff --git a/src/pgn.test.ts b/src/pgn.test.ts
@@ -6,9 +6,11 @@ import {
   ChildNode,
   defaultGame,
   emptyHeaders,
-  extendMainline,
+  extend,
   Game,
   isChildNode,
+  mainline,
+  mainlineEnd,
   makeComment,
   makePgn,
   Node,
@@ -83,8 +85,7 @@ test('make pgn', () => {
 
 test('extend mainline', () => {
   const game: Game<PgnNodeData> = defaultGame(emptyHeaders);
-  const mainline = 'e4 d5 a3 h6 Bg5'.split(' ').map(san => ({ san }));
-  extendMainline(game, mainline);
+  extend(mainlineEnd(game.moves), 'e4 d5 a3 h6 Bg5'.split(' ').map(san => ({ san })));
   expect(makePgn(game)).toEqual('1. e4 d5 2. a3 h6 3. Bg5 *\n');
 });
 
@@ -160,7 +161,7 @@ testPgnFile(
   },
   game => {
     expect(game.headers.get('Variant')).toBe('Antichess');
-    expect(Array.from(game.moves.mainline()).map(move => move.san)).toStrictEqual(['e3', 'e6', 'b4', 'Bxb4', 'Qg4']);
+    expect(Array.from(mainline(game.moves)).map(move => move.san)).toStrictEqual(['e3', 'e6', 'b4', 'Bxb4', 'Qg4']);
   },
 );
 testPgnFile(
@@ -187,12 +188,12 @@ testPgnFile(
     allValid: true,
   },
   game => {
-    expect(Array.from(game.moves.mainline()).map(move => move.san)).toStrictEqual(['e4', 'e5', 'Nf3', 'Nc6', 'Bb5']);
+    expect(Array.from(mainline(game.moves)).map(move => move.san)).toStrictEqual(['e4', 'e5', 'Nf3', 'Nc6', 'Bb5']);
   },
 );
 
 test('tricky tokens', () => {
-  const steps = Array.from(parsePgn('O-O-O !! 0-0-0# ??')[0].moves.mainline());
+  const steps = Array.from(mainline(parsePgn('O-O-O !! 0-0-0# ??')[0].moves));
   expect(steps[0].san).toBe('O-O-O');
   expect(steps[0].nags).toEqual([3]);
   expect(steps[1].san).toBe('O-O-O#');

diff --git a/src/pgn.ts b/src/pgn.ts
@@ -120,37 +120,41 @@ export const defaultGame = <T>(initHeaders: () => Map<string, string> = defaultH
 
 export class Node<T> {
   children: ChildNode<T>[] = [];
+}
 
-  *mainline(): Iterable<T> {
-    let node: Node<T> = this;
-    while (node.children.length) {
-      const child = node.children[0];
-      yield child.data;
-      node = child;
-    }
+export class ChildNode<T> extends Node<T> {
+  constructor(public data: T) {
+    super();
   }
 }
 
-export const extendMainline = <T>(game: Game<T>, data: T[]) => {
-  let node = game.moves;
+export const isChildNode = <T>(node: Node<T>): node is ChildNode<T> => node instanceof ChildNode;
+
+export function* mainlineNodes<T>(node: Node<T>): Iterable<ChildNode<T>> {
   while (node.children.length) {
     const child = node.children[0];
+    yield child;
     node = child;
   }
-  data.forEach(d => {
-    const newNode = new ChildNode(d);
-    node.children = [newNode];
-    node = newNode;
-  });
-};
+}
 
-export class ChildNode<T> extends Node<T> {
-  constructor(public data: T) {
-    super();
-  }
+export function* mainline<T>(node: Node<T>): Iterable<T> {
+  for (const child of mainlineNodes(node)) yield child.data;
 }
 
-export const isChildNode = <T>(node: Node<T>): node is ChildNode<T> => node instanceof ChildNode;
+export function mainlineEnd<T>(node: Node<T>): Node<T> {
+  while (node.children.length) node = node.children[0];
+  return node;
+}
+
+export const extend = <T>(node: Node<T>, data: T[]): Node<T> => {
+  for (const d of data) {
+    const child = new ChildNode(d);
+    node.children.push(child);
+    node = child;
+  }
+  return node;
+};
 
 export class Box<T> {
   constructor(public value: T) {}