rtp-packet-list-builder.rs

use bytes; // 1.5.0

pub struct RtpPacketBuilder<P> {
    timestamp: u16,
    payload: Option<P>,
}

impl<P> RtpPacketBuilder<P> {
    pub fn write<O>(self, vec: &mut (dyn RtpWriter<Payload = P, Output = O>)) -> Option<O> {
        vec.push(&[(self.timestamp >> 8) as u8]);
        vec.push(&[(self.timestamp & 0xff) as u8]);
        if let Some(payload) = self.payload {
            vec.push_payload(payload)
        }
        vec.finish()
    }
    
    pub fn new() -> Self {
        Self {
            timestamp: 0,
            payload: None,
        }
    }
}

#[derive(Debug, Clone)]
pub struct RtpPacketHeaderBuilder {
    timestamp: u16,
    // ...
}

impl RtpPacketHeaderBuilder {
    pub fn timestamp(&self) -> u16 {
        self.timestamp
    }
    pub fn set_timestamp(&mut self, timestamp: u16) {
        self.timestamp = timestamp;
    }
    // ...
}

pub struct RtpPacketListBuilder<'a, O, P> {
    hdr: RtpPacketHeaderBuilder,
    reader: Box<dyn FnMut(&mut RtpPacketHeaderBuilder) -> Option<P> + 'a>,
    writer: Box<dyn RtpWriter<Payload = P, Output = O> + 'a>,
}

impl<'a, O, P> RtpPacketListBuilder<'a, O, P> {
    pub fn new<R: FnMut(&mut RtpPacketHeaderBuilder) -> Option<P> + 'a, W: RtpWriter<Payload = P, Output = O> + 'a>(reader: R, writer: W) -> Self
    {
        Self {
            hdr: RtpPacketHeaderBuilder { timestamp: 0 },
            reader: Box::new(reader),
            writer: Box::new(writer),
        }
    }
}

impl<'a, O, P> Iterator for RtpPacketListBuilder<'a, O, P> {
    type Item = O;

    fn next(&mut self) -> Option<Self::Item> {
        let Some(payload) = (self.reader)(&mut self.hdr) else {
            return None;
        };
        let mut builder = RtpPacketBuilder::<P>::new();
        builder.payload = Some(payload);
        builder.timestamp = self.hdr.timestamp;
        let ret = builder.write(self.writer.as_mut());
        ret
    }
}

pub trait RtpWriter {
    type Output;
    type Payload;

    fn push(&mut self, data: &[u8]);
    fn push_payload(&mut self, payload: Self::Payload);
    fn finish(&mut self) -> Option<Self::Output>;
}

#[derive(Default, Debug)]
pub struct PacketVec<'a> {
    header: Option<Vec<u8>>,
    payload: Vec<&'a [u8]>,
}

impl<'a> RtpWriter for PacketVec<'a> {
    type Output = Vec<Vec<u8>>;
    type Payload = &'a [u8];

    fn push(&mut self, data: &[u8]) {
        let p = self.header.get_or_insert_with(|| Vec::new());
        p.extend_from_slice(data)
    }

    fn push_payload(&mut self, data: Self::Payload) {
        self.payload.push(data);
    }

    fn finish(&mut self) -> Option<Self::Output> {
        if self.header.is_none() && self.payload.is_empty() {
            return None;
        }
        let mut ret = vec![];
        if let Some(hdr) = self.header.take() {
            ret.push(hdr);
        }
        for payload in self.payload.drain(0..=self.payload.len() - 1) {
            ret.push(payload.to_vec());
        }
        Some(ret)
    }
}

#[derive(Default, Debug)]
pub struct PacketBytes {
    current_packet: Option<bytes::BytesMut>,
}

impl RtpWriter for PacketBytes {
    type Output = bytes::Bytes;
    type Payload = bytes::Bytes;

    fn push(&mut self, data: &[u8]) {
        let p = self.current_packet.get_or_insert_with(|| bytes::BytesMut::new());
        p.extend_from_slice(data)
    }

    fn push_payload(&mut self, data: Self::Payload) {
        let p = self.current_packet.get_or_insert_with(|| bytes::BytesMut::new());
        p.extend_from_slice(&data)
    }
    
    fn finish(&mut self) -> Option<Self::Output> {
        self.current_packet.take().map(|p| p.freeze())
    }
}

fn main() {
    let mut ts = 60;
    let mut i = 0;
    let data = bytes::Bytes::from_static([0x01, 0x02, 0x03, 0x04, 0x05, 0x06, 0x07, 0x08, 0x09].as_ref());
    
    let packet_list = PacketBytes::default();

    let pl = RtpPacketListBuilder::new(|hdr: &mut RtpPacketHeaderBuilder| {
        hdr.set_timestamp(ts);
        ts += 1;
        i += 1;
        if ts < 70 {
            println!("{i} {}", data.len());
            Some(data.slice(0..=data.len() - i))
        } else {
            None
        }
    }, packet_list);

    for packet in pl {
        println!("{packet:?}")
    }
}