<commit_before>#![feature(test)]

extern crate ipc_channel;
extern crate test;

/// Allows doing multiple inner iterations per bench.iter() run.
///
/// This is mostly to amortise the overhead of spawning a thread in the benchmark
/// when sending larger messages (that might be fragmented).
///
/// Note that you need to compensate the displayed results
/// for the proportionally longer runs yourself,
/// as the benchmark framework doesn't know about the inner iterations...
const ITERATIONS: usize = 1;

mod platform {
    extern crate crossbeam;

    use ipc_channel::platform;
    use ITERATIONS;

    use std::sync::{mpsc, Mutex};
    use test;

    #[bench]
    fn create_channel(b: &mut test::Bencher) {
        b.iter(|| {
            for _ in 0..ITERATIONS {
                platform::channel().unwrap();
            }
        });
    }

    fn bench_transfer_data(b: &mut test::Bencher, size: usize) {
        let data: Vec<u8> = (0..size).map(|i| (i % 251) as u8).collect();
        let (tx, rx) = platform::channel().unwrap();

        let (wait_tx, wait_rx) = mpsc::channel();
        let wait_rx = Mutex::new(wait_rx);

        if size > platform::OsIpcSender::get_max_fragment_size() {
            b.iter(|| {
                crossbeam::scope(|scope| {
                    let tx = tx.clone();
                    scope.spawn(|| {
                        let wait_rx = wait_rx.lock().unwrap();
                        let tx = tx;
                        for _ in 0..ITERATIONS {
                            tx.send(&data, vec![], vec![]).unwrap();
                            if ITERATIONS > 1 {
                                // Prevent beginning of the next send
                                // from overlapping with receive of last fragment,
                                // as otherwise results of runs with a large tail fragment
                                // are significantly skewed.
                                wait_rx.recv().unwrap();
                            }
                        }
                    });
                    for _ in 0..ITERATIONS {
                        rx.recv().unwrap();
                        if ITERATIONS > 1 {
                            wait_tx.send(()).unwrap();
                        }
                    }
                    // For reasons mysterious to me,
                    // not returning a value *from every branch*
                    // adds some 100 ns or so of overhead to all results --
                    // which is quite significant for very short tests...
                    0
                })
            });
        } else {
            b.iter(|| {
                for _ in 0..ITERATIONS {
                    tx.send(&data, vec![], vec![]).unwrap();
                    rx.recv().unwrap();
                }
                0
            });
        }
    }

    #[bench]
    fn transfer_data_00_1(b: &mut test::Bencher) {
        bench_transfer_data(b, 1);
    }
    #[bench]
    fn transfer_data_01_2(b: &mut test::Bencher) {
        bench_transfer_data(b, 2);
    }
    #[bench]
    fn transfer_data_02_4(b: &mut test::Bencher) {
        bench_transfer_data(b, 4);
    }
    #[bench]
    fn transfer_data_03_8(b: &mut test::Bencher) {
        bench_transfer_data(b, 8);
    }
    #[bench]
    fn transfer_data_04_16(b: &mut test::Bencher) {
        bench_transfer_data(b, 16);
    }
    #[bench]
    fn transfer_data_05_32(b: &mut test::Bencher) {
        bench_transfer_data(b, 32);
    }
    #[bench]
    fn transfer_data_06_64(b: &mut test::Bencher) {
        bench_transfer_data(b, 64);
    }
    #[bench]
    fn transfer_data_07_128(b: &mut test::Bencher) {
        bench_transfer_data(b, 128);
    }
    #[bench]
    fn transfer_data_08_256(b: &mut test::Bencher) {
        bench_transfer_data(b, 256);
    }
    #[bench]
    fn transfer_data_09_512(b: &mut test::Bencher) {
        bench_transfer_data(b, 512);
    }
    #[bench]
    fn transfer_data_10_1k(b: &mut test::Bencher) {
        bench_transfer_data(b, 1 * 1024);
    }
    #[bench]
    fn transfer_data_11_2k(b: &mut test::Bencher) {
        bench_transfer_data(b, 2 * 1024);
    }
    #[bench]
    fn transfer_data_12_4k(b: &mut test::Bencher) {
        bench_transfer_data(b, 4 * 1024);
    }
    #[bench]
    fn transfer_data_13_8k(b: &mut test::Bencher) {
        bench_transfer_data(b, 8 * 1024);
    }
    #[bench]
    fn transfer_data_14_16k(b: &mut test::Bencher) {
        bench_transfer_data(b, 16 * 1024);
    }
    #[bench]
    fn transfer_data_15_32k(b: &mut test::Bencher) {
        bench_transfer_data(b, 32 * 1024);
    }
    #[bench]
    fn transfer_data_16_64k(b: &mut test::Bencher) {
        bench_transfer_data(b, 64 * 1024);
    }
    #[bench]
    fn transfer_data_17_128k(b: &mut test::Bencher) {
        bench_transfer_data(b, 128 * 1024);
    }
    #[bench]
    fn transfer_data_18_256k(b: &mut test::Bencher) {
        bench_transfer_data(b, 256 * 1024);
    }
    #[bench]
    fn transfer_data_19_512k(b: &mut test::Bencher) {
        bench_transfer_data(b, 512 * 1024);
    }
    #[bench]
    fn transfer_data_20_1m(b: &mut test::Bencher) {
        bench_transfer_data(b, 1 * 1024 * 1024);
    }
    #[bench]
    fn transfer_data_21_2m(b: &mut test::Bencher) {
        bench_transfer_data(b, 2 * 1024 * 1024);
    }
    #[bench]
    fn transfer_data_22_4m(b: &mut test::Bencher) {
        bench_transfer_data(b, 4 * 1024 * 1024);
    }
    #[bench]
    fn transfer_data_23_8m(b: &mut test::Bencher) {
        bench_transfer_data(b, 8 * 1024 * 1024);
    }
}

mod ipc {
    mod receiver_set {
        use ipc_channel::ipc::{self, IpcReceiverSet};
        use test;

        // Benchmark selecting over a set of `n` receivers,
        // with `to_send` of them actually having pending data.
        fn bench_send_on_m_of_n(b: &mut test::Bencher, to_send: usize, n: usize) {
            let mut active = Vec::with_capacity(to_send);
            let mut dormant = Vec::with_capacity(n - to_send);
            let mut rx_set = IpcReceiverSet::new().unwrap();
            for _ in 0..to_send {
                let (tx, rx) = ipc::channel().unwrap();
                rx_set.add(rx).unwrap();
                active.push(tx);
            }
            for _ in to_send..n {
                let (tx, rx) = ipc::channel::<()>().unwrap();
                rx_set.add(rx).unwrap();
                dormant.push(tx);
            }
            b.iter(|| {
                for tx in active.iter() {
                    tx.send(()).unwrap();
                }
                let mut received = 0;
                while received < to_send {
                    for result in rx_set.select().unwrap().into_iter() {
                        let (_, _) = result.unwrap();
                        received += 1;
                    }
                }
            });
        }

        #[bench]
        fn send_on_1_of_1(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 1, 1);
        }

        #[bench]
        fn send_on_1_of_5(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 1, 5);
        }

        #[bench]
        fn send_on_2_of_5(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 2, 5);
        }

        #[bench]
        fn send_on_5_of_5(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 5, 5);
        }

        #[bench]
        fn send_on_1_of_20(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 1, 20);
        }

        #[bench]
        fn send_on_5_of_20(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 5, 20);
        }

        #[bench]
        fn send_on_20_of_20(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 20, 20);
        }

        #[bench]
        fn send_on_1_of_100(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 1, 100);
        }

        #[bench]
        fn send_on_5_of_100(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 5, 100);
        }
        #[bench]
        fn send_on_20_of_100(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 20, 100);
        }

        #[bench]
        fn send_on_100_of_100(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 100, 100);
        }

        fn create_set_of_n(n: usize) -> IpcReceiverSet {
            let mut rx_set = IpcReceiverSet::new().unwrap();
            for _ in 0..n {
                let (_, rx) = ipc::channel::<()>().unwrap();
                rx_set.add(rx).unwrap();
            }
            rx_set
        }

        #[bench]
        fn create_and_destroy_empty_set(b: &mut test::Bencher) {
            b.iter(|| {
                create_set_of_n(0);
            });
        }

        #[bench]
        fn create_and_destroy_set_of_1(b: &mut test::Bencher) {
            b.iter(|| {
                create_set_of_n(1);
            });
        }

        #[bench]
        fn create_and_destroy_set_of_10(b: &mut test::Bencher) {
            b.iter(|| {
                create_set_of_n(10);
            });
        }

        #[bench]
        fn create_and_destroy_set_of_100(b: &mut test::Bencher) {
            b.iter(|| {
                create_set_of_n(100);
            });
        }

        // Benchmark performance of removing closed receivers from set.
        // This also includes the time for adding receivers,
        // as there is no way to measure the removing in isolation.
        fn bench_remove_closed(b: &mut test::Bencher, n: usize) {
            b.iter(|| {
                let mut rx_set = create_set_of_n(n);

                let mut dropped_count = 0;
                while dropped_count < n {
                    // On `select()`, receivers with a "ClosedChannel" event will be closed,
                    // and automatically dropped from the set.
                    dropped_count += rx_set.select().unwrap().len();
                }
            });
        }

        #[bench]
        fn add_and_remove_1_closed_receivers(b: &mut test::Bencher) {
            bench_remove_closed(b, 1);
        }

        #[bench]
        fn add_and_remove_10_closed_receivers(b: &mut test::Bencher) {
            bench_remove_closed(b, 10);
        }

        #[bench]
        fn add_and_remove_100_closed_receivers(b: &mut test::Bencher) {
            bench_remove_closed(b, 100);
        }
    }
}
<commit_msg>Benchmarks: Simplify implementation of `bench_send_on_m_of_n()`<commit_after>#![feature(test)]

extern crate ipc_channel;
extern crate test;

/// Allows doing multiple inner iterations per bench.iter() run.
///
/// This is mostly to amortise the overhead of spawning a thread in the benchmark
/// when sending larger messages (that might be fragmented).
///
/// Note that you need to compensate the displayed results
/// for the proportionally longer runs yourself,
/// as the benchmark framework doesn't know about the inner iterations...
const ITERATIONS: usize = 1;

mod platform {
    extern crate crossbeam;

    use ipc_channel::platform;
    use ITERATIONS;

    use std::sync::{mpsc, Mutex};
    use test;

    #[bench]
    fn create_channel(b: &mut test::Bencher) {
        b.iter(|| {
            for _ in 0..ITERATIONS {
                platform::channel().unwrap();
            }
        });
    }

    fn bench_transfer_data(b: &mut test::Bencher, size: usize) {
        let data: Vec<u8> = (0..size).map(|i| (i % 251) as u8).collect();
        let (tx, rx) = platform::channel().unwrap();

        let (wait_tx, wait_rx) = mpsc::channel();
        let wait_rx = Mutex::new(wait_rx);

        if size > platform::OsIpcSender::get_max_fragment_size() {
            b.iter(|| {
                crossbeam::scope(|scope| {
                    let tx = tx.clone();
                    scope.spawn(|| {
                        let wait_rx = wait_rx.lock().unwrap();
                        let tx = tx;
                        for _ in 0..ITERATIONS {
                            tx.send(&data, vec![], vec![]).unwrap();
                            if ITERATIONS > 1 {
                                // Prevent beginning of the next send
                                // from overlapping with receive of last fragment,
                                // as otherwise results of runs with a large tail fragment
                                // are significantly skewed.
                                wait_rx.recv().unwrap();
                            }
                        }
                    });
                    for _ in 0..ITERATIONS {
                        rx.recv().unwrap();
                        if ITERATIONS > 1 {
                            wait_tx.send(()).unwrap();
                        }
                    }
                    // For reasons mysterious to me,
                    // not returning a value *from every branch*
                    // adds some 100 ns or so of overhead to all results --
                    // which is quite significant for very short tests...
                    0
                })
            });
        } else {
            b.iter(|| {
                for _ in 0..ITERATIONS {
                    tx.send(&data, vec![], vec![]).unwrap();
                    rx.recv().unwrap();
                }
                0
            });
        }
    }

    #[bench]
    fn transfer_data_00_1(b: &mut test::Bencher) {
        bench_transfer_data(b, 1);
    }
    #[bench]
    fn transfer_data_01_2(b: &mut test::Bencher) {
        bench_transfer_data(b, 2);
    }
    #[bench]
    fn transfer_data_02_4(b: &mut test::Bencher) {
        bench_transfer_data(b, 4);
    }
    #[bench]
    fn transfer_data_03_8(b: &mut test::Bencher) {
        bench_transfer_data(b, 8);
    }
    #[bench]
    fn transfer_data_04_16(b: &mut test::Bencher) {
        bench_transfer_data(b, 16);
    }
    #[bench]
    fn transfer_data_05_32(b: &mut test::Bencher) {
        bench_transfer_data(b, 32);
    }
    #[bench]
    fn transfer_data_06_64(b: &mut test::Bencher) {
        bench_transfer_data(b, 64);
    }
    #[bench]
    fn transfer_data_07_128(b: &mut test::Bencher) {
        bench_transfer_data(b, 128);
    }
    #[bench]
    fn transfer_data_08_256(b: &mut test::Bencher) {
        bench_transfer_data(b, 256);
    }
    #[bench]
    fn transfer_data_09_512(b: &mut test::Bencher) {
        bench_transfer_data(b, 512);
    }
    #[bench]
    fn transfer_data_10_1k(b: &mut test::Bencher) {
        bench_transfer_data(b, 1 * 1024);
    }
    #[bench]
    fn transfer_data_11_2k(b: &mut test::Bencher) {
        bench_transfer_data(b, 2 * 1024);
    }
    #[bench]
    fn transfer_data_12_4k(b: &mut test::Bencher) {
        bench_transfer_data(b, 4 * 1024);
    }
    #[bench]
    fn transfer_data_13_8k(b: &mut test::Bencher) {
        bench_transfer_data(b, 8 * 1024);
    }
    #[bench]
    fn transfer_data_14_16k(b: &mut test::Bencher) {
        bench_transfer_data(b, 16 * 1024);
    }
    #[bench]
    fn transfer_data_15_32k(b: &mut test::Bencher) {
        bench_transfer_data(b, 32 * 1024);
    }
    #[bench]
    fn transfer_data_16_64k(b: &mut test::Bencher) {
        bench_transfer_data(b, 64 * 1024);
    }
    #[bench]
    fn transfer_data_17_128k(b: &mut test::Bencher) {
        bench_transfer_data(b, 128 * 1024);
    }
    #[bench]
    fn transfer_data_18_256k(b: &mut test::Bencher) {
        bench_transfer_data(b, 256 * 1024);
    }
    #[bench]
    fn transfer_data_19_512k(b: &mut test::Bencher) {
        bench_transfer_data(b, 512 * 1024);
    }
    #[bench]
    fn transfer_data_20_1m(b: &mut test::Bencher) {
        bench_transfer_data(b, 1 * 1024 * 1024);
    }
    #[bench]
    fn transfer_data_21_2m(b: &mut test::Bencher) {
        bench_transfer_data(b, 2 * 1024 * 1024);
    }
    #[bench]
    fn transfer_data_22_4m(b: &mut test::Bencher) {
        bench_transfer_data(b, 4 * 1024 * 1024);
    }
    #[bench]
    fn transfer_data_23_8m(b: &mut test::Bencher) {
        bench_transfer_data(b, 8 * 1024 * 1024);
    }
}

mod ipc {
    mod receiver_set {
        use ipc_channel::ipc::{self, IpcReceiverSet};
        use test;

        // Benchmark selecting over a set of `n` receivers,
        // with `to_send` of them actually having pending data.
        fn bench_send_on_m_of_n(b: &mut test::Bencher, to_send: usize, n: usize) {
            let mut senders = Vec::with_capacity(n);
            let mut rx_set = IpcReceiverSet::new().unwrap();
            for _ in 0..n {
                let (tx, rx) = ipc::channel().unwrap();
                rx_set.add(rx).unwrap();
                senders.push(tx);
            }
            b.iter(|| {
                for tx in senders.iter().take(to_send) {
                    tx.send(()).unwrap();
                }
                let mut received = 0;
                while received < to_send {
                    received += rx_set.select().unwrap().len();
                }
            });
        }

        #[bench]
        fn send_on_1_of_1(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 1, 1);
        }

        #[bench]
        fn send_on_1_of_5(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 1, 5);
        }

        #[bench]
        fn send_on_2_of_5(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 2, 5);
        }

        #[bench]
        fn send_on_5_of_5(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 5, 5);
        }

        #[bench]
        fn send_on_1_of_20(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 1, 20);
        }

        #[bench]
        fn send_on_5_of_20(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 5, 20);
        }

        #[bench]
        fn send_on_20_of_20(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 20, 20);
        }

        #[bench]
        fn send_on_1_of_100(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 1, 100);
        }

        #[bench]
        fn send_on_5_of_100(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 5, 100);
        }
        #[bench]
        fn send_on_20_of_100(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 20, 100);
        }

        #[bench]
        fn send_on_100_of_100(b: &mut test::Bencher) {
            bench_send_on_m_of_n(b, 100, 100);
        }

        fn create_set_of_n(n: usize) -> IpcReceiverSet {
            let mut rx_set = IpcReceiverSet::new().unwrap();
            for _ in 0..n {
                let (_, rx) = ipc::channel::<()>().unwrap();
                rx_set.add(rx).unwrap();
            }
            rx_set
        }

        #[bench]
        fn create_and_destroy_empty_set(b: &mut test::Bencher) {
            b.iter(|| {
                create_set_of_n(0);
            });
        }

        #[bench]
        fn create_and_destroy_set_of_1(b: &mut test::Bencher) {
            b.iter(|| {
                create_set_of_n(1);
            });
        }

        #[bench]
        fn create_and_destroy_set_of_10(b: &mut test::Bencher) {
            b.iter(|| {
                create_set_of_n(10);
            });
        }

        #[bench]
        fn create_and_destroy_set_of_100(b: &mut test::Bencher) {
            b.iter(|| {
                create_set_of_n(100);
            });
        }

        // Benchmark performance of removing closed receivers from set.
        // This also includes the time for adding receivers,
        // as there is no way to measure the removing in isolation.
        fn bench_remove_closed(b: &mut test::Bencher, n: usize) {
            b.iter(|| {
                let mut rx_set = create_set_of_n(n);

                let mut dropped_count = 0;
                while dropped_count < n {
                    // On `select()`, receivers with a "ClosedChannel" event will be closed,
                    // and automatically dropped from the set.
                    dropped_count += rx_set.select().unwrap().len();
                }
            });
        }

        #[bench]
        fn add_and_remove_1_closed_receivers(b: &mut test::Bencher) {
            bench_remove_closed(b, 1);
        }

        #[bench]
        fn add_and_remove_10_closed_receivers(b: &mut test::Bencher) {
            bench_remove_closed(b, 10);
        }

        #[bench]
        fn add_and_remove_100_closed_receivers(b: &mut test::Bencher) {
            bench_remove_closed(b, 100);
        }
    }
}
<|endoftext|>


<commit_before>#![feature(test)]

extern crate test;
extern crate xmltree;

use test::Bencher;

use xmltree::Element;
use xmltree::ParseError;
use std::fs::File;
use std::io::Cursor;
use std::io::Read;

fn _parse_file(filename: &str) -> std::result::Result<xmltree::Element, xmltree::ParseError> {
    let mut file = File::open(filename).unwrap();
    let mut data = String::new();
    file.read_to_string(&mut data).unwrap();

    Element::parse(data.as_bytes())
}

#[bench]
fn bench_01(b: &mut Bencher) {
    b.iter(|| {
        let filename = "tests/data/01.xml";
        let e: Element = _parse_file(filename).unwrap();

        assert_eq!(e.name, "project");
        let e2: &Element = e.get_child("libraries")
            .expect("Missing libraries child element");
        assert_eq!(e2.name, "libraries");

        assert!(e.get_child("doesnotexist").is_none());

        let mut buf = Vec::new();
        e.write(&mut buf);

        let e2 = Element::parse(Cursor::new(buf)).unwrap();
        assert_eq!(e, e2);
    })
}

#[bench]
fn bench_02(b: &mut Bencher) {
    let filename = "tests/data/02.xml";
    b.iter(|| {
        let _ = _parse_file(filename).unwrap();
    });
}

#[bench]
fn bench_03(b: &mut Bencher) {
    let filename = "tests/data/03.xml";
    b.iter(|| {
        let _ = _parse_file(filename).unwrap();
    });
}

#[bench]
fn bench_04(b: &mut Bencher) {
    let filename = "tests/data/04.xml";
    b.iter(|| {
        let _ = _parse_file(filename).unwrap();
    });
}

#[bench]
fn bench_rw(b: &mut Bencher) {
    let filename = "tests/data/rw.xml";
    b.iter(|| {
        let e = _parse_file(filename).unwrap();

        let mut buf = Vec::new();
        e.write(&mut buf);

        let e2 = Element::parse(Cursor::new(buf)).unwrap();
        assert_eq!(e, e2);
    })
}

#[bench]
fn bench_mal_01(b: &mut Bencher) {
    let data = r##"
        <?xml version="1.0" encoding="utf-8" standalone="yes"?>
        <names>
            <name first="bob" last="jones />
            <name first="elizabeth" last="smith" />
        </names>
    "##;

    b.iter(|| {
        let names_element = Element::parse(data.as_bytes());
        if let Err(ParseError::MalformedXml(..)) = names_element {
            // OK
        } else {
            panic!("unexpected parse result");
        }
    });
}

#[bench]
fn bench_mal_02(b: &mut Bencher) {
    let data = r##"
            this is not even close
            to XML
    "##;

    b.iter(|| {
        let names_element = Element::parse(data.as_bytes());
        if let Err(ParseError::MalformedXml(..)) = names_element {
            // OK
        } else {
            panic!("unexpected parse result");
        }
    });
}

#[bench]
fn bench_mal_03(b: &mut Bencher) {
    let data = r##"
        <?xml version="1.0" encoding="utf-8" standalone="yes"?>
        <names>
            <name first="bob" last="jones"></badtag>
            <name first="elizabeth" last="smith" />
        </names>
    "##;

    b.iter(|| {
        let names_element = Element::parse(data.as_bytes());
        if let Err(ParseError::MalformedXml(..)) = names_element {
            // OK
        } else {
            panic!("unexpected parse result");
        }
    });
}

#[bench]
fn bench_new(b: &mut Bencher) {
    b.iter(|| {
        let e = Element::new("foo");
        assert_eq!(e.name.as_str(), "foo");
        assert_eq!(e.attributes.len(), 0);
        assert_eq!(e.children.len(), 0);
        assert_eq!(e.text, None);
    });
}

#[bench]
fn bench_take(b: &mut Bencher) {
    let data_xml_1 = r##"
        <?xml version="1.0" encoding="utf-8" standalone="yes"?>
        <names>
            <name first="bob" last="jones"></name>
            <name first="elizabeth" last="smith" />
            <remove_me key="value">
                <child />
            </remove_me>
        </names>
    "##;

    let data_xml_2 = r##"
        <?xml version="1.0" encoding="utf-8" standalone="yes"?>
        <names>
            <name first="bob" last="jones"></name>
            <name first="elizabeth" last="smith" />
        </names>
    "##;

    b.iter(|| {
        let mut data_1 = Element::parse(data_xml_1.as_bytes()).unwrap();
        let data_2 = Element::parse(data_xml_2.as_bytes()).unwrap();

        if let Some(removed) = data_1.take_child("remove_me") {
            assert_eq!(removed.children.len(), 1);
        } else {
            panic!("take_child failed");
        }

        assert_eq!(data_1, data_2);
    });
}

#[bench]
fn bench_ns1_rw(b: &mut Bencher) {
    let filename = "tests/data/ns2.xml";
    b.iter(|| {
        let e = _parse_file(filename).unwrap();
        let mut buf = Vec::new();
        e.write(&mut buf);
    });
}

#[bench]
fn bench_ns2_rw(b: &mut Bencher) {
    let filename = "tests/data/ns2.xml";
    b.iter(|| {
        let e = _parse_file(filename).unwrap();
        let mut buf = Vec::new();
        e.write(&mut buf);
    });
}
<commit_msg>add unwraps to benchs too<commit_after>#![feature(test)]

extern crate test;
extern crate xmltree;

use test::Bencher;

use xmltree::Element;
use xmltree::ParseError;
use std::fs::File;
use std::io::Cursor;
use std::io::Read;

fn _parse_file(filename: &str) -> std::result::Result<xmltree::Element, xmltree::ParseError> {
    let mut file = File::open(filename).unwrap();
    let mut data = String::new();
    file.read_to_string(&mut data).unwrap();

    Element::parse(data.as_bytes())
}

#[bench]
fn bench_01(b: &mut Bencher) {
    b.iter(|| {
        let filename = "tests/data/01.xml";
        let e: Element = _parse_file(filename).unwrap();

        assert_eq!(e.name, "project");
        let e2: &Element = e.get_child("libraries")
            .expect("Missing libraries child element");
        assert_eq!(e2.name, "libraries");

        assert!(e.get_child("doesnotexist").is_none());

        let mut buf = Vec::new();
        e.write(&mut buf).unwrap();

        let e2 = Element::parse(Cursor::new(buf)).unwrap();
        assert_eq!(e, e2);
    })
}

#[bench]
fn bench_02(b: &mut Bencher) {
    let filename = "tests/data/02.xml";
    b.iter(|| {
        let _ = _parse_file(filename).unwrap();
    });
}

#[bench]
fn bench_03(b: &mut Bencher) {
    let filename = "tests/data/03.xml";
    b.iter(|| {
        let _ = _parse_file(filename).unwrap();
    });
}

#[bench]
fn bench_04(b: &mut Bencher) {
    let filename = "tests/data/04.xml";
    b.iter(|| {
        let _ = _parse_file(filename).unwrap();
    });
}

#[bench]
fn bench_rw(b: &mut Bencher) {
    let filename = "tests/data/rw.xml";
    b.iter(|| {
        let e = _parse_file(filename).unwrap();

        let mut buf = Vec::new();
        e.write(&mut buf).unwrap();

        let e2 = Element::parse(Cursor::new(buf)).unwrap();
        assert_eq!(e, e2);
    })
}

#[bench]
fn bench_mal_01(b: &mut Bencher) {
    let data = r##"
        <?xml version="1.0" encoding="utf-8" standalone="yes"?>
        <names>
            <name first="bob" last="jones />
            <name first="elizabeth" last="smith" />
        </names>
    "##;

    b.iter(|| {
        let names_element = Element::parse(data.as_bytes());
        if let Err(ParseError::MalformedXml(..)) = names_element {
            // OK
        } else {
            panic!("unexpected parse result");
        }
    });
}

#[bench]
fn bench_mal_02(b: &mut Bencher) {
    let data = r##"
            this is not even close
            to XML
    "##;

    b.iter(|| {
        let names_element = Element::parse(data.as_bytes());
        if let Err(ParseError::MalformedXml(..)) = names_element {
            // OK
        } else {
            panic!("unexpected parse result");
        }
    });
}

#[bench]
fn bench_mal_03(b: &mut Bencher) {
    let data = r##"
        <?xml version="1.0" encoding="utf-8" standalone="yes"?>
        <names>
            <name first="bob" last="jones"></badtag>
            <name first="elizabeth" last="smith" />
        </names>
    "##;

    b.iter(|| {
        let names_element = Element::parse(data.as_bytes());
        if let Err(ParseError::MalformedXml(..)) = names_element {
            // OK
        } else {
            panic!("unexpected parse result");
        }
    });
}

#[bench]
fn bench_new(b: &mut Bencher) {
    b.iter(|| {
        let e = Element::new("foo");
        assert_eq!(e.name.as_str(), "foo");
        assert_eq!(e.attributes.len(), 0);
        assert_eq!(e.children.len(), 0);
        assert_eq!(e.text, None);
    });
}

#[bench]
fn bench_take(b: &mut Bencher) {
    let data_xml_1 = r##"
        <?xml version="1.0" encoding="utf-8" standalone="yes"?>
        <names>
            <name first="bob" last="jones"></name>
            <name first="elizabeth" last="smith" />
            <remove_me key="value">
                <child />
            </remove_me>
        </names>
    "##;

    let data_xml_2 = r##"
        <?xml version="1.0" encoding="utf-8" standalone="yes"?>
        <names>
            <name first="bob" last="jones"></name>
            <name first="elizabeth" last="smith" />
        </names>
    "##;

    b.iter(|| {
        let mut data_1 = Element::parse(data_xml_1.as_bytes()).unwrap();
        let data_2 = Element::parse(data_xml_2.as_bytes()).unwrap();

        if let Some(removed) = data_1.take_child("remove_me") {
            assert_eq!(removed.children.len(), 1);
        } else {
            panic!("take_child failed");
        }

        assert_eq!(data_1, data_2);
    });
}

#[bench]
fn bench_ns1_rw(b: &mut Bencher) {
    let filename = "tests/data/ns2.xml";
    b.iter(|| {
        let e = _parse_file(filename).unwrap();
        let mut buf = Vec::new();
        e.write(&mut buf).unwrap();
    });
}

#[bench]
fn bench_ns2_rw(b: &mut Bencher) {
    let filename = "tests/data/ns2.xml";
    b.iter(|| {
        let e = _parse_file(filename).unwrap();
        let mut buf = Vec::new();
        e.write(&mut buf).unwrap();
    });
}
<|endoftext|>


<commit_before>pub use self::StatusCode::*;

pub enum StatusCode {
    Informational(uint, &'static str),
    Continue,
    SwitchingProtocols,
    Successful(uint, &'static str),
    OK,
    Created,
    Accepted,
    NonAuthoritativeInformation,
    NoContent,
    ResetContent,
    PartialContent,
    Redirection(uint, &'static str),
    MultipleChoices,
    MovedPermanently,
    Found,
    SeeOther,
    NotModified,
    UseProxy,
    TemporaryRedirect,
    ClientError(uint, &'static str),
    BadRequest,
    Unauthorized,
    Forbidden,
    NotFound,
    MethodNotAllowed,
    NotAcceptable,
    ProxyAuthenticationRequired,
    RequestTimeout,
    Conflict,
    Gone,
    LengthRequired,
    PreconditionFailed,
    RequestEntityTooLarge,
    RequestUriTooLong,
    UnsupportedMediaType,
    RequestedRangeNotSatisfiable,
    ExpectationFailed,
    ServerError(uint, &'static str),
    InternalServerError,
    NotImplemented,
    BadGateway,
    ServiceUnavailable,
    GatewayTimeout,
    HttpVersionNotSupported,
    Other(uint, &'static str),
}

impl StatusCode {
    pub fn to_code(&self) -> (uint, &'static str) {
        match *self {
            Informational(num, string) |
                Successful(num, string) |
                Redirection(num, string) |
                ClientError(num, string) |
                ServerError(num, string) => (num, string),
            Continue => (100, "Continue"),
            SwitchingProtocols => (101, "Switching Protocols"),
            OK => (200, "OK"),
            Created => (201, "Created"),
            Accepted => (202, "Accepted"),
            NonAuthoritativeInformation => (203, "Non-Authoritative Information"),
            NoContent => (204, "No Content"),
            ResetContent => (205, "Reset Content"),
            PartialContent => (206, "Partial Content"),
            MultipleChoices => (300, "Multiple Choices"),
            MovedPermanently => (301, "Moved Permanently"),
            Found => (302, "Found"),
            SeeOther => (303, "See Other"),
            NotModified => (304, "Not Modified"),
            UseProxy => (305, "Use Proxy"),
            TemporaryRedirect => (307, "Temporary Redirect"),
            BadRequest => (400, "Bad Request"),
            Unauthorized => (401, "Unauthorized"),
            Forbidden => (403, "Forbidden"),
            NotFound => (404, "Not Found"),
            MethodNotAllowed => (405, "Method Not Allowed"),
            NotAcceptable => (406, "Not Acceptable"),
            ProxyAuthenticationRequired => (407, "Proxy Authentication Required"),
            RequestTimeout => (408, "Request Timeout"),
            Conflict => (409, "Conflict"),
            Gone => (410, "Gone"),
            LengthRequired => (411, "Length Required"),
            PreconditionFailed => (412, "Precondition Failed"),
            RequestEntityTooLarge => (413, "Request Entity Too Large"),
            RequestUriTooLong => (414, "Request-URI Too Long"),
            UnsupportedMediaType => (415, "Unsupported Media Type"),
            RequestedRangeNotSatisfiable => (416, "Requested Range Not Satisfiable"),
            ExpectationFailed => (417, "Expectation Failed"),
            InternalServerError => (500, "Internal Server Error"),
            NotImplemented => (501, "Not Implemented"),
            BadGateway => (502, "Bad Gateway"),
            ServiceUnavailable => (503, "Service Unavailable"),
            GatewayTimeout => (504, "Gateway Timeout"),
            HttpVersionNotSupported => (505, "HTTP Version Not Supported"),
            Other(n, s) => (n, s),
        }
    }
}

pub trait ToStatusCode {
    fn to_status(&self) -> Result<StatusCode, ()>;
}

impl ToStatusCode for StatusCode {
    fn to_status(&self) -> Result<StatusCode, ()> {
        Ok(*self)
    }
}

impl ToStatusCode for (uint, &'static str) {
    fn to_status(&self) -> Result<StatusCode, ()> {
        let (code, name) = *self;
        Ok(Other(code, name))
    }
}

impl ToStatusCode for int {
    fn to_status(&self) -> Result<StatusCode, ()> {
        (*self as uint).to_status()
    }
}

impl ToStatusCode for uint {
    fn to_status(&self) -> Result<StatusCode, ()> {
        match *self {
            num @ 102 ... 199 => Ok(Informational(num, "Informational")),
            num @ 207 ... 299 => Ok(Successful(num, "Successful")),
            num @ 306 | num @ 308 ... 399 => Ok(Redirection(num, "Redirection")),
            num @ 402 | num @ 418 ... 499 => Ok(ClientError(num, "Client Error")),
            num @ 506 ... 599 => Ok(ServerError(num, "Server Error")),
            100 => Ok(Continue),
            101 => Ok(SwitchingProtocols),
            200 => Ok(OK),
            201 => Ok(Created),
            202 => Ok(Accepted),
            203 => Ok(NonAuthoritativeInformation),
            204 => Ok(NoContent),
            205 => Ok(ResetContent),
            206 => Ok(PartialContent),
            300 => Ok(MultipleChoices),
            301 => Ok(MovedPermanently),
            302 => Ok(Found),
            303 => Ok(SeeOther),
            304 => Ok(NotModified),
            305 => Ok(UseProxy),
            307 => Ok(TemporaryRedirect),
            400 => Ok(BadRequest),
            401 => Ok(Unauthorized),
            403 => Ok(Forbidden),
            404 => Ok(NotFound),
            405 => Ok(MethodNotAllowed),
            406 => Ok(NotAcceptable),
            407 => Ok(ProxyAuthenticationRequired),
            408 => Ok(RequestTimeout),
            409 => Ok(Conflict),
            410 => Ok(Gone),
            411 => Ok(LengthRequired),
            412 => Ok(PreconditionFailed),
            413 => Ok(RequestEntityTooLarge),
            414 => Ok(RequestUriTooLong),
            415 => Ok(UnsupportedMediaType),
            416 => Ok(RequestedRangeNotSatisfiable),
            417 => Ok(ExpectationFailed),
            500 => Ok(InternalServerError),
            501 => Ok(NotImplemented),
            502 => Ok(BadGateway),
            503 => Ok(ServiceUnavailable),
            504 => Ok(GatewayTimeout),
            505 => Ok(HttpVersionNotSupported),
            _ => Err(())
        }
    }
}

impl ToStatusCode for () {
    fn to_status(&self) -> Result<StatusCode, ()> {
        Err(())
    }
}
<commit_msg>Update to rust master<commit_after>pub use self::StatusCode::*;

#[deriving(Copy)]
pub enum StatusCode {
    Informational(uint, &'static str),
    Continue,
    SwitchingProtocols,
    Successful(uint, &'static str),
    OK,
    Created,
    Accepted,
    NonAuthoritativeInformation,
    NoContent,
    ResetContent,
    PartialContent,
    Redirection(uint, &'static str),
    MultipleChoices,
    MovedPermanently,
    Found,
    SeeOther,
    NotModified,
    UseProxy,
    TemporaryRedirect,
    ClientError(uint, &'static str),
    BadRequest,
    Unauthorized,
    Forbidden,
    NotFound,
    MethodNotAllowed,
    NotAcceptable,
    ProxyAuthenticationRequired,
    RequestTimeout,
    Conflict,
    Gone,
    LengthRequired,
    PreconditionFailed,
    RequestEntityTooLarge,
    RequestUriTooLong,
    UnsupportedMediaType,
    RequestedRangeNotSatisfiable,
    ExpectationFailed,
    ServerError(uint, &'static str),
    InternalServerError,
    NotImplemented,
    BadGateway,
    ServiceUnavailable,
    GatewayTimeout,
    HttpVersionNotSupported,
    Other(uint, &'static str),
}

impl StatusCode {
    pub fn to_code(&self) -> (uint, &'static str) {
        match *self {
            Informational(num, string) |
                Successful(num, string) |
                Redirection(num, string) |
                ClientError(num, string) |
                ServerError(num, string) => (num, string),
            Continue => (100, "Continue"),
            SwitchingProtocols => (101, "Switching Protocols"),
            OK => (200, "OK"),
            Created => (201, "Created"),
            Accepted => (202, "Accepted"),
            NonAuthoritativeInformation => (203, "Non-Authoritative Information"),
            NoContent => (204, "No Content"),
            ResetContent => (205, "Reset Content"),
            PartialContent => (206, "Partial Content"),
            MultipleChoices => (300, "Multiple Choices"),
            MovedPermanently => (301, "Moved Permanently"),
            Found => (302, "Found"),
            SeeOther => (303, "See Other"),
            NotModified => (304, "Not Modified"),
            UseProxy => (305, "Use Proxy"),
            TemporaryRedirect => (307, "Temporary Redirect"),
            BadRequest => (400, "Bad Request"),
            Unauthorized => (401, "Unauthorized"),
            Forbidden => (403, "Forbidden"),
            NotFound => (404, "Not Found"),
            MethodNotAllowed => (405, "Method Not Allowed"),
            NotAcceptable => (406, "Not Acceptable"),
            ProxyAuthenticationRequired => (407, "Proxy Authentication Required"),
            RequestTimeout => (408, "Request Timeout"),
            Conflict => (409, "Conflict"),
            Gone => (410, "Gone"),
            LengthRequired => (411, "Length Required"),
            PreconditionFailed => (412, "Precondition Failed"),
            RequestEntityTooLarge => (413, "Request Entity Too Large"),
            RequestUriTooLong => (414, "Request-URI Too Long"),
            UnsupportedMediaType => (415, "Unsupported Media Type"),
            RequestedRangeNotSatisfiable => (416, "Requested Range Not Satisfiable"),
            ExpectationFailed => (417, "Expectation Failed"),
            InternalServerError => (500, "Internal Server Error"),
            NotImplemented => (501, "Not Implemented"),
            BadGateway => (502, "Bad Gateway"),
            ServiceUnavailable => (503, "Service Unavailable"),
            GatewayTimeout => (504, "Gateway Timeout"),
            HttpVersionNotSupported => (505, "HTTP Version Not Supported"),
            Other(n, s) => (n, s),
        }
    }
}

pub trait ToStatusCode {
    fn to_status(&self) -> Result<StatusCode, ()>;
}

impl ToStatusCode for StatusCode {
    fn to_status(&self) -> Result<StatusCode, ()> {
        Ok(*self)
    }
}

impl ToStatusCode for (uint, &'static str) {
    fn to_status(&self) -> Result<StatusCode, ()> {
        let (code, name) = *self;
        Ok(Other(code, name))
    }
}

impl ToStatusCode for int {
    fn to_status(&self) -> Result<StatusCode, ()> {
        (*self as uint).to_status()
    }
}

impl ToStatusCode for uint {
    fn to_status(&self) -> Result<StatusCode, ()> {
        match *self {
            num @ 102 ... 199 => Ok(Informational(num, "Informational")),
            num @ 207 ... 299 => Ok(Successful(num, "Successful")),
            num @ 306 | num @ 308 ... 399 => Ok(Redirection(num, "Redirection")),
            num @ 402 | num @ 418 ... 499 => Ok(ClientError(num, "Client Error")),
            num @ 506 ... 599 => Ok(ServerError(num, "Server Error")),
            100 => Ok(Continue),
            101 => Ok(SwitchingProtocols),
            200 => Ok(OK),
            201 => Ok(Created),
            202 => Ok(Accepted),
            203 => Ok(NonAuthoritativeInformation),
            204 => Ok(NoContent),
            205 => Ok(ResetContent),
            206 => Ok(PartialContent),
            300 => Ok(MultipleChoices),
            301 => Ok(MovedPermanently),
            302 => Ok(Found),
            303 => Ok(SeeOther),
            304 => Ok(NotModified),
            305 => Ok(UseProxy),
            307 => Ok(TemporaryRedirect),
            400 => Ok(BadRequest),
            401 => Ok(Unauthorized),
            403 => Ok(Forbidden),
            404 => Ok(NotFound),
            405 => Ok(MethodNotAllowed),
            406 => Ok(NotAcceptable),
            407 => Ok(ProxyAuthenticationRequired),
            408 => Ok(RequestTimeout),
            409 => Ok(Conflict),
            410 => Ok(Gone),
            411 => Ok(LengthRequired),
            412 => Ok(PreconditionFailed),
            413 => Ok(RequestEntityTooLarge),
            414 => Ok(RequestUriTooLong),
            415 => Ok(UnsupportedMediaType),
            416 => Ok(RequestedRangeNotSatisfiable),
            417 => Ok(ExpectationFailed),
            500 => Ok(InternalServerError),
            501 => Ok(NotImplemented),
            502 => Ok(BadGateway),
            503 => Ok(ServiceUnavailable),
            504 => Ok(GatewayTimeout),
            505 => Ok(HttpVersionNotSupported),
            _ => Err(())
        }
    }
}

impl ToStatusCode for () {
    fn to_status(&self) -> Result<StatusCode, ()> {
        Err(())
    }
}
<|endoftext|>


<commit_before>// Copyright 2019, The Gtk-rs Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

use ::ffi::{self, cairo_status_t};
use ::{Status, Surface, UserDataKey};

use libc::{c_void, c_double, c_uchar, c_uint};
use std::any::Any;
use std::cell::{Cell, RefCell, RefMut};
use std::io;
use std::panic::AssertUnwindSafe;
use std::rc::Rc;

macro_rules! for_stream_constructors {
    ($constructor_ffi: ident) => {
        /// Takes full ownership of the output stream,
        /// which is not allowed to borrow any lifetime shorter than `'static`.
        ///
        /// Because the underlying `cairo_surface_t` is reference-counted,
        /// a lifetime parameter in a Rust wrapper type would not be enough to track
        /// how long it can keep writing to the stream.
        pub fn for_stream<W: io::Write + 'static>(width: f64, height: f64, stream: W) -> Self {
            Self {
                inner: Surface::_for_stream(
                    ffi::$constructor_ffi,
                    width,
                    height,
                    stream,
                ),
            }
        }

        /// Allows writing to a borrowed stream. The lifetime of the borrow is not tracked.
        ///
        /// # Safety
        ///
        /// The value that `stream` points to must live at least until the underlying `cairo_surface_t`
        /// (which maybe be longer then the Rust `PdfSurface` wrapper, because of reference-counting),
        /// or until the output stream is removed from the surface with [`Surface::take_output_stream`].
        ///
        /// Since the former is hard to track for sure, the latter is strongly recommended.
        /// The concrete type behind the `Box<dyn Any>` value returned by `take_output_stream`
        /// is private, so you won’t be able to downcast it.
        /// But removing it anyway ensures that later writes do no go through a dangling pointer.
        pub unsafe fn for_raw_stream<W: io::Write + 'static>(width: f64, height: f64, stream: *mut W) -> Self {
            Self {
                inner: Surface::_for_raw_stream(
                    ffi::$constructor_ffi,
                    width,
                    height,
                    stream,
                ),
            }
        }
    };
}

impl Surface {
    pub(crate) fn _for_stream<W: io::Write + 'static>(
        constructor: Constructor,
        width: f64,
        height: f64,
        stream: W,
    ) -> Self {
        let env_rc = Rc::new(CallbackEnvironment {
            mutable: RefCell::new(MutableCallbackEnvironment {
                stream: Some(Box::new(stream)),
                io_error: None,
                unwind_payload: None,
            }),
            saw_already_borrowed: Cell::new(false),
        });
        let env: *const CallbackEnvironment = &*env_rc;
        unsafe {
            let ptr = constructor(Some(write_callback::<W>), env as *mut c_void, width, height);
            let surface = Surface::from_raw_full(ptr);
            surface.set_user_data(&STREAM_CALLBACK_ENVIRONMENT, env_rc);
            surface
        }
    }

    pub(crate) unsafe fn _for_raw_stream<W: io::Write + 'static>(
        constructor: Constructor,
        width: f64,
        height: f64,
        stream: *mut W,
    ) -> Self {
        Self::_for_stream(constructor, width, height, RawStream(stream))
    }

    fn with_stream_env<'surface, R, F>(&'surface self, f: F) -> Option<R>
        where F: FnOnce(RefMut<'surface, MutableCallbackEnvironment>) -> Option<R>
    {
        let env = self.get_user_data_ptr(&STREAM_CALLBACK_ENVIRONMENT)?;

        // Safety: since `STREAM_CALLBACK_ENVIRONMENT` is private and we never
        // call `set_user_data` again or `remove_user_data` with it,
        // the contract of `get_user_data_ptr` says that the user data entry
        // lives as long as the underlying `cairo_surface_t`
        // which is at least as long as this `Surface` wrapper.
        let env: &'surface _ = unsafe { &*env.as_ptr() };

        if env.saw_already_borrowed.get() {
            panic!("The output stream’s RefCell was already borrowed when cairo attempted a write")
        }

        let mut mutable = env.mutable.borrow_mut();
        if let Some(payload) = mutable.unwind_payload.take() {
            std::panic::resume_unwind(payload)
        }
        f(mutable)
    }

    /// Remove and return the output stream, if any.
    ///
    /// This is relevant for surfaces created for example with [`PdfSurface::for_stream`].
    /// Consider calling [`Surface::finish`] first,
    /// to ensure that all writes to the stream are done.
    ///
    /// Use [`Box::downcast`] to recover the concrete stream type.
    ///
    /// # Panics
    ///
    /// This panics if the output stream’s internal `RefCell` is already borrowed,
    /// or was already borrowed during a previous attempt to write to the stream.
    ///
    /// This can happen if the return value of `borrow_output_stream` or `borrow_io_error`
    /// was not dropped early enough, or if the stream’s `Write` implementation
    /// manipulates this `Surface` during a write.
    pub fn take_output_stream(&self) -> Option<Box<dyn Any>> {
        self.with_stream_env(|mut env| env.stream.take())
    }

    /// Remove and return the error that occurred while writing to the output stream, if any.
    ///
    /// # Panics
    ///
    /// This panics if the output stream’s internal `RefCell` is already borrowed,
    /// or was already borrowed during a previous attempt to write to the stream.
    ///
    /// This can happen if the return value of `borrow_output_stream` or `borrow_io_error`
    /// was not dropped early enough, or if the stream’s `Write` implementation
    /// manipulates this `Surface` during a write.
    pub fn take_io_error(&self) -> Result<(), io::Error> {
        match self.with_stream_env(|mut env| env.io_error.take()) {
            Some(error) => Err(error),
            None => Ok(()),
        }
    }

    /// Borrow the output stream, if any.
    ///
    /// This is relevant for surfaces created for example with [`PdfSurface::for_stream`].
    /// Consider calling [`Surface::finish`] first,
    /// to ensure that all writes to the stream are done.
    ///
    /// Use [`<dyn Any>::downcast_mut`] to recover the concrete stream type.
    ///
    /// # Panics
    ///
    /// This panics if the output stream’s internal `RefCell` is already borrowed,
    /// or was already borrowed during a previous attempt to write to the stream.
    ///
    /// This can happen if the return value of `borrow_output_stream` or `borrow_io_error`
    /// was not dropped early enough, or if the stream’s `Write` implementation
    /// manipulates this `Surface` during a write.
    pub fn borrow_output_stream<'surface>(&'surface self) -> Option<RefMut<'surface, dyn Any>> {
        self.with_stream_env(|env| {
            env.stream.as_ref()?;
            Some(RefMut::map(env, |env| &mut **env.stream.as_mut().unwrap()))
        })
    }

    /// Borrow the error that occurred while writing to the output stream, if any.
    ///
    /// # Panics
    ///
    /// This panics if the output stream’s internal `RefCell` is already borrowed,
    /// or was already borrowed during a previous attempt to write to the stream.
    ///
    /// This can happen if the return value of `borrow_output_stream` or `borrow_io_error`
    /// was not dropped early enough, or if the stream’s `Write` implementation
    /// manipulates this `Surface` during a write.
    pub fn borrow_io_error<'surface>(&'surface self) -> Option<RefMut<'surface, io::Error>> {
        self.with_stream_env(|env| {
            env.io_error.as_ref()?;
            Some(RefMut::map(env, |env| env.io_error.as_mut().unwrap()))
        })
    }
}

pub(crate) type Constructor = unsafe extern fn(
    ffi::cairo_write_func_t,
    *mut c_void,
    c_double,
    c_double,
) -> *mut ffi::cairo_surface_t;

static STREAM_CALLBACK_ENVIRONMENT: UserDataKey<CallbackEnvironment> =
    UserDataKey::new();

struct CallbackEnvironment {
    mutable: RefCell<MutableCallbackEnvironment>,
    saw_already_borrowed: Cell<bool>,
}

struct MutableCallbackEnvironment {
    stream: Option<Box<dyn Any>>,
    io_error: Option<io::Error>,
    unwind_payload: Option<Box<dyn Any + Send + 'static>>,
}

// Safety: unwinding into C is undefined behavior (https://github.com/rust-lang/rust/issues/58794)
// so code outside of the `catch_unwind` call must never panic.
extern "C" fn write_callback<W: io::Write + 'static>(
    env: *mut c_void,
    data: *mut c_uchar,
    length: c_uint,
) -> cairo_status_t {
    // This is consistent with the type of `env` in `Surface::_for_stream`.
    let env: *const CallbackEnvironment = env as _;

    // Safety: the user data entry keeps `Rc<CallbackEnvironment>` alive
    // until the surface is destroyed.
    // If this is called by cairo, the surface is still alive.
    let env: &CallbackEnvironment = unsafe { &*env };

    if let Ok(mut mutable) = env.mutable.try_borrow_mut() {
        if let MutableCallbackEnvironment {
            stream: Some(stream),
            // Don’t attempt another write if a previous one errored or panicked:
            io_error: io_error @ None,
            unwind_payload: unwind_payload @ None,
        } = &mut *mutable {
            // Safety: `write_callback<W>` was instanciated in `Surface::_for_stream`
            // with a W parameter consistent with the box that was unsized to `Box<dyn Any>`.
            let stream = unsafe {
                stream.downcast_mut_unchecked::<W>()
            };
            // Safety: this is the callback contract from cairo’s API
            let data = unsafe {
                std::slice::from_raw_parts(data, length as usize)
            };
            // Because `<W as Write>::write_all` is a generic,
            // we must conservatively assume that it can panic.
            let result = std::panic::catch_unwind(AssertUnwindSafe(|| stream.write_all(data)));
            match result {
                Ok(Ok(())) => {
                    return Status::Success.into()
                }
                Ok(Err(error)) => {
                    *io_error = Some(error);
                }
                Err(payload) => {
                    *unwind_payload = Some(payload);
                }
            }
        }
    } else {
        env.saw_already_borrowed.set(true)
    }
    Status::WriteError.into()
}

struct RawStream<W>(*mut W);

impl<W: io::Write> io::Write for RawStream<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> { unsafe { (*self.0).write(buf) }}
    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> { unsafe { (*self.0).write_all(buf) }}
    fn flush(&mut self) -> io::Result<()> { unsafe { (*self.0).flush() } }
}

trait AnyExt {
    /// Any::downcast_mut, but YOLO
    unsafe fn downcast_mut_unchecked<T>(&mut self) -> &mut T {
        let ptr = self as *mut Self as *mut T;
        &mut *ptr
    }
}
impl AnyExt for dyn Any {}
<commit_msg>Result over Option for errors<commit_after>// Copyright 2019, The Gtk-rs Project Developers.
// See the COPYRIGHT file at the top-level directory of this distribution.
// Licensed under the MIT license, see the LICENSE file or <http://opensource.org/licenses/MIT>

use ::ffi::{self, cairo_status_t};
use ::{Status, Surface, UserDataKey};

use libc::{c_void, c_double, c_uchar, c_uint};
use std::any::Any;
use std::cell::{Cell, RefCell, RefMut};
use std::io;
use std::panic::AssertUnwindSafe;
use std::rc::Rc;

macro_rules! for_stream_constructors {
    ($constructor_ffi: ident) => {
        /// Takes full ownership of the output stream,
        /// which is not allowed to borrow any lifetime shorter than `'static`.
        ///
        /// Because the underlying `cairo_surface_t` is reference-counted,
        /// a lifetime parameter in a Rust wrapper type would not be enough to track
        /// how long it can keep writing to the stream.
        pub fn for_stream<W: io::Write + 'static>(width: f64, height: f64, stream: W) -> Self {
            Self {
                inner: Surface::_for_stream(
                    ffi::$constructor_ffi,
                    width,
                    height,
                    stream,
                ),
            }
        }

        /// Allows writing to a borrowed stream. The lifetime of the borrow is not tracked.
        ///
        /// # Safety
        ///
        /// The value that `stream` points to must live at least until the underlying `cairo_surface_t`
        /// (which maybe be longer then the Rust `PdfSurface` wrapper, because of reference-counting),
        /// or until the output stream is removed from the surface with [`Surface::take_output_stream`].
        ///
        /// Since the former is hard to track for sure, the latter is strongly recommended.
        /// The concrete type behind the `Box<dyn Any>` value returned by `take_output_stream`
        /// is private, so you won’t be able to downcast it.
        /// But removing it anyway ensures that later writes do no go through a dangling pointer.
        pub unsafe fn for_raw_stream<W: io::Write + 'static>(width: f64, height: f64, stream: *mut W) -> Self {
            Self {
                inner: Surface::_for_raw_stream(
                    ffi::$constructor_ffi,
                    width,
                    height,
                    stream,
                ),
            }
        }
    };
}

impl Surface {
    pub(crate) fn _for_stream<W: io::Write + 'static>(
        constructor: Constructor,
        width: f64,
        height: f64,
        stream: W,
    ) -> Self {
        let env_rc = Rc::new(CallbackEnvironment {
            mutable: RefCell::new(MutableCallbackEnvironment {
                stream: Some(Box::new(stream)),
                io_error: None,
                unwind_payload: None,
            }),
            saw_already_borrowed: Cell::new(false),
        });
        let env: *const CallbackEnvironment = &*env_rc;
        unsafe {
            let ptr = constructor(Some(write_callback::<W>), env as *mut c_void, width, height);
            let surface = Surface::from_raw_full(ptr);
            surface.set_user_data(&STREAM_CALLBACK_ENVIRONMENT, env_rc);
            surface
        }
    }

    pub(crate) unsafe fn _for_raw_stream<W: io::Write + 'static>(
        constructor: Constructor,
        width: f64,
        height: f64,
        stream: *mut W,
    ) -> Self {
        Self::_for_stream(constructor, width, height, RawStream(stream))
    }

    fn stream_env<'s>(&'s self) -> Option<RefMut<'s, MutableCallbackEnvironment>>
    {
        let env = self.get_user_data_ptr(&STREAM_CALLBACK_ENVIRONMENT)?;

        // Safety: since `STREAM_CALLBACK_ENVIRONMENT` is private and we never
        // call `set_user_data` again or `remove_user_data` with it,
        // the contract of `get_user_data_ptr` says that the user data entry
        // lives as long as the underlying `cairo_surface_t`
        // which is at least as long as this `Surface` wrapper.
        let env: &'s CallbackEnvironment = unsafe { &*env.as_ptr() };

        if env.saw_already_borrowed.get() {
            panic!("The output stream’s RefCell was already borrowed when cairo attempted a write")
        }

        let mut mutable = env.mutable.borrow_mut();
        if let Some(payload) = mutable.unwind_payload.take() {
            std::panic::resume_unwind(payload)
        }
        Some(mutable)
    }

    /// Remove and return the output stream, if any.
    ///
    /// This is relevant for surfaces created for example with [`PdfSurface::for_stream`].
    /// Consider calling [`Surface::finish`] first,
    /// to ensure that all writes to the stream are done.
    ///
    /// Use [`Box::downcast`] to recover the concrete stream type.
    ///
    /// # Panics
    ///
    /// This panics if the output stream’s internal `RefCell` is already borrowed,
    /// or was already borrowed during a previous attempt to write to the stream.
    ///
    /// This can happen if the return value of `borrow_output_stream` or `borrow_io_error`
    /// was not dropped early enough, or if the stream’s `Write` implementation
    /// manipulates this `Surface` during a write.
    pub fn take_output_stream(&self) -> Option<Box<dyn Any>> {
        self.stream_env()?.stream.take()
    }

    /// Remove and return the error that occurred while writing to the output stream, if any.
    ///
    /// # Panics
    ///
    /// This panics if the output stream’s internal `RefCell` is already borrowed,
    /// or was already borrowed during a previous attempt to write to the stream.
    ///
    /// This can happen if the return value of `borrow_output_stream` or `borrow_io_error`
    /// was not dropped early enough, or if the stream’s `Write` implementation
    /// manipulates this `Surface` during a write.
    pub fn take_io_error(&self) -> Result<(), io::Error> {
        some_is_err(|| self.stream_env()?.io_error.take())
    }

    /// Borrow the output stream, if any.
    ///
    /// This is relevant for surfaces created for example with [`PdfSurface::for_stream`].
    /// Consider calling [`Surface::finish`] first,
    /// to ensure that all writes to the stream are done.
    ///
    /// Use [`<dyn Any>::downcast_mut`] to recover the concrete stream type.
    ///
    /// # Panics
    ///
    /// This panics if the output stream’s internal `RefCell` is already borrowed,
    /// or was already borrowed during a previous attempt to write to the stream.
    ///
    /// This can happen if the return value of `borrow_output_stream` or `borrow_io_error`
    /// was not dropped early enough, or if the stream’s `Write` implementation
    /// manipulates this `Surface` during a write.
    pub fn borrow_output_stream<'surface>(&'surface self) -> Option<RefMut<'surface, dyn Any>> {
        ref_mut_filter_map(self.stream_env()?, |env| env.stream.as_mut().map(|b| &mut **b))
    }

    /// Borrow the error that occurred while writing to the output stream, if any.
    ///
    /// # Panics
    ///
    /// This panics if the output stream’s internal `RefCell` is already borrowed,
    /// or was already borrowed during a previous attempt to write to the stream.
    ///
    /// This can happen if the return value of `borrow_output_stream` or `borrow_io_error`
    /// was not dropped early enough, or if the stream’s `Write` implementation
    /// manipulates this `Surface` during a write.
    pub fn borrow_io_error<'surface>(&'surface self) -> Result<(), RefMut<'surface, io::Error>> {
        some_is_err(|| ref_mut_filter_map(self.stream_env()?, |env| env.io_error.as_mut()))
    }
}

pub(crate) type Constructor = unsafe extern fn(
    ffi::cairo_write_func_t,
    *mut c_void,
    c_double,
    c_double,
) -> *mut ffi::cairo_surface_t;

static STREAM_CALLBACK_ENVIRONMENT: UserDataKey<CallbackEnvironment> =
    UserDataKey::new();

struct CallbackEnvironment {
    mutable: RefCell<MutableCallbackEnvironment>,
    saw_already_borrowed: Cell<bool>,
}

struct MutableCallbackEnvironment {
    stream: Option<Box<dyn Any>>,
    io_error: Option<io::Error>,
    unwind_payload: Option<Box<dyn Any + Send + 'static>>,
}

// Safety: unwinding into C is undefined behavior (https://github.com/rust-lang/rust/issues/58794)
// so code outside of the `catch_unwind` call must never panic.
extern "C" fn write_callback<W: io::Write + 'static>(
    env: *mut c_void,
    data: *mut c_uchar,
    length: c_uint,
) -> cairo_status_t {
    // This is consistent with the type of `env` in `Surface::_for_stream`.
    let env: *const CallbackEnvironment = env as _;

    // Safety: the user data entry keeps `Rc<CallbackEnvironment>` alive
    // until the surface is destroyed.
    // If this is called by cairo, the surface is still alive.
    let env: &CallbackEnvironment = unsafe { &*env };

    if let Ok(mut mutable) = env.mutable.try_borrow_mut() {
        if let MutableCallbackEnvironment {
            stream: Some(stream),
            // Don’t attempt another write if a previous one errored or panicked:
            io_error: io_error @ None,
            unwind_payload: unwind_payload @ None,
        } = &mut *mutable {
            // Safety: `write_callback<W>` was instanciated in `Surface::_for_stream`
            // with a W parameter consistent with the box that was unsized to `Box<dyn Any>`.
            let stream = unsafe {
                stream.downcast_mut_unchecked::<W>()
            };
            // Safety: this is the callback contract from cairo’s API
            let data = unsafe {
                std::slice::from_raw_parts(data, length as usize)
            };
            // Because `<W as Write>::write_all` is a generic,
            // we must conservatively assume that it can panic.
            let result = std::panic::catch_unwind(AssertUnwindSafe(|| stream.write_all(data)));
            match result {
                Ok(Ok(())) => {
                    return Status::Success.into()
                }
                Ok(Err(error)) => {
                    *io_error = Some(error);
                }
                Err(payload) => {
                    *unwind_payload = Some(payload);
                }
            }
        }
    } else {
        env.saw_already_borrowed.set(true)
    }
    Status::WriteError.into()
}

struct RawStream<W>(*mut W);

impl<W: io::Write> io::Write for RawStream<W> {
    fn write(&mut self, buf: &[u8]) -> io::Result<usize> { unsafe { (*self.0).write(buf) }}
    fn write_all(&mut self, buf: &[u8]) -> io::Result<()> { unsafe { (*self.0).write_all(buf) }}
    fn flush(&mut self) -> io::Result<()> { unsafe { (*self.0).flush() } }
}

trait AnyExt {
    /// Any::downcast_mut, but YOLO
    unsafe fn downcast_mut_unchecked<T>(&mut self) -> &mut T {
        let ptr = self as *mut Self as *mut T;
        &mut *ptr
    }
}
impl AnyExt for dyn Any {}

fn some_is_err<E>(f: impl FnOnce() -> Option<E>) -> Result<(), E> {
    match f() {
        Some(e) => Err(e),
        None => Ok(())
    }
}

fn ref_mut_filter_map<T, U:? Sized>(mut r: RefMut<T>, f: impl Fn(&mut T) -> Option<&mut U>)
    -> Option<RefMut<U>>
{
    f(&mut r)?;
    Some(RefMut::map(r, |value| f(value).unwrap()))
}
<|endoftext|>


<commit_before>//
// Copyright 2022 The Project Oak Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

//! Virtio vsock packet implementation.
//!
//! A packet consists of a header and an optional payload.
//!
//! See <https://docs.oasis-open.org/virtio/virtio/v1.1/csprd01/virtio-v1.1-csprd01.html#x1-3960006>.

use alloc::{vec, vec::Vec};
use anyhow::Context;
use bitflags::bitflags;
use core::mem::size_of;
use num_enum::{IntoPrimitive, TryFromPrimitive};

/// The size of the packet header in bytes.
const HEADER_SIZE: usize = 44;
/// The offset to the src_cid field.
const SRC_CID_OFFSET: usize = 0;
/// The offset to the dst_cid field.
const DST_CID_OFFSET: usize = 8;
/// The offset to the src_port field.
const SRC_PORT_OFFSET: usize = 16;
/// The offset to the dst_port field.
const DST_PORT_OFFSET: usize = 20;
/// The offset to the len field.
const LEN_OFFSET: usize = 24;
/// The offset ot the type field.
const TYPE_OFFSET: usize = 28;
/// The offset ot the op field.
const OP_OFFSET: usize = 30;
/// The offset ot the flags field.
const FLAGS_OFFSET: usize = 32;
/// The offset ot the buf_alloc field.
const BUF_ALLOC_OFFSET: usize = 36;
/// The offset ot the fwd_cnt field.
const FWD_CNT_OFFSET: usize = 40;

pub struct Packet {
    buffer: Vec<u8>,
}

impl Packet {
    /// Creates a new `Packet` from an existing buffer.
    pub fn new(buffer: Vec<u8>) -> anyhow::Result<Self> {
        if buffer.len() < HEADER_SIZE {
            anyhow::bail!("vsock packet too short");
        }
        Ok(Self { buffer })
    }

    /// Creates a new control `Packet` with only a header.
    pub fn new_header_only() -> Self {
        Self::new_with_buffer_size(HEADER_SIZE)
    }

    /// Creates a new data `Packet` with the given payload length.
    pub fn new_with_payload(payload_len: usize) -> Self {
        let mut result = Self::new_with_buffer_size(HEADER_SIZE + payload_len);
        result.set_len(payload_len as u32);
        result.set_op(VSockOp::Rw).unwrap();
        result
    }

    /// Creates a new `Packet` with the specified buffer size.
    fn new_with_buffer_size(buffer_len: usize) -> Self {
        let buffer = vec![0u8; buffer_len];
        let mut result = Self { buffer };
        result.set_type(VSockType::Stream);
        result
    }

    /// Gets the source CID.
    pub fn get_src_cid(&self) -> u64 {
        self.read_u64(SRC_CID_OFFSET)
    }

    /// Sets the source CID.
    pub fn set_src_cid(&mut self, src_cid: u64) {
        self.write_u64(SRC_CID_OFFSET, src_cid);
    }

    /// Gets the destination CID.
    pub fn get_dst_cid(&self) -> u64 {
        self.read_u64(DST_CID_OFFSET)
    }

    /// Sets the destination CID.
    pub fn set_dst_cid(&mut self, dst_cid: u64) {
        self.write_u64(DST_CID_OFFSET, dst_cid);
    }

    /// Gets the source port.
    pub fn get_src_port(&self) -> u32 {
        self.read_u32(SRC_PORT_OFFSET)
    }

    /// Sets the source port.
    pub fn set_src_port(&mut self, src_port: u32) {
        self.write_u32(SRC_PORT_OFFSET, src_port);
    }

    /// Gets the destination port.
    pub fn get_dst_port(&self) -> u32 {
        self.read_u32(DST_PORT_OFFSET)
    }

    /// Sets the destination port.
    pub fn set_dst_port(&mut self, dst_port: u32) {
        self.write_u32(DST_PORT_OFFSET, dst_port);
    }

    /// Gets the payload length.
    pub fn get_len(&self) -> u32 {
        self.read_u32(LEN_OFFSET)
    }

    /// Sets the payload length.
    fn set_len(&mut self, len: u32) {
        self.write_u32(LEN_OFFSET, len);
    }

    /// Gets the type of socket the packet is intended for.
    pub fn get_type(&self) -> anyhow::Result<VSockType> {
        VSockType::try_from_primitive(self.read_u16(TYPE_OFFSET))
            .map_err(anyhow::Error::msg)
            .context("invalid socket type")
    }

    /// Sets the type of socket the packet is intended for.
    fn set_type(&mut self, socket_type: VSockType) {
        self.write_u16(TYPE_OFFSET, socket_type.into())
    }

    /// Gets the op that the packet represents.
    pub fn get_op(&self) -> anyhow::Result<VSockOp> {
        VSockOp::try_from_primitive(self.read_u16(OP_OFFSET))
            .map_err(anyhow::Error::msg)
            .context("invalid op")
    }

    /// Sets the op that the packet represents.
    pub fn set_op(&mut self, op: VSockOp) -> anyhow::Result<()> {
        if self.get_len() > 0 && op != VSockOp::Rw {
            anyhow::bail!("non-empty payloads are only allowed with RW ops");
        }
        self.write_u16(OP_OFFSET, op.into());
        Ok(())
    }

    /// Gets the flags.
    ///
    /// Only currently used for connection shutdown requests.
    pub fn get_flags(&self) -> VSockFlags {
        VSockFlags::from_bits_truncate(self.read_u32(FLAGS_OFFSET))
    }

    /// Sets the flags.
    ///
    /// Only currently used for connection shutdown requests.
    pub fn set_flags(&mut self, flags: VSockFlags) {
        self.write_u32(FLAGS_OFFSET, flags.bits());
    }

    /// Gets the size of the peer's stream buffer.
    ///
    /// Used to facilitate flow-control calculations.
    pub fn get_buf_alloc(&self) -> u32 {
        self.read_u32(BUF_ALLOC_OFFSET)
    }

    /// Sets the size of the stream buffer.
    pub fn set_buf_alloc(&mut self, buf_alloc: u32) {
        self.write_u32(BUF_ALLOC_OFFSET, buf_alloc);
    }

    /// Gets the number of bytes that the peer has read out of the stream buffer (hence making space
    /// in the buffer).
    ///
    /// Used to facilitate flow-control calculations.
    pub fn get_fwd_cnt(&self) -> u32 {
        self.read_u32(FWD_CNT_OFFSET)
    }

    /// Sets the number of bytes read from the stream buffer.
    pub fn set_fwd_cnt(&mut self, fwd_cnt: u32) {
        self.write_u32(FWD_CNT_OFFSET, fwd_cnt);
    }

    /// Gets the payload.
    pub fn get_payload(&self) -> &[u8] {
        &self.buffer[HEADER_SIZE..(HEADER_SIZE + self.get_len() as usize)]
    }

    /// Sets the payload of a data packet from a slice.
    ///
    /// The length of the slice must match the packets configures payload length. Only usable if the
    /// packet's op is `VSockOp::Rw`.
    pub fn set_payload(&mut self, data: &[u8]) -> anyhow::Result<()> {
        if self.get_op()? != VSockOp::Rw {
            anyhow::bail!("non-empty payloads are only allowed with RW ops");
        }
        let len = self.get_len();
        if len as usize != data.len() {
            anyhow::bail!("data length does not match the paacket's configured payload length");
        }

        self.buffer[HEADER_SIZE..(HEADER_SIZE + len as usize)].copy_from_slice(data);
        Ok(())
    }

    fn read_u16(&self, offset: usize) -> u16 {
        let mut data = [0; size_of::<u16>()];
        data.copy_from_slice(&self.buffer[offset..(offset + size_of::<u16>())]);
        u16::from_le_bytes(data)
    }

    fn read_u32(&self, offset: usize) -> u32 {
        let mut data = [0; size_of::<u32>()];
        data.copy_from_slice(&self.buffer[offset..(offset + size_of::<u32>())]);
        u32::from_le_bytes(data)
    }

    fn read_u64(&self, offset: usize) -> u64 {
        let mut data = [0; size_of::<u64>()];
        data.copy_from_slice(&self.buffer[offset..(offset + size_of::<u64>())]);
        u64::from_le_bytes(data)
    }

    fn write_u16(&mut self, offset: usize, value: u16) {
        let dest = &mut self.buffer[offset..(offset + size_of::<u16>())];
        dest.copy_from_slice(&value.to_le_bytes()[..]);
    }

    fn write_u32(&mut self, offset: usize, value: u32) {
        let dest = &mut self.buffer[offset..(offset + size_of::<u32>())];
        dest.copy_from_slice(&value.to_le_bytes()[..]);
    }

    fn write_u64(&mut self, offset: usize, value: u64) {
        let dest = &mut self.buffer[offset..(offset + size_of::<u64>())];
        dest.copy_from_slice(&value.to_le_bytes()[..]);
    }
}

/// Vsock Ops.
#[derive(Debug, Eq, PartialEq, TryFromPrimitive, IntoPrimitive)]
#[repr(u16)]
pub enum VSockOp {
    Invalid = 0,
    /// Connection request.
    Request = 1,
    /// Connections accepted response.
    Response = 2,
    /// Connection reset, either in reponse to a shutdown request or invalid packets received.
    Rst = 3,
    /// Connection shutdown request.
    Shutdown = 4,
    /// Represents a data packet.
    ///
    /// Used to send payload. Only data packets with this op can contain a payload.
    Rw = 5,
    /// Give update on credit to support flow control, either in response to a credit request or at
    /// any time it might be useful to inform the peer of the state of the stream buffer.
    CreditUpdate = 6,
    /// Request for update on credit to calculate stream buffer availability.
    CreditRequest = 7,
}

bitflags! {
    /// Flags about a socket connection.
    ///
    /// Used in a connection shutdown request.
    /// See <https://docs.oasis-open.org/virtio/virtio/v1.1/csprd01/virtio-v1.1-csprd01.html#x1-4070005>.
    pub struct VSockFlags: u32 {
        /// Indicates that no more payload data will be received.
        const NO_RECEIVE = 1;
        /// Indicates that no more payload data will be senn.
        const NO_SEND = 2;
    }
}

/// Socket Type.
#[derive(Debug, Eq, PartialEq, TryFromPrimitive, IntoPrimitive)]
#[repr(u16)]
pub enum VSockType {
    /// Only stream sockets are currently supported in the Virtio spec.
    Stream = 1,
}
<commit_msg>Add tests for the vsock packet implementation (#2888)<commit_after>//
// Copyright 2022 The Project Oak Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

//! Virtio vsock packet implementation.
//!
//! A packet consists of a header and an optional payload.
//!
//! See <https://docs.oasis-open.org/virtio/virtio/v1.1/csprd01/virtio-v1.1-csprd01.html#x1-3960006>.

use alloc::{vec, vec::Vec};
use anyhow::Context;
use bitflags::bitflags;
use core::mem::size_of;
use num_enum::{IntoPrimitive, TryFromPrimitive};

/// The size of the packet header in bytes.
const HEADER_SIZE: usize = 44;
/// The offset to the src_cid field.
const SRC_CID_OFFSET: usize = 0;
/// The offset to the dst_cid field.
const DST_CID_OFFSET: usize = 8;
/// The offset to the src_port field.
const SRC_PORT_OFFSET: usize = 16;
/// The offset to the dst_port field.
const DST_PORT_OFFSET: usize = 20;
/// The offset to the len field.
const LEN_OFFSET: usize = 24;
/// The offset ot the type field.
const TYPE_OFFSET: usize = 28;
/// The offset ot the op field.
const OP_OFFSET: usize = 30;
/// The offset ot the flags field.
const FLAGS_OFFSET: usize = 32;
/// The offset ot the buf_alloc field.
const BUF_ALLOC_OFFSET: usize = 36;
/// The offset ot the fwd_cnt field.
const FWD_CNT_OFFSET: usize = 40;

pub struct Packet {
    buffer: Vec<u8>,
}

impl Packet {
    /// Creates a new `Packet` from an existing buffer.
    pub fn new(buffer: Vec<u8>) -> anyhow::Result<Self> {
        if buffer.len() < HEADER_SIZE {
            anyhow::bail!("vsock packet too short");
        }
        Ok(Self { buffer })
    }

    /// Creates a new control `Packet` with only a header.
    pub fn new_header_only() -> Self {
        Self::new_with_buffer_size(HEADER_SIZE)
    }

    /// Creates a new data `Packet` with the given payload length.
    pub fn new_with_payload(payload_len: usize) -> Self {
        let mut result = Self::new_with_buffer_size(HEADER_SIZE + payload_len);
        result.set_len(payload_len as u32);
        result.set_op(VSockOp::Rw).unwrap();
        result
    }

    /// Creates a new `Packet` with the specified buffer size.
    fn new_with_buffer_size(buffer_len: usize) -> Self {
        let buffer = vec![0u8; buffer_len];
        let mut result = Self { buffer };
        result.set_type(VSockType::Stream);
        result
    }

    /// Gets the source CID.
    pub fn get_src_cid(&self) -> u64 {
        self.read_u64(SRC_CID_OFFSET)
    }

    /// Sets the source CID.
    pub fn set_src_cid(&mut self, src_cid: u64) {
        self.write_u64(SRC_CID_OFFSET, src_cid);
    }

    /// Gets the destination CID.
    pub fn get_dst_cid(&self) -> u64 {
        self.read_u64(DST_CID_OFFSET)
    }

    /// Sets the destination CID.
    pub fn set_dst_cid(&mut self, dst_cid: u64) {
        self.write_u64(DST_CID_OFFSET, dst_cid);
    }

    /// Gets the source port.
    pub fn get_src_port(&self) -> u32 {
        self.read_u32(SRC_PORT_OFFSET)
    }

    /// Sets the source port.
    pub fn set_src_port(&mut self, src_port: u32) {
        self.write_u32(SRC_PORT_OFFSET, src_port);
    }

    /// Gets the destination port.
    pub fn get_dst_port(&self) -> u32 {
        self.read_u32(DST_PORT_OFFSET)
    }

    /// Sets the destination port.
    pub fn set_dst_port(&mut self, dst_port: u32) {
        self.write_u32(DST_PORT_OFFSET, dst_port);
    }

    /// Gets the payload length.
    pub fn get_len(&self) -> u32 {
        self.read_u32(LEN_OFFSET)
    }

    /// Sets the payload length.
    fn set_len(&mut self, len: u32) {
        self.write_u32(LEN_OFFSET, len);
    }

    /// Gets the type of socket the packet is intended for.
    pub fn get_type(&self) -> anyhow::Result<VSockType> {
        VSockType::try_from_primitive(self.read_u16(TYPE_OFFSET))
            .map_err(anyhow::Error::msg)
            .context("invalid socket type")
    }

    /// Sets the type of socket the packet is intended for.
    fn set_type(&mut self, socket_type: VSockType) {
        self.write_u16(TYPE_OFFSET, socket_type.into())
    }

    /// Gets the op that the packet represents.
    pub fn get_op(&self) -> anyhow::Result<VSockOp> {
        VSockOp::try_from_primitive(self.read_u16(OP_OFFSET))
            .map_err(anyhow::Error::msg)
            .context("invalid op")
    }

    /// Sets the op that the packet represents.
    pub fn set_op(&mut self, op: VSockOp) -> anyhow::Result<()> {
        if self.get_len() > 0 && op != VSockOp::Rw {
            anyhow::bail!("non-empty payloads are only allowed with RW ops");
        }
        self.write_u16(OP_OFFSET, op.into());
        Ok(())
    }

    /// Gets the flags.
    ///
    /// Only currently used for connection shutdown requests.
    pub fn get_flags(&self) -> VSockFlags {
        VSockFlags::from_bits_truncate(self.read_u32(FLAGS_OFFSET))
    }

    /// Sets the flags.
    ///
    /// Only currently used for connection shutdown requests.
    pub fn set_flags(&mut self, flags: VSockFlags) {
        self.write_u32(FLAGS_OFFSET, flags.bits());
    }

    /// Gets the size of the peer's stream buffer.
    ///
    /// Used to facilitate flow-control calculations.
    pub fn get_buf_alloc(&self) -> u32 {
        self.read_u32(BUF_ALLOC_OFFSET)
    }

    /// Sets the size of the stream buffer.
    pub fn set_buf_alloc(&mut self, buf_alloc: u32) {
        self.write_u32(BUF_ALLOC_OFFSET, buf_alloc);
    }

    /// Gets the number of bytes that the peer has read out of the stream buffer (hence making space
    /// in the buffer).
    ///
    /// Used to facilitate flow-control calculations.
    pub fn get_fwd_cnt(&self) -> u32 {
        self.read_u32(FWD_CNT_OFFSET)
    }

    /// Sets the number of bytes read from the stream buffer.
    pub fn set_fwd_cnt(&mut self, fwd_cnt: u32) {
        self.write_u32(FWD_CNT_OFFSET, fwd_cnt);
    }

    /// Gets the payload.
    pub fn get_payload(&self) -> &[u8] {
        &self.buffer[HEADER_SIZE..(HEADER_SIZE + self.get_len() as usize)]
    }

    /// Sets the payload of a data packet from a slice.
    ///
    /// The length of the slice must match the packets configures payload length. Only usable if the
    /// packet's op is `VSockOp::Rw`.
    pub fn set_payload(&mut self, data: &[u8]) -> anyhow::Result<()> {
        if self.get_op()? != VSockOp::Rw {
            anyhow::bail!("non-empty payloads are only allowed with RW ops");
        }
        let len = self.get_len();
        if len as usize != data.len() {
            anyhow::bail!("data length does not match the paacket's configured payload length");
        }

        self.buffer[HEADER_SIZE..(HEADER_SIZE + len as usize)].copy_from_slice(data);
        Ok(())
    }

    fn read_u16(&self, offset: usize) -> u16 {
        let mut data = [0; size_of::<u16>()];
        data.copy_from_slice(&self.buffer[offset..(offset + size_of::<u16>())]);
        u16::from_le_bytes(data)
    }

    fn read_u32(&self, offset: usize) -> u32 {
        let mut data = [0; size_of::<u32>()];
        data.copy_from_slice(&self.buffer[offset..(offset + size_of::<u32>())]);
        u32::from_le_bytes(data)
    }

    fn read_u64(&self, offset: usize) -> u64 {
        let mut data = [0; size_of::<u64>()];
        data.copy_from_slice(&self.buffer[offset..(offset + size_of::<u64>())]);
        u64::from_le_bytes(data)
    }

    fn write_u16(&mut self, offset: usize, value: u16) {
        let dest = &mut self.buffer[offset..(offset + size_of::<u16>())];
        dest.copy_from_slice(&value.to_le_bytes()[..]);
    }

    fn write_u32(&mut self, offset: usize, value: u32) {
        let dest = &mut self.buffer[offset..(offset + size_of::<u32>())];
        dest.copy_from_slice(&value.to_le_bytes()[..]);
    }

    fn write_u64(&mut self, offset: usize, value: u64) {
        let dest = &mut self.buffer[offset..(offset + size_of::<u64>())];
        dest.copy_from_slice(&value.to_le_bytes()[..]);
    }
}

/// Vsock Ops.
#[derive(Debug, Eq, PartialEq, TryFromPrimitive, IntoPrimitive)]
#[repr(u16)]
pub enum VSockOp {
    Invalid = 0,
    /// Connection request.
    Request = 1,
    /// Connections accepted response.
    Response = 2,
    /// Connection reset, either in reponse to a shutdown request or invalid packets received.
    Rst = 3,
    /// Connection shutdown request.
    Shutdown = 4,
    /// Represents a data packet.
    ///
    /// Used to send payload. Only data packets with this op can contain a payload.
    Rw = 5,
    /// Give update on credit to support flow control, either in response to a credit request or at
    /// any time it might be useful to inform the peer of the state of the stream buffer.
    CreditUpdate = 6,
    /// Request for update on credit to calculate stream buffer availability.
    CreditRequest = 7,
}

bitflags! {
    /// Flags about a socket connection.
    ///
    /// Used in a connection shutdown request.
    /// See <https://docs.oasis-open.org/virtio/virtio/v1.1/csprd01/virtio-v1.1-csprd01.html#x1-4070005>.
    pub struct VSockFlags: u32 {
        /// Indicates that no more payload data will be received.
        const NO_RECEIVE = 1;
        /// Indicates that no more payload data will be senn.
        const NO_SEND = 2;
    }
}

/// Socket Type.
#[derive(Debug, Eq, PartialEq, TryFromPrimitive, IntoPrimitive)]
#[repr(u16)]
pub enum VSockType {
    /// Only stream sockets are currently supported in the Virtio spec.
    Stream = 1,
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_set_and_get_all_header_fields() {
        let mut packet = Packet::new_header_only();
        packet.set_src_cid(1234);
        packet.set_dst_cid(2);
        packet.set_src_port(1023);
        packet.set_dst_port(8888);
        packet.set_flags(VSockFlags::all());
        packet.set_op(VSockOp::Shutdown).unwrap();
        packet.set_buf_alloc(4096);
        packet.set_fwd_cnt(12);

        assert_eq!(packet.get_src_cid(), 1234);
        assert_eq!(packet.get_dst_cid(), 2);
        assert_eq!(packet.get_src_port(), 1023);
        assert_eq!(packet.get_dst_port(), 8888);
        assert_eq!(packet.get_flags(), VSockFlags::all());
        assert_eq!(packet.get_type().unwrap(), VSockType::Stream);
        assert_eq!(packet.get_op().unwrap(), VSockOp::Shutdown);
        assert_eq!(packet.get_len(), 0);
        assert_eq!(packet.get_buf_alloc(), 4096);
        assert_eq!(packet.get_fwd_cnt(), 12);
    }

    #[test]
    fn test_invalid_payload() {
        let mut packet = Packet::new_with_payload(5);
        let result = packet.set_payload(&[1, 2, 3, 4]);
        assert!(result.is_err());
    }

    #[test]
    fn test_valid_payload() {
        let mut packet = Packet::new_with_payload(4);
        let result = packet.set_payload(&[1, 2, 3, 4]);
        assert!(result.is_ok());
    }

    #[test]
    fn test_new_from_buffer() {
        let mut packet = Packet::new_with_payload(4);
        packet.set_payload(&[1, 2, 3, 4]).unwrap();

        let packet = Packet::new(packet.buffer.clone()).unwrap();
        assert_eq!(packet.get_payload(), &[1, 2, 3, 4]);
        assert_eq!(packet.get_len(), 4);
        assert_eq!(packet.get_type().unwrap(), VSockType::Stream);
        assert_eq!(packet.get_op().unwrap(), VSockOp::Rw);
        assert_eq!(packet.get_flags(), VSockFlags::empty());
    }

    #[test]
    fn test_invalid_op() {
        let mut packet = Packet::new_with_payload(1);
        assert!(packet.set_op(VSockOp::Rst).is_err())
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate criterion;
#[macro_use]
extern crate serde_derive;

use std::collections::BTreeMap;
use std::fs::{create_dir_all, File};
use std::io::Write;
use std::path::Path;

use criterion::profiler::Profiler;
use criterion::Criterion;
use handlebars::{to_json, Handlebars, Template};
use pprof::protos::Message;
use pprof::ProfilerGuard;
use serde_json::value::Value as Json;

#[derive(Default)]
struct CpuProfiler<'a> {
    guard: Option<ProfilerGuard<'a>>,
}

impl<'a> Profiler for CpuProfiler<'a> {
    fn start_profiling(&mut self, _benchmark_id: &str, benchmark_dir: &Path) {
        create_dir_all(&benchmark_dir).unwrap();

        let guard = ProfilerGuard::new(100).unwrap();
        self.guard = Some(guard);
    }

    fn stop_profiling(&mut self, benchmark_id: &str, benchmark_dir: &Path) {
        if let Ok(ref report) = self.guard.as_ref().unwrap().report().build() {
            let fg_file_name = benchmark_dir.join(format!("{}.svg", benchmark_id));
            let fg_file = File::create(fg_file_name).unwrap();
            report.flamegraph(fg_file).unwrap();

            let pb_file_name = benchmark_dir.join(format!("{}.pb", benchmark_id));
            let mut pb_file = File::create(pb_file_name).unwrap();
            let profile = report.pprof().unwrap();

            let mut content = Vec::new();
            profile.encode(&mut content).unwrap();
            pb_file.write_all(&content).unwrap();
        };

        self.guard = None;
    }
}

fn profiled() -> Criterion {
    Criterion::default().with_profiler(CpuProfiler::default())
}

#[derive(Serialize)]
struct DataWrapper {
    v: String,
}

#[derive(Serialize)]
struct RowWrapper {
    real: Vec<DataWrapper>,
    dummy: Vec<DataWrapper>,
}

#[derive(Serialize)]
struct NestedRowWrapper {
    parent: Vec<Vec<DataWrapper>>,
}

static SOURCE: &'static str = "

    <title>{{year}}</title>

CSL {{year}}

    {{#each teams}}

      {{name}}: {{score}}

    {{/each}}

";

fn make_data() -> BTreeMap<String, Json> {
    let mut data = BTreeMap::new();

    data.insert("year".to_string(), to_json("2015"));

    let mut teams = Vec::new();

    for v in vec![
        ("Jiangsu", 43u16),
        ("Beijing", 27u16),
        ("Guangzhou", 22u16),
        ("Shandong", 12u16),
    ]
    .iter()
    {
        let (name, score) = *v;
        let mut t = BTreeMap::new();
        t.insert("name".to_string(), to_json(name));
        t.insert("score".to_string(), to_json(score));
        teams.push(t)
    }

    data.insert("teams".to_string(), to_json(&teams));
    data
}

fn parse_template(c: &mut Criterion) {
    c.bench_function("parse_template", move |b| {
        b.iter(|| Template::compile(SOURCE).ok().unwrap())
    });
}

fn render_template(c: &mut Criterion) {
    let mut handlebars = Handlebars::new();
    handlebars
        .register_template_string("table", SOURCE)
        .ok()
        .expect("Invalid template format");

    let data = make_data();
    c.bench_function("render_template", move |b| {
        b.iter(|| handlebars.render("table", &data).ok().unwrap())
    });
}

fn large_loop_helper(c: &mut Criterion) {
    let mut handlebars = Handlebars::new();
    handlebars
        .register_template_string("test", "BEFORE\n{{#each real}}{{this.v}}{{/each}}AFTER")
        .ok()
        .expect("Invalid template format");

    let real: Vec<DataWrapper> = (1..1000)
        .map(|i| DataWrapper {
            v: format!("n={}", i),
        })
        .collect();
    let dummy: Vec<DataWrapper> = (1..1000)
        .map(|i| DataWrapper {
            v: format!("n={}", i),
        })
        .collect();
    let rows = RowWrapper { real, dummy };

    c.bench_function("large_loop_helper", move |b| {
        b.iter(|| handlebars.render("test", &rows).ok().unwrap())
    });
}

fn large_nested_loop(c: &mut Criterion) {
    let mut handlebars = Handlebars::new();
    handlebars
        .register_template_string(
            "test",
            "BEFORE\n{{#each parent as |child|}}{{#each child}}{{this.v}}{{/each}}{{/each}}AFTER",
        )
        .ok()
        .expect("Invalid template format");

    let parent: Vec<Vec<DataWrapper>> = (1..100)
        .map(|_| {
            (1..10)
                .map(|v| DataWrapper {
                    v: format!("v={}", v),
                })
                .collect()
        })
        .collect();

    let rows = NestedRowWrapper { parent };

    c.bench_function("large_nested_loop", move |b| {
        b.iter(|| handlebars.render("test", &rows).ok().unwrap())
    });
}

criterion_group!(
    name = benches;
    config = profiled();
    targets = parse_template, render_template, large_loop_helper, large_nested_loop
);
criterion_main!(benches);
<commit_msg>(test) update benchmark code to use context based api<commit_after>#[macro_use]
extern crate criterion;
#[macro_use]
extern crate serde_derive;

use std::collections::BTreeMap;
use std::fs::{create_dir_all, File};
use std::io::Write;
use std::path::Path;

use criterion::profiler::Profiler;
use criterion::Criterion;
use handlebars::{to_json, Context, Handlebars, Template};
use pprof::protos::Message;
use pprof::ProfilerGuard;
use serde_json::value::Value as Json;

#[derive(Default)]
struct CpuProfiler<'a> {
    guard: Option<ProfilerGuard<'a>>,
}

impl<'a> Profiler for CpuProfiler<'a> {
    fn start_profiling(&mut self, _benchmark_id: &str, benchmark_dir: &Path) {
        create_dir_all(&benchmark_dir).unwrap();

        let guard = ProfilerGuard::new(100).unwrap();
        self.guard = Some(guard);
    }

    fn stop_profiling(&mut self, benchmark_id: &str, benchmark_dir: &Path) {
        if let Ok(ref report) = self.guard.as_ref().unwrap().report().build() {
            let fg_file_name = benchmark_dir.join(format!("{}.svg", benchmark_id));
            let fg_file = File::create(fg_file_name).unwrap();
            report.flamegraph(fg_file).unwrap();

            let pb_file_name = benchmark_dir.join(format!("{}.pb", benchmark_id));
            let mut pb_file = File::create(pb_file_name).unwrap();
            let profile = report.pprof().unwrap();

            let mut content = Vec::new();
            profile.encode(&mut content).unwrap();
            pb_file.write_all(&content).unwrap();
        };

        self.guard = None;
    }
}

fn profiled() -> Criterion {
    Criterion::default().with_profiler(CpuProfiler::default())
}

#[derive(Serialize)]
struct DataWrapper {
    v: String,
}

#[derive(Serialize)]
struct RowWrapper {
    real: Vec<DataWrapper>,
    dummy: Vec<DataWrapper>,
}

#[derive(Serialize)]
struct NestedRowWrapper {
    parent: Vec<Vec<DataWrapper>>,
}

static SOURCE: &'static str = "

    <title>{{year}}</title>

CSL {{year}}

    {{#each teams}}

      {{name}}: {{score}}

    {{/each}}

";

fn make_data() -> BTreeMap<String, Json> {
    let mut data = BTreeMap::new();

    data.insert("year".to_string(), to_json("2015"));

    let mut teams = Vec::new();

    for v in vec![
        ("Jiangsu", 43u16),
        ("Beijing", 27u16),
        ("Guangzhou", 22u16),
        ("Shandong", 12u16),
    ]
    .iter()
    {
        let (name, score) = *v;
        let mut t = BTreeMap::new();
        t.insert("name".to_string(), to_json(name));
        t.insert("score".to_string(), to_json(score));
        teams.push(t)
    }

    data.insert("teams".to_string(), to_json(&teams));
    data
}

fn parse_template(c: &mut Criterion) {
    c.bench_function("parse_template", move |b| {
        b.iter(|| Template::compile(SOURCE).ok().unwrap())
    });
}

fn render_template(c: &mut Criterion) {
    let mut handlebars = Handlebars::new();
    handlebars
        .register_template_string("table", SOURCE)
        .ok()
        .expect("Invalid template format");

    let ctx = Context::wraps(make_data()).unwrap();
    c.bench_function("render_template", move |b| {
        b.iter(|| handlebars.render_with_context("table", &ctx).ok().unwrap())
    });
}

fn large_loop_helper(c: &mut Criterion) {
    let mut handlebars = Handlebars::new();
    handlebars
        .register_template_string("test", "BEFORE\n{{#each real}}{{this.v}}{{/each}}AFTER")
        .ok()
        .expect("Invalid template format");

    let real: Vec<DataWrapper> = (1..1000)
        .map(|i| DataWrapper {
            v: format!("n={}", i),
        })
        .collect();
    let dummy: Vec<DataWrapper> = (1..1000)
        .map(|i| DataWrapper {
            v: format!("n={}", i),
        })
        .collect();
    let rows = RowWrapper { real, dummy };

    let ctx = Context::wraps(&rows).unwrap();
    c.bench_function("large_loop_helper", move |b| {
        b.iter(|| handlebars.render_with_context("test", &ctx).ok().unwrap())
    });
}

fn large_nested_loop(c: &mut Criterion) {
    let mut handlebars = Handlebars::new();
    handlebars
        .register_template_string(
            "test",
            "BEFORE\n{{#each parent as |child|}}{{#each child}}{{this.v}}{{/each}}{{/each}}AFTER",
        )
        .ok()
        .expect("Invalid template format");

    let parent: Vec<Vec<DataWrapper>> = (1..100)
        .map(|_| {
            (1..10)
                .map(|v| DataWrapper {
                    v: format!("v={}", v),
                })
                .collect()
        })
        .collect();

    let rows = NestedRowWrapper { parent };

    let ctx = Context::wraps(&rows).unwrap();
    c.bench_function("large_nested_loop", move |b| {
        b.iter(|| handlebars.render_with_context("test", &ctx).ok().unwrap())
    });
}

criterion_group!(
    name = benches;
    config = profiled();
    targets = parse_template, render_template, large_loop_helper, large_nested_loop
);
criterion_main!(benches);
<|endoftext|>