Add timestamp iterator

src/isobmff/isobmff-demuxer.ts

@@ -1305,18 +1305,12 @@ export class IsobmffDemuxer extends Demuxer {
 
 abstract class IsobmffTrackBacking<Chunk extends EncodedVideoChunk | EncodedAudioChunk> implements InputTrackBacking {
 	chunkToSampleIndex = new WeakMap<Chunk, number>();
-	sampleIndexToChunk = new Map<number, WeakRef<Chunk>>();
-
 	chunkToFragmentLocation = new WeakMap<Chunk, {
 		fragment: Fragment;
 		sampleIndex: number;
 	}>();
 
-	fragmentLocationToChunk = new Map<string, WeakRef<Chunk>>();
-
-	constructor(public internalTrack: InternalTrack) {
-
-	}
+	constructor(public internalTrack: InternalTrack) {}
 
 	getCodec(): Promise<MediaCodec> {
 		throw new Error('Not implemented on base class.');
@@ -1355,7 +1349,20 @@ abstract class IsobmffTrackBacking<Chunk extends EncodedVideoChunk | EncodedAudi
 		return this.fetchChunkForSampleIndex(0, options);
 	}
 
+	private roundToMicrosecond(timestamp: number) {
+		// We transform the timestamp so that chunk retrieval behaves expectedly: All chunks returned have a timestamp
+		// that's floored to the microseconds, and that timestamp may be before the actual timestamp. But since the
+		// actual timestamp is never communicated to the outside, chunk retrieval should work like the timestamp is
+		// exactly equal to its floored version. This means, when we retrieve the chunk for timestamp 0.333333, but the
+		// chunk's true, unrounded timestamp is 1/3, then we would not retrieve that chunk, despite the chunk having a
+		// floored timestamp of 0.333333. That's why we transform the search timestamp by first flooring it to the
+		// microsecond, and then adding "1-eps" to it to make sure get all chunks whose timestamps will round down to
+		// a value included by the search timestamp.
+		return (Math.floor(timestamp * 1e6) + 0.99999999) / 1e6;
+	}
+
 	async getChunk(timestamp: number, options: ChunkRetrievalOptions) {
+		timestamp = this.roundToMicrosecond(timestamp);
 		const timestampInTimescale = timestamp * this.internalTrack.timescale;
 
 		if (this.internalTrack.demuxer.isFragmented) {
@@ -1441,6 +1448,7 @@ abstract class IsobmffTrackBacking<Chunk extends EncodedVideoChunk | EncodedAudi
 	}
 
 	async getKeyChunk(timestamp: number, options: ChunkRetrievalOptions) {
+		timestamp = this.roundToMicrosecond(timestamp);
 		const timestampInTimescale = timestamp * this.internalTrack.timescale;
 
 		if (this.internalTrack.demuxer.isFragmented) {
@@ -1545,11 +1553,6 @@ abstract class IsobmffTrackBacking<Chunk extends EncodedVideoChunk | EncodedAudi
 			return null;
 		}
 
-		const existingChunk = this.sampleIndexToChunk.get(sampleIndex)?.deref();
-		if (existingChunk) {
-			return existingChunk;
-		}
-
 		const sampleTable = this.internalTrack.demuxer.getSampleTableForTrack(this.internalTrack);
 		const sampleInfo = getSampleInfo(sampleTable, sampleIndex);
 		if (!sampleInfo) {
@@ -1577,7 +1580,6 @@ abstract class IsobmffTrackBacking<Chunk extends EncodedVideoChunk | EncodedAudi
 		const chunk = this.createChunk(data, timestamp, duration, sampleInfo.isKeyFrame);
 
 		this.chunkToSampleIndex.set(chunk, sampleIndex);
-		this.sampleIndexToChunk.set(sampleIndex, new WeakRef(chunk));
 
 		return chunk;
 	}
@@ -1587,12 +1589,6 @@ abstract class IsobmffTrackBacking<Chunk extends EncodedVideoChunk | EncodedAudi
 			return null;
 		}
 
-		const compositeKey = `${fragment.moofOffset}:${sampleIndex}`;
-		const existingChunk = this.fragmentLocationToChunk.get(compositeKey)?.deref();
-		if (existingChunk) {
-			return existingChunk;
-		}
-
 		const trackData = fragment.trackData.get(this.internalTrack.id)!;
 		const sample = trackData.samples[sampleIndex];
 		assert(sample);
@@ -1615,7 +1611,6 @@ abstract class IsobmffTrackBacking<Chunk extends EncodedVideoChunk | EncodedAudi
 		const chunk = this.createChunk(data, timestamp, duration, sample.isKeyFrame);
 
 		this.chunkToFragmentLocation.set(chunk, { fragment, sampleIndex });
-		this.fragmentLocationToChunk.set(compositeKey, new WeakRef(chunk));
 
 		return chunk;
 	}

src/media-drain.ts

@@ -1,5 +1,5 @@
 import { ChunkRetrievalOptions, InputAudioTrack, InputVideoTrack } from './input-track';
-import { assert, promiseWithResolvers } from './misc';
+import { AnyIterable, assert, promiseWithResolvers, toAsyncIterator } from './misc';
 
 abstract class BaseChunkDrain<Chunk extends EncodedVideoChunk | EncodedAudioChunk> {
 	abstract getFirstChunk(options?: ChunkRetrievalOptions): Promise<Chunk | null>;
@@ -78,71 +78,154 @@ abstract class BaseMediaFrameDrain<
 	abstract createDecoder(onMedia: (media: MediaFrame) => unknown): Promise<VideoDecoder | AudioDecoder>;
 	abstract createChunkDrain(): BaseChunkDrain<Chunk>;
 
-	protected async getKeyMediaFrame(timestamp: number): Promise<MediaFrame | null> {
-		let result: MediaFrame | null = null;
-
-		const decoder = await this.createDecoder(frame => result = frame);
-		const chunkDrain = this.createChunkDrain();
-		const chunk = await chunkDrain.getKeyChunk(timestamp);
-		if (!chunk) {
-			return null;
-		}
-
-		decoder.decode(chunk);
-
-		await decoder.flush();
-		decoder.close();
-
-		return result;
+	private duplicateFrame(frame: MediaFrame) {
+		return structuredClone(frame);
 	}
 
-	protected async getMediaFrame(timestamp: number): Promise<MediaFrame | null> {
-		let result: MediaFrame | null = null;
+	protected async* mediaFramesAtTimestamps(timestamps: AnyIterable<number>) {
+		const timestampIterator = toAsyncIterator(timestamps);
+		const timestampsOfInterest: number[] = [];
+
+		const frameQueue: (MediaFrame | null)[] = [];
+		let { promise: queueNotEmpty, resolve: onQueueNotEmpty } = promiseWithResolvers();
+		let { promise: queueDequeue, resolve: onQueueDequeue } = promiseWithResolvers();
+		let decoderIsFlushed = false;
+		let ended = false;
+
+		const MAX_QUEUE_SIZE = 8;
+
+		let lastUsedFrame: MediaFrame | null = null;
+		const pushToQueue = (frame: MediaFrame | null) => {
+			frameQueue.push(frame);
+			onQueueNotEmpty();
+			({ promise: queueNotEmpty, resolve: onQueueNotEmpty } = promiseWithResolvers());
+		};
 
 		const decoder = await this.createDecoder((frame) => {
-			if (frame.timestamp / 1e6 <= timestamp) {
-				result?.close();
-				result = frame;
+			onQueueDequeue();
+
+			if (ended) {
+				frame.close();
+				return;
+			}
+
+			let frameUsed = false;
+			while (timestampsOfInterest.length > 0 && timestampsOfInterest[0] === frame.timestamp) {
+				pushToQueue(this.duplicateFrame(frame));
+				timestampsOfInterest.shift();
+				frameUsed = true;
+			}
+
+			if (frameUsed) {
+				lastUsedFrame?.close();
+				lastUsedFrame = frame;
 			} else {
 				frame.close();
 			}
 		});
-		const chunkDrain = this.createChunkDrain();
-		const keyChunk = await chunkDrain.getKeyChunk(timestamp);
-		if (!keyChunk) {
-			return null;
-		}
 
-		const targetChunk = await chunkDrain.getChunk(timestamp);
-		assert(targetChunk);
+		// The following is the "pump" process that keeps pumping chunks into the decoder
+		void (async () => {
+			const chunkDrain = this.createChunkDrain();
+			let lastKeyChunk: Chunk | null = null;
+			let lastChunk: Chunk | null = null;
 
-		decoder.decode(keyChunk);
+			for await (const timestamp of timestampIterator) {
+				while (frameQueue.length + decoder.decodeQueueSize > MAX_QUEUE_SIZE) {
+					({ promise: queueDequeue, resolve: onQueueDequeue } = promiseWithResolvers());
+					await queueDequeue;
+				}
 
-		let currentChunk = keyChunk;
-		while (currentChunk !== targetChunk) {
-			const nextChunk = await chunkDrain.getNextChunk(currentChunk);
-			assert(nextChunk);
+				if (ended) {
+					break;
+				}
 
-			currentChunk = nextChunk;
-			decoder.decode(nextChunk);
+				const targetChunk = await chunkDrain.getChunk(timestamp);
+				if (!targetChunk) {
+					pushToQueue(null);
+					continue;
+				}
 
-			if (decoder.decodeQueueSize >= 10) {
-				await new Promise(resolve => decoder.addEventListener('dequeue', resolve, { once: true }));
+				const keyChunk = await chunkDrain.getKeyChunk(timestamp);
+				if (!keyChunk) {
+					pushToQueue(null);
+					continue;
+				}
+
+				timestampsOfInterest.push(targetChunk.timestamp);
+
+				if (
+					lastKeyChunk
+					&& keyChunk.timestamp === lastKeyChunk.timestamp
+					&& targetChunk.timestamp >= lastChunk!.timestamp
+				) {
+					assert(lastChunk);
+
+					if (targetChunk.timestamp === lastChunk.timestamp && timestampsOfInterest.length === 1) {
+						// Special case: We have a repeat chunk, but the frame for that chunk has already been decoded.
+						// Therefore, we need to push the frame here instead of in the decoder callback.
+						if (lastUsedFrame) {
+							pushToQueue(this.duplicateFrame(lastUsedFrame));
+						}
+						timestampsOfInterest.shift();
+					}
+				} else {
+					lastKeyChunk = keyChunk;
+					lastChunk = keyChunk;
+					decoder.decode(keyChunk);
+				}
+
+				while (lastChunk.timestamp !== targetChunk.timestamp) {
+					const nextChunk = await chunkDrain.getNextChunk(lastChunk);
+					assert(nextChunk);
+
+					lastChunk = nextChunk;
+					decoder.decode(nextChunk);
+				}
+
+				if (decoder.decodeQueueSize >= 10) {
+					await new Promise(resolve => decoder.addEventListener('dequeue', resolve, { once: true }));
+				}
 			}
+
+			await decoder.flush();
+			decoder.close();
+
+			decoderIsFlushed = true;
+			onQueueNotEmpty(); // To unstuck the generator
+		})();
+
+		try {
+			while (true) {
+				if (frameQueue.length > 0) {
+					const nextFrame = frameQueue.shift();
+					assert(nextFrame !== undefined);
+					yield nextFrame;
+					onQueueDequeue();
+				} else if (!decoderIsFlushed) {
+					await queueNotEmpty;
+				} else {
+					break;
+				}
+			}
+		} finally {
+			ended = true;
+			onQueueDequeue();
+
+			for (const frame of frameQueue) {
+				frame?.close();
+			}
+			(lastUsedFrame as MediaFrame | null)?.close();
 		}
-
-		await decoder.flush();
-		decoder.close();
-
-		return result;
 	}
 
-	protected async* mediaFrames(startTimestamp = 0, endTimestamp = Infinity) {
+	protected async* mediaFramesInRange(startTimestamp = 0, endTimestamp = Infinity) {
 		const frameQueue: MediaFrame[] = [];
 		let firstFrameQueued = false;
 		let lastFrame: MediaFrame | null = null;
 		let { promise: queueNotEmpty, resolve: onQueueNotEmpty } = promiseWithResolvers();
 		let { promise: queueDequeue, resolve: onQueueDequeue } = promiseWithResolvers();
+		let decoderIsFlushed = false;
 		let ended = false;
 
 		const MAX_QUEUE_SIZE = 8;
@@ -192,8 +275,6 @@ abstract class BaseMediaFrameDrain<
 			return;
 		}
 
-		let decoderIsFlushed = false;
-
 		// The following is the "pump" process that keeps pumping chunks into the decoder
 		void (async () => {
 			let currentChunk: Chunk | null = keyChunk;
@@ -320,16 +401,19 @@ export class VideoFrameDrain extends BaseMediaFrameDrain<EncodedVideoChunk, Vide
 		return new EncodedVideoChunkDrain(this.videoTrack);
 	}
 
-	getKeyFrame(timestamp: number) {
-		return this.getKeyMediaFrame(timestamp);
-	}
-
-	getFrame(timestamp: number) {
-		return this.getMediaFrame(timestamp);
+	async getFrame(timestamp: number) {
+		for await (const frame of this.mediaFramesAtTimestamps([timestamp])) {
+			return frame;
+		}
+		throw new Error('Internal error: Iterator returned nothing.');
 	}
 
 	frames(startTimestamp = 0, endTimestamp = Infinity) {
-		return this.mediaFrames(startTimestamp, endTimestamp);
+		return this.mediaFramesInRange(startTimestamp, endTimestamp);
+	}
+
+	framesAtTimestamps(timestamps: AnyIterable<number>) {
+		return this.mediaFramesAtTimestamps(timestamps);
 	}
 }
 
@@ -386,6 +470,12 @@ export class CanvasDrain {
 			yield this.videoFrameToWrappedCanvas(frame);
 		}
 	}
+
+	async* canvasesAtTimestamps(timestamps: AnyIterable<number>) {
+		for await (const frame of this.videoFrameDrain.framesAtTimestamps(timestamps)) {
+			yield frame && this.videoFrameToWrappedCanvas(frame);
+		}
+	}
 }
 
 export class EncodedAudioChunkDrain extends BaseChunkDrain<EncodedAudioChunk> {
@@ -437,16 +527,19 @@ export class AudioDataDrain extends BaseMediaFrameDrain<EncodedAudioChunk, Audio
 		return new EncodedAudioChunkDrain(this.audioTrack);
 	}
 
-	getKeyData(timestamp: number) {
-		return this.getKeyMediaFrame(timestamp);
-	}
-
-	getData(timestamp: number) {
-		return this.getMediaFrame(timestamp);
+	async getData(timestamp: number) {
+		for await (const data of this.mediaFramesAtTimestamps([timestamp])) {
+			return data;
+		}
+		throw new Error('Internal error: Iterator returned nothing.');
 	}
 
 	data(startTimestamp = 0, endTimestamp = Infinity) {
-		return this.mediaFrames(startTimestamp, endTimestamp);
+		return this.mediaFramesInRange(startTimestamp, endTimestamp);
+	}
+
+	dataAtTimestamps(timestamps: AnyIterable<number>) {
+		return this.mediaFramesAtTimestamps(timestamps);
 	}
 }
 
@@ -499,4 +592,10 @@ export class AudioBufferDrain {
 			yield this.audioDataToWrappedArrayBuffer(data);
 		}
 	}
+
+	async* buffersAtTimestamps(timestamps: AnyIterable<number>) {
+		for await (const data of this.audioDataDrain.dataAtTimestamps(timestamps)) {
+			yield data && this.audioDataToWrappedArrayBuffer(data);
+		}
+	}
 }

src/misc.ts

@@ -165,7 +165,7 @@ export const binarySearchExact = <T>(arr: T[], key: number, valueGetter: (x: T)
 
 		if (midVal === key) {
 			res = mid;
-			high = mid - 1; // continue searching left to find the lowest index
+			high = mid - 1; // Continue searching left to find the lowest index
 		} else if (midVal < key) {
 			low = mid + 1;
 		} else {
@@ -213,3 +213,17 @@ export const removeItem = <T>(arr: T[], item: T) => {
 		arr.splice(index, 1);
 	}
 };
+
+export type AnyIterable<T> =
+	| Iterable<T>
+	| AsyncIterable<T>;
+
+export const toAsyncIterator = async function* <T>(source: AnyIterable<T>): AsyncGenerator<T, void, unknown> {
+	if (Symbol.iterator in source) {
+		// @ts-expect-error Trust me
+		yield* source[Symbol.iterator]();
+	} else {
+		// @ts-expect-error Trust me
+		yield* source[Symbol.asyncIterator]();
+	}
+};