small improvement for video data demux

ScuffleCloud · Jan 13, 2025 · 711b1d5 · 711b1d5
1 parent 8b4d1b0
commit 711b1d5
Showing 1 changed file with 76 additions and 59 deletions.
diff --git a/crates/flv/src/video.rs b/crates/flv/src/video.rs
@@ -35,14 +35,12 @@ impl VideoData {
     pub fn demux(reader: &mut io::Cursor<Bytes>) -> io::Result<Self> {
         let byte = reader.read_u8()?;
         let enhanced = (byte & 0b1000_0000) != 0;
-        let frame_type = FrameType::from((byte >> 4) & 0b0111);
-        let codec_id = byte & 0b0000_1111;
-        let body = if enhanced {
-            VideoDataBody::demux_enhanced(EnhancedPacketType::from(codec_id), reader)?
-        } else {
-            VideoDataBody::demux(VideoCodecId::from(codec_id), reader)?
-        };
-        Ok(VideoData { frame_type, body })
+        let frame_type_byte = (byte >> 4) & 0b0111;
+        let packet_type_byte = byte & 0b0000_1111;
+        Ok(VideoData {
+            frame_type: FrameType::from(frame_type_byte),
+            body: VideoDataBody::demux(VideoPacketType::new(packet_type_byte, enhanced), reader)?,
+        })
     }
 }
 
@@ -75,63 +73,82 @@ pub enum VideoDataBody {
     Unknown { codec_id: VideoCodecId, data: Bytes },
 }
 
-impl VideoDataBody {
-    pub fn demux(codec_id: VideoCodecId, reader: &mut io::Cursor<Bytes>) -> io::Result<Self> {
-        match codec_id {
-            VideoCodecId::Avc => {
-                let avc_packet_type = AvcPacketType::from(reader.read_u8()?);
-                Ok(VideoDataBody::Avc(AvcPacket::demux(avc_packet_type, reader)?))
-            }
-            _ => Ok(VideoDataBody::Unknown {
-                codec_id,
-                data: reader.extract_remaining(),
-            }),
+#[derive(Debug, Clone, PartialEq, Copy, Eq, PartialOrd, Ord, Hash)]
+pub enum VideoPacketType {
+    CodecId(VideoCodecId),
+    Enhanced(EnhancedPacketType),
+}
+
+impl VideoPacketType {
+    pub fn new(byte: u8, enhanced: bool) -> Self {
+        if enhanced {
+            Self::Enhanced(EnhancedPacketType::from(byte))
+        } else {
+            Self::CodecId(VideoCodecId::from(byte))
         }
     }
+}
 
-    pub fn demux_enhanced(packet_type: EnhancedPacketType, reader: &mut io::Cursor<Bytes>) -> io::Result<Self> {
-        let mut video_codec = [0; 4];
-        reader.read_exact(&mut video_codec)?;
-        let video_codec = VideoFourCC::from(video_codec);
-
+impl VideoDataBody {
+    pub fn demux(packet_type: VideoPacketType, reader: &mut io::Cursor<Bytes>) -> io::Result<Self> {
         match packet_type {
-            EnhancedPacketType::SequenceEnd => return Ok(VideoDataBody::Enhanced(EnhancedPacket::SequenceEnd)),
-            EnhancedPacketType::Metadata => {
-                return Ok(VideoDataBody::Enhanced(EnhancedPacket::Metadata(reader.extract_remaining())))
-            }
-            _ => {}
-        }
-
-        match (video_codec, packet_type) {
-            (VideoFourCC::Av1, EnhancedPacketType::SequenceStart) => Ok(VideoDataBody::Enhanced(EnhancedPacket::Av1(
-                Av1Packet::SequenceStart(AV1CodecConfigurationRecord::demux(reader)?),
-            ))),
-            (VideoFourCC::Av1, EnhancedPacketType::Mpeg2SequenceStart) => Ok(VideoDataBody::Enhanced(EnhancedPacket::Av1(
-                Av1Packet::SequenceStart(AV1VideoDescriptor::demux(reader)?.codec_configuration_record),
-            ))),
-            (VideoFourCC::Av1, EnhancedPacketType::CodedFrames) => Ok(VideoDataBody::Enhanced(EnhancedPacket::Av1(
-                Av1Packet::Raw(reader.extract_remaining()),
-            ))),
-            (VideoFourCC::Hevc, EnhancedPacketType::SequenceStart) => Ok(VideoDataBody::Enhanced(EnhancedPacket::Hevc(
-                HevcPacket::SequenceStart(HEVCDecoderConfigurationRecord::demux(reader)?),
-            ))),
-            (VideoFourCC::Hevc, EnhancedPacketType::CodedFrames) => {
-                Ok(VideoDataBody::Enhanced(EnhancedPacket::Hevc(HevcPacket::Nalu {
-                    composition_time: Some(reader.read_i24::<BigEndian>()?),
+            VideoPacketType::CodecId(codec_id) => match codec_id {
+                VideoCodecId::Avc => {
+                    let avc_packet_type = AvcPacketType::from(reader.read_u8()?);
+                    Ok(VideoDataBody::Avc(AvcPacket::demux(avc_packet_type, reader)?))
+                }
+                _ => Ok(VideoDataBody::Unknown {
+                    codec_id,
                     data: reader.extract_remaining(),
-                })))
-            }
-            (VideoFourCC::Hevc, EnhancedPacketType::CodedFramesX) => {
-                Ok(VideoDataBody::Enhanced(EnhancedPacket::Hevc(HevcPacket::Nalu {
-                    composition_time: None,
-                    data: reader.extract_remaining(),
-                })))
+                }),
+            },
+            VideoPacketType::Enhanced(packet_type) => {
+                let mut video_codec = [0; 4];
+                reader.read_exact(&mut video_codec)?;
+                let video_codec = VideoFourCC::from(video_codec);
+
+                match packet_type {
+                    EnhancedPacketType::SequenceEnd => return Ok(VideoDataBody::Enhanced(EnhancedPacket::SequenceEnd)),
+                    EnhancedPacketType::Metadata => {
+                        return Ok(VideoDataBody::Enhanced(EnhancedPacket::Metadata(reader.extract_remaining())))
+                    }
+                    _ => {}
+                }
+
+                match (video_codec, packet_type) {
+                    (VideoFourCC::Av1, EnhancedPacketType::SequenceStart) => Ok(VideoDataBody::Enhanced(
+                        EnhancedPacket::Av1(Av1Packet::SequenceStart(AV1CodecConfigurationRecord::demux(reader)?)),
+                    )),
+                    (VideoFourCC::Av1, EnhancedPacketType::Mpeg2SequenceStart) => {
+                        Ok(VideoDataBody::Enhanced(EnhancedPacket::Av1(Av1Packet::SequenceStart(
+                            AV1VideoDescriptor::demux(reader)?.codec_configuration_record,
+                        ))))
+                    }
+                    (VideoFourCC::Av1, EnhancedPacketType::CodedFrames) => Ok(VideoDataBody::Enhanced(EnhancedPacket::Av1(
+                        Av1Packet::Raw(reader.extract_remaining()),
+                    ))),
+                    (VideoFourCC::Hevc, EnhancedPacketType::SequenceStart) => Ok(VideoDataBody::Enhanced(
+                        EnhancedPacket::Hevc(HevcPacket::SequenceStart(HEVCDecoderConfigurationRecord::demux(reader)?)),
+                    )),
+                    (VideoFourCC::Hevc, EnhancedPacketType::CodedFrames) => {
+                        Ok(VideoDataBody::Enhanced(EnhancedPacket::Hevc(HevcPacket::Nalu {
+                            composition_time: Some(reader.read_i24::<BigEndian>()?),
+                            data: reader.extract_remaining(),
+                        })))
+                    }
+                    (VideoFourCC::Hevc, EnhancedPacketType::CodedFramesX) => {
+                        Ok(VideoDataBody::Enhanced(EnhancedPacket::Hevc(HevcPacket::Nalu {
+                            composition_time: None,
+                            data: reader.extract_remaining(),
+                        })))
+                    }
+                    _ => Ok(VideoDataBody::Enhanced(EnhancedPacket::Unknown {
+                        packet_type,
+                        video_codec,
+                        data: reader.extract_remaining(),
+                    })),
+                }
             }
-            _ => Ok(VideoDataBody::Enhanced(EnhancedPacket::Unknown {
-                packet_type,
-                video_codec,
-                data: reader.extract_remaining(),
-            })),
         }
     }
 }