vk-video 0.3.1

[DEPRECATED] Renamed to `gpu-video`
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86

use std::sync::LazyLock;

use bytes::{BufMut, BytesMut};
use memchr::memmem::Finder;

#[derive(Debug, Default)]
pub(crate) struct NALUSplitter {
    buffer: BytesMut,
    pts: Option<u64>,
    previous_search_end: usize,
}

fn find_start_of_next_nalu(buf: &[u8]) -> Option<usize> {
    static FINDER: LazyLock<Finder> = LazyLock::new(|| Finder::new(&[0, 0, 1]));

    if buf.len() < 4 {
        return None;
    };

    // If the start code is at the beginning of the buffer, we need to skip it, because this means we
    // would give the parser only the start code, without a nal unit before it.
    //
    // The code can be either 0, 0, 0, 1 or 0, 0, 1.
    // We're looking for the sequence 0, 0, 1. If we find it starting at the second byte of input,
    // this could mean either that it's the longer 0, 0, 0, 1 code, or that there is one byte of
    // the previous nalu and the shorter code. This is what is checked here.
    if buf[0] != 0 && buf[1..4] == [0, 0, 1] {
        return Some(5);
    }

    FINDER.find(&buf[2..]).map(|i| i + 5)
}

impl NALUSplitter {
    pub(crate) fn push(
        &mut self,
        bytestream: &[u8],
        pts: Option<u64>,
    ) -> Vec<(Vec<u8>, Option<u64>)> {
        let mut output_pts = if self.buffer.is_empty() {
            pts
        } else {
            self.pts
        };

        self.buffer.put(bytestream);
        let mut result = Vec::new();

        while let Some(i) = find_start_of_next_nalu(&self.buffer[self.previous_search_end..]) {
            let nalu = self.buffer.split_to(self.previous_search_end + i);
            self.previous_search_end = 0;
            result.push((nalu.to_vec(), output_pts));
            output_pts = pts;
        }

        // This will cause the whole start code to be reprocessed when the beginning of the next start code
        // is at the end of current buffer.
        self.previous_search_end = self.buffer.len().saturating_sub(3);

        self.pts = pts;

        result
    }

    pub(crate) fn flush(&mut self) -> Vec<(Vec<u8>, Option<u64>)> {
        if self.buffer.is_empty() {
            return vec![];
        }

        let mut result = Vec::new();
        while let Some(i) = find_start_of_next_nalu(&self.buffer[self.previous_search_end..]) {
            let nalu = self.buffer.split_to(self.previous_search_end + i);
            self.previous_search_end = 0;
            result.push((nalu.to_vec(), self.pts));
        }

        result.push((
            [self.buffer.to_vec(), [0, 0, 1].to_vec()].concat(),
            self.pts,
        ));
        self.buffer = BytesMut::new();
        self.previous_search_end = 0;

        result
    }
}