avi.rs

#[macro_use] extern crate macro_attr;
#[macro_use] extern crate bmp;
use std::fs;
use std::mem;
use std::io::{BufWriter, Write};
use std::iter;
use bmp::{Image, Pixel};

trait Serialize where Self: Sized {
    fn serialize(&self) -> Vec<u8>;

    fn sizeof(&self) -> u32 {
        mem::size_of::<Self>() as u32
    }
}

macro_rules! Serialize {
    (() $(pub)* struct $name:ident { $($i:ident : $t:ty ,)* }) => { Serialize! { @impl $name $( $i )* } };

    (@impl $name:ident $( $i:ident )*) => {
        impl Serialize for $name {
            fn serialize(&self) -> Vec<u8> {
                let mut vec = Vec::new();
                $(
                    vec.extend(self.$i.serialize());
                )*
                vec
            }

            fn sizeof(&self) -> u32 {
                0 $( + self.$i.sizeof() )*
            }
        }
    };
}

impl Serialize for u8 {
    fn serialize(&self) -> Vec<u8> {
        unsafe { mem::transmute::<u8, [u8; 1]>(*self).to_vec() }
    }
}

impl Serialize for i16 {
    fn serialize(&self) -> Vec<u8> {
        unsafe { mem::transmute::<i16, [u8; 2]>(*self).to_vec() }
    }
}

impl Serialize for u16 {
    fn serialize(&self) -> Vec<u8> {
        unsafe { mem::transmute::<u16, [u8; 2]>(*self).to_vec() }
    }
}

impl Serialize for i32 {
    fn serialize(&self) -> Vec<u8> {
        unsafe { mem::transmute::<i32, [u8; 4]>(*self).to_vec() }
    }
}

impl Serialize for u32 {
    fn serialize(&self) -> Vec<u8> {
        unsafe { mem::transmute::<u32, [u8; 4]>(*self).to_vec() }
    }
}

impl Serialize for [u32; 4] {
    fn serialize(&self) -> Vec<u8> {
        unsafe { mem::transmute::<[u32; 4], [u8; 16]>(*self).to_vec() }
    }
}

impl<A: Serialize> Serialize for Vec<A> {
    fn serialize(&self) -> Vec<u8> {
        self.iter().map(|i| i.serialize()).flatten().collect()
    }

    fn sizeof(&self) -> u32 {
        self.iter().map(|i| i.sizeof()).sum()
    }
}

// ------

struct FourCC([u8;4]);

impl<'a> From<&'a str> for FourCC {
    fn from(s: &'a str) -> FourCC {
        assert_eq!(s.len(), 4);
        let s = s.clone().as_bytes();

        FourCC([s[0], s[1], s[2], s[3]])
    }
}

impl Serialize for FourCC {
    fn serialize(&self) -> Vec<u8> {
        self.0.to_vec()
    }
}

enum RiffData {
    Chunk(Chunk),
    List(List),
}

impl RiffData {
    fn chunk<D: Serialize>(ck_id_str: &str, ck_data: D) -> RiffData {
        RiffData::Chunk(Chunk::new(FourCC::from(ck_id_str), ck_data))
    }

    fn list(list_type_str: &str, list_data: Vec<RiffData>) -> RiffData {
        RiffData::List(List::new(FourCC::from(list_type_str), list_data))
    }
}

impl Serialize for RiffData {
    fn serialize(&self) -> Vec<u8> {
        match self {
            RiffData::Chunk(c) => c.serialize(),
            RiffData::List(c) => c.serialize(),
        }
    }

    fn sizeof(&self) -> u32 {
        match self {
            RiffData::Chunk(c) => c.sizeof(),
            RiffData::List(c) => c.sizeof(),
        }
    }
}

macro_attr! {
    #[derive(Serialize!)]
    struct Chunk {
        ck_id: FourCC,
        ck_size: u32,
        ck_data: Vec<u8>,
    }
}

impl Chunk {
    fn new<S: Serialize>(ck_id: FourCC, ck_data: S) -> Chunk {
        Chunk {
            ck_id: ck_id,
            ck_size: ck_data.sizeof(),
            ck_data: ck_data.serialize(),
        }
    }
}

macro_attr! {
    #[derive(Serialize!)]
    struct List {
        list_id: FourCC,
        list_size: u32,
        list_type: FourCC,
        list_data: Vec<RiffData>,
    }
}

impl List {
    fn new(list_type: FourCC, list_data: Vec<RiffData>) -> List {
        List {
            list_id: FourCC::from("LIST"),
            list_size: list_type.sizeof() + list_data.sizeof(),
            list_type: list_type,
            list_data: list_data,
        }
    }
}

macro_attr!{
    #[derive(Serialize!)]
    struct AviMainHeader {
        dw_micro_sec_per_frame: u32,
        dw_max_bytes_per_sec: u32,
        dw_padding_granularity: u32,
        dw_flags: u32,
        dw_total_frames: u32,
        dw_initial_frames: u32,
        dw_streams: u32,
        dw_suggested_buffer_size: u32,
        dw_width: u32,
        dw_height: u32,
        dw_reserved: [u32; 4],
    }
}

impl AviMainHeader {
    fn new(
        micro_per_frame: u32,
        max_bytes_per_sec: u32,
        padding_granularity: u32,
        flags: u32,
        total_frames: u32,
        width: u32,
        height: u32
    ) -> AviMainHeader {
        AviMainHeader {
            dw_micro_sec_per_frame: micro_per_frame,
            dw_max_bytes_per_sec: max_bytes_per_sec,
            dw_padding_granularity: padding_granularity,
            dw_flags: flags,
            dw_total_frames: total_frames,
            dw_initial_frames: 0, // not for interleaved videos
            dw_streams: 1, // for video and audio media
            dw_suggested_buffer_size: 5 * 30 * width * height,
            dw_width: width,
            dw_height: height,
            dw_reserved: [0; 4],
        }
    }
}

macro_attr! {
    #[derive(Serialize!)]
    struct RcFrame {
        left: i16,
        top: i16,
        right: i16,
        bottom: i16,
    }
}

enum FccType {
    Auds,
    Mids,
    Txts,
    Vids,
}

impl FccType {
    fn to_fcc(&self) -> FourCC {
        let fcc = match self {
            FccType::Auds => "auds",
            FccType::Mids => "mids",
            FccType::Txts => "txts",
            FccType::Vids => "vids",
        };

        FourCC::from(fcc)
    }
}

macro_attr! {
    #[derive(Serialize!)]
    struct AviStreamHeader {
        fcc_type: FourCC,
        fcc_handler: FourCC,
        dw_flags: u32,
        w_priority: u16,
        w_language: u16,
        dw_initial_frames: u32,
        dw_scale: u32,
        dw_rate: u32,
        dw_start: u32,
        dw_length: u32,
        dw_suggested_buffer_size: u32,
        dw_quality: u32,
        dw_sample_size: u32,
        rc_frame: RcFrame,
    }
}

impl AviStreamHeader {
    fn new(
        fcc_type: FccType,
        fcc_handler: FourCC,
        rate: u32,
        length: u32,
        width: i16,
        height: i16,
    ) -> AviStreamHeader {
        AviStreamHeader {
            fcc_type: fcc_type.to_fcc(),
            fcc_handler: fcc_handler,
            dw_flags: 0,
            w_priority: 0,
            w_language: 0,
            dw_initial_frames: 0,
            dw_scale: 1,
            dw_rate: rate,
            dw_start: 0,
            dw_length: length,
            dw_suggested_buffer_size: 0,
            dw_quality: (0 - 1) as u32,
            dw_sample_size: 0,
            rc_frame: RcFrame {
                left: 0,
                top: 0,
                right: width,
                bottom: height,
            },
        }
    }
}

macro_attr! {
    #[derive(Serialize!)]
    struct BitmapInfoHeader {
        width: i32,
        height: i32,
        planes: u16,
        bit_count: u16,
        compression: u32,
        size_image: u32,
        x_pels_per_meter: i32,
        y_pels_per_meter: i32,
        clr_used: u32,
        clr_important: u32,
    }
}

macro_attr! {
    #[derive(Serialize!)]
    struct RGBQuad {
        blue: u8,
        green: u8,
        red: u8,
        reserved: u8,
    }
}

impl RGBQuad {
    fn bgr(blue: u8, green: u8, red: u8) -> RGBQuad {
        RGBQuad {
            blue: blue,
            green: green,
            red: red,
            reserved: 0,
        }
    }
}

macro_attr! {
    #[derive(Serialize!)]
    struct BitmapInfo {
        bmi_header: BitmapInfoHeader,
        bmi_colors: RGBQuad,
    }
}

enum StreamFormat {
    Bitmap(BitmapInfo),
    WaveFormat(u8),
}

impl Serialize for StreamFormat {
    fn serialize(&self) -> Vec<u8> {
        match self {
            StreamFormat::Bitmap(c) => c.serialize(),
            StreamFormat::WaveFormat(c) => c.serialize(),
        }
    }

    fn sizeof(&self) -> u32 {
        match self {
            StreamFormat::Bitmap(c) => c.sizeof(),
            StreamFormat::WaveFormat(c) => c.sizeof(),
        }
    }
}

macro_attr! {
    #[derive(Serialize!)]
    struct Riff {
        header_type: FourCC,
        file_size: u32,
        file_type: FourCC,
        data: Vec<RiffData>,
    }
}

impl Riff {
    fn avi(hdrl: Vec<RiffData>, movi: Vec<RiffData>) -> Riff {
        let _hdrl = RiffData::list("hdrl", hdrl);
        let _movi = RiffData::list("movi", movi);
        let _hdrl_sizeof = _hdrl.sizeof() as usize;

        Riff {
            header_type: FourCC::from("RIFF"),
            file_size: 4 + _hdrl.sizeof() + _movi.sizeof(),
            file_type: FourCC::from("AVI "),
            data: vec![
                _hdrl,
                RiffData::chunk("JUNK", iter::repeat(0).take((2048 - (_hdrl_sizeof + 32) % 2048) % 2048).collect::<Vec<u8>>()),
                _movi,
            ],
        }
    }
}

fn main() {
    let width = 256;
    let height = 240;

    let avi = Riff::avi(
        vec![
            RiffData::chunk(
                "avih",
                AviMainHeader::new(
                    1_000_000, // 1 frame micro sec
                    100,
                    10,
                    0,
                    1,
                    width,
                    height,
                )
            ),
            RiffData::list("strl", vec![
                RiffData::chunk(
                    "strh",
                    AviStreamHeader::new(
                        FccType::Vids,
                        FourCC::from("cvid"),
                        1,
                        1,
                        width as i16,
                        height as i16,
                    ),
                ),
                RiffData::chunk(
                    "strf",
                    StreamFormat::Bitmap(BitmapInfo {
                        bmi_header: BitmapInfoHeader {
                            width: width as i32,
                            height: height as i32,
                            planes: 1,
                            bit_count: 32,
                            compression: 0,
                            size_image: width * height * 32,
                            x_pels_per_meter: 0,
                            y_pels_per_meter: 0,
                            clr_used: 0,
                            clr_important: 0,
                        },
                        bmi_colors: RGBQuad::bgr(255, 255, 255),
                    }),
                ),
            ]),
        ],
        vec![
            RiffData::chunk(
                "00db",
                (0..width*height).map(|_| RGBQuad::bgr(128, 128, 128)).collect::<Vec<RGBQuad>>(),
            ),
        ]
    );

    let mut file = BufWriter::new(fs::File::create("test.avi").unwrap());
    file.write(avi.serialize().as_slice()).unwrap();
}