Refactor Target and Writer, add single-file HLS writing mode

src/target.ts

@@ -7,7 +7,6 @@
  */
 
 import type { FileHandle } from 'node:fs/promises';
-import { BufferTargetWriter, NullTargetWriter, StreamTargetWriter, Writer } from './writer';
 import { Output } from './output';
 import * as nodeAlias from './node';
 import { assert } from './misc';
@@ -26,7 +25,18 @@ export abstract class Target {
 	_output: Output | null = null;
 
 	/** @internal */
-	abstract _createWriter(): Writer;
+	_ensureMonotonicity = false;
+
+	/** @internal */
+	abstract _start(): void;
+	/** @internal */
+	abstract _write(data: Uint8Array, pos: number): void;
+	/** @internal */
+	abstract _flush(): Promise<void>;
+	/** @internal */
+	abstract _finalize(): Promise<void>;
+	/** @internal */
+	abstract _close(): Promise<void>;
 
 	/**
 	 * Called each time data is written to the target. Will be called with the byte range into which data was written.
@@ -37,8 +47,19 @@ export abstract class Target {
 	onwrite: ((start: number, end: number) => unknown) | null = null;
 
 	onfinalized: (() => unknown) | null = null;
+
+	slice(offset: number) {
+		if (!Number.isInteger(offset) && offset < 0) {
+			throw new TypeError('offset must be a non-negative integer.');
+		}
+
+		return new RangedTarget(this, offset);
+	}
 }
 
+const ARRAY_BUFFER_INITIAL_SIZE = 2 ** 16;
+const ARRAY_BUFFER_MAX_SIZE = 2 ** 32;
+
 /**
  * A target that writes data directly into an ArrayBuffer in memory. Great for performance, but not suitable for very
  * large files. The buffer will be available once the output has been finalized.
@@ -50,8 +71,92 @@ export class BufferTarget extends Target {
 	buffer: ArrayBuffer | null = null;
 
 	/** @internal */
-	_createWriter() {
-		return new BufferTargetWriter(this);
+	_buffer: ArrayBuffer;
+	/** @internal */
+	_bytes: Uint8Array;
+	/** @internal */
+	_maxPos = 0;
+	/** @internal */
+	_supportsResize: boolean;
+
+	constructor() {
+		super();
+
+		this._supportsResize = 'resize' in new ArrayBuffer(0);
+		if (this._supportsResize) {
+			try {
+				// @ts-expect-error Don't want to bump "lib" in tsconfig
+				this._buffer = new ArrayBuffer(ARRAY_BUFFER_INITIAL_SIZE, { maxByteLength: ARRAY_BUFFER_MAX_SIZE });
+			} catch {
+				this._buffer = new ArrayBuffer(ARRAY_BUFFER_INITIAL_SIZE);
+				this._supportsResize = false;
+			}
+		} else {
+			this._buffer = new ArrayBuffer(ARRAY_BUFFER_INITIAL_SIZE);
+		}
+
+		this._bytes = new Uint8Array(this._buffer);
+	}
+
+	/** @internal */
+	_ensureSize(size: number) {
+		let newLength = this._buffer.byteLength;
+		while (newLength < size) newLength *= 2;
+
+		if (newLength === this._buffer.byteLength) return;
+
+		if (newLength > ARRAY_BUFFER_MAX_SIZE) {
+			throw new Error(
+				`ArrayBuffer exceeded maximum size of ${ARRAY_BUFFER_MAX_SIZE} bytes. Please consider using another`
+				+ ` target.`,
+			);
+		}
+
+		if (this._supportsResize) {
+			// Use resize if it exists
+			// @ts-expect-error Don't want to bump "lib" in tsconfig
+			// eslint-disable-next-line @typescript-eslint/no-unsafe-call
+			this._buffer.resize(newLength);
+			// The Uint8Array scales automatically
+		} else {
+			const newBuffer = new ArrayBuffer(newLength);
+			const newBytes = new Uint8Array(newBuffer);
+			newBytes.set(this._bytes, 0);
+
+			this._buffer = newBuffer;
+			this._bytes = newBytes;
+		}
+	}
+
+	/** @internal */
+	_start() {}
+
+	/** @internal */
+	_write(data: Uint8Array, pos: number) {
+		this._ensureSize(pos + data.byteLength);
+
+		this._bytes.set(data, pos);
+
+		this._maxPos = Math.max(this._maxPos, pos + data.byteLength);
+
+		this.onwrite?.(pos, pos + data.byteLength);
+	}
+
+	/** @internal */
+	async _flush() {}
+
+	/** @internal */
+	async _finalize() {
+		this.buffer = this._buffer.slice(0, this._maxPos);
+		this.onfinalized?.();
+	}
+
+	/** @internal */
+	async _close() {}
+
+	/** @internal */
+	_getSlice(start: number, end: number) {
+		return this._bytes.slice(start, end);
 	}
 }
 
@@ -85,6 +190,21 @@ export type StreamTargetOptions = {
 	chunkSize?: number;
 };
 
+const DEFAULT_CHUNK_SIZE = 2 ** 24;
+const MAX_CHUNKS_AT_ONCE = 2;
+
+type Chunk = {
+	start: number;
+	written: ChunkSection[];
+	data: Uint8Array<ArrayBuffer>;
+	shouldFlush: boolean;
+};
+
+type ChunkSection = {
+	start: number;
+	end: number;
+};
+
 /**
  * This target writes data to a [`WritableStream`](https://developer.mozilla.org/en-US/docs/Web/API/WritableStream),
  * making it a general-purpose target for writing data anywhere. It is also compatible with
@@ -100,6 +220,33 @@ export class StreamTarget extends Target {
 	/** @internal */
 	_options: StreamTargetOptions;
 
+	/** @internal */
+	_sections: {
+		data: Uint8Array;
+		start: number;
+	}[] = [];
+
+	/** @internal */
+	_lastWriteEnd = 0;
+	/** @internal */
+	_lastFlushEnd = 0;
+	/** @internal */
+	_streamWriter: WritableStreamDefaultWriter<StreamTargetChunk> | null = null;
+	/** @internal */
+	_writeError: unknown = null;
+
+	// These variables regard chunked mode:
+	/** @internal */
+	_chunked: boolean;
+	/** @internal */
+	_chunkSize: number;
+	/**
+	 * The data is divided up into fixed-size chunks, whose contents are first filled in RAM and then flushed out.
+	 * A chunk is flushed if all of its contents have been written.
+	 */
+	/** @internal */
+	_chunks: Chunk[] = [];
+
 	/** Creates a new {@link StreamTarget} which writes to the specified `writable`. */
 	constructor(
 		writable: WritableStream<StreamTargetChunk>,
@@ -122,11 +269,242 @@ export class StreamTarget extends Target {
 
 		this._writable = writable;
 		this._options = options;
+
+		this._chunked = options.chunked ?? false;
+		this._chunkSize = options.chunkSize ?? DEFAULT_CHUNK_SIZE;
 	}
 
 	/** @internal */
-	_createWriter() {
-		return new StreamTargetWriter(this);
+	_start() {
+		this._streamWriter = this._writable.getWriter();
+	}
+
+	/** @internal */
+	_write(data: Uint8Array, pos: number) {
+		if (pos > this._lastWriteEnd) {
+			const paddingBytesNeeded = pos - this._lastWriteEnd;
+			this._write(new Uint8Array(paddingBytesNeeded), this._lastWriteEnd);
+		}
+
+		this._sections.push({
+			data: data.slice(),
+			start: pos,
+		});
+
+		this._lastWriteEnd = Math.max(this._lastWriteEnd, pos + data.byteLength);
+
+		this.onwrite?.(pos, pos + data.byteLength);
+	}
+
+	/** @internal */
+	async _flush() {
+		if (this._writeError !== null) {
+			// eslint-disable-next-line @typescript-eslint/only-throw-error
+			throw this._writeError;
+		}
+
+		assert(this._streamWriter);
+		if (this._sections.length === 0) return;
+
+		const chunks: {
+			start: number;
+			size: number;
+			data?: Uint8Array<ArrayBuffer>;
+		}[] = [];
+		const sorted = [...this._sections].sort((a, b) => a.start - b.start);
+
+		chunks.push({
+			start: sorted[0]!.start,
+			size: sorted[0]!.data.byteLength,
+		});
+
+		// Figure out how many contiguous chunks we have
+		for (let i = 1; i < sorted.length; i++) {
+			const lastChunk = chunks[chunks.length - 1]!;
+			const section = sorted[i]!;
+
+			if (section.start <= lastChunk.start + lastChunk.size) {
+				lastChunk.size = Math.max(lastChunk.size, section.start + section.data.byteLength - lastChunk.start);
+			} else {
+				chunks.push({
+					start: section.start,
+					size: section.data.byteLength,
+				});
+			}
+		}
+
+		for (const chunk of chunks) {
+			chunk.data = new Uint8Array(chunk.size);
+
+			// Make sure to write the data in the correct order for correct overwriting
+			for (const section of this._sections) {
+				// Check if the section is in the chunk
+				if (chunk.start <= section.start && section.start < chunk.start + chunk.size) {
+					chunk.data.set(section.data, section.start - chunk.start);
+				}
+			}
+
+			if (this._streamWriter.desiredSize !== null && this._streamWriter.desiredSize <= 0) {
+				await this._streamWriter.ready; // Allow the writer to apply backpressure
+			}
+
+			if (this._chunked) {
+				// Let's first gather the data into bigger chunks before writing it
+				this._writeDataIntoChunks(chunk.data, chunk.start);
+				this._tryToFlushChunks();
+			} else {
+				if (this._ensureMonotonicity && chunk.start !== this._lastFlushEnd) {
+					throw new Error('Internal error: Monotonicity violation.');
+				}
+
+				void this._streamWriter.write({
+					type: 'write',
+					data: chunk.data,
+					position: chunk.start,
+				}).catch((error) => {
+					this._writeError ??= error;
+				});
+
+				this._lastFlushEnd = chunk.start + chunk.data.byteLength;
+			}
+		}
+
+		this._sections.length = 0;
+	}
+
+	/** @internal */
+	_writeDataIntoChunks(data: Uint8Array, position: number) {
+		// First, find the chunk to write the data into, or create one if none exists
+		let chunkIndex = this._chunks.findIndex(x => x.start <= position && position < x.start + this._chunkSize);
+		if (chunkIndex === -1) chunkIndex = this._createChunk(position);
+		const chunk = this._chunks[chunkIndex]!;
+
+		// Figure out how much to write to the chunk, and then write to the chunk
+		const relativePosition = position - chunk.start;
+		const toWrite = data.subarray(0, Math.min(this._chunkSize - relativePosition, data.byteLength));
+		chunk.data.set(toWrite, relativePosition);
+
+		// Create a section describing the region of data that was just written to
+		const section: ChunkSection = {
+			start: relativePosition,
+			end: relativePosition + toWrite.byteLength,
+		};
+		this._insertSectionIntoChunk(chunk, section);
+
+		// Queue chunk for flushing to target if it has been fully written to
+		if (chunk.written[0]!.start === 0 && chunk.written[0]!.end === this._chunkSize) {
+			chunk.shouldFlush = true;
+		}
+
+		// Make sure we don't hold too many chunks in memory at once to keep memory usage down
+		if (this._chunks.length > MAX_CHUNKS_AT_ONCE) {
+			// Flush all but the last chunk
+			for (let i = 0; i < this._chunks.length - 1; i++) {
+				this._chunks[i]!.shouldFlush = true;
+			}
+			this._tryToFlushChunks();
+		}
+
+		// If the data didn't fit in one chunk, recurse with the remaining data
+		if (toWrite.byteLength < data.byteLength) {
+			this._writeDataIntoChunks(data.subarray(toWrite.byteLength), position + toWrite.byteLength);
+		}
+	}
+
+	/** @internal */
+	_insertSectionIntoChunk(chunk: Chunk, section: ChunkSection) {
+		let low = 0;
+		let high = chunk.written.length - 1;
+		let index = -1;
+
+		// Do a binary search to find the last section with a start not larger than `section`'s start
+		while (low <= high) {
+			const mid = Math.floor(low + (high - low + 1) / 2);
+
+			if (chunk.written[mid]!.start <= section.start) {
+				low = mid + 1;
+				index = mid;
+			} else {
+				high = mid - 1;
+			}
+		}
+
+		// Insert the new section
+		chunk.written.splice(index + 1, 0, section);
+		if (index === -1 || chunk.written[index]!.end < section.start) index++;
+
+		// Merge overlapping sections
+		while (index < chunk.written.length - 1 && chunk.written[index]!.end >= chunk.written[index + 1]!.start) {
+			chunk.written[index]!.end = Math.max(chunk.written[index]!.end, chunk.written[index + 1]!.end);
+			chunk.written.splice(index + 1, 1);
+		}
+	}
+
+	/** @internal */
+	_createChunk(includesPosition: number) {
+		const start = Math.floor(includesPosition / this._chunkSize) * this._chunkSize;
+		const chunk: Chunk = {
+			start,
+			data: new Uint8Array(this._chunkSize),
+			written: [],
+			shouldFlush: false,
+		};
+		this._chunks.push(chunk);
+		this._chunks.sort((a, b) => a.start - b.start);
+
+		return this._chunks.indexOf(chunk);
+	}
+
+	/** @internal */
+	_tryToFlushChunks(force = false) {
+		assert(this._streamWriter);
+
+		for (let i = 0; i < this._chunks.length; i++) {
+			const chunk = this._chunks[i]!;
+			if (!chunk.shouldFlush && !force) continue;
+
+			for (const section of chunk.written) {
+				const position = chunk.start + section.start;
+				if (this._ensureMonotonicity && position !== this._lastFlushEnd) {
+					throw new Error('Internal error: Monotonicity violation.');
+				}
+
+				void this._streamWriter.write({
+					type: 'write',
+					data: chunk.data.subarray(section.start, section.end),
+					position,
+				}).catch((error) => {
+					this._writeError ??= error;
+				});
+
+				this._lastFlushEnd = chunk.start + section.end;
+			}
+
+			this._chunks.splice(i--, 1);
+		}
+	}
+
+	/** @internal */
+	async _finalize() {
+		if (this._chunked) {
+			this._tryToFlushChunks(true);
+		}
+
+		if (this._writeError !== null) {
+			// eslint-disable-next-line @typescript-eslint/only-throw-error
+			throw this._writeError;
+		}
+
+		assert(this._streamWriter);
+		await this._streamWriter.ready;
+		await this._streamWriter.close();
+
+		this.onfinalized?.();
+	}
+
+	/** @internal */
+	async _close() {
+		return this._streamWriter?.close();
 	}
 }
 
@@ -187,8 +565,30 @@ export class FilePathTarget extends Target {
 	}
 
 	/** @internal */
-	_createWriter(): Writer {
-		return this._streamTarget._createWriter();
+	_start() {
+		this._streamTarget._start();
+	}
+
+	/** @internal */
+	_write(data: Uint8Array, pos: number) {
+		this._streamTarget._write(data, pos);
+		this.onwrite?.(pos, pos + data.byteLength);
+	}
+
+	/** @internal */
+	async _flush() {
+		return this._streamTarget._flush();
+	}
+
+	/** @internal */
+	async _finalize() {
+		await this._streamTarget._finalize();
+		this.onfinalized?.();
+	}
+
+	/** @internal */
+	async _close() {
+		return this._streamTarget._close();
 	}
 }
 
@@ -200,7 +600,59 @@ export class FilePathTarget extends Target {
  */
 export class NullTarget extends Target {
 	/** @internal */
-	_createWriter() {
-		return new NullTargetWriter(this);
+	_start() {}
+
+	/** @internal */
+
+	_write(data: Uint8Array, pos: number) {
+		this.onwrite?.(pos, pos + data.byteLength);
 	}
+
+	/** @internal */
+	async _flush() {}
+
+	/** @internal */
+	async _finalize() {
+		this.onfinalized?.();
+	}
+
+	/** @internal */
+	async _close() {}
+}
+
+export class RangedTarget extends Target {
+	/** @internal */
+	_baseTarget: Target;
+	/** @internal */
+	_offset: number;
+
+	/** @internal */
+	constructor(baseTarget: Target, offset: number) {
+		super();
+
+		this._baseTarget = baseTarget;
+		this._offset = offset;
+	}
+
+	/** @internal */
+	_start() {}
+
+	/** @internal */
+	_write(data: Uint8Array, pos: number): void {
+		this._baseTarget._write(data, this._offset + pos);
+		this.onwrite?.(pos, pos + data.byteLength);
+	}
+
+	/** @internal */
+	_flush() {
+		return this._baseTarget._flush();
+	}
+
+	/** @internal */
+	async _finalize() {
+		this.onfinalized?.();
+	}
+
+	/** @internal */
+	async _close() {}
 }