feat: Add context buffer for more efficient loading

Author: effigies

src/schema/walk.ts

@@ -3,6 +3,7 @@ import { BIDSFile, FileOpener, FileTree } from '../types/filetree.ts'
 import type { DatasetIssues } from '../issues/datasetIssues.ts'
 import { loadTSV } from '../files/tsv.ts'
 import { loadJSON } from '../files/json.ts'
+import { queuedAsyncIterator } from '../utils/queue.ts'
 
 function* quickWalk(dir: FileTree): Generator<BIDSFile> {
   for (const file of dir.files) {
@@ -38,11 +39,13 @@ function pseudoFile(dir: FileTree, opaque: boolean): BIDSFile {
   )
 }
 
+type CleanupFunction = () => void
+
 /** Recursive algorithm for visiting each file in the dataset, creating a context */
 async function* _walkFileTree(
   fileTree: FileTree,
   dsContext: BIDSContextDataset,
-): AsyncIterable<BIDSContext> {
+): AsyncIterable<BIDSContext | CleanupFunction> {
   for (const file of fileTree.files) {
     yield new BIDSContext(file, dsContext)
   }
@@ -56,13 +59,23 @@ async function* _walkFileTree(
       yield* _walkFileTree(dir, dsContext)
     }
   }
-  loadTSV.cache.delete(fileTree.path)
-  loadJSON.cache.delete(fileTree.path)
+  yield () => {
+    loadTSV.cache.delete(fileTree.path)
+    loadJSON.cache.delete(fileTree.path)
+  }
 }
 
 /** Walk all files in the dataset and construct a context for each one */
 export async function* walkFileTree(
   dsContext: BIDSContextDataset,
+  bufferSize: number = 1,
 ): AsyncIterable<BIDSContext> {
-  yield* _walkFileTree(dsContext.tree, dsContext)
+  for await (
+    await using context of queuedAsyncIterator(
+      _walkFileTree(dsContext.tree, dsContext),
+      bufferSize,
+    )
+  ) {
+    yield context
+  }
 }

src/utils/queue.ts

@@ -0,0 +1,43 @@
+export async function* bufferAsyncIterator<T>(
+  asyncIterator: AsyncIterable<T>,
+  bufferSize: number = 5,
+): AsyncGenerator<T, void, undefined> {
+  const iterator = asyncIterator[Symbol.asyncIterator]()
+  const promises: Promise<IteratorResult<T>>[] = []
+
+  // Initialize buffer
+  for (let i = 0; i < bufferSize; i++) {
+    promises.push(iterator.next())
+  }
+
+  while (promises.length > 0) {
+    const result = await promises.shift()!
+
+    if (!result.done) {
+      yield result.value
+      // Keep buffer full
+      promises.push(iterator.next())
+    }
+  }
+}
+
+type CleanupFunction = () => void
+
+export async function* cleanupAsyncIterator<T extends object>(
+  asyncIterator: AsyncIterable<T | CleanupFunction>,
+): AsyncGenerator<T, void, undefined> {
+  for await (const item of asyncIterator) {
+    if (typeof item === 'function') {
+      item()
+    } else {
+      yield item
+    }
+  }
+}
+
+export async function* queuedAsyncIterator<T extends object>(
+  asyncIterator: AsyncIterable<T | CleanupFunction>,
+  bufferSize: number = 5,
+): AsyncGenerator<T, void, undefined> {
+  yield* cleanupAsyncIterator(bufferAsyncIterator(asyncIterator, bufferSize))
+}

src/validators/bids.ts

@@ -108,7 +108,7 @@ export async function validate(
     return false
   })
 
-  for await (const context of walkFileTree(dsContext)) {
+  for await (const context of walkFileTree(dsContext, 20)) {
     // TODO - Skip ignored files for now (some tests may reference ignored files)
     if (context.file.ignored) {
       continue