<commit_before>use crate::schema::*;
use diesel::result::Error;
use diesel::*;

#[test]
#[cfg(not(feature = "sqlite"))] // FIXME: This test is only valid when operating on a file and not :memory:
fn transaction_executes_fn_in_a_sql_transaction() {
    const TEST_NAME: &'static str = "transaction_executes_fn_in_a_sql_transaction";
    let conn1 = &mut connection_without_transaction();
    let conn2 = &mut connection_without_transaction();
    setup_test_table(conn1, TEST_NAME);

    fn get_count(conn: &mut TestConnection) -> i64 {
        count_test_table(conn, TEST_NAME)
    }

    conn1
        .transaction::<_, Error, _>(|conn1| {
            assert_eq!(0, get_count(conn1));
            assert_eq!(0, get_count(conn2));
            conn1.execute(&format!("INSERT INTO {} DEFAULT VALUES", TEST_NAME))?;
            assert_eq!(1, get_count(conn1));
            assert_eq!(0, get_count(conn2));
            Ok(())
        })
        .unwrap();

    assert_eq!(1, get_count(conn1));
    assert_eq!(1, get_count(conn2));

    drop_test_table(conn1, TEST_NAME);
}

#[test]
fn transaction_returns_the_returned_value() {
    let conn1 = &mut connection_without_transaction();

    assert_eq!(Ok(1), conn1.transaction::<_, Error, _>(|_| Ok(1)));
}

#[test]
fn transaction_is_rolled_back_when_returned_an_error() {
    let connection = &mut connection_without_transaction();
    let test_name = "transaction_is_rolled_back_when_returned_an_error";
    setup_test_table(connection, test_name);

    let _ = connection.transaction::<(), _, _>(|connection| {
        connection
            .execute(&format!("INSERT INTO {} DEFAULT VALUES", test_name))
            .unwrap();
        Err(Error::RollbackTransaction)
    });
    assert_eq!(0, count_test_table(connection, test_name));

    drop_test_table(connection, test_name);
}

#[test]
#[cfg(feature = "sqlite")]
fn transaction_rollback_returns_error() {
    let connection = &mut connection_without_transaction();
    let test_name = "transaction_rollback_returns_error";
    setup_test_table(connection, test_name);

    // Create a transaction that will fail to rollback.
    let r = connection.transaction::<usize, _, _>(|connection| {
        connection
            .execute(&format!("INSERT INTO {} DEFAULT VALUES", test_name))
            .unwrap();

        // This rollback would succeed, and cause any rollback later to fail.
        connection.execute("ROLLBACK").unwrap();

        // Return any error to trigger a rollback that fails in this case.
        Err(Error::NotFound)
    });

    // Verify that the transaction failed with `RollbackError`.
    assert!(r.is_err());
    assert!(matches!(r.unwrap_err(), Error::RollbackError(_)));

    assert_eq!(0, count_test_table(connection, test_name));
    drop_test_table(connection, test_name);
}

#[test]
fn transactions_can_be_nested() {
    let connection = &mut connection_without_transaction();
    const TEST_NAME: &'static str = "transactions_can_be_nested";
    setup_test_table(connection, TEST_NAME);
    fn get_count(connection: &mut TestConnection) -> i64 {
        count_test_table(connection, TEST_NAME)
    }

    let _ = connection.transaction::<(), _, _>(|connection| {
        connection
            .execute(&format!("INSERT INTO {} DEFAULT VALUES", TEST_NAME))
            .unwrap();
        assert_eq!(1, get_count(connection));
        let _ = connection.transaction::<(), _, _>(|connection| {
            connection
                .execute(&format!("INSERT INTO {} DEFAULT VALUES", TEST_NAME))
                .unwrap();
            assert_eq!(2, get_count(connection));
            Err(Error::RollbackTransaction)
        });
        assert_eq!(1, get_count(connection));
        let _ = connection.transaction::<(), Error, _>(|connection| {
            connection
                .execute(&format!("INSERT INTO {} DEFAULT VALUES", TEST_NAME))
                .unwrap();
            assert_eq!(2, get_count(connection));
            Ok(())
        });
        assert_eq!(2, get_count(connection));
        Err(Error::RollbackTransaction)
    });
    assert_eq!(0, get_count(connection));

    drop_test_table(connection, TEST_NAME);
}

#[test]
fn test_transaction_always_rolls_back() {
    let connection = &mut connection_without_transaction();
    let test_name = "test_transaction_always_rolls_back";
    setup_test_table(connection, test_name);

    let result = connection.test_transaction::<_, Error, _>(|connection| {
        connection.execute(&format!("INSERT INTO {} DEFAULT VALUES", test_name))?;
        assert_eq!(1, count_test_table(connection, test_name));
        Ok("success")
    });
    assert_eq!(0, count_test_table(connection, test_name));
    assert_eq!("success", result);

    drop_test_table(connection, test_name);
}

#[test]
#[should_panic(expected = "Transaction did not succeed")]
fn test_transaction_panics_on_error() {
    let connection = &mut connection_without_transaction();
    connection.test_transaction::<(), _, _>(|_| Err(()));
}

fn setup_test_table(connection: &mut TestConnection, table_name: &str) {
    use crate::schema_dsl::*;
    create_table(table_name, (integer("id").primary_key().auto_increment(),))
        .execute(connection)
        .unwrap();
}

fn drop_test_table(connection: &mut TestConnection, table_name: &str) {
    connection
        .execute(&format!("DROP TABLE {}", table_name))
        .unwrap();
}

fn count_test_table(connection: &mut TestConnection, table_name: &str) -> i64 {
    use diesel::dsl::sql;
    select(sql::<sql_types::BigInt>(&format!(
        "COUNT(*) FROM {}",
        table_name
    )))
    .first(connection)
    .unwrap()
}

#[test]
#[cfg(feature = "postgres")]
fn regression_test_for_2123() {
    let conn = &mut connection_without_transaction();
    // fail once
    let ret = conn.transaction(|conn| {
        let _ = conn.transaction(|conn| {
            // handling error
            match conn.execute("SELECT foo") {
                // do nothing
                Ok(_) => unreachable!("This query should fail"),
                // ignore the error
                Err(e) => eprintln!("error occurred: {}", e),
            };
            Ok::<_, Error>(())
        });

        conn.transaction(|conn| {
            let ret = conn.execute("SELECT 1");
            assert_eq!(Ok(1), ret);
            Ok::<_, Error>(())
        })
    });
    println!("{:?}", ret);
    // other transaction
    let ret = conn
        .build_transaction()
        .serializable()
        .run(|conn| conn.execute("SELECT 1"));
    // must be Ok(1), but get Err(AlreadyInTransaction)
    println!("{:?}", ret);
    assert_eq!(Ok(1), ret);
}
<commit_msg>add comments to address review comments<commit_after>use crate::schema::*;
use diesel::result::Error;
use diesel::*;

#[test]
#[cfg(not(feature = "sqlite"))] // FIXME: This test is only valid when operating on a file and not :memory:
fn transaction_executes_fn_in_a_sql_transaction() {
    const TEST_NAME: &'static str = "transaction_executes_fn_in_a_sql_transaction";
    let conn1 = &mut connection_without_transaction();
    let conn2 = &mut connection_without_transaction();
    setup_test_table(conn1, TEST_NAME);

    fn get_count(conn: &mut TestConnection) -> i64 {
        count_test_table(conn, TEST_NAME)
    }

    conn1
        .transaction::<_, Error, _>(|conn1| {
            assert_eq!(0, get_count(conn1));
            assert_eq!(0, get_count(conn2));
            conn1.execute(&format!("INSERT INTO {} DEFAULT VALUES", TEST_NAME))?;
            assert_eq!(1, get_count(conn1));
            assert_eq!(0, get_count(conn2));
            Ok(())
        })
        .unwrap();

    assert_eq!(1, get_count(conn1));
    assert_eq!(1, get_count(conn2));

    drop_test_table(conn1, TEST_NAME);
}

#[test]
fn transaction_returns_the_returned_value() {
    let conn1 = &mut connection_without_transaction();

    assert_eq!(Ok(1), conn1.transaction::<_, Error, _>(|_| Ok(1)));
}

#[test]
fn transaction_is_rolled_back_when_returned_an_error() {
    let connection = &mut connection_without_transaction();
    let test_name = "transaction_is_rolled_back_when_returned_an_error";
    setup_test_table(connection, test_name);

    let _ = connection.transaction::<(), _, _>(|connection| {
        connection
            .execute(&format!("INSERT INTO {} DEFAULT VALUES", test_name))
            .unwrap();
        Err(Error::RollbackTransaction)
    });
    assert_eq!(0, count_test_table(connection, test_name));

    drop_test_table(connection, test_name);
}

// This test uses a SQLite3 fact to generate a rollback error,
// so that we can verify `Error::RollbackError`. Reference:
// https://www.sqlite.org/lang_transaction.html
//
// The same trick cannot be used for PostgreSQL as it generates
// warning, but not error if a rollback is called twice. Reference:
// https://www.postgresql.org/docs/9.4/sql-rollback.html
//
// The same trick seems to work for MySQL as well based on the
// test result, but I cannot find a document support yet. Hence
// this test is marked for "sqlite" only as this moment. FIXME.
#[test]
#[cfg(feature = "sqlite")]
fn transaction_rollback_returns_error() {
    let connection = &mut connection_without_transaction();
    let test_name = "transaction_rollback_returns_error";
    setup_test_table(connection, test_name);

    // Create a transaction that will fail to rollback.
    let r = connection.transaction::<usize, _, _>(|connection| {
        connection
            .execute(&format!("INSERT INTO {} DEFAULT VALUES", test_name))
            .unwrap();

        // This rollback would succeed, and cause any rollback later to fail.
        connection.execute("ROLLBACK").unwrap();

        // Return any error to trigger a rollback that fails in this case.
        Err(Error::NotFound)
    });

    // Verify that the transaction failed with `RollbackError`.
    assert!(r.is_err());
    assert!(matches!(r.unwrap_err(), Error::RollbackError(_)));

    assert_eq!(0, count_test_table(connection, test_name));
    drop_test_table(connection, test_name);
}

#[test]
fn transactions_can_be_nested() {
    let connection = &mut connection_without_transaction();
    const TEST_NAME: &'static str = "transactions_can_be_nested";
    setup_test_table(connection, TEST_NAME);
    fn get_count(connection: &mut TestConnection) -> i64 {
        count_test_table(connection, TEST_NAME)
    }

    let _ = connection.transaction::<(), _, _>(|connection| {
        connection
            .execute(&format!("INSERT INTO {} DEFAULT VALUES", TEST_NAME))
            .unwrap();
        assert_eq!(1, get_count(connection));
        let _ = connection.transaction::<(), _, _>(|connection| {
            connection
                .execute(&format!("INSERT INTO {} DEFAULT VALUES", TEST_NAME))
                .unwrap();
            assert_eq!(2, get_count(connection));
            Err(Error::RollbackTransaction)
        });
        assert_eq!(1, get_count(connection));
        let _ = connection.transaction::<(), Error, _>(|connection| {
            connection
                .execute(&format!("INSERT INTO {} DEFAULT VALUES", TEST_NAME))
                .unwrap();
            assert_eq!(2, get_count(connection));
            Ok(())
        });
        assert_eq!(2, get_count(connection));
        Err(Error::RollbackTransaction)
    });
    assert_eq!(0, get_count(connection));

    drop_test_table(connection, TEST_NAME);
}

#[test]
fn test_transaction_always_rolls_back() {
    let connection = &mut connection_without_transaction();
    let test_name = "test_transaction_always_rolls_back";
    setup_test_table(connection, test_name);

    let result = connection.test_transaction::<_, Error, _>(|connection| {
        connection.execute(&format!("INSERT INTO {} DEFAULT VALUES", test_name))?;
        assert_eq!(1, count_test_table(connection, test_name));
        Ok("success")
    });
    assert_eq!(0, count_test_table(connection, test_name));
    assert_eq!("success", result);

    drop_test_table(connection, test_name);
}

#[test]
#[should_panic(expected = "Transaction did not succeed")]
fn test_transaction_panics_on_error() {
    let connection = &mut connection_without_transaction();
    connection.test_transaction::<(), _, _>(|_| Err(()));
}

fn setup_test_table(connection: &mut TestConnection, table_name: &str) {
    use crate::schema_dsl::*;
    create_table(table_name, (integer("id").primary_key().auto_increment(),))
        .execute(connection)
        .unwrap();
}

fn drop_test_table(connection: &mut TestConnection, table_name: &str) {
    connection
        .execute(&format!("DROP TABLE {}", table_name))
        .unwrap();
}

fn count_test_table(connection: &mut TestConnection, table_name: &str) -> i64 {
    use diesel::dsl::sql;
    select(sql::<sql_types::BigInt>(&format!(
        "COUNT(*) FROM {}",
        table_name
    )))
    .first(connection)
    .unwrap()
}

#[test]
#[cfg(feature = "postgres")]
fn regression_test_for_2123() {
    let conn = &mut connection_without_transaction();
    // fail once
    let ret = conn.transaction(|conn| {
        let _ = conn.transaction(|conn| {
            // handling error
            match conn.execute("SELECT foo") {
                // do nothing
                Ok(_) => unreachable!("This query should fail"),
                // ignore the error
                Err(e) => eprintln!("error occurred: {}", e),
            };
            Ok::<_, Error>(())
        });

        conn.transaction(|conn| {
            let ret = conn.execute("SELECT 1");
            assert_eq!(Ok(1), ret);
            Ok::<_, Error>(())
        })
    });
    println!("{:?}", ret);
    // other transaction
    let ret = conn
        .build_transaction()
        .serializable()
        .run(|conn| conn.execute("SELECT 1"));
    // must be Ok(1), but get Err(AlreadyInTransaction)
    println!("{:?}", ret);
    assert_eq!(Ok(1), ret);
}
<|endoftext|>


<commit_before>use common::slice::GetSlice;

use alloc::boxed::Box;

use arch::memory::Memory;

use collections::slice;
use collections::string::{String, ToString};
use collections::vec::Vec;

use common::debug;

use core::cmp;

use disk::Disk;
use disk::ide::Extent;

use fs::redoxfs::{FileSystem, Node, NodeData};

use fs::{KScheme, Resource, ResourceSeek, Url, VecResource};

use syscall::{Error, Result, O_CREAT, ENOENT, EIO};

/// A file resource
pub struct FileResource {
    pub scheme: *mut FileScheme,
    pub node: Node,
    pub vec: Vec<u8>,
    pub seek: usize,
    pub dirty: bool,
}

impl Resource for FileResource {
    fn dup(&self) -> Result<Box<Resource>> {
        Ok(box FileResource {
            scheme: self.scheme,
            node: self.node.clone(),
            vec: self.vec.clone(),
            seek: self.seek,
            dirty: self.dirty,
        })
    }

    fn path(&self, buf: &mut [u8]) -> Result<usize> {
        let path_a = b"file:/";
        for (b, p) in buf.iter_mut().zip(path_a.iter()) {
            *b = *p;
        }

        let path_b = self.node.name.as_bytes();
        for (b, p) in buf.iter_mut().skip(path_a.len()).zip(path_b.iter()) {
            *b = *p;
        }

        Ok(cmp::min(buf.len(), path_a.len() + path_b.len()))
    }

    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        let mut i = 0;
        while i < buf.len() && self.seek < self.vec.len() {
            match self.vec.get(self.seek) {
                Some(b) => buf[i] = *b,
                None => (),
            }
            self.seek += 1;
            i += 1;
        }
        Ok(i)
    }

    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        let mut i = 0;
        while i < buf.len() && self.seek < self.vec.len() {
            self.vec[self.seek] = buf[i];
            self.seek += 1;
            i += 1;
        }
        while i < buf.len() {
            self.vec.push(buf[i]);
            self.seek += 1;
            i += 1;
        }
        if i > 0 {
            self.dirty = true;
        }
        Ok(i)
    }

    fn seek(&mut self, pos: ResourceSeek) -> Result<usize> {
        match pos {
            ResourceSeek::Start(offset) => self.seek = offset,
            ResourceSeek::Current(offset) =>
                self.seek = cmp::max(0, self.seek as isize + offset) as usize,
            ResourceSeek::End(offset) =>
                self.seek = cmp::max(0, self.vec.len() as isize + offset) as usize,
        }
        while self.vec.len() < self.seek {
            self.vec.push(0);
        }
        Ok(self.seek)
    }

    // TODO: Check to make sure proper amount of bytes written. See Disk::write
    // TODO: Allow reallocation
    fn sync(&mut self) -> Result<()> {
        if self.dirty {
            let mut node_dirty = false;
            let mut pos = 0;
            let mut remaining = self.vec.len() as isize;
            for ref mut extent in &mut self.node.extents {
                if remaining > 0 && extent.empty() {
                    debug::d("Reallocate file, extra: ");
                    debug::ds(remaining);
                    debug::dl();

                    unsafe {
                        let sectors = ((remaining + 511) / 512) as u64;
                        if (*self.scheme).fs.header.free_space.length >= sectors * 512 {
                            extent.block = (*self.scheme).fs.header.free_space.block;
                            extent.length = remaining as u64;
                            (*self.scheme).fs.header.free_space.block = (*self.scheme)
                                                                            .fs
                                                                            .header
                                                                            .free_space
                                                                            .block +
                                                                        sectors;
                            (*self.scheme).fs.header.free_space.length = (*self.scheme)
                                                                             .fs
                                                                             .header
                                                                             .free_space
                                                                             .length -
                                                                         sectors * 512;

                            node_dirty = true;
                        }
                    }
                }

                // Make sure it is a valid extent
                if !extent.empty() {
                    let current_sectors = (extent.length as usize + 511) / 512;
                    let max_size = current_sectors * 512;

                    let size = cmp::min(remaining as usize, max_size);

                    if size as u64 != extent.length {
                        extent.length = size as u64;
                        node_dirty = true;
                    }

                    while self.vec.len() < pos + max_size {
                        self.vec.push(0);
                    }

                    unsafe {
                        let _ = (*self.scheme).fs.disk.write(extent.block, &self.vec[pos .. pos + max_size]);
                    }

                    self.vec.truncate(pos + size);

                    pos += size;
                    remaining -= size as isize;
                }
            }

            if node_dirty {
                debug::d("Node dirty, rewrite\n");

                if self.node.block > 0 {
                    unsafe {
                        if let Some(mut node_data) = Memory::<NodeData>::new(1) {
                            node_data.write(0, self.node.data());

                            let mut buffer = slice::from_raw_parts(node_data.address() as *mut u8, 512);
                            let _ = (*self.scheme).fs.disk.write(self.node.block, &mut buffer);

                            debug::d("Renode\n");

                            for mut node in (*self.scheme).fs.nodes.iter_mut() {
                                if node.block == self.node.block {
                                    *node = self.node.clone();
                                }
                            }
                        }
                    }
                } else {
                    debug::d("Need to place Node block\n");
                }
            }

            self.dirty = false;

            if remaining > 0 {
                debug::d("Need to defragment file, extra: ");
                debug::ds(remaining);
                debug::dl();
                return Err(Error::new(EIO));
            }
        }
        Ok(())
    }

    fn truncate(&mut self, len: usize) -> Result<()> {
        while len > self.vec.len() {
            self.vec.push(0);
        }
        self.vec.truncate(len);
        self.seek = cmp::min(self.seek, self.vec.len());
        self.dirty = true;
        Ok(())
    }
}

impl Drop for FileResource {
    fn drop(&mut self) {
        let _ = self.sync();
    }
}

/// A file scheme (pci + fs)
pub struct FileScheme {
    fs: FileSystem,
}

impl FileScheme {
    /// Create a new file scheme from an array of Disks
    pub fn new(mut disks: Vec<Box<Disk>>) -> Option<Box<Self>> {
        while ! disks.is_empty() {
            if let Some(fs) = FileSystem::from_disk(disks.remove(0)) {
                return Some(box FileScheme { fs: fs });
            }
        }

        None
    }
}

impl KScheme for FileScheme {
    fn on_irq(&mut self, _irq: u8) {
        /*if irq == self.fs.disk.irq {
        }*/
    }

    fn scheme(&self) -> &str {
        "file"
    }

    fn open(&mut self, url: &Url, flags: usize) -> Result<Box<Resource>> {
        let mut path = url.reference();
        while path.starts_with('/') {
            path = &path[1..];
        }
        if path.is_empty() || path.ends_with('/') {
            let mut list = String::new();
            let mut dirs: Vec<String> = Vec::new();

            for file in self.fs.list(path).iter() {
                let mut line = String::new();
                match file.find('/') {
                    Some(index) => {
                        let dirname = file.get_slice(..index + 1).to_string();
                        let mut found = false;
                        for dir in dirs.iter() {
                            if dirname == *dir {
                                found = true;
                                break;
                            }
                        }
                        if found {
                            line.clear();
                        } else {
                            line = dirname.clone();
                            dirs.push(dirname);
                        }
                    }
                    None => line = file.clone(),
                }
                if !line.is_empty() {
                    if !list.is_empty() {
                        list = list + "\n" + &line;
                    } else {
                        list = line;
                    }
                }
            }

            if list.len() > 0 {
                Ok(box VecResource::new(&url.string, list.into_bytes()))
            } else {
                Err(Error::new(ENOENT))
            }
        } else {
            match self.fs.node(path) {
                Some(node) => {
                    let mut vec: Vec<u8> = Vec::new();
                    for extent in &node.extents {
                        if extent.block > 0 && extent.length > 0 {
                            let current_sectors = (extent.length as usize + 511) / 512;
                            let max_size = current_sectors * 512;

                            let size = cmp::min(extent.length as usize, max_size);

                            let pos = vec.len();

                            while vec.len() < pos + max_size {
                                vec.push(0);
                            }

                            let _ = self.fs.disk.read(extent.block, &mut vec[pos..pos + max_size]);

                            vec.truncate(pos + size);
                        }
                    }

                    Ok(box FileResource {
                        scheme: self,
                        node: node,
                        vec: vec,
                        seek: 0,
                        dirty: false,
                    })
                }
                None => {
                    if flags & O_CREAT == O_CREAT {
                        // TODO: Create file
                        let mut node = Node {
                            block: 0,
                            name: path.to_string(),
                            extents: [Extent {
                                block: 0,
                                length: 0,
                            }; 16],
                        };

                        if self.fs.header.free_space.length >= 512 {
                            node.block = self.fs.header.free_space.block;
                            self.fs.header.free_space.block = self.fs.header.free_space.block + 1;
                            self.fs.header.free_space.length = self.fs.header.free_space.length -
                                                               512;
                        }

                        self.fs.nodes.push(node.clone());

                        Ok(box FileResource {
                            scheme: self,
                            node: node,
                            vec: Vec::new(),
                            seek: 0,
                            dirty: false,
                        })
                    } else {
                        Err(Error::new(ENOENT))
                    }
                }
            }
        }
    }

    fn unlink(&mut self, url: &Url) -> Result<()> {
        let mut ret = Err(Error::new(ENOENT));

        let mut path = url.reference();
        while path.starts_with('/') {
            path = &path[1..];
        }

        let mut i = 0;
        while i < self.fs.nodes.len() {
            let mut remove = false;

            if let Some(node) = self.fs.nodes.get(i) {
                remove = node.name == path;
            }

            if remove {
                self.fs.nodes.remove(i);
                ret = Ok(());
            } else {
                i += 1;
            }
        }

        ret
    }
}
<commit_msg>Chain both `path_a` and `path_b` iterators to combine both loops<commit_after>use common::slice::GetSlice;

use alloc::boxed::Box;

use arch::memory::Memory;

use collections::slice;
use collections::string::{String, ToString};
use collections::vec::Vec;

use common::debug;

use core::cmp;

use disk::Disk;
use disk::ide::Extent;

use fs::redoxfs::{FileSystem, Node, NodeData};

use fs::{KScheme, Resource, ResourceSeek, Url, VecResource};

use syscall::{Error, Result, O_CREAT, ENOENT, EIO};

/// A file resource
pub struct FileResource {
    pub scheme: *mut FileScheme,
    pub node: Node,
    pub vec: Vec<u8>,
    pub seek: usize,
    pub dirty: bool,
}

impl Resource for FileResource {
    fn dup(&self) -> Result<Box<Resource>> {
        Ok(box FileResource {
            scheme: self.scheme,
            node: self.node.clone(),
            vec: self.vec.clone(),
            seek: self.seek,
            dirty: self.dirty,
        })
    }

    fn path(&self, buf: &mut [u8]) -> Result<usize> {
        let path_a = b"file:/";
        let path_b = self.node.name.as_bytes();
        for (b, p) in buf.iter_mut().zip(path_a.iter().chain(path_b.iter())) {
            *b = *p;
        }

        Ok(cmp::min(buf.len(), path_a.len() + path_b.len()))
    }

    fn read(&mut self, buf: &mut [u8]) -> Result<usize> {
        let mut i = 0;
        while i < buf.len() && self.seek < self.vec.len() {
            match self.vec.get(self.seek) {
                Some(b) => buf[i] = *b,
                None => (),
            }
            self.seek += 1;
            i += 1;
        }
        Ok(i)
    }

    fn write(&mut self, buf: &[u8]) -> Result<usize> {
        let mut i = 0;
        while i < buf.len() && self.seek < self.vec.len() {
            self.vec[self.seek] = buf[i];
            self.seek += 1;
            i += 1;
        }
        while i < buf.len() {
            self.vec.push(buf[i]);
            self.seek += 1;
            i += 1;
        }
        if i > 0 {
            self.dirty = true;
        }
        Ok(i)
    }

    fn seek(&mut self, pos: ResourceSeek) -> Result<usize> {
        match pos {
            ResourceSeek::Start(offset) => self.seek = offset,
            ResourceSeek::Current(offset) =>
                self.seek = cmp::max(0, self.seek as isize + offset) as usize,
            ResourceSeek::End(offset) =>
                self.seek = cmp::max(0, self.vec.len() as isize + offset) as usize,
        }
        while self.vec.len() < self.seek {
            self.vec.push(0);
        }
        Ok(self.seek)
    }

    // TODO: Check to make sure proper amount of bytes written. See Disk::write
    // TODO: Allow reallocation
    fn sync(&mut self) -> Result<()> {
        if self.dirty {
            let mut node_dirty = false;
            let mut pos = 0;
            let mut remaining = self.vec.len() as isize;
            for ref mut extent in &mut self.node.extents {
                if remaining > 0 && extent.empty() {
                    debug::d("Reallocate file, extra: ");
                    debug::ds(remaining);
                    debug::dl();

                    unsafe {
                        let sectors = ((remaining + 511) / 512) as u64;
                        if (*self.scheme).fs.header.free_space.length >= sectors * 512 {
                            extent.block = (*self.scheme).fs.header.free_space.block;
                            extent.length = remaining as u64;
                            (*self.scheme).fs.header.free_space.block = (*self.scheme)
                                                                            .fs
                                                                            .header
                                                                            .free_space
                                                                            .block +
                                                                        sectors;
                            (*self.scheme).fs.header.free_space.length = (*self.scheme)
                                                                             .fs
                                                                             .header
                                                                             .free_space
                                                                             .length -
                                                                         sectors * 512;

                            node_dirty = true;
                        }
                    }
                }

                // Make sure it is a valid extent
                if !extent.empty() {
                    let current_sectors = (extent.length as usize + 511) / 512;
                    let max_size = current_sectors * 512;

                    let size = cmp::min(remaining as usize, max_size);

                    if size as u64 != extent.length {
                        extent.length = size as u64;
                        node_dirty = true;
                    }

                    while self.vec.len() < pos + max_size {
                        self.vec.push(0);
                    }

                    unsafe {
                        let _ = (*self.scheme).fs.disk.write(extent.block, &self.vec[pos .. pos + max_size]);
                    }

                    self.vec.truncate(pos + size);

                    pos += size;
                    remaining -= size as isize;
                }
            }

            if node_dirty {
                debug::d("Node dirty, rewrite\n");

                if self.node.block > 0 {
                    unsafe {
                        if let Some(mut node_data) = Memory::<NodeData>::new(1) {
                            node_data.write(0, self.node.data());

                            let mut buffer = slice::from_raw_parts(node_data.address() as *mut u8, 512);
                            let _ = (*self.scheme).fs.disk.write(self.node.block, &mut buffer);

                            debug::d("Renode\n");

                            for mut node in (*self.scheme).fs.nodes.iter_mut() {
                                if node.block == self.node.block {
                                    *node = self.node.clone();
                                }
                            }
                        }
                    }
                } else {
                    debug::d("Need to place Node block\n");
                }
            }

            self.dirty = false;

            if remaining > 0 {
                debug::d("Need to defragment file, extra: ");
                debug::ds(remaining);
                debug::dl();
                return Err(Error::new(EIO));
            }
        }
        Ok(())
    }

    fn truncate(&mut self, len: usize) -> Result<()> {
        while len > self.vec.len() {
            self.vec.push(0);
        }
        self.vec.truncate(len);
        self.seek = cmp::min(self.seek, self.vec.len());
        self.dirty = true;
        Ok(())
    }
}

impl Drop for FileResource {
    fn drop(&mut self) {
        let _ = self.sync();
    }
}

/// A file scheme (pci + fs)
pub struct FileScheme {
    fs: FileSystem,
}

impl FileScheme {
    /// Create a new file scheme from an array of Disks
    pub fn new(mut disks: Vec<Box<Disk>>) -> Option<Box<Self>> {
        while ! disks.is_empty() {
            if let Some(fs) = FileSystem::from_disk(disks.remove(0)) {
                return Some(box FileScheme { fs: fs });
            }
        }

        None
    }
}

impl KScheme for FileScheme {
    fn on_irq(&mut self, _irq: u8) {
        /*if irq == self.fs.disk.irq {
        }*/
    }

    fn scheme(&self) -> &str {
        "file"
    }

    fn open(&mut self, url: &Url, flags: usize) -> Result<Box<Resource>> {
        let mut path = url.reference();
        while path.starts_with('/') {
            path = &path[1..];
        }
        if path.is_empty() || path.ends_with('/') {
            let mut list = String::new();
            let mut dirs: Vec<String> = Vec::new();

            for file in self.fs.list(path).iter() {
                let mut line = String::new();
                match file.find('/') {
                    Some(index) => {
                        let dirname = file.get_slice(..index + 1).to_string();
                        let mut found = false;
                        for dir in dirs.iter() {
                            if dirname == *dir {
                                found = true;
                                break;
                            }
                        }
                        if found {
                            line.clear();
                        } else {
                            line = dirname.clone();
                            dirs.push(dirname);
                        }
                    }
                    None => line = file.clone(),
                }
                if !line.is_empty() {
                    if !list.is_empty() {
                        list = list + "\n" + &line;
                    } else {
                        list = line;
                    }
                }
            }

            if list.len() > 0 {
                Ok(box VecResource::new(&url.string, list.into_bytes()))
            } else {
                Err(Error::new(ENOENT))
            }
        } else {
            match self.fs.node(path) {
                Some(node) => {
                    let mut vec: Vec<u8> = Vec::new();
                    for extent in &node.extents {
                        if extent.block > 0 && extent.length > 0 {
                            let current_sectors = (extent.length as usize + 511) / 512;
                            let max_size = current_sectors * 512;

                            let size = cmp::min(extent.length as usize, max_size);

                            let pos = vec.len();

                            while vec.len() < pos + max_size {
                                vec.push(0);
                            }

                            let _ = self.fs.disk.read(extent.block, &mut vec[pos..pos + max_size]);

                            vec.truncate(pos + size);
                        }
                    }

                    Ok(box FileResource {
                        scheme: self,
                        node: node,
                        vec: vec,
                        seek: 0,
                        dirty: false,
                    })
                }
                None => {
                    if flags & O_CREAT == O_CREAT {
                        // TODO: Create file
                        let mut node = Node {
                            block: 0,
                            name: path.to_string(),
                            extents: [Extent {
                                block: 0,
                                length: 0,
                            }; 16],
                        };

                        if self.fs.header.free_space.length >= 512 {
                            node.block = self.fs.header.free_space.block;
                            self.fs.header.free_space.block = self.fs.header.free_space.block + 1;
                            self.fs.header.free_space.length = self.fs.header.free_space.length -
                                                               512;
                        }

                        self.fs.nodes.push(node.clone());

                        Ok(box FileResource {
                            scheme: self,
                            node: node,
                            vec: Vec::new(),
                            seek: 0,
                            dirty: false,
                        })
                    } else {
                        Err(Error::new(ENOENT))
                    }
                }
            }
        }
    }

    fn unlink(&mut self, url: &Url) -> Result<()> {
        let mut ret = Err(Error::new(ENOENT));

        let mut path = url.reference();
        while path.starts_with('/') {
            path = &path[1..];
        }

        let mut i = 0;
        while i < self.fs.nodes.len() {
            let mut remove = false;

            if let Some(node) = self.fs.nodes.get(i) {
                remove = node.name == path;
            }

            if remove {
                self.fs.nodes.remove(i);
                ret = Ok(());
            } else {
                i += 1;
            }
        }

        ret
    }
}
<|endoftext|>


<commit_before>use graph::Graph;
use adjacency_matrix::AdjacencyMatrix;

use std::usize;

// maximum_flow returns the maximum flow from source to sink in g.
pub fn maximum_flow(g: &Graph, source: usize, sink: usize) -> u32 {
    // r is the residual flow graph of g. Initially, that's just the same as g.
    let mut r = AdjacencyMatrix::from_graph(g);

    let mut flow = 0;
    loop {
        let path_capacity = r.add_augmenting_path(source, sink);
        if path_capacity > 0 {
            flow += path_capacity;
        } else {
            return flow;
        }
    }
}

impl AdjacencyMatrix {
    // addAugmentingPath looks for an augmenting path from source to sink in the residual flow
    // graph.
    // If an augmenting path is found, addAugmentingPath updates the residual flow graph to reflect
    // the addition of the augmenting path, and returns the capacity of the augmenting path.
    // If no augmenting path is found, addAugmentingPath returns 0.
    fn add_augmenting_path(&mut self, source: usize, sink: usize) -> u32 {
        let mut parents = vec![usize::MAX; self.size()];

        self.breadth_first_search(source, |edge| {
            parents[edge.y] = edge.x;
            // Return true to abort the BFS.
            edge.y == sink
        });

        self.substract_path(&parents, sink)
    }

    // substractPath substracts the path described by parents and end from the residual flow graph,
    // and returns the capacity of the path.
    fn substract_path(&mut self, parents: &Vec<usize>, end: usize) -> u32 {
        let capacity = self.path_capacity(parents, end);

        path(parents, end, &mut |x, y| {
            self[x][y] -= capacity;
            self[y][x] += capacity;
        });

        capacity
    }

    // pathCapacity returns the capacity of the path described by parents and end in the residual
    // flow graph.
    fn path_capacity(&self, parents: &Vec<usize>, end: usize) -> u32 {
        let mut capacity = 0;

        path(parents, end, &mut |x, y| {
            let c = self[x][y];
            if capacity == 0 || capacity > c {
                capacity = c
            }
        });

        capacity
    }
}

// path calls callback with each segment of the path described by parents and end.
fn path<C>(parents: &Vec<usize>, end: usize, callback: &mut C) where C: FnMut(usize, usize) {
    let p = parents[end];
    if p < usize::MAX {
        path(parents, p, callback);
        callback(p, end);
    }
}
<commit_msg>[edge_connectivity/rust] Replace callback-style iteration by classic vector iteration<commit_after>use graph::Graph;
use adjacency_matrix::AdjacencyMatrix;

use std::usize;

// maximum_flow returns the maximum flow from source to sink in g.
pub fn maximum_flow(g: &Graph, source: usize, sink: usize) -> u32 {
    // r is the residual flow graph of g. Initially, that's just the same as g.
    let mut r = AdjacencyMatrix::from_graph(g);

    let mut flow = 0;
    loop {
        let path_capacity = r.add_augmenting_path(source, sink);
        if path_capacity > 0 {
            flow += path_capacity;
        } else {
            return flow;
        }
    }
}

impl AdjacencyMatrix {
    // addAugmentingPath looks for an augmenting path from source to sink in the residual flow
    // graph.
    // If an augmenting path is found, addAugmentingPath updates the residual flow graph to reflect
    // the addition of the augmenting path, and returns the capacity of the augmenting path.
    // If no augmenting path is found, addAugmentingPath returns 0.
    fn add_augmenting_path(&mut self, source: usize, sink: usize) -> u32 {
        let mut parents = vec![usize::MAX; self.size()];

        self.breadth_first_search(source, |edge| {
            parents[edge.y] = edge.x;
            // Return true to abort the BFS.
            edge.y == sink
        });

        self.substract_path(&parents, sink)
    }

    // substractPath substracts the path described by parents and end from the residual flow graph,
    // and returns the capacity of the path.
    fn substract_path(&mut self, parents: &Vec<usize>, end: usize) -> u32 {
        let capacity = self.path_capacity(parents, end);

        for (x, y) in edges(parents, end) {
            self[x][y] -= capacity;
            self[y][x] += capacity;
        }

        capacity
    }

    // pathCapacity returns the capacity of the path described by parents and end in the residual
    // flow graph.
    fn path_capacity(&self, parents: &Vec<usize>, end: usize) -> u32 {
        let mut capacity = 0;

        for (x, y) in edges(parents, end) {
            let c = self[x][y];
            if capacity == 0 || capacity > c {
                capacity = c
            }
        }

        capacity
    }
}

// edges returns the edges of the path described by parents and end.
fn edges(parents: &Vec<usize>, mut end: usize) -> Vec<(usize, usize)> {
    let mut edges = vec![];

    loop {
        let p = parents[end];
        if p == usize::MAX {
            edges.reverse();
            return edges
        }
        edges.push((p, end));
        end = p
    }
}
<|endoftext|>


<commit_before>use crate::*;
use bls::create_proof_of_possession;

/// Builds a `BeaconState` for use in testing or benchmarking.
///
/// Building the `BeaconState` is a three step processes:
///
/// 1. Create a new `BeaconStateBuilder`.
/// 2. Call `Self::build()` or `Self::build_fast()` generate a  `BeaconState`.
/// 3. (Optional) Use builder functions to modify the `BeaconState`.
/// 4. Call `Self::cloned_state()` to obtain a `BeaconState` cloned from this struct.
///
/// Step (2) happens prior to step (3) because some functionality requires an existing
/// `BeaconState`.
///
/// Step (4) produces a clone of the BeaconState and doesn't consume the `BeaconStateBuilder` to
/// allow access to `self.keypairs` and `self.spec`.
pub struct BeaconStateBuilder {
    pub validator_count: usize,
    pub state: Option<BeaconState>,
    pub genesis_time: u64,
    pub latest_eth1_data: Eth1Data,
    pub spec: ChainSpec,
    pub keypairs: Vec<Keypair>,
}

impl BeaconStateBuilder {
    /// Create a new builder with the given number of validators.
    pub fn new(validator_count: usize) -> Self {
        let genesis_time = 10_000_000;

        let latest_eth1_data = Eth1Data {
            deposit_root: Hash256::zero(),
            block_hash: Hash256::zero(),
        };

        let spec = ChainSpec::foundation();

        Self {
            validator_count,
            state: None,
            genesis_time,
            latest_eth1_data,
            spec,
            keypairs: vec![],
        }
    }

    pub fn build(&mut self) -> Result<(), BeaconStateError> {
        self.keypairs = (0..self.validator_count)
            .collect::<Vec<usize>>()
            .iter()
            .map(|_| Keypair::random())
            .collect();

        let initial_validator_deposits = self.keypairs
            .iter()
            .map(|keypair| Deposit {
                branch: vec![], // branch verification is not specified.
                index: 0,       // index verification is not specified.
                deposit_data: DepositData {
                    amount: 32_000_000_000, // 32 ETH (in Gwei)
                    timestamp: self.genesis_time - 1,
                    deposit_input: DepositInput {
                        pubkey: keypair.pk.clone(),
                        withdrawal_credentials: Hash256::zero(), // Withdrawal not possible.
                        proof_of_possession: create_proof_of_possession(&keypair),
                    },
                },
            })
            .collect();

        let state = BeaconState::genesis(
            self.genesis_time,
            initial_validator_deposits,
            self.latest_eth1_data.clone(),
            &self.spec,
        )?;

        self.state = Some(state);

        Ok(())
    }

    pub fn build_fast(&mut self) -> Result<(), BeaconStateError> {
        let common_keypair = Keypair::random();

        let mut validator_registry = Vec::with_capacity(self.validator_count);
        let mut validator_balances = Vec::with_capacity(self.validator_count);
        self.keypairs = Vec::with_capacity(self.validator_count);

        for _ in 0..self.validator_count {
            self.keypairs.push(common_keypair.clone());
            validator_balances.push(32_000_000_000);
            validator_registry.push(Validator {
                pubkey: common_keypair.pk.clone(),
                withdrawal_credentials: Hash256::zero(),
                activation_epoch: self.spec.genesis_epoch,
                ..Validator::default()
            })
        }

        let state = BeaconState {
            validator_registry,
            validator_balances,
            ..BeaconState::genesis(
                self.genesis_time,
                vec![],
                self.latest_eth1_data.clone(),
                &self.spec,
            )?
        };

        self.state = Some(state);

        Ok(())
    }

    /// Sets the `BeaconState` to be in the last slot of the given epoch.
    ///
    /// Sets all justification/finalization parameters to be be as "perfect" as possible (i.e.,
    /// highest justified and finalized slots, full justification bitfield, etc).
    pub fn teleport_to_end_of_epoch(&mut self, epoch: Epoch) {
        let state = self.state.as_mut().expect("Genesis required");

        let slot = epoch.end_slot(self.spec.epoch_length);

        state.slot = slot;
        state.validator_registry_update_epoch = epoch - 1;

        state.previous_calculation_epoch = epoch - 1;
        state.current_calculation_epoch = epoch;

        state.previous_epoch_seed = Hash256::from(&b"previous_seed"[..]);
        state.current_epoch_seed = Hash256::from(&b"current_seed"[..]);

        state.previous_justified_epoch = epoch - 2;
        state.justified_epoch = epoch - 1;
        state.justification_bitfield = u64::max_value();
        state.finalized_epoch = epoch - 1;
    }

    /// Creates a full set of attestations for the `BeaconState`. Each attestation has full
    /// participation from its committee and references the expected beacon_block hashes.
    ///
    /// These attestations should be fully conducive to justification and finalization.
    pub fn insert_attestations(&mut self) {
        let state = self.state.as_mut().expect("Genesis required");

        state
            .build_epoch_cache(RelativeEpoch::Previous, &self.spec)
            .unwrap();
        state
            .build_epoch_cache(RelativeEpoch::Current, &self.spec)
            .unwrap();

        let current_epoch = state.current_epoch(&self.spec);
        let previous_epoch = state.previous_epoch(&self.spec);
        let current_epoch_depth =
            (state.slot - current_epoch.end_slot(self.spec.epoch_length)).as_usize();

        let previous_epoch_slots = previous_epoch.slot_iter(self.spec.epoch_length);
        let current_epoch_slots = current_epoch
            .slot_iter(self.spec.epoch_length)
            .take(current_epoch_depth);

        for slot in previous_epoch_slots.chain(current_epoch_slots) {
            let committees = state
                .get_crosslink_committees_at_slot(slot, &self.spec)
                .unwrap()
                .clone();

            for (committee, shard) in committees {
                state
                    .latest_attestations
                    .push(committee_to_pending_attestation(
                        state, &committee, shard, slot, &self.spec,
                    ))
            }
        }
    }

    /// Returns a cloned `BeaconState`.
    pub fn cloned_state(&self) -> BeaconState {
        self.state.as_ref().expect("Genesis required").clone()
    }
}

/// Builds a valid PendingAttestation with full participation for some committee.
fn committee_to_pending_attestation(
    state: &BeaconState,
    committee: &[usize],
    shard: u64,
    slot: Slot,
    spec: &ChainSpec,
) -> PendingAttestation {
    let current_epoch = state.current_epoch(spec);
    let previous_epoch = state.previous_epoch(spec);

    let mut aggregation_bitfield = Bitfield::new();
    let mut custody_bitfield = Bitfield::new();

    for (i, _) in committee.iter().enumerate() {
        aggregation_bitfield.set(i, true);
        custody_bitfield.set(i, true);
    }

    let is_previous_epoch = state.slot.epoch(spec.epoch_length) != slot.epoch(spec.epoch_length);

    let justified_epoch = if is_previous_epoch {
        state.previous_justified_epoch
    } else {
        state.justified_epoch
    };

    let epoch_boundary_root = if is_previous_epoch {
        *state
            .get_block_root(previous_epoch.start_slot(spec.epoch_length), spec)
            .unwrap()
    } else {
        *state
            .get_block_root(current_epoch.start_slot(spec.epoch_length), spec)
            .unwrap()
    };

    let justified_block_root = *state
        .get_block_root(justified_epoch.start_slot(spec.epoch_length), &spec)
        .unwrap();

    PendingAttestation {
        aggregation_bitfield,
        data: AttestationData {
            slot,
            shard,
            beacon_block_root: *state.get_block_root(slot, spec).unwrap(),
            epoch_boundary_root,
            shard_block_root: Hash256::zero(),
            latest_crosslink: Crosslink {
                epoch: slot.epoch(spec.epoch_length),
                shard_block_root: Hash256::zero(),
            },
            justified_epoch,
            justified_block_root,
        },
        custody_bitfield,
        inclusion_slot: slot,
    }
}
<commit_msg>Add docstrings to `BeaconStateBuilder`<commit_after>use crate::*;
use bls::create_proof_of_possession;

/// Builds a `BeaconState` for use in testing or benchmarking.
///
/// Building the `BeaconState` is a three step processes:
///
/// 1. Create a new `BeaconStateBuilder`.
/// 2. Call `Self::build()` or `Self::build_fast()` generate a  `BeaconState`.
/// 3. (Optional) Use builder functions to modify the `BeaconState`.
/// 4. Call `Self::cloned_state()` to obtain a `BeaconState` cloned from this struct.
///
/// Step (2) happens prior to step (3) because some functionality requires an existing
/// `BeaconState`.
///
/// Step (4) produces a clone of the BeaconState and doesn't consume the `BeaconStateBuilder` to
/// allow access to `self.keypairs` and `self.spec`.
pub struct BeaconStateBuilder {
    pub validator_count: usize,
    pub state: Option<BeaconState>,
    pub genesis_time: u64,
    pub latest_eth1_data: Eth1Data,
    pub spec: ChainSpec,
    pub keypairs: Vec<Keypair>,
}

impl BeaconStateBuilder {
    /// Create a new builder with the given number of validators.
    pub fn new(validator_count: usize) -> Self {
        let genesis_time = 10_000_000;

        let latest_eth1_data = Eth1Data {
            deposit_root: Hash256::zero(),
            block_hash: Hash256::zero(),
        };

        let spec = ChainSpec::foundation();

        Self {
            validator_count,
            state: None,
            genesis_time,
            latest_eth1_data,
            spec,
            keypairs: vec![],
        }
    }

    /// Builds a `BeaconState` using the `BeaconState::genesis(..)` function.
    ///
    /// Each validator is assigned a unique, randomly-generated keypair and all
    /// proof-of-possessions are verified during genesis.
    pub fn build(&mut self) -> Result<(), BeaconStateError> {
        self.keypairs = (0..self.validator_count)
            .collect::<Vec<usize>>()
            .iter()
            .map(|_| Keypair::random())
            .collect();

        let initial_validator_deposits = self
            .keypairs
            .iter()
            .map(|keypair| Deposit {
                branch: vec![], // branch verification is not specified.
                index: 0,       // index verification is not specified.
                deposit_data: DepositData {
                    amount: 32_000_000_000, // 32 ETH (in Gwei)
                    timestamp: self.genesis_time - 1,
                    deposit_input: DepositInput {
                        pubkey: keypair.pk.clone(),
                        withdrawal_credentials: Hash256::zero(), // Withdrawal not possible.
                        proof_of_possession: create_proof_of_possession(&keypair),
                    },
                },
            })
            .collect();

        let state = BeaconState::genesis(
            self.genesis_time,
            initial_validator_deposits,
            self.latest_eth1_data.clone(),
            &self.spec,
        )?;

        self.state = Some(state);

        Ok(())
    }

    /// Builds a `BeaconState` using the `BeaconState::genesis(..)` function, without supplying any
    /// validators. Instead validators are added to the state post-genesis.
    ///
    /// One keypair is randomly generated and all validators are assigned this same keypair.
    /// Proof-of-possessions are not created (or validated).
    ///
    /// This function runs orders of magnitude faster than `Self::build()`, however it will be
    /// erroneous for functions which use a validators public key as an identifier (e.g.,
    /// deposits).
    /// proof-of-possessions are verified during genesis.
    pub fn build_fast(&mut self) -> Result<(), BeaconStateError> {
        let common_keypair = Keypair::random();

        let mut validator_registry = Vec::with_capacity(self.validator_count);
        let mut validator_balances = Vec::with_capacity(self.validator_count);
        self.keypairs = Vec::with_capacity(self.validator_count);

        for _ in 0..self.validator_count {
            self.keypairs.push(common_keypair.clone());
            validator_balances.push(32_000_000_000);
            validator_registry.push(Validator {
                pubkey: common_keypair.pk.clone(),
                withdrawal_credentials: Hash256::zero(),
                activation_epoch: self.spec.genesis_epoch,
                ..Validator::default()
            })
        }

        let state = BeaconState {
            validator_registry,
            validator_balances,
            ..BeaconState::genesis(
                self.genesis_time,
                vec![],
                self.latest_eth1_data.clone(),
                &self.spec,
            )?
        };

        self.state = Some(state);

        Ok(())
    }

    /// Sets the `BeaconState` to be in the last slot of the given epoch.
    ///
    /// Sets all justification/finalization parameters to be be as "perfect" as possible (i.e.,
    /// highest justified and finalized slots, full justification bitfield, etc).
    pub fn teleport_to_end_of_epoch(&mut self, epoch: Epoch) {
        let state = self.state.as_mut().expect("Genesis required");

        let slot = epoch.end_slot(self.spec.epoch_length);

        state.slot = slot;
        state.validator_registry_update_epoch = epoch - 1;

        state.previous_calculation_epoch = epoch - 1;
        state.current_calculation_epoch = epoch;

        state.previous_epoch_seed = Hash256::from(&b"previous_seed"[..]);
        state.current_epoch_seed = Hash256::from(&b"current_seed"[..]);

        state.previous_justified_epoch = epoch - 2;
        state.justified_epoch = epoch - 1;
        state.justification_bitfield = u64::max_value();
        state.finalized_epoch = epoch - 1;
    }

    /// Creates a full set of attestations for the `BeaconState`. Each attestation has full
    /// participation from its committee and references the expected beacon_block hashes.
    ///
    /// These attestations should be fully conducive to justification and finalization.
    pub fn insert_attestations(&mut self) {
        let state = self.state.as_mut().expect("Genesis required");

        state
            .build_epoch_cache(RelativeEpoch::Previous, &self.spec)
            .unwrap();
        state
            .build_epoch_cache(RelativeEpoch::Current, &self.spec)
            .unwrap();

        let current_epoch = state.current_epoch(&self.spec);
        let previous_epoch = state.previous_epoch(&self.spec);
        let current_epoch_depth =
            (state.slot - current_epoch.end_slot(self.spec.epoch_length)).as_usize();

        let previous_epoch_slots = previous_epoch.slot_iter(self.spec.epoch_length);
        let current_epoch_slots = current_epoch
            .slot_iter(self.spec.epoch_length)
            .take(current_epoch_depth);

        for slot in previous_epoch_slots.chain(current_epoch_slots) {
            let committees = state
                .get_crosslink_committees_at_slot(slot, &self.spec)
                .unwrap()
                .clone();

            for (committee, shard) in committees {
                state
                    .latest_attestations
                    .push(committee_to_pending_attestation(
                        state, &committee, shard, slot, &self.spec,
                    ))
            }
        }
    }

    /// Returns a cloned `BeaconState`.
    pub fn cloned_state(&self) -> BeaconState {
        self.state.as_ref().expect("Genesis required").clone()
    }
}

/// Builds a valid PendingAttestation with full participation for some committee.
fn committee_to_pending_attestation(
    state: &BeaconState,
    committee: &[usize],
    shard: u64,
    slot: Slot,
    spec: &ChainSpec,
) -> PendingAttestation {
    let current_epoch = state.current_epoch(spec);
    let previous_epoch = state.previous_epoch(spec);

    let mut aggregation_bitfield = Bitfield::new();
    let mut custody_bitfield = Bitfield::new();

    for (i, _) in committee.iter().enumerate() {
        aggregation_bitfield.set(i, true);
        custody_bitfield.set(i, true);
    }

    let is_previous_epoch = state.slot.epoch(spec.epoch_length) != slot.epoch(spec.epoch_length);

    let justified_epoch = if is_previous_epoch {
        state.previous_justified_epoch
    } else {
        state.justified_epoch
    };

    let epoch_boundary_root = if is_previous_epoch {
        *state
            .get_block_root(previous_epoch.start_slot(spec.epoch_length), spec)
            .unwrap()
    } else {
        *state
            .get_block_root(current_epoch.start_slot(spec.epoch_length), spec)
            .unwrap()
    };

    let justified_block_root = *state
        .get_block_root(justified_epoch.start_slot(spec.epoch_length), &spec)
        .unwrap();

    PendingAttestation {
        aggregation_bitfield,
        data: AttestationData {
            slot,
            shard,
            beacon_block_root: *state.get_block_root(slot, spec).unwrap(),
            epoch_boundary_root,
            shard_block_root: Hash256::zero(),
            latest_crosslink: Crosslink {
                epoch: slot.epoch(spec.epoch_length),
                shard_block_root: Hash256::zero(),
            },
            justified_epoch,
            justified_block_root,
        },
        custody_bitfield,
        inclusion_slot: slot,
    }
}
<|endoftext|>


<commit_before>mod support;

use glutin::event::{
    ElementState, Event, KeyboardInput, VirtualKeyCode, WindowEvent,
};
use glutin::event_loop::{ControlFlow, EventLoop};
use glutin::monitor::{MonitorHandle, VideoMode};
use glutin::window::{Fullscreen, WindowBuilder};
use std::io::{stdin, stdout, Write};

fn main() {
    let el = EventLoop::new();

    print!(
        "Please choose the fullscreen mode: (1) exclusive, (2) borderless: "
    );
    stdout().flush().unwrap();

    let mut num = String::new();
    stdin().read_line(&mut num).unwrap();
    let num = num.trim().parse().ok().expect("Please enter a number");

    let fullscreen = Some(match num {
        1 => Fullscreen::Exclusive(prompt_for_video_mode(&prompt_for_monitor(
            &el,
        ))),
        2 => Fullscreen::Borderless(prompt_for_monitor(&el)),
        _ => panic!("Please enter a valid number"),
    });

    let mut is_maximized = false;
    let mut decorations = true;

    let wb = WindowBuilder::new()
        .with_title("Hello world!")
        .with_fullscreen(fullscreen.clone());
    let windowed_context = glutin::ContextBuilder::new()
        .build_windowed(wb, &el)
        .unwrap();

    let windowed_context = unsafe { windowed_context.make_current().unwrap() };

    let gl = support::load(&windowed_context.context());

    el.run(move |event, _, control_flow| {
        *control_flow = ControlFlow::Wait;

        match event {
            Event::WindowEvent { event, .. } => match event {
                WindowEvent::CloseRequested => {
                    *control_flow = ControlFlow::Exit
                }
                WindowEvent::KeyboardInput {
                    input:
                        KeyboardInput {
                            virtual_keycode: Some(virtual_code),
                            state,
                            ..
                        },
                    ..
                } => match (virtual_code, state) {
                    (VirtualKeyCode::Escape, _) => {
                        *control_flow = ControlFlow::Exit
                    }
                    (VirtualKeyCode::F, ElementState::Pressed) => {
                        if windowed_context.window().fullscreen().is_some() {
                            windowed_context.window().set_fullscreen(None);
                        } else {
                            windowed_context
                                .window()
                                .set_fullscreen(fullscreen.clone());
                        }
                    }
                    (VirtualKeyCode::S, ElementState::Pressed) => {
                        println!(
                            "window.fullscreen {:?}",
                            windowed_context.window().fullscreen()
                        );
                    }
                    (VirtualKeyCode::M, ElementState::Pressed) => {
                        is_maximized = !is_maximized;
                        windowed_context.window().set_maximized(is_maximized);
                    }
                    (VirtualKeyCode::D, ElementState::Pressed) => {
                        decorations = !decorations;
                        windowed_context.window().set_decorations(decorations);
                    }
                    _ => (),
                },
                _ => (),
            },
            Event::RedrawRequested(_) => {
                gl.draw_frame([1.0, 0.5, 0.7, 1.0]);
                windowed_context.swap_buffers().unwrap();
            }
            _ => {}
        }
    });
}

// Enumerate monitors and prompt user to choose one
fn prompt_for_monitor(el: &EventLoop<()>) -> MonitorHandle {
    for (num, monitor) in el.available_monitors().enumerate() {
        println!("Monitor #{}: {:?}", num, monitor.name());
    }

    print!("Please write the number of the monitor to use: ");
    stdout().flush().unwrap();

    let mut num = String::new();
    stdin().read_line(&mut num).unwrap();
    let num = num.trim().parse().ok().expect("Please enter a number");
    let monitor = el
        .available_monitors()
        .nth(num)
        .expect("Please enter a valid ID");

    println!("Using {:?}", monitor.name());

    monitor
}

fn prompt_for_video_mode(monitor: &MonitorHandle) -> VideoMode {
    for (i, video_mode) in monitor.video_modes().enumerate() {
        println!("Video mode #{}: {}", i, video_mode);
    }

    print!("Please write the number of the video mode to use: ");
    stdout().flush().unwrap();

    let mut num = String::new();
    stdin().read_line(&mut num).unwrap();
    let num = num.trim().parse().ok().expect("Please enter a number");
    let video_mode = monitor
        .video_modes()
        .nth(num)
        .expect("Please enter a valid ID");

    println!("Using {}", video_mode);

    video_mode
}
<commit_msg>Support resize for fullscreen example<commit_after>mod support;

use glutin::event::{
    ElementState, Event, KeyboardInput, VirtualKeyCode, WindowEvent,
};
use glutin::event_loop::{ControlFlow, EventLoop};
use glutin::monitor::{MonitorHandle, VideoMode};
use glutin::window::{Fullscreen, WindowBuilder};
use std::io::{stdin, stdout, Write};

fn main() {
    let el = EventLoop::new();

    print!(
        "Please choose the fullscreen mode: (1) exclusive, (2) borderless: "
    );
    stdout().flush().unwrap();

    let mut num = String::new();
    stdin().read_line(&mut num).unwrap();
    let num = num.trim().parse().ok().expect("Please enter a number");

    let fullscreen = Some(match num {
        1 => Fullscreen::Exclusive(prompt_for_video_mode(&prompt_for_monitor(
            &el,
        ))),
        2 => Fullscreen::Borderless(prompt_for_monitor(&el)),
        _ => panic!("Please enter a valid number"),
    });

    println!("Press (F) to toggle fullscreen, (D) to toggle window decorations, and (M) to toggle maximized/minimized.");

    let mut is_maximized = false;
    let mut decorations = true;

    let wb = WindowBuilder::new()
        .with_title("Hello world!")
        .with_fullscreen(fullscreen.clone());
    let windowed_context = glutin::ContextBuilder::new()
        .build_windowed(wb, &el)
        .unwrap();

    let windowed_context = unsafe { windowed_context.make_current().unwrap() };

    let gl = support::load(&windowed_context.context());

    el.run(move |event, _, control_flow| {
        *control_flow = ControlFlow::Wait;

        match event {
            Event::WindowEvent { event, .. } => match event {
                WindowEvent::CloseRequested => {
                    *control_flow = ControlFlow::Exit
                }
                WindowEvent::Resized(physical_size) => {
                    windowed_context.resize(physical_size);
                }
                WindowEvent::KeyboardInput {
                    input:
                        KeyboardInput {
                            virtual_keycode: Some(virtual_code),
                            state,
                            ..
                        },
                    ..
                } => match (virtual_code, state) {
                    (VirtualKeyCode::Escape, _) => {
                        *control_flow = ControlFlow::Exit
                    }
                    (VirtualKeyCode::F, ElementState::Pressed) => {
                        if windowed_context.window().fullscreen().is_some() {
                            windowed_context.window().set_fullscreen(None);
                        } else {
                            windowed_context
                                .window()
                                .set_fullscreen(fullscreen.clone());
                        }
                    }
                    (VirtualKeyCode::S, ElementState::Pressed) => {
                        println!(
                            "window.fullscreen {:?}",
                            windowed_context.window().fullscreen()
                        );
                    }
                    (VirtualKeyCode::M, ElementState::Pressed) => {
                        is_maximized = !is_maximized;
                        windowed_context.window().set_maximized(is_maximized);
                    }
                    (VirtualKeyCode::D, ElementState::Pressed) => {
                        decorations = !decorations;
                        windowed_context.window().set_decorations(decorations);
                    }
                    _ => (),
                },
                _ => (),
            },
            Event::RedrawRequested(_) => {
                gl.draw_frame([1.0, 0.5, 0.7, 1.0]);
                windowed_context.swap_buffers().unwrap();
            }
            _ => {}
        }
    });
}

// Enumerate monitors and prompt user to choose one
fn prompt_for_monitor(el: &EventLoop<()>) -> MonitorHandle {
    for (num, monitor) in el.available_monitors().enumerate() {
        println!("Monitor #{}: {:?}", num, monitor.name());
    }

    print!("Please write the number of the monitor to use: ");
    stdout().flush().unwrap();

    let mut num = String::new();
    stdin().read_line(&mut num).unwrap();
    let num = num.trim().parse().ok().expect("Please enter a number");
    let monitor = el
        .available_monitors()
        .nth(num)
        .expect("Please enter a valid ID");

    println!("Using {:?}", monitor.name());

    monitor
}

fn prompt_for_video_mode(monitor: &MonitorHandle) -> VideoMode {
    for (i, video_mode) in monitor.video_modes().enumerate() {
        println!("Video mode #{}: {}", i, video_mode);
    }

    print!("Please write the number of the video mode to use: ");
    stdout().flush().unwrap();

    let mut num = String::new();
    stdin().read_line(&mut num).unwrap();
    let num = num.trim().parse().ok().expect("Please enter a number");
    let video_mode = monitor
        .video_modes()
        .nth(num)
        .expect("Please enter a valid ID");

    println!("Using {}", video_mode);

    video_mode
}
<|endoftext|>


<commit_before>// Copyright 2020 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Data structures for commands of virtio video devices.

use remain::sorted;
use std::convert::{TryFrom, TryInto};
use std::io;
use thiserror::Error as ThisError;

use base::error;
use data_model::Le32;
use enumn::N;

use crate::virtio::video::control::*;
use crate::virtio::video::format::*;
use crate::virtio::video::params::Params;
use crate::virtio::video::protocol::*;
use crate::virtio::video::resource::{ResourceType, UnresolvedGuestResource};
use crate::virtio::Reader;

/// An error indicating a failure while reading a request from the guest.
#[sorted]
#[derive(Debug, ThisError)]
pub enum ReadCmdError {
    /// Invalid argument is passed.
    #[error("invalid argument passed to command")]
    InvalidArgument,
    /// The type of the command was invalid.
    #[error("invalid command type: {0}")]
    InvalidCmdType(u32),
    /// Failed to read an object.
    #[error("failed to read object: {0}")]
    IoError(#[from] io::Error),
    /// The type of the requested control was unsupported.
    #[error("unsupported control type: {0}")]
    UnsupportedCtrlType(u32),
}

#[derive(PartialEq, Eq, PartialOrd, Ord, N, Clone, Copy, Debug)]
#[repr(u32)]
pub enum QueueType {
    Input = VIRTIO_VIDEO_QUEUE_TYPE_INPUT,
    Output = VIRTIO_VIDEO_QUEUE_TYPE_OUTPUT,
}
impl_try_from_le32_for_enumn!(QueueType, "queue_type");

#[derive(Debug)]
pub enum VideoCmd {
    QueryCapability {
        queue_type: QueueType,
    },
    StreamCreate {
        stream_id: u32,
        coded_format: Format,
        input_resource_type: ResourceType,
        output_resource_type: ResourceType,
    },
    StreamDestroy {
        stream_id: u32,
    },
    StreamDrain {
        stream_id: u32,
    },
    ResourceCreate {
        stream_id: u32,
        queue_type: QueueType,
        resource_id: u32,
        plane_offsets: Vec<u32>,
        resource: UnresolvedGuestResource,
    },
    ResourceQueue {
        stream_id: u32,
        queue_type: QueueType,
        resource_id: u32,
        timestamp: u64,
        data_sizes: Vec<u32>,
    },
    ResourceDestroyAll {
        stream_id: u32,
        queue_type: QueueType,
    },
    QueueClear {
        stream_id: u32,
        queue_type: QueueType,
    },
    GetParams {
        stream_id: u32,
        queue_type: QueueType,
    },
    SetParams {
        stream_id: u32,
        queue_type: QueueType,
        params: Params,
    },
    QueryControl {
        query_ctrl_type: QueryCtrlType,
    },
    GetControl {
        stream_id: u32,
        ctrl_type: CtrlType,
    },
    SetControl {
        stream_id: u32,
        ctrl_val: CtrlVal,
    },
}

impl<'a> VideoCmd {
    /// Reads a request on virtqueue and construct a VideoCmd value.
    pub fn from_reader(r: &'a mut Reader) -> Result<Self, ReadCmdError> {
        use self::ReadCmdError::*;
        use self::VideoCmd::*;

        // Unlike structs in virtio_video.h in the kernel, our command structs in protocol.rs don't
        // have a field of `struct virtio_video_cmd_hdr`. So, we first read the header here and
        // a body below.
        let hdr = r.read_obj::<virtio_video_cmd_hdr>()?;

        Ok(match hdr.type_.into() {
            VIRTIO_VIDEO_CMD_QUERY_CAPABILITY => {
                let virtio_video_query_capability { queue_type, .. } = r.read_obj()?;
                QueryCapability {
                    queue_type: queue_type.try_into()?,
                }
            }
            VIRTIO_VIDEO_CMD_STREAM_CREATE => {
                let virtio_video_stream_create {
                    in_mem_type,
                    out_mem_type,
                    coded_format,
                    ..
                } = r.read_obj()?;

                let input_resource_type = match in_mem_type.into() {
                    VIRTIO_VIDEO_MEM_TYPE_VIRTIO_OBJECT => ResourceType::Object,
                    _ => {
                        error!("mem_type must be VIRTIO_OBJECT");
                        return Err(InvalidArgument);
                    }
                };

                let output_resource_type = match out_mem_type.into() {
                    VIRTIO_VIDEO_MEM_TYPE_VIRTIO_OBJECT => ResourceType::Object,
                    _ => {
                        error!("mem_type must be VIRTIO_OBJECT");
                        return Err(InvalidArgument);
                    }
                };

                StreamCreate {
                    stream_id: hdr.stream_id.into(),
                    coded_format: coded_format.try_into()?,
                    input_resource_type,
                    output_resource_type,
                }
            }
            VIRTIO_VIDEO_CMD_STREAM_DESTROY => {
                let virtio_video_stream_destroy { .. } = r.read_obj()?;
                StreamDestroy {
                    stream_id: hdr.stream_id.into(),
                }
            }
            VIRTIO_VIDEO_CMD_STREAM_DRAIN => {
                let virtio_video_stream_drain { .. } = r.read_obj()?;
                StreamDrain {
                    stream_id: hdr.stream_id.into(),
                }
            }
            VIRTIO_VIDEO_CMD_RESOURCE_CREATE => {
                let virtio_video_resource_create {
                    queue_type,
                    resource_id,
                    planes_layout,
                    num_planes,
                    plane_offsets,
                    num_entries,
                    ..
                } = r.read_obj()?;

                // Assume ChromeOS-specific requirements.
                if Into::<u32>::into(planes_layout) != VIRTIO_VIDEO_PLANES_LAYOUT_SINGLE_BUFFER {
                    error!(
                        "each buffer must be a single DMAbuf: {}",
                        Into::<u32>::into(planes_layout),
                    );
                    return Err(InvalidArgument);
                }

                let num_planes: u32 = num_planes.into();
                if num_planes as usize > plane_offsets.len() {
                    error!(
                        "num_planes must not exceed {} but {}",
                        plane_offsets.len(),
                        num_planes
                    );
                    return Err(InvalidArgument);
                }
                // We only support single-entry objects for now.
                if num_planes != 1 || num_entries[0] != 1 {
                    error!("the number of entries must be equal to 1");
                    return Err(InvalidArgument);
                }

                let plane_offsets = plane_offsets[0..num_planes as usize]
                    .iter()
                    .map(|x| Into::<u32>::into(*x))
                    .collect::<Vec<u32>>();

                let resource = r.read_obj()?;

                ResourceCreate {
                    stream_id: hdr.stream_id.into(),
                    queue_type: queue_type.try_into()?,
                    resource_id: resource_id.into(),
                    plane_offsets,
                    resource,
                }
            }
            VIRTIO_VIDEO_CMD_RESOURCE_QUEUE => {
                let virtio_video_resource_queue {
                    queue_type,
                    resource_id,
                    timestamp,
                    num_data_sizes,
                    data_sizes,
                    ..
                } = r.read_obj()?;

                let num_data_sizes: u32 = num_data_sizes.into();
                if num_data_sizes as usize > data_sizes.len() {
                    return Err(InvalidArgument);
                }
                let data_sizes = data_sizes[0..num_data_sizes as usize]
                    .iter()
                    .map(|x| Into::<u32>::into(*x))
                    .collect::<Vec<u32>>();
                ResourceQueue {
                    stream_id: hdr.stream_id.into(),
                    queue_type: queue_type.try_into()?,
                    resource_id: resource_id.into(),
                    timestamp: timestamp.into(),
                    data_sizes,
                }
            }
            VIRTIO_VIDEO_CMD_RESOURCE_DESTROY_ALL => {
                let virtio_video_resource_destroy_all { queue_type, .. } = r.read_obj()?;
                ResourceDestroyAll {
                    stream_id: hdr.stream_id.into(),
                    queue_type: queue_type.try_into()?,
                }
            }
            VIRTIO_VIDEO_CMD_QUEUE_CLEAR => {
                let virtio_video_queue_clear { queue_type, .. } = r.read_obj()?;
                QueueClear {
                    stream_id: hdr.stream_id.into(),
                    queue_type: queue_type.try_into()?,
                }
            }
            VIRTIO_VIDEO_CMD_GET_PARAMS => {
                let virtio_video_get_params { queue_type, .. } = r.read_obj()?;
                GetParams {
                    stream_id: hdr.stream_id.into(),
                    queue_type: queue_type.try_into()?,
                }
            }
            VIRTIO_VIDEO_CMD_SET_PARAMS => {
                let virtio_video_set_params { params } = r.read_obj()?;
                SetParams {
                    stream_id: hdr.stream_id.into(),
                    queue_type: params.queue_type.try_into()?,
                    params: params.try_into()?,
                }
            }
            VIRTIO_VIDEO_CMD_QUERY_CONTROL => {
                let body = r.read_obj::<virtio_video_query_control>()?;
                let query_ctrl_type = match body.control.into() {
                    VIRTIO_VIDEO_CONTROL_PROFILE => QueryCtrlType::Profile(
                        r.read_obj::<virtio_video_query_control_profile>()?
                            .format
                            .try_into()?,
                    ),
                    VIRTIO_VIDEO_CONTROL_LEVEL => QueryCtrlType::Level(
                        r.read_obj::<virtio_video_query_control_level>()?
                            .format
                            .try_into()?,
                    ),
                    t => {
                        return Err(ReadCmdError::UnsupportedCtrlType(t));
                    }
                };
                QueryControl { query_ctrl_type }
            }
            VIRTIO_VIDEO_CMD_GET_CONTROL => {
                let virtio_video_get_control { control, .. } = r.read_obj()?;
                let ctrl_type = match control.into() {
                    VIRTIO_VIDEO_CONTROL_BITRATE => CtrlType::Bitrate,
                    VIRTIO_VIDEO_CONTROL_BITRATE_PEAK => CtrlType::BitratePeak,
                    VIRTIO_VIDEO_CONTROL_BITRATE_MODE => CtrlType::BitrateMode,
                    VIRTIO_VIDEO_CONTROL_PROFILE => CtrlType::Profile,
                    VIRTIO_VIDEO_CONTROL_LEVEL => CtrlType::Level,
                    VIRTIO_VIDEO_CONTROL_FORCE_KEYFRAME => CtrlType::ForceKeyframe,
                    VIRTIO_VIDEO_CONTROL_PREPEND_SPSPPS_TO_IDR => CtrlType::PrependSpsPpsToIdr,
                    t => {
                        return Err(ReadCmdError::UnsupportedCtrlType(t));
                    }
                };
                GetControl {
                    stream_id: hdr.stream_id.into(),
                    ctrl_type,
                }
            }
            VIRTIO_VIDEO_CMD_SET_CONTROL => {
                let virtio_video_set_control { control, .. } = r.read_obj()?;
                let ctrl_val = match control.into() {
                    VIRTIO_VIDEO_CONTROL_BITRATE => CtrlVal::Bitrate(
                        r.read_obj::<virtio_video_control_val_bitrate>()?
                            .bitrate
                            .into(),
                    ),
                    VIRTIO_VIDEO_CONTROL_BITRATE_PEAK => CtrlVal::BitratePeak(
                        r.read_obj::<virtio_video_control_val_bitrate_peak>()?
                            .bitrate_peak
                            .into(),
                    ),
                    VIRTIO_VIDEO_CONTROL_BITRATE_MODE => CtrlVal::BitrateMode(
                        r.read_obj::<virtio_video_control_val_bitrate_mode>()?
                            .bitrate_mode
                            .try_into()?,
                    ),
                    VIRTIO_VIDEO_CONTROL_PROFILE => CtrlVal::Profile(
                        r.read_obj::<virtio_video_control_val_profile>()?
                            .profile
                            .try_into()?,
                    ),
                    VIRTIO_VIDEO_CONTROL_LEVEL => CtrlVal::Level(
                        r.read_obj::<virtio_video_control_val_level>()?
                            .level
                            .try_into()?,
                    ),
                    VIRTIO_VIDEO_CONTROL_FORCE_KEYFRAME => CtrlVal::ForceKeyframe,
                    VIRTIO_VIDEO_CONTROL_PREPEND_SPSPPS_TO_IDR => CtrlVal::PrependSpsPpsToIdr(
                        r.read_obj::<virtio_video_control_val_prepend_spspps_to_idr>()?
                            .prepend_spspps_to_idr
                            != 0,
                    ),
                    t => {
                        return Err(ReadCmdError::UnsupportedCtrlType(t));
                    }
                };
                SetControl {
                    stream_id: hdr.stream_id.into(),
                    ctrl_val,
                }
            }
            _ => return Err(ReadCmdError::InvalidCmdType(hdr.type_.into())),
        })
    }
}
<commit_msg>Revert "virtio: video: reject resources with more than one entry"<commit_after>// Copyright 2020 The Chromium OS Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Data structures for commands of virtio video devices.

use remain::sorted;
use std::convert::{TryFrom, TryInto};
use std::io;
use thiserror::Error as ThisError;

use base::error;
use data_model::Le32;
use enumn::N;

use crate::virtio::video::control::*;
use crate::virtio::video::format::*;
use crate::virtio::video::params::Params;
use crate::virtio::video::protocol::*;
use crate::virtio::video::resource::{ResourceType, UnresolvedGuestResource};
use crate::virtio::Reader;

/// An error indicating a failure while reading a request from the guest.
#[sorted]
#[derive(Debug, ThisError)]
pub enum ReadCmdError {
    /// Invalid argument is passed.
    #[error("invalid argument passed to command")]
    InvalidArgument,
    /// The type of the command was invalid.
    #[error("invalid command type: {0}")]
    InvalidCmdType(u32),
    /// Failed to read an object.
    #[error("failed to read object: {0}")]
    IoError(#[from] io::Error),
    /// The type of the requested control was unsupported.
    #[error("unsupported control type: {0}")]
    UnsupportedCtrlType(u32),
}

#[derive(PartialEq, Eq, PartialOrd, Ord, N, Clone, Copy, Debug)]
#[repr(u32)]
pub enum QueueType {
    Input = VIRTIO_VIDEO_QUEUE_TYPE_INPUT,
    Output = VIRTIO_VIDEO_QUEUE_TYPE_OUTPUT,
}
impl_try_from_le32_for_enumn!(QueueType, "queue_type");

#[derive(Debug)]
pub enum VideoCmd {
    QueryCapability {
        queue_type: QueueType,
    },
    StreamCreate {
        stream_id: u32,
        coded_format: Format,
        input_resource_type: ResourceType,
        output_resource_type: ResourceType,
    },
    StreamDestroy {
        stream_id: u32,
    },
    StreamDrain {
        stream_id: u32,
    },
    ResourceCreate {
        stream_id: u32,
        queue_type: QueueType,
        resource_id: u32,
        plane_offsets: Vec<u32>,
        resource: UnresolvedGuestResource,
    },
    ResourceQueue {
        stream_id: u32,
        queue_type: QueueType,
        resource_id: u32,
        timestamp: u64,
        data_sizes: Vec<u32>,
    },
    ResourceDestroyAll {
        stream_id: u32,
        queue_type: QueueType,
    },
    QueueClear {
        stream_id: u32,
        queue_type: QueueType,
    },
    GetParams {
        stream_id: u32,
        queue_type: QueueType,
    },
    SetParams {
        stream_id: u32,
        queue_type: QueueType,
        params: Params,
    },
    QueryControl {
        query_ctrl_type: QueryCtrlType,
    },
    GetControl {
        stream_id: u32,
        ctrl_type: CtrlType,
    },
    SetControl {
        stream_id: u32,
        ctrl_val: CtrlVal,
    },
}

impl<'a> VideoCmd {
    /// Reads a request on virtqueue and construct a VideoCmd value.
    pub fn from_reader(r: &'a mut Reader) -> Result<Self, ReadCmdError> {
        use self::ReadCmdError::*;
        use self::VideoCmd::*;

        // Unlike structs in virtio_video.h in the kernel, our command structs in protocol.rs don't
        // have a field of `struct virtio_video_cmd_hdr`. So, we first read the header here and
        // a body below.
        let hdr = r.read_obj::<virtio_video_cmd_hdr>()?;

        Ok(match hdr.type_.into() {
            VIRTIO_VIDEO_CMD_QUERY_CAPABILITY => {
                let virtio_video_query_capability { queue_type, .. } = r.read_obj()?;
                QueryCapability {
                    queue_type: queue_type.try_into()?,
                }
            }
            VIRTIO_VIDEO_CMD_STREAM_CREATE => {
                let virtio_video_stream_create {
                    in_mem_type,
                    out_mem_type,
                    coded_format,
                    ..
                } = r.read_obj()?;

                let input_resource_type = match in_mem_type.into() {
                    VIRTIO_VIDEO_MEM_TYPE_VIRTIO_OBJECT => ResourceType::Object,
                    _ => {
                        error!("mem_type must be VIRTIO_OBJECT");
                        return Err(InvalidArgument);
                    }
                };

                let output_resource_type = match out_mem_type.into() {
                    VIRTIO_VIDEO_MEM_TYPE_VIRTIO_OBJECT => ResourceType::Object,
                    _ => {
                        error!("mem_type must be VIRTIO_OBJECT");
                        return Err(InvalidArgument);
                    }
                };

                StreamCreate {
                    stream_id: hdr.stream_id.into(),
                    coded_format: coded_format.try_into()?,
                    input_resource_type,
                    output_resource_type,
                }
            }
            VIRTIO_VIDEO_CMD_STREAM_DESTROY => {
                let virtio_video_stream_destroy { .. } = r.read_obj()?;
                StreamDestroy {
                    stream_id: hdr.stream_id.into(),
                }
            }
            VIRTIO_VIDEO_CMD_STREAM_DRAIN => {
                let virtio_video_stream_drain { .. } = r.read_obj()?;
                StreamDrain {
                    stream_id: hdr.stream_id.into(),
                }
            }
            VIRTIO_VIDEO_CMD_RESOURCE_CREATE => {
                let virtio_video_resource_create {
                    queue_type,
                    resource_id,
                    planes_layout,
                    num_planes,
                    plane_offsets,
                    ..
                } = r.read_obj()?;

                // Assume ChromeOS-specific requirements.
                if Into::<u32>::into(planes_layout) != VIRTIO_VIDEO_PLANES_LAYOUT_SINGLE_BUFFER {
                    error!(
                        "each buffer must be a single DMAbuf: {}",
                        Into::<u32>::into(planes_layout),
                    );
                    return Err(InvalidArgument);
                }

                let num_planes: u32 = num_planes.into();
                if num_planes as usize > plane_offsets.len() {
                    error!(
                        "num_planes must not exceed {} but {}",
                        plane_offsets.len(),
                        num_planes
                    );
                    return Err(InvalidArgument);
                }
                let plane_offsets = plane_offsets[0..num_planes as usize]
                    .iter()
                    .map(|x| Into::<u32>::into(*x))
                    .collect::<Vec<u32>>();

                let resource = r.read_obj()?;

                ResourceCreate {
                    stream_id: hdr.stream_id.into(),
                    queue_type: queue_type.try_into()?,
                    resource_id: resource_id.into(),
                    plane_offsets,
                    resource,
                }
            }
            VIRTIO_VIDEO_CMD_RESOURCE_QUEUE => {
                let virtio_video_resource_queue {
                    queue_type,
                    resource_id,
                    timestamp,
                    num_data_sizes,
                    data_sizes,
                    ..
                } = r.read_obj()?;

                let num_data_sizes: u32 = num_data_sizes.into();
                if num_data_sizes as usize > data_sizes.len() {
                    return Err(InvalidArgument);
                }
                let data_sizes = data_sizes[0..num_data_sizes as usize]
                    .iter()
                    .map(|x| Into::<u32>::into(*x))
                    .collect::<Vec<u32>>();
                ResourceQueue {
                    stream_id: hdr.stream_id.into(),
                    queue_type: queue_type.try_into()?,
                    resource_id: resource_id.into(),
                    timestamp: timestamp.into(),
                    data_sizes,
                }
            }
            VIRTIO_VIDEO_CMD_RESOURCE_DESTROY_ALL => {
                let virtio_video_resource_destroy_all { queue_type, .. } = r.read_obj()?;
                ResourceDestroyAll {
                    stream_id: hdr.stream_id.into(),
                    queue_type: queue_type.try_into()?,
                }
            }
            VIRTIO_VIDEO_CMD_QUEUE_CLEAR => {
                let virtio_video_queue_clear { queue_type, .. } = r.read_obj()?;
                QueueClear {
                    stream_id: hdr.stream_id.into(),
                    queue_type: queue_type.try_into()?,
                }
            }
            VIRTIO_VIDEO_CMD_GET_PARAMS => {
                let virtio_video_get_params { queue_type, .. } = r.read_obj()?;
                GetParams {
                    stream_id: hdr.stream_id.into(),
                    queue_type: queue_type.try_into()?,
                }
            }
            VIRTIO_VIDEO_CMD_SET_PARAMS => {
                let virtio_video_set_params { params } = r.read_obj()?;
                SetParams {
                    stream_id: hdr.stream_id.into(),
                    queue_type: params.queue_type.try_into()?,
                    params: params.try_into()?,
                }
            }
            VIRTIO_VIDEO_CMD_QUERY_CONTROL => {
                let body = r.read_obj::<virtio_video_query_control>()?;
                let query_ctrl_type = match body.control.into() {
                    VIRTIO_VIDEO_CONTROL_PROFILE => QueryCtrlType::Profile(
                        r.read_obj::<virtio_video_query_control_profile>()?
                            .format
                            .try_into()?,
                    ),
                    VIRTIO_VIDEO_CONTROL_LEVEL => QueryCtrlType::Level(
                        r.read_obj::<virtio_video_query_control_level>()?
                            .format
                            .try_into()?,
                    ),
                    t => {
                        return Err(ReadCmdError::UnsupportedCtrlType(t));
                    }
                };
                QueryControl { query_ctrl_type }
            }
            VIRTIO_VIDEO_CMD_GET_CONTROL => {
                let virtio_video_get_control { control, .. } = r.read_obj()?;
                let ctrl_type = match control.into() {
                    VIRTIO_VIDEO_CONTROL_BITRATE => CtrlType::Bitrate,
                    VIRTIO_VIDEO_CONTROL_BITRATE_PEAK => CtrlType::BitratePeak,
                    VIRTIO_VIDEO_CONTROL_BITRATE_MODE => CtrlType::BitrateMode,
                    VIRTIO_VIDEO_CONTROL_PROFILE => CtrlType::Profile,
                    VIRTIO_VIDEO_CONTROL_LEVEL => CtrlType::Level,
                    VIRTIO_VIDEO_CONTROL_FORCE_KEYFRAME => CtrlType::ForceKeyframe,
                    VIRTIO_VIDEO_CONTROL_PREPEND_SPSPPS_TO_IDR => CtrlType::PrependSpsPpsToIdr,
                    t => {
                        return Err(ReadCmdError::UnsupportedCtrlType(t));
                    }
                };
                GetControl {
                    stream_id: hdr.stream_id.into(),
                    ctrl_type,
                }
            }
            VIRTIO_VIDEO_CMD_SET_CONTROL => {
                let virtio_video_set_control { control, .. } = r.read_obj()?;
                let ctrl_val = match control.into() {
                    VIRTIO_VIDEO_CONTROL_BITRATE => CtrlVal::Bitrate(
                        r.read_obj::<virtio_video_control_val_bitrate>()?
                            .bitrate
                            .into(),
                    ),
                    VIRTIO_VIDEO_CONTROL_BITRATE_PEAK => CtrlVal::BitratePeak(
                        r.read_obj::<virtio_video_control_val_bitrate_peak>()?
                            .bitrate_peak
                            .into(),
                    ),
                    VIRTIO_VIDEO_CONTROL_BITRATE_MODE => CtrlVal::BitrateMode(
                        r.read_obj::<virtio_video_control_val_bitrate_mode>()?
                            .bitrate_mode
                            .try_into()?,
                    ),
                    VIRTIO_VIDEO_CONTROL_PROFILE => CtrlVal::Profile(
                        r.read_obj::<virtio_video_control_val_profile>()?
                            .profile
                            .try_into()?,
                    ),
                    VIRTIO_VIDEO_CONTROL_LEVEL => CtrlVal::Level(
                        r.read_obj::<virtio_video_control_val_level>()?
                            .level
                            .try_into()?,
                    ),
                    VIRTIO_VIDEO_CONTROL_FORCE_KEYFRAME => CtrlVal::ForceKeyframe,
                    VIRTIO_VIDEO_CONTROL_PREPEND_SPSPPS_TO_IDR => CtrlVal::PrependSpsPpsToIdr(
                        r.read_obj::<virtio_video_control_val_prepend_spspps_to_idr>()?
                            .prepend_spspps_to_idr
                            != 0,
                    ),
                    t => {
                        return Err(ReadCmdError::UnsupportedCtrlType(t));
                    }
                };
                SetControl {
                    stream_id: hdr.stream_id.into(),
                    ctrl_val,
                }
            }
            _ => return Err(ReadCmdError::InvalidCmdType(hdr.type_.into())),
        })
    }
}
<|endoftext|>