MPEG-TS demuxer optimizations: faster Annex B segmentation, better NALU iteration

dev/demux.html

@@ -11,17 +11,14 @@
 		const file = fileInput.files[0];
 		const input = new Mediabunny.Input({
 			formats: Mediabunny.ALL_FORMATS,
-			source: new Mediabunny.BlobSource(file),
+			source: new Mediabunny.BufferSource(await file.arrayBuffer()),
 		});
 
-		const track = await input.getPrimaryVideoTrack();
+		const track = await input.getPrimaryAudioTrack();
 		const sink = new Mediabunny.EncodedPacketSink(track);
 
-		let currentPacket = await sink.getFirstPacket();
-		while (currentPacket) {
-			console.log(currentPacket);
-			currentPacket = await sink.getNextPacket(currentPacket);
-		}
+		console.log("Done")
+
 		/*
 		
 		//const secondPacket = await sink.getNextPacket(firstPacket);

src/codec-data.ts

@@ -36,6 +36,10 @@ import { MetadataTags } from './metadata';
 // https://stackoverflow.com/questions/24884827
 
 export enum AvcNalUnitType {
+	NON_IDR_SLICE = 1,
+	SLICE_DPA = 2,
+	SLICE_DPB = 3,
+	SLICE_DPC = 4,
 	IDR = 5,
 	SEI = 6,
 	SPS = 7,
@@ -57,68 +61,68 @@ export enum HevcNalUnitType {
 	SUFFIX_SEI_NUT = 40,
 }
 
-/** Finds all NAL units in an AVC packet in Annex B format. */
-export const findNalUnitsInAnnexB = (packetData: Uint8Array) => {
-	const nalUnits: Uint8Array[] = [];
+export type NalUnitLocation = {
+	offset: number;
+	length: number;
+};
+
+export const iterateNalUnitsInAnnexB = function* (packetData: Uint8Array): Generator<NalUnitLocation> {
 	let i = 0;
+	let nalStart = -1;
 
-	while (i < packetData.length) {
-		let startCodePos = -1;
-		let startCodeLength = 0;
-
-		for (let j = i; j < packetData.length - 3; j++) {
-			// Check for 3-byte start code (0x000001)
-			if (packetData[j] === 0 && packetData[j + 1] === 0 && packetData[j + 2] === 1) {
-				startCodePos = j;
-				startCodeLength = 3;
+	while (i < packetData.length - 2) {
+		const zeroIndex = packetData.indexOf(0, i);
+		if (zeroIndex === -1 || zeroIndex >= packetData.length - 2) {
 			break;
 		}
+		i = zeroIndex;
+
+		let startCodeLength = 0;
 
 		// Check for 4-byte start code (0x00000001)
 		if (
-				j < packetData.length - 4
-				&& packetData[j] === 0
-				&& packetData[j + 1] === 0
-				&& packetData[j + 2] === 0
-				&& packetData[j + 3] === 1
+			i + 3 < packetData.length
+			&& packetData[i + 1] === 0
+			&& packetData[i + 2] === 0
+			&& packetData[i + 3] === 1
 		) {
-				startCodePos = j;
 			startCodeLength = 4;
-				break;
-			}
+		} else if (packetData[i + 1] === 0 && packetData[i + 2] === 1) {
+			// Check for 3-byte start code (0x000001)
+			startCodeLength = 3;
 		}
 
-		if (startCodePos === -1) {
-			break; // No more start codes found
+		if (startCodeLength === 0) {
+			i++;
+			continue;
 		}
 
-		// If this isn't the first start code, extract the previous NAL unit
-		if (i > 0 && startCodePos > i) {
-			const nalData = packetData.subarray(i, startCodePos);
-			if (nalData.length > 0) {
-				nalUnits.push(nalData);
-			}
+		// If we had a previous NAL unit, yield it
+		if (nalStart !== -1 && i > nalStart) {
+			yield {
+				offset: nalStart,
+				length: i - nalStart,
+			};
 		}
 
-		i = startCodePos + startCodeLength;
+		nalStart = i + startCodeLength;
+		i = nalStart;
 	}
 
-	// Extract the last NAL unit if there is one
-	if (i < packetData.length) {
-		const nalData = packetData.subarray(i);
-		if (nalData.length > 0) {
-			nalUnits.push(nalData);
+	// Yield the last NAL unit if there is one
+	if (nalStart !== -1 && nalStart < packetData.length) {
+		yield {
+			offset: nalStart,
+			length: packetData.length - nalStart,
+		};
 	}
-	}
-
-	return nalUnits;
 };
 
-/** Finds all NAL units in an AVC packet in length-prefixed format. */
-const findNalUnitsInLengthPrefixed = (packetData: Uint8Array, lengthSize: 1 | 2 | 3 | 4) => {
-	const nalUnits: Uint8Array[] = [];
+const iterateNalUnitsInLengthPrefixed = function* (
+	packetData: Uint8Array,
+	lengthSize: 1 | 2 | 3 | 4,
+): Generator<NalUnitLocation> {
 	let offset = 0;
-
 	const dataView = new DataView(packetData.buffer, packetData.byteOffset, packetData.byteLength);
 
 	while (offset + lengthSize <= packetData.length) {
@@ -129,22 +133,40 @@ const findNalUnitsInLengthPrefixed = (packetData: Uint8Array, lengthSize: 1 | 2
 			nalUnitLength = dataView.getUint16(offset, false);
 		} else if (lengthSize === 3) {
 			nalUnitLength = getUint24(dataView, offset, false);
-		} else if (lengthSize === 4) {
-			nalUnitLength = dataView.getUint32(offset, false);
 		} else {
-			assertNever(lengthSize);
-			assert(false);
+			assert(lengthSize === 4);
+			nalUnitLength = dataView.getUint32(offset, false);
 		}
 
 		offset += lengthSize;
 
-		const nalUnit = packetData.subarray(offset, offset + nalUnitLength);
-		nalUnits.push(nalUnit);
+		yield {
+			offset,
+			length: nalUnitLength,
+		};
 
 		offset += nalUnitLength;
 	}
+};
 
-	return nalUnits;
+export const iterateAvcNalUnits = (packetData: Uint8Array, decoderConfig: VideoDecoderConfig) => {
+	if (decoderConfig.description) {
+		const bytes = toUint8Array(decoderConfig.description);
+		const lengthSizeMinusOne = bytes[4]! & 0b11;
+		const lengthSize = (lengthSizeMinusOne + 1) as 1 | 2 | 3 | 4;
+
+		return iterateNalUnitsInLengthPrefixed(packetData, lengthSize);
+	} else {
+		return iterateNalUnitsInAnnexB(packetData);
+	}
+};
+
+export const iterateAvcNalUnitsAnnexB = function* (packetData: Uint8Array): Generator<NalUnitLocation> {
+	yield* iterateNalUnitsInAnnexB(packetData);
+};
+
+export const extractNalUnitTypeForAvc = (byte: number) => {
+	return byte & 0x1F;
 };
 
 const removeEmulationPreventionBytes = (data: Uint8Array) => {
@@ -231,21 +253,6 @@ export type AvcDecoderConfigurationRecord = {
 	sequenceParameterSetExt: Uint8Array[] | null;
 };
 
-export const extractAvcNalUnits = (packetData: Uint8Array, decoderConfig: VideoDecoderConfig) => {
-	if (decoderConfig.description) {
-		// Stream is length-prefixed. Let's extract the size of the length prefix from the decoder config
-
-		const bytes = toUint8Array(decoderConfig.description);
-		const lengthSizeMinusOne = bytes[4]! & 0b11;
-		const lengthSize = (lengthSizeMinusOne + 1) as 1 | 2 | 3 | 4;
-
-		return findNalUnitsInLengthPrefixed(packetData, lengthSize);
-	} else {
-		// Stream is in Annex B format
-		return findNalUnitsInAnnexB(packetData);
-	}
-};
-
 export const concatAvcNalUnits = (nalUnits: Uint8Array[], decoderConfig: VideoDecoderConfig) => {
 	if (decoderConfig.description) {
 		// Stream is length-prefixed. Let's extract the size of the length prefix from the decoder config
@@ -261,18 +268,25 @@ export const concatAvcNalUnits = (nalUnits: Uint8Array[], decoderConfig: VideoDe
 	}
 };
 
-export const extractNalUnitTypeForAvc = (data: Uint8Array) => {
-	return data[0]! & 0x1F;
-};
-
 /** Builds an AvcDecoderConfigurationRecord from an AVC packet in Annex B format. */
 export const extractAvcDecoderConfigurationRecord = (packetData: Uint8Array): AvcDecoderConfigurationRecord | null => {
 	try {
-		const nalUnits = findNalUnitsInAnnexB(packetData);
+		const spsUnits: Uint8Array[] = [];
+		const ppsUnits: Uint8Array[] = [];
+		const spsExtUnits: Uint8Array[] = [];
 
-		const spsUnits = nalUnits.filter(unit => extractNalUnitTypeForAvc(unit) === AvcNalUnitType.SPS);
-		const ppsUnits = nalUnits.filter(unit => extractNalUnitTypeForAvc(unit) === AvcNalUnitType.PPS);
-		const spsExtUnits = nalUnits.filter(unit => extractNalUnitTypeForAvc(unit) === AvcNalUnitType.SPS_EXT);
+		for (const loc of iterateAvcNalUnitsAnnexB(packetData)) {
+			const nalUnit = packetData.subarray(loc.offset, loc.offset + loc.length);
+			const type = extractNalUnitTypeForAvc(nalUnit[0]!);
+
+			if (type === AvcNalUnitType.SPS) {
+				spsUnits.push(nalUnit);
+			} else if (type === AvcNalUnitType.PPS) {
+				ppsUnits.push(nalUnit);
+			} else if (type === AvcNalUnitType.SPS_EXT) {
+				spsExtUnits.push(nalUnit);
+			}
+		}
 
 		if (spsUnits.length === 0) {
 			return null;
@@ -823,23 +837,24 @@ export type HevcSpsInfo = {
 	minSpatialSegmentationIdc: number;
 };
 
-export const extractHevcNalUnits = (packetData: Uint8Array, decoderConfig: VideoDecoderConfig) => {
+export const iterateHevcNalUnits = (packetData: Uint8Array, decoderConfig: VideoDecoderConfig) => {
 	if (decoderConfig.description) {
-		// Stream is length-prefixed. Let's extract the size of the length prefix from the decoder config
-
 		const bytes = toUint8Array(decoderConfig.description);
 		const lengthSizeMinusOne = bytes[21]! & 0b11;
 		const lengthSize = (lengthSizeMinusOne + 1) as 1 | 2 | 3 | 4;
 
-		return findNalUnitsInLengthPrefixed(packetData, lengthSize);
+		return iterateNalUnitsInLengthPrefixed(packetData, lengthSize);
 	} else {
-		// Stream is in Annex B format
-		return findNalUnitsInAnnexB(packetData);
+		return iterateNalUnitsInAnnexB(packetData);
 	}
 };
 
-export const extractNalUnitTypeForHevc = (data: Uint8Array) => {
-	return (data[0]! >> 1) & 0x3F;
+export const iterateHevcNalUnitsAnnexB = function* (packetData: Uint8Array): Generator<NalUnitLocation> {
+	yield* iterateNalUnitsInAnnexB(packetData);
+};
+
+export const extractNalUnitTypeForHevc = (byte: number) => {
+	return (byte >> 1) & 0x3F;
 };
 
 /** Parses an HEVC SPS (Sequence Parameter Set) to extract video information. */
@@ -994,15 +1009,25 @@ export const parseHevcSps = (sps: Uint8Array): HevcSpsInfo | null => {
 /** Builds a HevcDecoderConfigurationRecord from an HEVC packet in Annex B format. */
 export const extractHevcDecoderConfigurationRecord = (packetData: Uint8Array) => {
 	try {
-		const nalUnits = findNalUnitsInAnnexB(packetData);
+		const vpsUnits: Uint8Array[] = [];
+		const spsUnits: Uint8Array[] = [];
+		const ppsUnits: Uint8Array[] = [];
+		const seiUnits: Uint8Array[] = [];
 
-		const vpsUnits = nalUnits.filter(unit => extractNalUnitTypeForHevc(unit) === HevcNalUnitType.VPS_NUT);
-		const spsUnits = nalUnits.filter(unit => extractNalUnitTypeForHevc(unit) === HevcNalUnitType.SPS_NUT);
-		const ppsUnits = nalUnits.filter(unit => extractNalUnitTypeForHevc(unit) === HevcNalUnitType.PPS_NUT);
-		const seiUnits = nalUnits.filter(
-			unit => extractNalUnitTypeForHevc(unit) === HevcNalUnitType.PREFIX_SEI_NUT
-				|| extractNalUnitTypeForHevc(unit) === HevcNalUnitType.SUFFIX_SEI_NUT,
-		);
+		for (const loc of iterateHevcNalUnitsAnnexB(packetData)) {
+			const nalUnit = packetData.subarray(loc.offset, loc.offset + loc.length);
+			const type = extractNalUnitTypeForHevc(nalUnit[0]!);
+
+			if (type === HevcNalUnitType.VPS_NUT) {
+				vpsUnits.push(nalUnit);
+			} else if (type === HevcNalUnitType.SPS_NUT) {
+				spsUnits.push(nalUnit);
+			} else if (type === HevcNalUnitType.PPS_NUT) {
+				ppsUnits.push(nalUnit);
+			} else if (type === HevcNalUnitType.PREFIX_SEI_NUT || type === HevcNalUnitType.SUFFIX_SEI_NUT) {
+				seiUnits.push(nalUnit);
+			}
+		}
 
 		if (spsUnits.length === 0 || ppsUnits.length === 0) return null;
 
@@ -1078,7 +1103,7 @@ export const extractHevcDecoderConfigurationRecord = (packetData: Uint8Array) =>
 				? [
 						{
 							arrayCompleteness: 1,
-							nalUnitType: extractNalUnitTypeForHevc(seiUnits[0]!),
+							nalUnitType: extractNalUnitTypeForHevc(seiUnits[0]![0]!),
 							nalUnits: seiUnits,
 						},
 					]
@@ -2041,20 +2066,23 @@ export const determineVideoPacketType = (
 ): PacketType | null => {
 	switch (codec) {
 		case 'avc': {
-			const nalUnits = extractAvcNalUnits(packetData, decoderConfig);
-			let isKeyframe = nalUnits.some(x => extractNalUnitTypeForAvc(x) === AvcNalUnitType.IDR);
+			for (const loc of iterateAvcNalUnits(packetData, decoderConfig)) {
+				const nalTypeByte = packetData[loc.offset]!;
+				const type = extractNalUnitTypeForAvc(nalTypeByte);
+
+				if (type >= AvcNalUnitType.NON_IDR_SLICE && type <= AvcNalUnitType.SLICE_DPC) {
+					return 'delta';
+				}
+
+				if (type === AvcNalUnitType.IDR) {
+					return 'key';
+				}
 
-			if (!isKeyframe && (!isChromium() || getChromiumVersion()! >= 144)) {
 				// In addition to IDR, Recovery Point SEI also counts as a valid H.264 keyframe by current consensus.
 				// See https://github.com/w3c/webcodecs/issues/650 for the relevant discussion. WebKit and Firefox have
 				// always supported them, but Chromium hasn't, therefore the (admittedly dirty) version check.
-
-				for (const nalUnit of nalUnits) {
-					const type = extractNalUnitTypeForAvc(nalUnit);
-					if (type !== AvcNalUnitType.SEI) {
-						continue;
-					}
-
+				if (type === AvcNalUnitType.SEI && (!isChromium() || getChromiumVersion()! >= 144)) {
+					const nalUnit = packetData.subarray(loc.offset, loc.offset + loc.length);
 					const bytes = removeEmulationPreventionBytes(nalUnit);
 					let pos = 1; // Skip NALU header
 
@@ -2095,8 +2123,7 @@ export const determineVideoPacketType = (
 							if (recoveryFrameCount === 0 && exactMatchFlag === 1) {
 								// https://github.com/w3c/webcodecs/pull/910
 								// "recovery_frame_cnt == 0 and exact_match_flag=1 in the SEI recovery payload"
-								isKeyframe = true;
-								break;
+								return 'key';
 							}
 						}
 
@@ -2105,17 +2132,22 @@ export const determineVideoPacketType = (
 				}
 			}
 
-			return isKeyframe ? 'key' : 'delta';
+			return 'delta';
 		};
 
 		case 'hevc': {
-			const nalUnits = extractHevcNalUnits(packetData, decoderConfig);
-			const isKeyframe = nalUnits.some((x) => {
-				const type = extractNalUnitTypeForHevc(x);
-				return HevcNalUnitType.BLA_W_LP <= type && type <= HevcNalUnitType.RSV_IRAP_VCL23;
-			});
+			for (const loc of iterateHevcNalUnits(packetData, decoderConfig)) {
+				const type = extractNalUnitTypeForHevc(packetData[loc.offset]!);
+				if (type < HevcNalUnitType.BLA_W_LP) {
+					return 'delta';
+				}
 
-			return isKeyframe ? 'key' : 'delta';
+				if (type <= HevcNalUnitType.RSV_IRAP_VCL23) {
+					return 'key';
+				}
+			}
+
+			return 'delta';
 		};
 
 		case 'vp8': {

src/isobmff/isobmff-muxer.ts

@@ -28,7 +28,7 @@ import {
 	concatNalUnitsInLengthPrefixed,
 	extractAvcDecoderConfigurationRecord,
 	extractHevcDecoderConfigurationRecord,
-	findNalUnitsInAnnexB,
+	iterateNalUnitsInAnnexB,
 	serializeAvcDecoderConfigurationRecord,
 	serializeHevcDecoderConfigurationRecord,
 } from '../codec-data';
@@ -465,7 +465,8 @@ export class IsobmffMuxer extends Muxer {
 
 			let packetData = packet.data;
 			if (trackData.info.requiresAnnexBTransformation) {
-				const nalUnits = findNalUnitsInAnnexB(packetData);
+				const nalUnits = [...iterateNalUnitsInAnnexB(packetData)]
+					.map(loc => packetData.subarray(loc.offset, loc.offset + loc.length));
 				if (nalUnits.length === 0) {
 					// It's not valid Annex B data
 					throw new Error(

src/media-sink.ts

@@ -11,11 +11,11 @@ import {
 	concatAvcNalUnits,
 	deserializeAvcDecoderConfigurationRecord,
 	determineVideoPacketType,
-	extractAvcNalUnits,
-	extractHevcNalUnits,
 	extractNalUnitTypeForAvc,
 	extractNalUnitTypeForHevc,
 	HevcNalUnitType,
+	iterateAvcNalUnits,
+	iterateHevcNalUnits,
 	parseAvcSps,
 } from './codec-data';
 import { CustomVideoDecoder, customVideoDecoders, CustomAudioDecoder, customAudioDecoders } from './custom-coder';
@@ -945,12 +945,15 @@ class VideoDecoderWrapper extends DecoderWrapper<VideoSample> {
 
 			// Workaround for https://issues.chromium.org/issues/470109459
 			if (isChromium() && this.currentPacketIndex === 0 && this.codec === 'avc') {
-				const nalUnits = extractAvcNalUnits(packet.data, this.decoderConfig);
-				const filteredNalUnits = nalUnits.filter((x) => {
-					const type = extractNalUnitTypeForAvc(x);
+				const filteredNalUnits: Uint8Array[] = [];
+
+				for (const loc of iterateAvcNalUnits(packet.data, this.decoderConfig)) {
+					const type = extractNalUnitTypeForAvc(packet.data[loc.offset]!);
 					// These trip up Chromium's key frame detection, so let's strip them
-					return !(type >= 20 && type <= 31);
-				});
+					if (!(type >= 20 && type <= 31)) {
+						filteredNalUnits.push(packet.data.subarray(loc.offset, loc.offset + loc.length));
+					}
+				}
 
 				const newData = concatAvcNalUnits(filteredNalUnits, this.decoderConfig);
 				packet = new EncodedPacket(newData, packet.type, packet.timestamp, packet.duration);
@@ -1081,11 +1084,14 @@ class VideoDecoderWrapper extends DecoderWrapper<VideoSample> {
 	 * and causes bugs upstream. So, let's take the dropping into our own hands.
 	 */
 	hasHevcRaslPicture(packetData: Uint8Array) {
-		const nalUnits = extractHevcNalUnits(packetData, this.decoderConfig);
-		return nalUnits.some((x) => {
-			const type = extractNalUnitTypeForHevc(x);
-			return type === HevcNalUnitType.RASL_N || type === HevcNalUnitType.RASL_R;
-		});
+		for (const loc of iterateHevcNalUnits(packetData, this.decoderConfig)) {
+			const type = extractNalUnitTypeForHevc(packetData[loc.offset]!);
+			if (type === HevcNalUnitType.RASL_N || type === HevcNalUnitType.RASL_R) {
+				return true;
+			}
+		}
+
+		return false;
 	}
 
 	/** Handler for the WebCodecs VideoDecoder for ironing out browser differences. */

src/mpeg-ts/mpeg-ts-demuxer.ts

@@ -583,12 +583,14 @@ export class MpegTsDemuxer extends Demuxer {
 			return null;
 		}
 
-		const syncByte = readU8(slice);
+		const bytes = readBytes(slice, TS_PACKET_SIZE);
+
+		const syncByte = bytes[0]!;
 		if (syncByte !== 0x47) {
 			throw new Error('Invalid TS packet sync byte. Likely an internal bug, please report this file.');
 		}
 
-		const nextTwoBytes = readU16Be(slice);
+		const nextTwoBytes = (bytes[1]! << 8) + bytes[2]!;
 		// eslint-disable-next-line @typescript-eslint/no-unused-vars
 		const transportErrorIndicator = nextTwoBytes >> 15;
 		const payloadUnitStartIndicator = (nextTwoBytes >> 14) & 0x1;
@@ -596,7 +598,7 @@ export class MpegTsDemuxer extends Demuxer {
 		const transportPriority = (nextTwoBytes >> 13) & 0x1;
 		const pid = nextTwoBytes & 0x1FFF;
 
-		const nextByte = readU8(slice);
+		const nextByte = bytes[3]!;
 		// eslint-disable-next-line @typescript-eslint/no-unused-vars
 		const transportScramblingControl = nextByte >> 6;
 		const adaptationFieldControl = (nextByte >> 4) & 0x3;
@@ -607,7 +609,7 @@ export class MpegTsDemuxer extends Demuxer {
 			payloadUnitStartIndicator,
 			pid,
 			adaptationFieldControl,
-			body: readBytes(slice, TS_PACKET_SIZE - 4),
+			body: bytes.subarray(4),
 		};
 	}
 }
@@ -1063,25 +1065,6 @@ export abstract class MpegTsTrackBacking implements InputTrackBacking {
 
 		release();
 
-		const pesPacketHasKeyframe = async (sectionStartPos: number) => {
-			const section = await demuxer.readSection(sectionStartPos, true);
-			assert(section);
-			assert(section.pid === this.elementaryStream.pid);
-
-			const fullPesPacket = readPesPacket(section);
-			assert(fullPesPacket);
-
-			// Only mark the first packet
-			const context = new PacketReadingContext(this, fullPesPacket, false);
-			await this.markNextPacket(context);
-
-			if (!context.suppliedPacket) {
-				return false;
-			}
-
-			return this.getPacketType(context.suppliedPacket.data) === 'key';
-		};
-
 		// Starting from the binary search guess, let's now find the moment where the packet timestamps cross the
 		// search timestamp. This point will then be used as the center around which we search.
 		outer:
@@ -1226,36 +1209,39 @@ export abstract class MpegTsTrackBacking implements InputTrackBacking {
 				continue;
 			}
 
-			if (!(await pesPacketHasKeyframe(searchPos))) {
+			const section = await demuxer.readSection(searchPos, true);
+			assert(section);
+
+			const pesPacket = readPesPacket(section);
+			if (!pesPacket) {
+				throw new Error(MISSING_PES_PACKET_ERROR);
+			}
+
+			const context = new PacketReadingContext(this, pesPacket, false);
+			await this.markNextPacket(context);
+
+			if (!context.suppliedPacket) {
 				continue;
 			}
 
-			// Found a PES packet with a keyframe. Set up a PacketBuffer and pull until we get the keyframe.
-			const keySection = await demuxer.readSection(searchPos, true);
-			assert(keySection);
+			// Check if this packet is a keyframe
+			if (this.getPacketType(context.suppliedPacket.data) !== 'key') {
+				continue;
+			}
 
-			const keyPesPacket = readPesPacket(keySection);
-			assert(keyPesPacket);
+			const buffer = new PacketBuffer(this, context);
 
-			const keyContext = new PacketReadingContext(this, keyPesPacket, true);
-			const keyBuffer = new PacketBuffer(this, keyContext);
-
-			// Pull until we get a keyframe
-			while (true) {
-				const result = await keyBuffer.readNext();
+			const result = await buffer.readNext();
 			assert(result); // How else?
 
-				if (this.getPacketType(result.packet.data) === 'key') {
 			const packet = this.createEncodedPacket(result.packet, result.duration, options);
-					this.packetBuffers.set(packet, keyBuffer);
+			this.packetBuffers.set(packet, buffer);
 			this.packetSectionStarts.set(packet, result.packet.sectionStartPos);
 
 			return packet;
 		}
 	}
 }
-	}
-}
 
 class MpegTsVideoTrackBacking extends MpegTsTrackBacking implements InputVideoTrackBacking {
 	override elementaryStream: ElementaryVideoStream;
@@ -1324,7 +1310,7 @@ class MpegTsVideoTrackBacking extends MpegTsTrackBacking implements InputVideoTr
 		assert(!context.suppliedPacket);
 
 		const codec = this.elementaryStream.info.codec;
-		const CHUNK_SIZE = 128;
+		const CHUNK_SIZE = 1024;
 
 		let packetStartPos: number | null = null;
 
@@ -1332,38 +1318,35 @@ class MpegTsVideoTrackBacking extends MpegTsTrackBacking implements InputVideoTr
 			let remaining = context.ensureBuffered(CHUNK_SIZE);
 			if (remaining instanceof Promise) remaining = await remaining;
 
-			const startPos = context.currentPos;
-
-			while (context.currentPos - startPos < remaining) {
-				const byte = context.readU8();
-
-				// Look for 0x00 as potential start of a start code
-				if (byte !== 0x00) {
-					continue;
+			if (remaining === 0) {
+				break;
 			}
 
+			const chunkStartPos = context.currentPos;
+			const chunk = context.readBytes(remaining);
+			const length = chunk.byteLength;
+
+			let i = 0;
+			while (i < length) {
+				const zeroIndex = chunk.indexOf(0, i);
+				if (zeroIndex === -1 || zeroIndex >= length) {
+					break;
+				}
+				i = zeroIndex;
+
 				// Check if we have enough bytes to identify a start code
-				const posBeforeZero = context.currentPos - 1;
+				const posBeforeZero = chunkStartPos + i;
 
-				let remaining = context.ensureBuffered(4);
-				if (remaining instanceof Promise) remaining = await remaining;
-
-				if (remaining < 4) {
-					// Not enough data left
-					if (packetStartPos !== null) {
-						// Return what we have
-						const packetLength = context.endPos - packetStartPos;
-						context.seekTo(packetStartPos);
-						return context.supplyPacket(packetLength, 0);
+				// Need at least 4 more bytes after the 0x00 to check for start code + NAL type
+				if (i + 4 >= length) {
+					// Not enough data in current chunk, seek back and let the next iteration handle it
+					context.seekTo(posBeforeZero);
+					break;
 				}
 
-					return;
-				}
-
-				// Read potential start code bytes
-				const b1 = context.readU8();
-				const b2 = context.readU8();
-				const b3 = context.readU8();
+				const b1 = chunk[i + 1]!;
+				const b2 = chunk[i + 2]!;
+				const b3 = chunk[i + 3]!;
 
 				let startCodeLength = 0;
 				let nalUnitTypeByte: number | null = null;
@@ -1371,7 +1354,7 @@ class MpegTsVideoTrackBacking extends MpegTsTrackBacking implements InputVideoTr
 				// Check for 4-byte start code (0x00000001)
 				if (b1 === 0x00 && b2 === 0x00 && b3 === 0x01) {
 					startCodeLength = 4;
-					nalUnitTypeByte = context.readU8();
+					nalUnitTypeByte = chunk[i + 4]!;
 				} else if (b1 === 0x00 && b2 === 0x01) {
 					// 3-byte start code (0x000001)
 					startCodeLength = 3;
@@ -1379,8 +1362,8 @@ class MpegTsVideoTrackBacking extends MpegTsTrackBacking implements InputVideoTr
 				}
 
 				if (startCodeLength === 0) {
-					// Not a start code, rewind and continue
-					context.seekTo(posBeforeZero + 1);
+					// Not a start code, continue
+					i++;
 					continue;
 				}
 
@@ -1389,14 +1372,15 @@ class MpegTsVideoTrackBacking extends MpegTsTrackBacking implements InputVideoTr
 				if (packetStartPos === null) {
 					// This is our first start code, mark packet start
 					packetStartPos = startCodePos;
+					i += startCodeLength;
 					continue;
 				}
 
 				// We have a second start code. Check if it's an AUD.
 				if (nalUnitTypeByte !== null) {
 					const nalUnitType = codec === 'avc'
-						? extractNalUnitTypeForAvc(new Uint8Array([nalUnitTypeByte]))
-						: extractNalUnitTypeForHevc(new Uint8Array([nalUnitTypeByte]));
+						? extractNalUnitTypeForAvc(nalUnitTypeByte)
+						: extractNalUnitTypeForHevc(nalUnitTypeByte);
 					const isAud = codec === 'avc'
 						? nalUnitType === AvcNalUnitType.AUD
 						: nalUnitType === HevcNalUnitType.AUD_NUT;
@@ -1410,6 +1394,7 @@ class MpegTsVideoTrackBacking extends MpegTsTrackBacking implements InputVideoTr
 				}
 
 				// Not an AUD, continue searching
+				i += startCodeLength;
 			}
 
 			if (remaining < CHUNK_SIZE) {
@@ -1687,15 +1672,14 @@ class PacketReadingContext {
 		while (true) {
 			this.advanceCurrentPacket();
 			const currentPesPacket = this.getCurrentPesPacket();
-			const relativeStartOffset = 0;
 			const relativeEndOffset = length - offset;
 
 			if (relativeEndOffset <= currentPesPacket.data.byteLength) {
-				result.set(currentPesPacket.data.subarray(relativeStartOffset, relativeEndOffset), offset);
+				result.set(currentPesPacket.data.subarray(0, relativeEndOffset), offset);
 				break;
 			}
 
-			result.set(currentPesPacket.data.subarray(relativeStartOffset), offset);
+			result.set(currentPesPacket.data, offset);
 			offset += currentPesPacket.data.byteLength;
 		}
 
@@ -1864,18 +1848,25 @@ class PacketBuffer {
 			return false;
 		}
 
-		this.context.suppliedPacket = null;
+		let suppliedPacket: SuppliedPacket | null;
+		if (this.context.suppliedPacket) {
+			// Small optimization: there was already a supplied packet in the context, so let's first use that one
+			suppliedPacket = this.context.suppliedPacket;
+		} else {
 			await this.backing.markNextPacket(this.context);
+			suppliedPacket = this.context.suppliedPacket;
+		}
+		this.context.suppliedPacket = null;
 
-		if (!this.context.suppliedPacket) {
+		if (!suppliedPacket) {
 			this.reachedEnd = true;
 			this.flushReorderBuffer();
 
 			return false;
 		}
 
-		this.decodeOrderPackets.push(this.context.suppliedPacket);
-		this.processPacketThroughReorderBuffer(this.context.suppliedPacket);
+		this.decodeOrderPackets.push(suppliedPacket);
+		this.processPacketThroughReorderBuffer(suppliedPacket);
 
 		return true;
 	}

test/node/annex-b-conversion.test.ts

@@ -8,7 +8,7 @@ import { Mp4OutputFormat } from '../../src/output-format.js';
 import { BufferTarget } from '../../src/target.js';
 import { Conversion } from '../../src/conversion.js';
 import { EncodedPacketSink } from '../../src/media-sink.js';
-import { extractAvcNalUnits } from '../../src/codec-data.js';
+import { iterateAvcNalUnits } from '../../src/codec-data.js';
 
 const __dirname = new URL('.', import.meta.url).pathname;
 
@@ -26,7 +26,8 @@ test('Annex B to length-prefixed conversion, MP4', async () => {
 	const originalFirstPacket = await originalSink.getFirstPacket();
 	expect([...originalFirstPacket!.data.slice(0, 4)]).toEqual([0, 0, 0, 1]);
 
-	const originalNalUnits = extractAvcNalUnits(originalFirstPacket!.data, originalDecoderConfig);
+	const originalNalUnits = [...iterateAvcNalUnits(originalFirstPacket!.data, originalDecoderConfig)]
+		.map(loc => originalFirstPacket!.data.subarray(loc.offset, loc.offset + loc.length));
 
 	const output = new Output({
 		format: new Mp4OutputFormat(),
@@ -49,6 +50,7 @@ test('Annex B to length-prefixed conversion, MP4', async () => {
 	const newFirstPacket = await newSink.getFirstPacket();
 	expect([...newFirstPacket!.data.slice(0, 4)]).not.toEqual([0, 0, 0, 1]); // Successfully converted
 
-	const newNalUnits = extractAvcNalUnits(newFirstPacket!.data, newDecoderConfig);
+	const newNalUnits = [...iterateAvcNalUnits(newFirstPacket!.data, newDecoderConfig)]
+		.map(loc => newFirstPacket!.data.subarray(loc.offset, loc.offset + loc.length));
 	expect(newNalUnits).toEqual(originalNalUnits); // Content is the same though
 });

test/node/mpeg-ts-demuxing.test.ts

@@ -385,7 +385,7 @@ test('MPEG-TS video key packets', async () => {
 
 test('MPEG-TS audio key packets', async () => {
 	using input = new Input({
-		source: new FilePathSource(path.join(__dirname, '../public/193039199_mp4_h264_aac_fhd_7.ts')),
+		source: new FilePathSource(path.join(__dirname, '../public/0.ts')),
 		formats: ALL_FORMATS,
 	});
 
@@ -513,9 +513,12 @@ test('MPEG-TS with HEVC video', async () => {
 
 	const sink = new EncodedPacketSink(videoTrack);
 
+	let i = 0;
 	for await (const packet of sink.packets()) {
 		expect(packet.data.slice(0, 4)).toEqual(new Uint8Array([0, 0, 0, 1])); // Annex B
 		expect(packet.duration).toBeCloseTo(0.04166666666);
+		expect(packet.type).toBe(i > 0 ? 'delta' : 'key');
+		i++;
 	}
 });
 
@@ -543,13 +546,10 @@ test('MPEG-TS with MP3 audio', async () => {
 	const firstPacket = await sink.getFirstPacket();
 	assert(firstPacket);
 
-	expect(firstPacket.data[0]).toBe(0xff); // MP3 sync byte
-	expect(firstPacket.type).toBe('key');
-	expect(firstPacket.duration).toBeGreaterThan(0);
-
 	let count = 0;
 	for await (const packet of sink.packets()) {
 		expect(packet.data[0]).toBe(0xff);
+		expect(packet.type).toBe('key');
 		count++;
 	}