media_codec/
packet.rs

1use std::{
2    io::{self, Read, Write},
3    sync::Arc,
4};
5
6use bitflags::bitflags;
7use media_core::{buffer::Buffer, error::Error, invalid_param_error, Result};
8use num_rational::Rational64;
9
10bitflags! {
11    #[repr(transparent)]
12    #[derive(Clone, Copy, Debug, Eq, PartialEq)]
13    pub struct PacketFlags: u32 {
14        const Key = 1;
15        const Corrupt = 2;
16    }
17}
18
19#[derive(Clone)]
20enum PacketData<'a> {
21    Borrowed(&'a [u8]),
22    Owned(Vec<u8>),
23    Buffer(Arc<Buffer>),
24}
25
26impl PacketData<'_> {
27    fn from_slice(slice: &[u8]) -> PacketData<'_> {
28        PacketData::Borrowed(slice)
29    }
30
31    fn from_vec(vec: Vec<u8>) -> PacketData<'static> {
32        PacketData::Owned(vec)
33    }
34
35    fn from_buffer(buffer: Arc<Buffer>) -> PacketData<'static> {
36        PacketData::Buffer(buffer)
37    }
38
39    fn as_slice(&self) -> &[u8] {
40        match self {
41            PacketData::Borrowed(slice) => slice,
42            PacketData::Owned(vec) => vec.as_slice(),
43            PacketData::Buffer(buffer) => buffer.data(),
44        }
45    }
46
47    fn as_mut_slice(&mut self) -> Option<&mut [u8]> {
48        match self {
49            PacketData::Borrowed(_) => None,
50            PacketData::Owned(ref mut vec) => Some(vec.as_mut_slice()),
51            PacketData::Buffer(buffer) => {
52                let mut_buffer = Arc::get_mut(buffer)?;
53                Some(mut_buffer.data_mut())
54            }
55        }
56    }
57
58    fn len(&self) -> usize {
59        match self {
60            PacketData::Borrowed(slice) => slice.len(),
61            PacketData::Owned(vec) => vec.len(),
62            PacketData::Buffer(buffer) => buffer.len(),
63        }
64    }
65
66    fn is_empty(&self) -> bool {
67        self.len() == 0
68    }
69
70    fn capacity(&self) -> usize {
71        match &self {
72            PacketData::Borrowed(slice) => slice.len(),
73            PacketData::Owned(vec) => vec.capacity(),
74            PacketData::Buffer(buffer) => buffer.capacity(),
75        }
76    }
77
78    fn into_owned(self) -> PacketData<'static> {
79        match self {
80            PacketData::Borrowed(slice) => PacketData::Owned(slice.to_vec()),
81            PacketData::Owned(vec) => PacketData::Owned(vec),
82            PacketData::Buffer(buffer) => PacketData::Buffer(buffer),
83        }
84    }
85}
86
87#[derive(Clone)]
88pub struct Packet<'a> {
89    pub pts: Option<i64>,
90    pub dts: Option<i64>,
91    pub duration: Option<i64>,
92    pub time_base: Option<Rational64>,
93    pub flags: PacketFlags,
94    pub pos: Option<usize>,
95    pub stream_index: Option<usize>,
96    data: PacketData<'a>,
97}
98
99impl<'a> Packet<'a> {
100    fn from_data(data: PacketData<'a>) -> Self {
101        Self {
102            pts: None,
103            dts: None,
104            duration: None,
105            time_base: None,
106            flags: PacketFlags::empty(),
107            pos: None,
108            stream_index: None,
109            data,
110        }
111    }
112
113    pub fn new(size: usize) -> Self {
114        Self::from_data(PacketData::from_vec(vec![0; size]))
115    }
116
117    pub fn with_capacity(capacity: usize) -> Self {
118        Self::from_data(PacketData::from_vec(Vec::with_capacity(capacity)))
119    }
120
121    pub fn from_slice(data: &'a [u8]) -> Self {
122        Self::from_data(PacketData::from_slice(data))
123    }
124
125    pub fn from_buffer(buffer: Arc<Buffer>) -> Self {
126        Self::from_data(PacketData::from_buffer(buffer))
127    }
128
129    pub fn into_owned(self) -> Packet<'static> {
130        Packet {
131            pts: self.pts,
132            dts: self.dts,
133            duration: self.duration,
134            time_base: self.time_base,
135            flags: self.flags,
136            pos: self.pos,
137            stream_index: self.stream_index,
138            data: self.data.into_owned(),
139        }
140    }
141
142    pub fn len(&self) -> usize {
143        self.data.len()
144    }
145
146    pub fn is_empty(&self) -> bool {
147        self.data.is_empty()
148    }
149
150    pub fn capacity(&self) -> usize {
151        self.data.capacity()
152    }
153
154    pub fn data(&self) -> &[u8] {
155        self.data.as_slice()
156    }
157
158    pub fn data_mut(&mut self) -> Option<&mut [u8]> {
159        self.data.as_mut_slice()
160    }
161
162    pub fn truncate(&mut self, len: usize) -> Result<()> {
163        let current_len = self.data.len();
164        if len > current_len {
165            return Err(invalid_param_error!(len));
166        }
167
168        match &mut self.data {
169            PacketData::Borrowed(slice) => {
170                self.data = PacketData::Borrowed(&slice[..len]);
171            }
172            PacketData::Owned(vec) => {
173                vec.truncate(len);
174            }
175            PacketData::Buffer(buffer) => {
176                let buffer = Arc::get_mut(buffer).ok_or_else(|| Error::Invalid("buffer is shared".to_string()))?;
177                buffer.resize(len);
178            }
179        }
180
181        Ok(())
182    }
183}
184
185pub trait ReadPacket: Read {
186    fn read_packet(&mut self, size: usize) -> io::Result<Packet<'_>> {
187        let mut packet = Packet::new(size);
188
189        if let Some(data_mut) = packet.data_mut() {
190            self.read_exact(data_mut)?;
191        } else {
192            // Packet::new always creates an owned packet
193            unreachable!()
194        }
195
196        Ok(packet)
197    }
198}
199
200pub trait WritePacket: Write {
201    fn write_packet(&mut self, packet: &Packet) -> io::Result<()> {
202        self.write_all(packet.data())
203    }
204}
205
206impl<T: Read + ?Sized> ReadPacket for T {}
207impl<T: Write + ?Sized> WritePacket for T {}