1#![cfg_attr(
28 feature = "backtrace",
29 feature(error_generic_member_access, provide_any)
30)]
31
32use std::{
33 mem::MaybeUninit,
34 os::raw::{c_int, c_uint, c_ulong},
35};
36
37#[cfg(feature = "backtrace")]
38use std::backtrace::Backtrace;
39use std::{ptr, slice};
40
41use thiserror::Error;
42
43#[cfg(feature = "vp9")]
44use vpx_sys::vp8e_enc_control_id::*;
45use vpx_sys::vpx_codec_cx_pkt_kind::VPX_CODEC_CX_FRAME_PKT;
46use vpx_sys::*;
47
48#[derive(Copy, Clone, Debug, PartialEq, Eq, Hash)]
49pub enum VideoCodecId {
50 VP8,
51 #[cfg(feature = "vp9")]
52 VP9,
53}
54
55impl Default for VideoCodecId {
56 #[cfg(not(feature = "vp9"))]
57 fn default() -> VideoCodecId {
58 VideoCodecId::VP8
59 }
60
61 #[cfg(feature = "vp9")]
62 fn default() -> VideoCodecId {
63 VideoCodecId::VP9
64 }
65}
66
67pub struct Encoder {
68 ctx: vpx_codec_ctx_t,
69 width: usize,
70 height: usize,
71}
72
73#[derive(Debug, Error)]
74#[error("VPX encode error: {msg}")]
75pub struct Error {
76 msg: String,
77 #[cfg(feature = "backtrace")]
78 #[backtrace]
79 backtrace: Backtrace,
80}
81
82impl From<String> for Error {
83 fn from(msg: String) -> Self {
84 Self {
85 msg,
86 #[cfg(feature = "backtrace")]
87 backtrace: Backtrace::capture(),
88 }
89 }
90}
91
92pub type Result<T> = std::result::Result<T, Error>;
93
94macro_rules! call_vpx {
95 ($x:expr) => {{
96 let result = unsafe { $x }; let result_int = unsafe { std::mem::transmute::<_, i32>(result) };
98 if result_int != 0 {
100 return Err(Error::from(format!(
101 "Function call failed (error code {}).",
102 result_int
103 )));
104 }
105 result
106 }};
107}
108
109macro_rules! call_vpx_ptr {
110 ($x:expr) => {{
111 let result = unsafe { $x }; if result.is_null() {
113 return Err(Error::from("Bad pointer.".to_string()));
114 }
115 result
116 }};
117}
118
119impl Encoder {
120 pub fn new(config: Config) -> Result<Self> {
121 let i = match config.codec {
122 VideoCodecId::VP8 => call_vpx_ptr!(vpx_codec_vp8_cx()),
123 #[cfg(feature = "vp9")]
124 VideoCodecId::VP9 => call_vpx_ptr!(vpx_codec_vp9_cx()),
125 };
126
127 if config.width % 2 != 0 {
128 return Err(Error::from("Width must be divisible by 2".to_string()));
129 }
130 if config.height % 2 != 0 {
131 return Err(Error::from("Height must be divisible by 2".to_string()));
132 }
133
134 let c = MaybeUninit::zeroed();
135 let mut c = unsafe { c.assume_init() };
136 call_vpx!(vpx_codec_enc_config_default(i, &mut c, 0));
137
138 c.g_w = config.width;
139 c.g_h = config.height;
140 c.g_timebase.num = config.timebase[0];
141 c.g_timebase.den = config.timebase[1];
142 c.rc_target_bitrate = config.bitrate;
143
144 c.g_threads = 8;
145 c.g_error_resilient = VPX_ERROR_RESILIENT_DEFAULT;
146
147 let ctx = MaybeUninit::zeroed();
148 let mut ctx = unsafe { ctx.assume_init() };
149
150 match config.codec {
151 VideoCodecId::VP8 => {
152 call_vpx!(vpx_codec_enc_init_ver(
153 &mut ctx,
154 i,
155 &c,
156 0,
157 vpx_sys::VPX_ENCODER_ABI_VERSION as i32
158 ));
159 }
160 #[cfg(feature = "vp9")]
161 VideoCodecId::VP9 => {
162 call_vpx!(vpx_codec_enc_init_ver(
163 &mut ctx,
164 i,
165 &c,
166 0,
167 vpx_sys::VPX_ENCODER_ABI_VERSION as i32
168 ));
169 call_vpx!(vpx_codec_control_(
171 &mut ctx,
172 VP8E_SET_CPUUSED as _,
173 6 as c_int
174 ));
175 call_vpx!(vpx_codec_control_(
177 &mut ctx,
178 VP9E_SET_ROW_MT as _,
179 1 as c_int
180 ));
181 }
182 };
183
184 Ok(Self {
185 ctx,
186 width: config.width as usize,
187 height: config.height as usize,
188 })
189 }
190
191 pub fn encode(&mut self, pts: i64, data: &[u8]) -> Result<Packets> {
192 assert!(2 * data.len() >= 3 * self.width * self.height);
193
194 let image = MaybeUninit::zeroed();
195 let mut image = unsafe { image.assume_init() };
196
197 call_vpx_ptr!(vpx_img_wrap(
198 &mut image,
199 vpx_img_fmt::VPX_IMG_FMT_I420,
200 self.width as _,
201 self.height as _,
202 1,
203 data.as_ptr() as _,
204 ));
205
206 call_vpx!(vpx_codec_encode(
207 &mut self.ctx,
208 &image,
209 pts,
210 1, 0, vpx_sys::VPX_DL_REALTIME as c_ulong,
213 ));
214
215 Ok(Packets {
216 ctx: &mut self.ctx,
217 iter: ptr::null(),
218 })
219 }
220
221 pub fn finish(mut self) -> Result<Finish> {
222 call_vpx!(vpx_codec_encode(
223 &mut self.ctx,
224 ptr::null(),
225 -1, 1, 0, vpx_sys::VPX_DL_REALTIME as c_ulong,
229 ));
230
231 Ok(Finish {
232 enc: self,
233 iter: ptr::null(),
234 })
235 }
236}
237
238impl Drop for Encoder {
239 fn drop(&mut self) {
240 unsafe {
241 let result = vpx_codec_destroy(&mut self.ctx);
242 if result != vpx_sys::VPX_CODEC_OK {
243 eprintln!("failed to destroy vpx codec: {result:?}");
244 }
245 }
246 }
247}
248
249#[derive(Clone, Copy, Debug)]
250pub struct Frame<'a> {
251 pub data: &'a [u8],
253 pub key: bool,
255 pub pts: i64,
257}
258
259#[derive(Clone, Copy, Debug)]
260pub struct Config {
261 pub width: c_uint,
263 pub height: c_uint,
265 pub timebase: [c_int; 2],
267 pub bitrate: c_uint,
269 pub codec: VideoCodecId,
271}
272
273pub struct Packets<'a> {
274 ctx: &'a mut vpx_codec_ctx_t,
275 iter: vpx_codec_iter_t,
276}
277
278impl<'a> Iterator for Packets<'a> {
279 type Item = Frame<'a>;
280 fn next(&mut self) -> Option<Self::Item> {
281 loop {
282 unsafe {
283 let pkt = vpx_codec_get_cx_data(self.ctx, &mut self.iter);
284 if pkt.is_null() {
285 return None;
286 } else if (*pkt).kind == VPX_CODEC_CX_FRAME_PKT {
287 let f = &(*pkt).data.frame;
288 return Some(Frame {
289 data: slice::from_raw_parts(f.buf as _, f.sz as usize),
290 key: (f.flags & VPX_FRAME_IS_KEY) != 0,
291 pts: f.pts,
292 });
293 } else {
294 }
296 }
297 }
298 }
299}
300
301pub struct Finish {
302 enc: Encoder,
303 iter: vpx_codec_iter_t,
304}
305
306impl Finish {
307 pub fn next(&mut self) -> Result<Option<Frame>> {
308 let mut tmp = Packets {
309 ctx: &mut self.enc.ctx,
310 iter: self.iter,
311 };
312
313 if let Some(packet) = tmp.next() {
314 self.iter = tmp.iter;
315 Ok(Some(packet))
316 } else {
317 call_vpx!(vpx_codec_encode(
318 tmp.ctx,
319 ptr::null(),
320 -1, 1, 0, vpx_sys::VPX_DL_REALTIME as c_ulong,
324 ));
325
326 tmp.iter = ptr::null();
327 if let Some(packet) = tmp.next() {
328 self.iter = tmp.iter;
329 Ok(Some(packet))
330 } else {
331 Ok(None)
332 }
333 }
334 }
335}