<commit_before>//! A control flow graph represented as mappings of extended basic blocks to their predecessors
//! and successors.
//!
//! Successors are represented as extended basic blocks while predecessors are represented by basic
//! blocks. Basic blocks are denoted by tuples of EBB and branch/jump instructions. Each
//! predecessor tuple corresponds to the end of a basic block.
//!
//! ```c
//!     Ebb0:
//!         ...          ; beginning of basic block
//!
//!         ...
//!
//!         brz vx, Ebb1 ; end of basic block
//!
//!         ...          ; beginning of basic block
//!
//!         ...
//!
//!         jmp Ebb2     ; end of basic block
//! ```
//!
//! Here `Ebb1` and `Ebb2` would each have a single predecessor denoted as `(Ebb0, brz)`
//! and `(Ebb0, jmp Ebb2)` respectively.

use ir::{Function, Inst, Ebb};
use ir::instructions::BranchInfo;
use entity_map::{EntityMap, Keys};
use std::collections::HashSet;
use std::mem;

/// A basic block denoted by its enclosing Ebb and last instruction.
pub type BasicBlock = (Ebb, Inst);

/// A container for the successors and predecessors of some Ebb.
#[derive(Debug, Clone, Default)]
pub struct CFGNode {
    /// EBBs that are the targets of branches and jumps in this EBB.
    pub successors: Vec<Ebb>,
    /// Basic blocks that can branch or jump to this EBB.
    pub predecessors: Vec<BasicBlock>,
}

/// The Control Flow Graph maintains a mapping of ebbs to their predecessors
/// and successors where predecessors are basic blocks and successors are
/// extended basic blocks.
#[derive(Debug)]
pub struct ControlFlowGraph {
    entry_block: Option<Ebb>,
    data: EntityMap<Ebb, CFGNode>,
}

impl ControlFlowGraph {
    /// Allocate a new blank control flow graph.
    pub fn new() -> ControlFlowGraph {
        ControlFlowGraph {
            entry_block: None,
            data: EntityMap::new(),
        }
    }

    /// Allocate and compute the control flow graph for `func`.
    pub fn with_function(func: &Function) -> ControlFlowGraph {
        let mut cfg = ControlFlowGraph::new();
        cfg.compute(func);
        cfg
    }

    /// Compute the control flow graph of `func`.
    ///
    /// This will clear and overwrite any information already stored in this data structure.
    pub fn compute(&mut self, func: &Function) {
        self.entry_block = func.layout.entry_block();
        self.data.clear();
        self.data.resize(func.dfg.num_ebbs());

        for ebb in &func.layout {
            self.compute_ebb(func, ebb);
        }
    }

    fn compute_ebb(&mut self, func: &Function, ebb: Ebb) {
        for inst in func.layout.ebb_insts(ebb) {
            match func.dfg[inst].analyze_branch(&func.dfg.value_lists) {
                BranchInfo::SingleDest(dest, _) => {
                    self.add_edge((ebb, inst), dest);
                }
                BranchInfo::Table(jt) => {
                    for (_, dest) in func.jump_tables[jt].entries() {
                        self.add_edge((ebb, inst), dest);
                    }
                }
                BranchInfo::NotABranch => {}
            }
        }
    }

    fn invalidate_ebb_successors(&mut self, ebb: Ebb) {
        // Temporarily take ownership because we need mutable access to self.data inside the loop.
        // Unfortunately borrowck cannot see that our mut accesses to predecessors don't alias
        // our iteration over successors.
        let mut successors = mem::replace(&mut self.data[ebb].successors, Vec::new());
        for suc in successors.iter().cloned() {
            self.data[suc].predecessors.retain(|&(e, _)| e != ebb);
        }
        successors.clear();
        self.data[ebb].successors = successors;
    }

    /// Recompute the control flow graph of `ebb`.
    ///
    /// This is for use after modifying instructions within a specific EBB. It recomputes all edges
    /// from `ebb` while leaving edges to `ebb` intact. Its functionality a subset of that of the
    /// more expensive `compute`, and should be used when we know we don't need to recompute the CFG
    /// from scratch, but rather that our changes have been restricted to specific EBBs.
    pub fn recompute_ebb(&mut self, func: &Function, ebb: Ebb) {
        self.invalidate_ebb_successors(ebb);
        self.compute_ebb(func, ebb);
    }

    fn add_edge(&mut self, from: BasicBlock, to: Ebb) {
        self.data[from.0].successors.push(to);
        self.data[to].predecessors.push(from);
    }

    /// Get the CFG predecessor basic blocks to `ebb`.
    pub fn get_predecessors(&self, ebb: Ebb) -> &[BasicBlock] {
        &self.data[ebb].predecessors
    }

    /// Get the CFG successors to `ebb`.
    pub fn get_successors(&self, ebb: Ebb) -> &[Ebb] {
        &self.data[ebb].successors
    }

    /// Return [reachable] ebbs in post-order.
    pub fn postorder_ebbs(&self) -> Vec<Ebb> {
        let entry_block = match self.entry_block {
            None => {
                return Vec::new();
            }
            Some(eb) => eb,
        };

        let mut grey = HashSet::new();
        let mut black = HashSet::new();
        let mut stack = vec![entry_block.clone()];
        let mut postorder = Vec::new();

        while !stack.is_empty() {
            let node = stack.pop().unwrap();
            if !grey.contains(&node) {
                // This is a white node. Mark it as gray.
                grey.insert(node);
                stack.push(node);
                // Get any children we've never seen before.
                for child in self.get_successors(node) {
                    if !grey.contains(child) {
                        stack.push(child.clone());
                    }
                }
            } else if !black.contains(&node) {
                postorder.push(node.clone());
                black.insert(node.clone());
            }
        }
        postorder
    }

    /// An iterator across all of the ebbs stored in the CFG.
    pub fn ebbs(&self) -> Keys<Ebb> {
        self.data.keys()
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ir::{Function, InstBuilder, Cursor, types};

    #[test]
    fn empty() {
        let func = Function::new();
        let cfg = ControlFlowGraph::with_function(&func);
        assert_eq!(None, cfg.ebbs().next());
    }

    #[test]
    fn no_predecessors() {
        let mut func = Function::new();
        let ebb0 = func.dfg.make_ebb();
        let ebb1 = func.dfg.make_ebb();
        let ebb2 = func.dfg.make_ebb();
        func.layout.append_ebb(ebb0);
        func.layout.append_ebb(ebb1);
        func.layout.append_ebb(ebb2);

        let cfg = ControlFlowGraph::with_function(&func);
        let nodes = cfg.ebbs().collect::<Vec<_>>();
        assert_eq!(nodes.len(), 3);

        let mut fun_ebbs = func.layout.ebbs();
        for ebb in nodes {
            assert_eq!(ebb, fun_ebbs.next().unwrap());
            assert_eq!(cfg.get_predecessors(ebb).len(), 0);
            assert_eq!(cfg.get_successors(ebb).len(), 0);
        }
    }

    #[test]
    fn branches_and_jumps() {
        let mut func = Function::new();
        let ebb0 = func.dfg.make_ebb();
        let cond = func.dfg.append_ebb_arg(ebb0, types::I32);
        let ebb1 = func.dfg.make_ebb();
        let ebb2 = func.dfg.make_ebb();

        let br_ebb0_ebb2;
        let br_ebb1_ebb1;
        let jmp_ebb0_ebb1;
        let jmp_ebb1_ebb2;

        {
            let dfg = &mut func.dfg;
            let cur = &mut Cursor::new(&mut func.layout);

            cur.insert_ebb(ebb0);
            br_ebb0_ebb2 = dfg.ins(cur).brnz(cond, ebb2, &[]);
            jmp_ebb0_ebb1 = dfg.ins(cur).jump(ebb1, &[]);

            cur.insert_ebb(ebb1);
            br_ebb1_ebb1 = dfg.ins(cur).brnz(cond, ebb1, &[]);
            jmp_ebb1_ebb2 = dfg.ins(cur).jump(ebb2, &[]);

            cur.insert_ebb(ebb2);
        }

        let mut cfg = ControlFlowGraph::with_function(&func);

        {
            let ebb0_predecessors = cfg.get_predecessors(ebb0);
            let ebb1_predecessors = cfg.get_predecessors(ebb1);
            let ebb2_predecessors = cfg.get_predecessors(ebb2);

            let ebb0_successors = cfg.get_successors(ebb0);
            let ebb1_successors = cfg.get_successors(ebb1);
            let ebb2_successors = cfg.get_successors(ebb2);

            assert_eq!(ebb0_predecessors.len(), 0);
            assert_eq!(ebb1_predecessors.len(), 2);
            assert_eq!(ebb2_predecessors.len(), 2);

            assert_eq!(ebb1_predecessors.contains(&(ebb0, jmp_ebb0_ebb1)), true);
            assert_eq!(ebb1_predecessors.contains(&(ebb1, br_ebb1_ebb1)), true);
            assert_eq!(ebb2_predecessors.contains(&(ebb0, br_ebb0_ebb2)), true);
            assert_eq!(ebb2_predecessors.contains(&(ebb1, jmp_ebb1_ebb2)), true);

            assert_eq!(ebb0_successors.len(), 2);
            assert_eq!(ebb1_successors.len(), 2);
            assert_eq!(ebb2_successors.len(), 0);

            assert_eq!(ebb0_successors.contains(&ebb1), true);
            assert_eq!(ebb0_successors.contains(&ebb2), true);
            assert_eq!(ebb1_successors.contains(&ebb1), true);
            assert_eq!(ebb1_successors.contains(&ebb2), true);
        }

        // Change some instructions and recompute ebb0
        func.dfg.replace(br_ebb0_ebb2).brnz(cond, ebb1, &[]);
        func.dfg.replace(jmp_ebb0_ebb1).return_(&[]);
        cfg.recompute_ebb(&mut func, ebb0);
        let br_ebb0_ebb1 = br_ebb0_ebb2;

        {
            let ebb0_predecessors = cfg.get_predecessors(ebb0);
            let ebb1_predecessors = cfg.get_predecessors(ebb1);
            let ebb2_predecessors = cfg.get_predecessors(ebb2);

            let ebb0_successors = cfg.get_successors(ebb0);
            let ebb1_successors = cfg.get_successors(ebb1);
            let ebb2_successors = cfg.get_successors(ebb2);

            assert_eq!(ebb0_predecessors.len(), 0);
            assert_eq!(ebb1_predecessors.len(), 2);
            assert_eq!(ebb2_predecessors.len(), 1);

            assert_eq!(ebb1_predecessors.contains(&(ebb0, br_ebb0_ebb1)), true);
            assert_eq!(ebb1_predecessors.contains(&(ebb1, br_ebb1_ebb1)), true);
            assert_eq!(ebb2_predecessors.contains(&(ebb0, br_ebb0_ebb2)), false);
            assert_eq!(ebb2_predecessors.contains(&(ebb1, jmp_ebb1_ebb2)), true);

            assert_eq!(ebb0_successors.len(), 1);
            assert_eq!(ebb1_successors.len(), 2);
            assert_eq!(ebb2_successors.len(), 0);

            assert_eq!(ebb0_successors.contains(&ebb1), true);
            assert_eq!(ebb0_successors.contains(&ebb2), false);
            assert_eq!(ebb1_successors.contains(&ebb1), true);
            assert_eq!(ebb1_successors.contains(&ebb2), true);
        }
    }
}
<commit_msg>Remove cfg.postorder_ebbs().<commit_after>//! A control flow graph represented as mappings of extended basic blocks to their predecessors
//! and successors.
//!
//! Successors are represented as extended basic blocks while predecessors are represented by basic
//! blocks. Basic blocks are denoted by tuples of EBB and branch/jump instructions. Each
//! predecessor tuple corresponds to the end of a basic block.
//!
//! ```c
//!     Ebb0:
//!         ...          ; beginning of basic block
//!
//!         ...
//!
//!         brz vx, Ebb1 ; end of basic block
//!
//!         ...          ; beginning of basic block
//!
//!         ...
//!
//!         jmp Ebb2     ; end of basic block
//! ```
//!
//! Here `Ebb1` and `Ebb2` would each have a single predecessor denoted as `(Ebb0, brz)`
//! and `(Ebb0, jmp Ebb2)` respectively.

use ir::{Function, Inst, Ebb};
use ir::instructions::BranchInfo;
use entity_map::EntityMap;
use std::mem;

/// A basic block denoted by its enclosing Ebb and last instruction.
pub type BasicBlock = (Ebb, Inst);

/// A container for the successors and predecessors of some Ebb.
#[derive(Debug, Clone, Default)]
pub struct CFGNode {
    /// EBBs that are the targets of branches and jumps in this EBB.
    pub successors: Vec<Ebb>,
    /// Basic blocks that can branch or jump to this EBB.
    pub predecessors: Vec<BasicBlock>,
}

/// The Control Flow Graph maintains a mapping of ebbs to their predecessors
/// and successors where predecessors are basic blocks and successors are
/// extended basic blocks.
#[derive(Debug)]
pub struct ControlFlowGraph {
    entry_block: Option<Ebb>,
    data: EntityMap<Ebb, CFGNode>,
}

impl ControlFlowGraph {
    /// Allocate a new blank control flow graph.
    pub fn new() -> ControlFlowGraph {
        ControlFlowGraph {
            entry_block: None,
            data: EntityMap::new(),
        }
    }

    /// Allocate and compute the control flow graph for `func`.
    pub fn with_function(func: &Function) -> ControlFlowGraph {
        let mut cfg = ControlFlowGraph::new();
        cfg.compute(func);
        cfg
    }

    /// Compute the control flow graph of `func`.
    ///
    /// This will clear and overwrite any information already stored in this data structure.
    pub fn compute(&mut self, func: &Function) {
        self.entry_block = func.layout.entry_block();
        self.data.clear();
        self.data.resize(func.dfg.num_ebbs());

        for ebb in &func.layout {
            self.compute_ebb(func, ebb);
        }
    }

    fn compute_ebb(&mut self, func: &Function, ebb: Ebb) {
        for inst in func.layout.ebb_insts(ebb) {
            match func.dfg[inst].analyze_branch(&func.dfg.value_lists) {
                BranchInfo::SingleDest(dest, _) => {
                    self.add_edge((ebb, inst), dest);
                }
                BranchInfo::Table(jt) => {
                    for (_, dest) in func.jump_tables[jt].entries() {
                        self.add_edge((ebb, inst), dest);
                    }
                }
                BranchInfo::NotABranch => {}
            }
        }
    }

    fn invalidate_ebb_successors(&mut self, ebb: Ebb) {
        // Temporarily take ownership because we need mutable access to self.data inside the loop.
        // Unfortunately borrowck cannot see that our mut accesses to predecessors don't alias
        // our iteration over successors.
        let mut successors = mem::replace(&mut self.data[ebb].successors, Vec::new());
        for suc in successors.iter().cloned() {
            self.data[suc].predecessors.retain(|&(e, _)| e != ebb);
        }
        successors.clear();
        self.data[ebb].successors = successors;
    }

    /// Recompute the control flow graph of `ebb`.
    ///
    /// This is for use after modifying instructions within a specific EBB. It recomputes all edges
    /// from `ebb` while leaving edges to `ebb` intact. Its functionality a subset of that of the
    /// more expensive `compute`, and should be used when we know we don't need to recompute the CFG
    /// from scratch, but rather that our changes have been restricted to specific EBBs.
    pub fn recompute_ebb(&mut self, func: &Function, ebb: Ebb) {
        self.invalidate_ebb_successors(ebb);
        self.compute_ebb(func, ebb);
    }

    fn add_edge(&mut self, from: BasicBlock, to: Ebb) {
        self.data[from.0].successors.push(to);
        self.data[to].predecessors.push(from);
    }

    /// Get the CFG predecessor basic blocks to `ebb`.
    pub fn get_predecessors(&self, ebb: Ebb) -> &[BasicBlock] {
        &self.data[ebb].predecessors
    }

    /// Get the CFG successors to `ebb`.
    pub fn get_successors(&self, ebb: Ebb) -> &[Ebb] {
        &self.data[ebb].successors
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use ir::{Function, InstBuilder, Cursor, types};

    #[test]
    fn empty() {
        let func = Function::new();
        ControlFlowGraph::with_function(&func);
    }

    #[test]
    fn no_predecessors() {
        let mut func = Function::new();
        let ebb0 = func.dfg.make_ebb();
        let ebb1 = func.dfg.make_ebb();
        let ebb2 = func.dfg.make_ebb();
        func.layout.append_ebb(ebb0);
        func.layout.append_ebb(ebb1);
        func.layout.append_ebb(ebb2);

        let cfg = ControlFlowGraph::with_function(&func);

        let mut fun_ebbs = func.layout.ebbs();
        for ebb in func.layout.ebbs() {
            assert_eq!(ebb, fun_ebbs.next().unwrap());
            assert_eq!(cfg.get_predecessors(ebb).len(), 0);
            assert_eq!(cfg.get_successors(ebb).len(), 0);
        }
    }

    #[test]
    fn branches_and_jumps() {
        let mut func = Function::new();
        let ebb0 = func.dfg.make_ebb();
        let cond = func.dfg.append_ebb_arg(ebb0, types::I32);
        let ebb1 = func.dfg.make_ebb();
        let ebb2 = func.dfg.make_ebb();

        let br_ebb0_ebb2;
        let br_ebb1_ebb1;
        let jmp_ebb0_ebb1;
        let jmp_ebb1_ebb2;

        {
            let dfg = &mut func.dfg;
            let cur = &mut Cursor::new(&mut func.layout);

            cur.insert_ebb(ebb0);
            br_ebb0_ebb2 = dfg.ins(cur).brnz(cond, ebb2, &[]);
            jmp_ebb0_ebb1 = dfg.ins(cur).jump(ebb1, &[]);

            cur.insert_ebb(ebb1);
            br_ebb1_ebb1 = dfg.ins(cur).brnz(cond, ebb1, &[]);
            jmp_ebb1_ebb2 = dfg.ins(cur).jump(ebb2, &[]);

            cur.insert_ebb(ebb2);
        }

        let mut cfg = ControlFlowGraph::with_function(&func);

        {
            let ebb0_predecessors = cfg.get_predecessors(ebb0);
            let ebb1_predecessors = cfg.get_predecessors(ebb1);
            let ebb2_predecessors = cfg.get_predecessors(ebb2);

            let ebb0_successors = cfg.get_successors(ebb0);
            let ebb1_successors = cfg.get_successors(ebb1);
            let ebb2_successors = cfg.get_successors(ebb2);

            assert_eq!(ebb0_predecessors.len(), 0);
            assert_eq!(ebb1_predecessors.len(), 2);
            assert_eq!(ebb2_predecessors.len(), 2);

            assert_eq!(ebb1_predecessors.contains(&(ebb0, jmp_ebb0_ebb1)), true);
            assert_eq!(ebb1_predecessors.contains(&(ebb1, br_ebb1_ebb1)), true);
            assert_eq!(ebb2_predecessors.contains(&(ebb0, br_ebb0_ebb2)), true);
            assert_eq!(ebb2_predecessors.contains(&(ebb1, jmp_ebb1_ebb2)), true);

            assert_eq!(ebb0_successors.len(), 2);
            assert_eq!(ebb1_successors.len(), 2);
            assert_eq!(ebb2_successors.len(), 0);

            assert_eq!(ebb0_successors.contains(&ebb1), true);
            assert_eq!(ebb0_successors.contains(&ebb2), true);
            assert_eq!(ebb1_successors.contains(&ebb1), true);
            assert_eq!(ebb1_successors.contains(&ebb2), true);
        }

        // Change some instructions and recompute ebb0
        func.dfg.replace(br_ebb0_ebb2).brnz(cond, ebb1, &[]);
        func.dfg.replace(jmp_ebb0_ebb1).return_(&[]);
        cfg.recompute_ebb(&mut func, ebb0);
        let br_ebb0_ebb1 = br_ebb0_ebb2;

        {
            let ebb0_predecessors = cfg.get_predecessors(ebb0);
            let ebb1_predecessors = cfg.get_predecessors(ebb1);
            let ebb2_predecessors = cfg.get_predecessors(ebb2);

            let ebb0_successors = cfg.get_successors(ebb0);
            let ebb1_successors = cfg.get_successors(ebb1);
            let ebb2_successors = cfg.get_successors(ebb2);

            assert_eq!(ebb0_predecessors.len(), 0);
            assert_eq!(ebb1_predecessors.len(), 2);
            assert_eq!(ebb2_predecessors.len(), 1);

            assert_eq!(ebb1_predecessors.contains(&(ebb0, br_ebb0_ebb1)), true);
            assert_eq!(ebb1_predecessors.contains(&(ebb1, br_ebb1_ebb1)), true);
            assert_eq!(ebb2_predecessors.contains(&(ebb0, br_ebb0_ebb2)), false);
            assert_eq!(ebb2_predecessors.contains(&(ebb1, jmp_ebb1_ebb2)), true);

            assert_eq!(ebb0_successors.len(), 1);
            assert_eq!(ebb1_successors.len(), 2);
            assert_eq!(ebb2_successors.len(), 0);

            assert_eq!(ebb0_successors.contains(&ebb1), true);
            assert_eq!(ebb0_successors.contains(&ebb2), false);
            assert_eq!(ebb1_successors.contains(&ebb1), true);
            assert_eq!(ebb1_successors.contains(&ebb2), true);
        }
    }
}
<|endoftext|>


<commit_before>use crate::support::{database, project};

#[test]
fn migration_revert_runs_the_last_migration_down() {
    let p = project("migration_revert_runs_the_last_migration_down")
        .folder("migrations")
        .build();
    let db = database(&p.database_url());

    p.create_migration(
        "12345_create_users_table",
        "CREATE TABLE users ( id INTEGER )",
        "DROP TABLE users",
    );

    // Make sure the project is setup
    p.command("setup").run();

    assert!(db.table_exists("users"));

    let result = p.command("migration").arg("revert").run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);
    assert!(
        result.stdout().contains("Rolling back migration 12345"),
        "Unexpected stdout {}",
        result.stdout()
    );
    assert!(!db.table_exists("users"));
}

#[test]
fn migration_revert_respects_migration_dir_var() {
    let p = project("migration_revert_var").folder("foo").build();
    let db = database(&p.database_url());

    p.create_migration_in_directory(
        "foo",
        "12345_create_users_table",
        "CREATE TABLE users ( id INTEGER )",
        "DROP TABLE users",
    );

    // Make sure the project is setup.
    p.command("setup").arg("--migration-dir=foo").run();

    assert!(db.table_exists("users"));

    let result = p
        .command("migration")
        .arg("revert")
        .arg("--migration-dir=foo")
        .run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);
    assert!(
        result.stdout().contains("Rolling back migration 12345"),
        "Unexpected stdout {}",
        result.stdout()
    );
    assert!(!db.table_exists("users"));
}

#[test]
fn migration_revert_respects_migration_dir_env() {
    let p = project("migration_revert_env").folder("bar").build();
    let db = database(&p.database_url());

    p.create_migration_in_directory(
        "bar",
        "12345_create_users_table",
        "CREATE TABLE users ( id INTEGER )",
        "DROP TABLE users",
    );

    // Make sure the project is setup.
    p.command("setup").arg("--migration-dir=bar").run();

    assert!(db.table_exists("users"));

    let result = p
        .command("migration")
        .arg("revert")
        .env("MIGRATION_DIRECTORY", "bar")
        .run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);
    assert!(
        result.stdout().contains("Rolling back migration 12345"),
        "Unexpected stdout {}",
        result.stdout()
    );
    assert!(!db.table_exists("users"));
}

#[test]
fn migration_revert_respects_migration_dir_from_diesel_toml() {
    let p = project("migration_revert_respects_migration_dir_from_diesel_toml")
        .folder("custom_migrations")
        .file(
            "diesel.toml",
            r#"
            [migrations_directory]
            dir = "custom_migrations"
            "#,
        )
        .build();

    let db = database(&p.database_url());

    p.create_migration_in_directory(
        "custom_migrations",
        "12345_create_users_table",
        "CREATE TABLE users ( id INTEGER )",
        "DROP TABLE users",
    );

    // Make sure the project is setup.
    p.command("setup").run();

    assert!(db.table_exists("users"));

    let result = p.command("migration").arg("revert").run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);
    assert!(
        result.stdout().contains("Rolling back migration 12345"),
        "Unexpected stdout {}",
        result.stdout()
    );
    assert!(!db.table_exists("users"));
}

#[test]
fn migration_revert_runs_the_last_two_migration_down() {
    let p = project("migration_revert_runs_the_last_two_migration_down")
        .folder("migrations")
        .build();
    let db = database(&p.database_url());

    p.create_migration(
        "2017-08-31-210424_create_customers",
        "CREATE TABLE customers ( id INTEGER PRIMARY KEY )",
        "DROP TABLE customers",
    );

    p.create_migration(
        "2017-09-03-210424_create_contracts",
        "CREATE TABLE contracts ( id INTEGER PRIMARY KEY )",
        "DROP TABLE contracts",
    );

    p.create_migration(
        "2017-09-12-210424_create_bills",
        "CREATE TABLE bills ( id INTEGER PRIMARY KEY )",
        "DROP TABLE bills",
    );

    // Make sure the project is setup
    p.command("setup").run();

    assert!(db.table_exists("customers"));
    assert!(db.table_exists("contracts"));
    assert!(db.table_exists("bills"));

    // Reverts the last two migration files. The `contracts` and `bills` tables should be dropped.
    // The `customers` table shouldn't be reverted.
    let result = p
        .command("migration")
        .arg("revert")
        .arg("-n")
        .arg("2")
        .run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);
    assert!(
        result.stdout().contains("Rolling back migration 2017-09-12-210424_create_bills\nRolling back migration 2017-09-03-210424_create_contracts"),
        "Unexpected stdout {}",
        result.stdout()
    );

    assert!(db.table_exists("customers"));
    assert!(!db.table_exists("contracts"));
    assert!(!db.table_exists("bills"));
}

#[test]
fn migration_revert_all_runs_the_migrations_down() {
    let p = project("migration_revert_all_runs_the_migrations_down")
        .folder("migrations")
        .build();
    let db = database(&p.database_url());

    p.create_migration(
        "2017-08-31-210424_create_customers",
        "CREATE TABLE customers ( id INTEGER PRIMARY KEY )",
        "DROP TABLE customers",
    );

    p.create_migration(
        "2017-09-03-210424_create_contracts",
        "CREATE TABLE contracts ( id INTEGER PRIMARY KEY )",
        "DROP TABLE contracts",
    );

    p.create_migration(
        "2017-09-12-210424_create_bills",
        "CREATE TABLE bills ( id INTEGER PRIMARY KEY )",
        "DROP TABLE bills",
    );

    // Make sure the project is setup
    p.command("setup").run();

    assert!(db.table_exists("customers"));
    assert!(db.table_exists("contracts"));
    assert!(db.table_exists("bills"));

    // Reverts all migrations. The migrations `contracts`, `bills`, `customers` and
    // `00000000000000_diesel_initial_setup` should be reverted.
    let result = p
        .command("migration")
        .arg("revert")
        .arg("-n")
        .arg("all")
        .run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);

    assert!(
        result.stdout()
            == "Rolling back migration 2017-09-12-210424_create_bills\n\
                Rolling back migration 2017-09-03-210424_create_contracts\n\
                Rolling back migration 2017-08-31-210424_create_customers\n\
                Rolling back migration 00000000000000_diesel_initial_setup\n",
        "Unexpected stdout : {}",
        result.stdout()
    );

    assert!(!db.table_exists("customers"));
    assert!(!db.table_exists("contracts"));
    assert!(!db.table_exists("bills"));
    assert!(!db.table_exists("00000000000000_diesel_initial_setup"));
}

#[test]
fn migration_revert_with_zero_should_throw_an_error() {
    let p = project("migration_revert_with_zero_should_throw_an_error")
        .folder("migrations")
        .build();

    // Make sure the project is setup
    p.command("setup").run();

    // Should not revert any migration.
    let result = p
        .command("migration")
        .arg("revert")
        .arg("-n")
        .arg("0")
        .run();

    assert!(!result.is_success(), "Result was unsuccessful {:?}", result);

    assert!(
        result.stderr() == "error: Invalid value for '--number <REVERT_NUMBER>': <REVERT_NUMBER> must be at least equal to 1.\n",
        "Unexpected stderr : {}",
        result.stderr()
    );
}

#[test]
fn migration_revert_with_an_invalid_input_should_throw_an_error() {
    let p = project("migration_revert_with_an_invalid_input_should_throw_an_error")
        .folder("migrations")
        .build();

    // Make sure the project is setup
    p.command("setup").run();

    // Should not revert any migration.
    let result = p
        .command("migration")
        .arg("revert")
        .arg("-n")
        .arg("infinite")
        .run();

    assert!(!result.is_success(), "Result was unsuccessful {:?}", result);

    assert!(
        result.stderr() == "error: Invalid value for '--number <REVERT_NUMBER>': Cannot parse <REVERT_NUMBER>. The input must be an integer or 'all'.\n",
        "Unexpected stderr : {}",
        result.stderr()
    );
}
<commit_msg>Update tests for the revert feature<commit_after>use crate::support::{database, project};

#[test]
fn migration_revert_runs_the_last_migration_down() {
    let p = project("migration_revert_runs_the_last_migration_down")
        .folder("migrations")
        .build();
    let db = database(&p.database_url());

    p.create_migration(
        "12345_create_users_table",
        "CREATE TABLE users ( id INTEGER )",
        "DROP TABLE users",
    );

    // Make sure the project is setup
    p.command("setup").run();

    assert!(db.table_exists("users"));

    let result = p.command("migration").arg("revert").run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);
    assert!(
        result.stdout().contains("Rolling back migration 12345"),
        "Unexpected stdout {}",
        result.stdout()
    );
    assert!(!db.table_exists("users"));
}

#[test]
fn migration_revert_respects_migration_dir_var() {
    let p = project("migration_revert_var").folder("foo").build();
    let db = database(&p.database_url());

    p.create_migration_in_directory(
        "foo",
        "12345_create_users_table",
        "CREATE TABLE users ( id INTEGER )",
        "DROP TABLE users",
    );

    // Make sure the project is setup.
    p.command("setup").arg("--migration-dir=foo").run();

    assert!(db.table_exists("users"));

    let result = p
        .command("migration")
        .arg("revert")
        .arg("--migration-dir=foo")
        .run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);
    assert!(
        result.stdout().contains("Rolling back migration 12345"),
        "Unexpected stdout {}",
        result.stdout()
    );
    assert!(!db.table_exists("users"));
}

#[test]
fn migration_revert_respects_migration_dir_env() {
    let p = project("migration_revert_env").folder("bar").build();
    let db = database(&p.database_url());

    p.create_migration_in_directory(
        "bar",
        "12345_create_users_table",
        "CREATE TABLE users ( id INTEGER )",
        "DROP TABLE users",
    );

    // Make sure the project is setup.
    p.command("setup").arg("--migration-dir=bar").run();

    assert!(db.table_exists("users"));

    let result = p
        .command("migration")
        .arg("revert")
        .env("MIGRATION_DIRECTORY", "bar")
        .run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);
    assert!(
        result.stdout().contains("Rolling back migration 12345"),
        "Unexpected stdout {}",
        result.stdout()
    );
    assert!(!db.table_exists("users"));
}

#[test]
fn migration_revert_respects_migration_dir_from_diesel_toml() {
    let p = project("migration_revert_respects_migration_dir_from_diesel_toml")
        .folder("custom_migrations")
        .file(
            "diesel.toml",
            r#"
            [migrations_directory]
            dir = "custom_migrations"
            "#,
        )
        .build();

    let db = database(&p.database_url());

    p.create_migration_in_directory(
        "custom_migrations",
        "12345_create_users_table",
        "CREATE TABLE users ( id INTEGER )",
        "DROP TABLE users",
    );

    // Make sure the project is setup.
    p.command("setup").run();

    assert!(db.table_exists("users"));

    let result = p.command("migration").arg("revert").run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);
    assert!(
        result.stdout().contains("Rolling back migration 12345"),
        "Unexpected stdout {}",
        result.stdout()
    );
    assert!(!db.table_exists("users"));
}

#[test]
fn migration_revert_runs_the_last_two_migration_down() {
    let p = project("migration_revert_runs_the_last_two_migration_down")
        .folder("migrations")
        .build();
    let db = database(&p.database_url());

    p.create_migration(
        "2017-08-31-210424_create_customers",
        "CREATE TABLE customers ( id INTEGER PRIMARY KEY )",
        "DROP TABLE customers",
    );

    p.create_migration(
        "2017-09-03-210424_create_contracts",
        "CREATE TABLE contracts ( id INTEGER PRIMARY KEY )",
        "DROP TABLE contracts",
    );

    p.create_migration(
        "2017-09-12-210424_create_bills",
        "CREATE TABLE bills ( id INTEGER PRIMARY KEY )",
        "DROP TABLE bills",
    );

    // Make sure the project is setup
    p.command("setup").run();

    assert!(db.table_exists("customers"));
    assert!(db.table_exists("contracts"));
    assert!(db.table_exists("bills"));

    // Reverts the last two migration files. The `contracts` and `bills` tables should be dropped.
    // The `customers` table shouldn't be reverted.
    let result = p
        .command("migration")
        .arg("revert")
        .arg("-n")
        .arg("2")
        .run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);
    assert!(
        result.stdout().contains("Rolling back migration 2017-09-12-210424_create_bills\nRolling back migration 2017-09-03-210424_create_contracts"),
        "Unexpected stdout {}",
        result.stdout()
    );

    assert!(db.table_exists("customers"));
    assert!(!db.table_exists("contracts"));
    assert!(!db.table_exists("bills"));
}

#[test]
fn migration_revert_all_runs_the_migrations_down() {
    let p = project("migration_revert_all_runs_the_migrations_down")
        .folder("migrations")
        .build();
    let db = database(&p.database_url());

    p.create_migration(
        "2017-08-31-210424_create_customers",
        "CREATE TABLE customers ( id INTEGER PRIMARY KEY )",
        "DROP TABLE customers",
    );

    p.create_migration(
        "2017-09-03-210424_create_contracts",
        "CREATE TABLE contracts ( id INTEGER PRIMARY KEY )",
        "DROP TABLE contracts",
    );

    p.create_migration(
        "2017-09-12-210424_create_bills",
        "CREATE TABLE bills ( id INTEGER PRIMARY KEY )",
        "DROP TABLE bills",
    );

    // Make sure the project is setup
    p.command("setup").run();

    assert!(db.table_exists("customers"));
    assert!(db.table_exists("contracts"));
    assert!(db.table_exists("bills"));

    // Reverts all migrations. The migrations `contracts`, `bills`, `customers` and
    // `00000000000000_diesel_initial_setup` should be reverted.
    let result = p
        .command("migration")
        .arg("revert")
        .arg("-n")
        .arg("all")
        .run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);

    assert!(
        result.stdout()
            == "Rolling back migration 2017-09-12-210424_create_bills\n\
                Rolling back migration 2017-09-03-210424_create_contracts\n\
                Rolling back migration 2017-08-31-210424_create_customers\n\
                Rolling back migration 00000000000000_diesel_initial_setup\n",
        "Unexpected stdout : {}",
        result.stdout()
    );

    assert!(!db.table_exists("customers"));
    assert!(!db.table_exists("contracts"));
    assert!(!db.table_exists("bills"));
}

#[test]
fn migration_revert_with_zero_should_throw_an_error() {
    let p = project("migration_revert_with_zero_should_throw_an_error")
        .folder("migrations")
        .build();

    // Make sure the project is setup
    p.command("setup").run();

    // Should not revert any migration.
    let result = p
        .command("migration")
        .arg("revert")
        .arg("-n")
        .arg("0")
        .run();

    assert!(!result.is_success(), "Result was unsuccessful {:?}", result);

    assert!(
        result.stderr() == "error: Invalid value for '--number <REVERT_NUMBER>': <REVERT_NUMBER> must be at least equal to 1.\n",
        "Unexpected stderr : {}",
        result.stderr()
    );
}

#[test]
fn migration_revert_with_an_invalid_input_should_throw_an_error() {
    let p = project("migration_revert_with_an_invalid_input_should_throw_an_error")
        .folder("migrations")
        .build();

    // Make sure the project is setup
    p.command("setup").run();

    // Should not revert any migration.
    let result = p
        .command("migration")
        .arg("revert")
        .arg("-n")
        .arg("infinite")
        .run();

    assert!(!result.is_success(), "Result was unsuccessful {:?}", result);

    assert!(
        result.stderr() == "error: Invalid value for '--number <REVERT_NUMBER>': Cannot parse <REVERT_NUMBER>. The input must be an integer or 'all'.\n",
        "Unexpected stderr : {}",
        result.stderr()
    );
}

#[test]
fn migration_revert_with_more_than_the_number_of_migrations_should_run_all_the_migrations_down() {
    let p = project("migration_revert_with_more_than_the_number_of_migrations_should_run_all_the_migrations_down")
        .folder("migrations")
        .build();
    let db = database(&p.database_url());

    p.create_migration(
        "2017-08-31-210424_create_customers",
        "CREATE TABLE customers ( id INTEGER PRIMARY KEY )",
        "DROP TABLE customers",
    );

    p.create_migration(
        "2017-09-03-210424_create_contracts",
        "CREATE TABLE contracts ( id INTEGER PRIMARY KEY )",
        "DROP TABLE contracts",
    );

    p.create_migration(
        "2017-09-12-210424_create_bills",
        "CREATE TABLE bills ( id INTEGER PRIMARY KEY )",
        "DROP TABLE bills",
    );

    // Make sure the project is setup
    p.command("setup").run();

    assert!(db.table_exists("customers"));
    assert!(db.table_exists("contracts"));
    assert!(db.table_exists("bills"));

    // Reverts all migrations. The migrations `contracts`, `bills`, `customers` and
    // `00000000000000_diesel_initial_setup` should be reverted.
    let result = p
        .command("migration")
        .arg("revert")
        .arg("-n")
        .arg("1000")
        .run();

    assert!(result.is_success(), "Result was unsuccessful {:?}", result);

    assert!(
        result.stdout()
            == "Rolling back migration 2017-09-12-210424_create_bills\n\
                Rolling back migration 2017-09-03-210424_create_contracts\n\
                Rolling back migration 2017-08-31-210424_create_customers\n\
                Rolling back migration 00000000000000_diesel_initial_setup\n",
        "Unexpected stdout : {}",
        result.stdout()
    );

    assert!(!db.table_exists("customers"));
    assert!(!db.table_exists("contracts"));
    assert!(!db.table_exists("bills"));
}
<|endoftext|>


<commit_before>use diesel::*;
use helpers::*;

type Backend = <TestConnection as Connection>::Backend;

#[test]
fn simple_belongs_to() {
    table! {
        users {
            id -> Integer,
            name -> Text,
        }
    }

    table! {
        posts {
            id -> Integer,
            user_id -> Integer,
            title -> Text,
        }
    }

    allow_tables_to_appear_in_same_query!(users, posts);

    #[derive(Identifiable)]
    pub struct User {
        id: i32,
    }

    #[derive(Associations, Identifiable)]
    #[belongs_to(User)]
    pub struct Post {
        id: i32,
        user_id: i32,
    }

    joinable!(posts -> users(user_id));

    let _can_join_tables = posts::table
        .inner_join(users::table)
        .select((users::id, users::name, posts::id))
        .filter(
            posts::id
                .eq(1)
                .and(posts::user_id.eq(2))
                .and(posts::title.eq("Bar")),
        );

    let _can_reverse_join_tables = users::table
        .inner_join(posts::table)
        .select((posts::id, posts::user_id, posts::title))
        .filter(users::id.eq(1).and(users::name.eq("Sean")));

    let t = User { id: 42 };

    let belong_to = Post::belonging_to(&t);
    let filter = posts::table.filter(posts::user_id.eq(42));

    assert_eq!(
        debug_query::<Backend, _>(&belong_to).to_string(),
        debug_query::<Backend, _>(&filter).to_string()
    );
}

#[test]
fn custom_foreign_key() {
    table! {
        users {
            id -> Integer,
            name -> Text,
        }
    }

    table! {
        posts {
            id -> Integer,
            belongs_to_user -> Integer,
            title -> Text,
        }
    }

    allow_tables_to_appear_in_same_query!(users, posts);

    #[derive(Identifiable)]
    pub struct User {
        id: i32,
    }

    #[derive(Associations, Identifiable)]
    #[belongs_to(User, foreign_key = "belongs_to_user")]
    pub struct Post {
        id: i32,
        belongs_to_user: i32,
    }

    joinable!(posts -> users(belongs_to_user));

    let _can_join_tables = posts::table
        .inner_join(users::table)
        .select((users::id, users::name))
        .filter(
            posts::id
                .eq(1)
                .and(posts::belongs_to_user.eq(2))
                .and(posts::title.eq("Bar")),
        );

    let _can_reverse_join_tables = users::table
        .inner_join(posts::table)
        .select((posts::id, posts::belongs_to_user, posts::title))
        .filter(users::id.eq(1).and(users::name.eq("Sean")));

    let t = User { id: 42 };

    let belong_to = Post::belonging_to(&t);
    let filter = posts::table.filter(posts::belongs_to_user.eq(42));

    assert_eq!(
        debug_query::<Backend, _>(&belong_to).to_string(),
        debug_query::<Backend, _>(&filter).to_string()
    );
}

#[test]
fn self_referential() {
    table! {
        trees {
            id -> Integer,
            parent_id -> Nullable<Integer>,
        }
    }

    #[derive(Associations, Identifiable)]
    #[belongs_to(Tree, foreign_key = "parent_id")]
    pub struct Tree {
        id: i32,
        parent_id: Option<i32>,
    }
    let t = Tree {
        id: 42,
        parent_id: None,
    };

    let belong_to = Tree::belonging_to(&t);
    let filter = trees::table.filter(trees::parent_id.eq(42));
    assert_eq!(
        debug_query::<Backend, _>(&belong_to).to_string(),
        debug_query::<Backend, _>(&filter).to_string()
    );
}
<commit_msg>Add some missing test cases for `#[belongs_to]`<commit_after>use diesel::*;
use helpers::*;

type Backend = <TestConnection as Connection>::Backend;

#[test]
fn simple_belongs_to() {
    table! {
        users {
            id -> Integer,
            name -> Text,
        }
    }

    table! {
        posts {
            id -> Integer,
            user_id -> Integer,
            title -> Text,
        }
    }

    allow_tables_to_appear_in_same_query!(users, posts);

    #[derive(Identifiable)]
    pub struct User {
        id: i32,
    }

    #[derive(Associations, Identifiable)]
    #[belongs_to(User)]
    pub struct Post {
        id: i32,
        user_id: i32,
    }

    joinable!(posts -> users(user_id));

    let _can_join_tables = posts::table
        .inner_join(users::table)
        .select((users::id, users::name, posts::id))
        .filter(
            posts::id
                .eq(1)
                .and(posts::user_id.eq(2))
                .and(posts::title.eq("Bar")),
        );

    let _can_reverse_join_tables = users::table
        .inner_join(posts::table)
        .select((posts::id, posts::user_id, posts::title))
        .filter(users::id.eq(1).and(users::name.eq("Sean")));

    let t = User { id: 42 };

    let belong_to = Post::belonging_to(&t);
    let filter = posts::table.filter(posts::user_id.eq(42));

    assert_eq!(
        debug_query::<Backend, _>(&belong_to).to_string(),
        debug_query::<Backend, _>(&filter).to_string()
    );
}

#[test]
fn custom_foreign_key() {
    table! {
        users {
            id -> Integer,
            name -> Text,
        }
    }

    table! {
        posts {
            id -> Integer,
            belongs_to_user -> Integer,
            title -> Text,
        }
    }

    allow_tables_to_appear_in_same_query!(users, posts);

    #[derive(Identifiable)]
    pub struct User {
        id: i32,
    }

    #[derive(Associations, Identifiable)]
    #[belongs_to(User, foreign_key = "belongs_to_user")]
    pub struct Post {
        id: i32,
        belongs_to_user: i32,
    }

    joinable!(posts -> users(belongs_to_user));

    let _can_join_tables = posts::table
        .inner_join(users::table)
        .select((users::id, users::name))
        .filter(
            posts::id
                .eq(1)
                .and(posts::belongs_to_user.eq(2))
                .and(posts::title.eq("Bar")),
        );

    let _can_reverse_join_tables = users::table
        .inner_join(posts::table)
        .select((posts::id, posts::belongs_to_user, posts::title))
        .filter(users::id.eq(1).and(users::name.eq("Sean")));

    let t = User { id: 42 };

    let belong_to = Post::belonging_to(&t);
    let filter = posts::table.filter(posts::belongs_to_user.eq(42));

    assert_eq!(
        debug_query::<Backend, _>(&belong_to).to_string(),
        debug_query::<Backend, _>(&filter).to_string()
    );
}

#[test]
fn self_referential() {
    table! {
        trees {
            id -> Integer,
            parent_id -> Nullable<Integer>,
        }
    }

    #[derive(Associations, Identifiable)]
    #[belongs_to(Tree, foreign_key = "parent_id")]
    pub struct Tree {
        id: i32,
        parent_id: Option<i32>,
    }
    let t = Tree {
        id: 42,
        parent_id: None,
    };

    let belong_to = Tree::belonging_to(&t);
    let filter = trees::table.filter(trees::parent_id.eq(42));
    assert_eq!(
        debug_query::<Backend, _>(&belong_to).to_string(),
        debug_query::<Backend, _>(&filter).to_string()
    );
}

#[test]
fn multiple_associations() {
    table! {
        users {
            id -> Integer,
        }
    }

    table! {
        posts {
            id -> Integer,
        }
    }

    table! {
        comments {
            id -> Integer,
            user_id -> Integer,
            post_id -> Integer,
        }
    }

    #[derive(Identifiable)]
    struct User {
        id: i32,
    }

    #[derive(Identifiable)]
    struct Post {
        id: i32,
    }

    #[derive(Identifiable, Associations)]
    #[belongs_to(User)]
    #[belongs_to(Post)]
    struct Comment {
        id: i32,
        user_id: i32,
        post_id: i32,
    }

    let user = User { id: 1 };
    let post = Post { id: 2 };

    let query = Comment::belonging_to(&user);
    let expected = comments::table.filter(comments::user_id.eq(1));
    assert_eq!(
        debug_query::<Backend, _>(&query).to_string(),
        debug_query::<Backend, _>(&expected).to_string()
    );
    let query = Comment::belonging_to(&post);
    let expected = comments::table.filter(comments::post_id.eq(2));
    assert_eq!(
        debug_query::<Backend, _>(&query).to_string(),
        debug_query::<Backend, _>(&expected).to_string()
    );
}

#[test]
fn foreign_key_field_with_column_rename() {
    table! {
        users {
            id -> Integer,
        }
    }

    table! {
        posts {
            id -> Integer,
            user_id -> Integer,
        }
    }

    #[derive(Identifiable, Clone, Copy)]
    pub struct User {
        id: i32,
    }

    #[derive(Associations, Identifiable, Clone, Copy, PartialEq, Debug)]
    #[belongs_to(User)]
    pub struct Post {
        id: i32,
        #[column_name = "user_id"]
        author_id: i32,
    }

    let user1 = User { id: 1 };
    let user2 = User { id: 2 };
    let post1 = Post {
        id: 1,
        author_id: 2,
    };
    let post2 = Post {
        id: 2,
        author_id: 1,
    };

    let query = Post::belonging_to(&user1);
    let expected = posts::table.filter(posts::user_id.eq(1));
    assert_eq!(
        debug_query::<Backend, _>(&query).to_string(),
        debug_query::<Backend, _>(&expected).to_string()
    );

    let users = vec![user1, user2];
    let posts = vec![post1, post2].grouped_by(&users);
    assert_eq!(vec![vec![post2], vec![post1]], posts);
}

#[test]
fn tuple_struct() {
    table! {
        users {
            id -> Integer,
        }
    }

    table! {
        posts {
            id -> Integer,
            user_id -> Integer,
        }
    }

    #[derive(Identifiable)]
    pub struct User {
        id: i32,
    }

    #[derive(Associations, Identifiable)]
    #[belongs_to(User)]
    pub struct Post(#[column_name = "id"] i32, #[column_name = "user_id"] i32);

    let user = User { id: 1 };

    let query = Post::belonging_to(&user);
    let expected = posts::table.filter(posts::user_id.eq(1));
    assert_eq!(
        debug_query::<Backend, _>(&query).to_string(),
        debug_query::<Backend, _>(&expected).to_string()
    );
}
<|endoftext|>


<commit_before>//! `cpuid` intrinsics
#![allow(clippy::module_name_repetitions)]

use crate::arch::asm;
#[cfg(test)]
use stdarch_test::assert_instr;

/// Result of the `cpuid` instruction.
#[allow(clippy::missing_inline_in_public_items)]
// ^^ the derived impl of Debug for CpuidResult is not #[inline] and that's OK.
#[derive(Copy, Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
#[stable(feature = "simd_x86", since = "1.27.0")]
pub struct CpuidResult {
    /// EAX register.
    #[stable(feature = "simd_x86", since = "1.27.0")]
    pub eax: u32,
    /// EBX register.
    #[stable(feature = "simd_x86", since = "1.27.0")]
    pub ebx: u32,
    /// ECX register.
    #[stable(feature = "simd_x86", since = "1.27.0")]
    pub ecx: u32,
    /// EDX register.
    #[stable(feature = "simd_x86", since = "1.27.0")]
    pub edx: u32,
}

/// Returns the result of the `cpuid` instruction for a given `leaf` (`EAX`)
/// and
/// `sub_leaf` (`ECX`).
///
/// The highest-supported leaf value is returned by the first tuple argument of
/// [`__get_cpuid_max(0)`](fn.__get_cpuid_max.html). For leaves containung
/// sub-leaves, the second tuple argument returns the highest-supported
/// sub-leaf
/// value.
///
/// The [CPUID Wikipedia page][wiki_cpuid] contains how to query which
/// information using the `EAX` and `ECX` registers, and the interpretation of
/// the results returned in `EAX`, `EBX`, `ECX`, and `EDX`.
///
/// The references are:
/// - [Intel 64 and IA-32 Architectures Software Developer's Manual Volume 2:
///   Instruction Set Reference, A-Z][intel64_ref].
/// - [AMD64 Architecture Programmer's Manual, Volume 3: General-Purpose and
///   System Instructions][amd64_ref].
///
/// [wiki_cpuid]: https://en.wikipedia.org/wiki/CPUID
/// [intel64_ref]: http://www.intel.de/content/dam/www/public/us/en/documents/manuals/64-ia-32-architectures-software-developer-instruction-set-reference-manual-325383.pdf
/// [amd64_ref]: http://support.amd.com/TechDocs/24594.pdf
#[inline]
#[cfg_attr(test, assert_instr(cpuid))]
#[stable(feature = "simd_x86", since = "1.27.0")]
pub unsafe fn __cpuid_count(leaf: u32, sub_leaf: u32) -> CpuidResult {
    let eax;
    let ebx;
    let ecx;
    let edx;

    // LLVM sometimes reserves `ebx` for its internal use, we so we need to use
    // a scratch register for it instead.
    #[cfg(target_arch = "x86")]
    {
        asm!(
            "movl {0}, ebx",
            "cpuid",
            "xchgl {0}, ebx",
            out(reg) ebx,
            inout("eax") leaf => eax,
            inout("ecx") sub_leaf => ecx,
            out("edx") edx,
            options(nostack, preserves_flags),
        );
    }
    #[cfg(target_arch = "x86_64")]
    {
        asm!(
            "movq {0:r}, rbx",
            "cpuid",
            "xchgq {0:r}, rbx",
            out(reg) ebx,
            inout("eax") leaf => eax,
            inout("ecx") sub_leaf => ecx,
            out("edx") edx,
            options(nostack, preserves_flags),
        );
    }
    CpuidResult { eax, ebx, ecx, edx }
}

/// See [`__cpuid_count`](fn.__cpuid_count.html).
#[inline]
#[cfg_attr(test, assert_instr(cpuid))]
#[stable(feature = "simd_x86", since = "1.27.0")]
pub unsafe fn __cpuid(leaf: u32) -> CpuidResult {
    __cpuid_count(leaf, 0)
}

/// Does the host support the `cpuid` instruction?
#[inline]
pub fn has_cpuid() -> bool {
    #[cfg(target_env = "sgx")]
    {
        false
    }
    #[cfg(all(not(target_env = "sgx"), target_arch = "x86_64"))]
    {
        true
    }
    #[cfg(all(not(target_env = "sgx"), target_arch = "x86"))]
    {
        // Optimization for i586 and i686 Rust targets which SSE enabled
        // and support cpuid:
        #[cfg(target_feature = "sse")]
        {
            true
        }

        // If SSE is not enabled, detect whether cpuid is available:
        #[cfg(not(target_feature = "sse"))]
        unsafe {
            // On `x86` the `cpuid` instruction is not always available.
            // This follows the approach indicated in:
            // http://wiki.osdev.org/CPUID#Checking_CPUID_availability
            // https://software.intel.com/en-us/articles/using-cpuid-to-detect-the-presence-of-sse-41-and-sse-42-instruction-sets/
            // which detects whether `cpuid` is available by checking whether
            // the 21st bit of the EFLAGS register is modifiable or not.
            // If it is, then `cpuid` is available.
            let result: u32;
            asm!(
                // Read eflags and save a copy of it
                "pushfd",
                "pop {result}",
                "mov {result}, {saved_flags}",
                // Flip 21st bit of the flags
                "xor $0x200000, {result}",
                // Load the modified flags and read them back.
                // Bit 21 can only be modified if cpuid is available.
                "push {result}",
                "popfd",
                "pushfd",
                "pop {result}",
                // Use xor to find out whether bit 21 has changed
                "xor {saved_flags}, {result}",
                result = out(reg) result,
                saved_flags = out(reg) _,
                options(nomem, att_syntax),
            );
            // There is a race between popfd (A) and pushfd (B)
            // where other bits beyond 21st may have been modified due to
            // interrupts, a debugger stepping through the asm, etc.
            //
            // Therefore, explicitly check whether the 21st bit
            // was modified or not.
            //
            // If the result is zero, the cpuid bit was not modified.
            // If the result is `0x200000` (non-zero), then the cpuid
            // was correctly modified and the CPU supports the cpuid
            // instruction:
            (result & 0x200000) != 0
        }
    }
}

/// Returns the highest-supported `leaf` (`EAX`) and sub-leaf (`ECX`) `cpuid`
/// values.
///
/// If `cpuid` is supported, and `leaf` is zero, then the first tuple argument
/// contains the highest `leaf` value that `cpuid` supports. For `leaf`s
/// containing sub-leafs, the second tuple argument contains the
/// highest-supported sub-leaf value.
///
/// See also [`__cpuid`](fn.__cpuid.html) and
/// [`__cpuid_count`](fn.__cpuid_count.html).
#[inline]
#[stable(feature = "simd_x86", since = "1.27.0")]
pub unsafe fn __get_cpuid_max(leaf: u32) -> (u32, u32) {
    let CpuidResult { eax, ebx, .. } = __cpuid(leaf);
    (eax, ebx)
}

#[cfg(test)]
mod tests {
    use crate::core_arch::x86::*;

    #[test]
    fn test_always_has_cpuid() {
        // all currently-tested targets have the instruction
        // FIXME: add targets without `cpuid` to CI
        assert!(cpuid::has_cpuid());
    }

    #[test]
    fn test_has_cpuid_idempotent() {
        assert_eq!(cpuid::has_cpuid(), cpuid::has_cpuid());
    }
}
<commit_msg>Use mov and xchg instead of movl(q) and xchgl(q)<commit_after>//! `cpuid` intrinsics
#![allow(clippy::module_name_repetitions)]

use crate::arch::asm;
#[cfg(test)]
use stdarch_test::assert_instr;

/// Result of the `cpuid` instruction.
#[allow(clippy::missing_inline_in_public_items)]
// ^^ the derived impl of Debug for CpuidResult is not #[inline] and that's OK.
#[derive(Copy, Clone, Debug, Eq, Ord, PartialEq, PartialOrd)]
#[stable(feature = "simd_x86", since = "1.27.0")]
pub struct CpuidResult {
    /// EAX register.
    #[stable(feature = "simd_x86", since = "1.27.0")]
    pub eax: u32,
    /// EBX register.
    #[stable(feature = "simd_x86", since = "1.27.0")]
    pub ebx: u32,
    /// ECX register.
    #[stable(feature = "simd_x86", since = "1.27.0")]
    pub ecx: u32,
    /// EDX register.
    #[stable(feature = "simd_x86", since = "1.27.0")]
    pub edx: u32,
}

/// Returns the result of the `cpuid` instruction for a given `leaf` (`EAX`)
/// and
/// `sub_leaf` (`ECX`).
///
/// The highest-supported leaf value is returned by the first tuple argument of
/// [`__get_cpuid_max(0)`](fn.__get_cpuid_max.html). For leaves containung
/// sub-leaves, the second tuple argument returns the highest-supported
/// sub-leaf
/// value.
///
/// The [CPUID Wikipedia page][wiki_cpuid] contains how to query which
/// information using the `EAX` and `ECX` registers, and the interpretation of
/// the results returned in `EAX`, `EBX`, `ECX`, and `EDX`.
///
/// The references are:
/// - [Intel 64 and IA-32 Architectures Software Developer's Manual Volume 2:
///   Instruction Set Reference, A-Z][intel64_ref].
/// - [AMD64 Architecture Programmer's Manual, Volume 3: General-Purpose and
///   System Instructions][amd64_ref].
///
/// [wiki_cpuid]: https://en.wikipedia.org/wiki/CPUID
/// [intel64_ref]: http://www.intel.de/content/dam/www/public/us/en/documents/manuals/64-ia-32-architectures-software-developer-instruction-set-reference-manual-325383.pdf
/// [amd64_ref]: http://support.amd.com/TechDocs/24594.pdf
#[inline]
#[cfg_attr(test, assert_instr(cpuid))]
#[stable(feature = "simd_x86", since = "1.27.0")]
pub unsafe fn __cpuid_count(leaf: u32, sub_leaf: u32) -> CpuidResult {
    let eax;
    let ebx;
    let ecx;
    let edx;

    // LLVM sometimes reserves `ebx` for its internal use, we so we need to use
    // a scratch register for it instead.
    #[cfg(target_arch = "x86")]
    {
        asm!(
            "mov {0}, ebx",
            "cpuid",
            "xchg {0}, ebx",
            out(reg) ebx,
            inout("eax") leaf => eax,
            inout("ecx") sub_leaf => ecx,
            out("edx") edx,
            options(nostack, preserves_flags),
        );
    }
    #[cfg(target_arch = "x86_64")]
    {
        asm!(
            "mov {0:r}, rbx",
            "cpuid",
            "xchg {0:r}, rbx",
            out(reg) ebx,
            inout("eax") leaf => eax,
            inout("ecx") sub_leaf => ecx,
            out("edx") edx,
            options(nostack, preserves_flags),
        );
    }
    CpuidResult { eax, ebx, ecx, edx }
}

/// See [`__cpuid_count`](fn.__cpuid_count.html).
#[inline]
#[cfg_attr(test, assert_instr(cpuid))]
#[stable(feature = "simd_x86", since = "1.27.0")]
pub unsafe fn __cpuid(leaf: u32) -> CpuidResult {
    __cpuid_count(leaf, 0)
}

/// Does the host support the `cpuid` instruction?
#[inline]
pub fn has_cpuid() -> bool {
    #[cfg(target_env = "sgx")]
    {
        false
    }
    #[cfg(all(not(target_env = "sgx"), target_arch = "x86_64"))]
    {
        true
    }
    #[cfg(all(not(target_env = "sgx"), target_arch = "x86"))]
    {
        // Optimization for i586 and i686 Rust targets which SSE enabled
        // and support cpuid:
        #[cfg(target_feature = "sse")]
        {
            true
        }

        // If SSE is not enabled, detect whether cpuid is available:
        #[cfg(not(target_feature = "sse"))]
        unsafe {
            // On `x86` the `cpuid` instruction is not always available.
            // This follows the approach indicated in:
            // http://wiki.osdev.org/CPUID#Checking_CPUID_availability
            // https://software.intel.com/en-us/articles/using-cpuid-to-detect-the-presence-of-sse-41-and-sse-42-instruction-sets/
            // which detects whether `cpuid` is available by checking whether
            // the 21st bit of the EFLAGS register is modifiable or not.
            // If it is, then `cpuid` is available.
            let result: u32;
            asm!(
                // Read eflags and save a copy of it
                "pushfd",
                "pop {result}",
                "mov {result}, {saved_flags}",
                // Flip 21st bit of the flags
                "xor $0x200000, {result}",
                // Load the modified flags and read them back.
                // Bit 21 can only be modified if cpuid is available.
                "push {result}",
                "popfd",
                "pushfd",
                "pop {result}",
                // Use xor to find out whether bit 21 has changed
                "xor {saved_flags}, {result}",
                result = out(reg) result,
                saved_flags = out(reg) _,
                options(nomem, att_syntax),
            );
            // There is a race between popfd (A) and pushfd (B)
            // where other bits beyond 21st may have been modified due to
            // interrupts, a debugger stepping through the asm, etc.
            //
            // Therefore, explicitly check whether the 21st bit
            // was modified or not.
            //
            // If the result is zero, the cpuid bit was not modified.
            // If the result is `0x200000` (non-zero), then the cpuid
            // was correctly modified and the CPU supports the cpuid
            // instruction:
            (result & 0x200000) != 0
        }
    }
}

/// Returns the highest-supported `leaf` (`EAX`) and sub-leaf (`ECX`) `cpuid`
/// values.
///
/// If `cpuid` is supported, and `leaf` is zero, then the first tuple argument
/// contains the highest `leaf` value that `cpuid` supports. For `leaf`s
/// containing sub-leafs, the second tuple argument contains the
/// highest-supported sub-leaf value.
///
/// See also [`__cpuid`](fn.__cpuid.html) and
/// [`__cpuid_count`](fn.__cpuid_count.html).
#[inline]
#[stable(feature = "simd_x86", since = "1.27.0")]
pub unsafe fn __get_cpuid_max(leaf: u32) -> (u32, u32) {
    let CpuidResult { eax, ebx, .. } = __cpuid(leaf);
    (eax, ebx)
}

#[cfg(test)]
mod tests {
    use crate::core_arch::x86::*;

    #[test]
    fn test_always_has_cpuid() {
        // all currently-tested targets have the instruction
        // FIXME: add targets without `cpuid` to CI
        assert!(cpuid::has_cpuid());
    }

    #[test]
    fn test_has_cpuid_idempotent() {
        assert_eq!(cpuid::has_cpuid(), cpuid::has_cpuid());
    }
}
<|endoftext|>


<commit_before>use actix::prelude::*;
use godcoin::prelude::*;

mod common;

pub use common::*;

#[test]
fn fresh_blockchain() {
    System::run(|| {
        let minter = TestMinter::new();
        let chain = minter.chain();
        assert!(chain.get_block(0).is_some());
        assert_eq!(chain.get_chain_height(), 0);

        let owner = chain.get_owner();
        assert_eq!(owner.minter, minter.genesis_info().minter_key.0);
        assert_eq!(
            owner.script,
            script::Builder::new().push(OpFrame::False).build()
        );
        assert_eq!(owner.wallet, (&minter.genesis_info().script).into());

        assert!(chain.get_block(1).is_none());
        System::current().stop();
    })
    .unwrap();
}

#[test]
fn mint_tx_verification() {
    System::run(|| {
        let minter = TestMinter::new();
        let chain = minter.chain();
        let config = verify::Config::strict();

        let create_tx = |fee: &str| {
            let mut tx = MintTx {
                base: create_tx(TxType::MINT, fee),
                to: (&minter.genesis_info().script).into(),
                amount: Balance::default(),
                script: minter.genesis_info().script.clone(),
            };
            tx.append_sign(&minter.genesis_info().wallet_keys[3]);
            tx.append_sign(&minter.genesis_info().wallet_keys[0]);
            TxVariant::MintTx(tx)
        };

        let tx = create_tx("0 GOLD");
        assert!(chain.verify_tx(&tx, &[], config).is_ok());

        let tx = create_tx("1 GOLD");
        assert_eq!(
            chain.verify_tx(&tx, &[], config).unwrap_err(),
            verify::TxErr::InsufficientFeeAmount
        );

        let mut tx = create_tx("0 GOLD");
        tx.signature_pairs.remove(1);
        assert!(check_sigs(&tx));
        assert_eq!(
            chain.verify_tx(&tx, &[], config).unwrap_err(),
            verify::TxErr::ScriptRetFalse
        );

        let mut tx = create_tx("0 GOLD");
        tx.signature_pairs.clear();
        assert!(check_sigs(&tx));
        assert_eq!(
            chain.verify_tx(&tx, &[], config).unwrap_err(),
            verify::TxErr::ScriptRetFalse
        );

        System::current().stop();
    })
    .unwrap();
}

#[test]
fn mint_tx_updates_balances() {
    System::run(|| {
        let minter = TestMinter::new();

        let mut tx = MintTx {
            base: create_tx(TxType::MINT, "0 GOLD"),
            to: (&minter.genesis_info().script).into(),
            amount: get_balance("10.0 GOLD", "1000 SILVER"),
            script: minter.genesis_info().script.clone(),
        };

        tx.append_sign(&minter.genesis_info().wallet_keys[1]);
        tx.append_sign(&minter.genesis_info().wallet_keys[0]);

        let tx = TxVariant::MintTx(tx);
        let fut = minter.request(MsgRequest::Broadcast(tx));
        System::current().arbiter().send(
            fut.then(move |res| {
                let res = res.unwrap();
                assert!(!res.is_err(), format!("{:?}", res));

                minter.produce_block().map(|_| minter)
            })
            .then(|res| {
                let minter = res.unwrap();
                let chain = minter.chain();
                let props = chain.get_properties();
                let expected_bal = get_balance("10 GOLD", "1000 SILVER");
                assert_eq!(props.token_supply, expected_bal);

                let bal = chain.get_balance(&(&minter.genesis_info().script).into());
                assert_eq!(bal, expected_bal);

                System::current().stop();
                Ok(())
            }),
        );
    })
    .unwrap();
}
<commit_msg>Add test with invalid signature for mint tx<commit_after>use actix::prelude::*;
use godcoin::{prelude::*, crypto::Signature};

mod common;

pub use common::*;

#[test]
fn fresh_blockchain() {
    System::run(|| {
        let minter = TestMinter::new();
        let chain = minter.chain();
        assert!(chain.get_block(0).is_some());
        assert_eq!(chain.get_chain_height(), 0);

        let owner = chain.get_owner();
        assert_eq!(owner.minter, minter.genesis_info().minter_key.0);
        assert_eq!(
            owner.script,
            script::Builder::new().push(OpFrame::False).build()
        );
        assert_eq!(owner.wallet, (&minter.genesis_info().script).into());

        assert!(chain.get_block(1).is_none());
        System::current().stop();
    })
    .unwrap();
}

#[test]
fn mint_tx_verification() {
    System::run(|| {
        let minter = TestMinter::new();
        let chain = minter.chain();
        let config = verify::Config::strict();

        let create_tx = |fee: &str| {
            let mut tx = MintTx {
                base: create_tx(TxType::MINT, fee),
                to: (&minter.genesis_info().script).into(),
                amount: Balance::default(),
                script: minter.genesis_info().script.clone(),
            };
            tx.append_sign(&minter.genesis_info().wallet_keys[3]);
            tx.append_sign(&minter.genesis_info().wallet_keys[0]);
            TxVariant::MintTx(tx)
        };

        let tx = create_tx("0 GOLD");
        assert!(chain.verify_tx(&tx, &[], config).is_ok());

        let tx = create_tx("1 GOLD");
        assert_eq!(
            chain.verify_tx(&tx, &[], config).unwrap_err(),
            verify::TxErr::InsufficientFeeAmount
        );

        let mut tx = create_tx("0 GOLD");
        tx.signature_pairs.remove(1);
        assert!(check_sigs(&tx));
        assert_eq!(
            chain.verify_tx(&tx, &[], config).unwrap_err(),
            verify::TxErr::ScriptRetFalse
        );

        let mut tx = create_tx("0 GOLD");
        tx.signature_pairs.clear();
        assert!(check_sigs(&tx));
        assert_eq!(
            chain.verify_tx(&tx, &[], config).unwrap_err(),
            verify::TxErr::ScriptRetFalse
        );

        let mut tx = create_tx("0 GOLD");
        tx.signature_pairs.clear();
        tx.signature_pairs.push(SigPair {
            pub_key: minter.genesis_info().wallet_keys[0].0.clone(),
            signature: Signature::from_slice(&[0; 64]).unwrap(),
        });
        assert!(!check_sigs(&tx));
        assert_eq!(
            chain.verify_tx(&tx, &[], config).unwrap_err(),
            verify::TxErr::ScriptRetFalse
        );

        System::current().stop();
    })
    .unwrap();
}

#[test]
fn mint_tx_updates_balances() {
    System::run(|| {
        let minter = TestMinter::new();

        let mut tx = MintTx {
            base: create_tx(TxType::MINT, "0 GOLD"),
            to: (&minter.genesis_info().script).into(),
            amount: get_balance("10.0 GOLD", "1000 SILVER"),
            script: minter.genesis_info().script.clone(),
        };

        tx.append_sign(&minter.genesis_info().wallet_keys[1]);
        tx.append_sign(&minter.genesis_info().wallet_keys[0]);

        let tx = TxVariant::MintTx(tx);
        let fut = minter.request(MsgRequest::Broadcast(tx));
        System::current().arbiter().send(
            fut.then(move |res| {
                let res = res.unwrap();
                assert!(!res.is_err(), format!("{:?}", res));

                minter.produce_block().map(|_| minter)
            })
            .then(|res| {
                let minter = res.unwrap();
                let chain = minter.chain();
                let props = chain.get_properties();
                let expected_bal = get_balance("10 GOLD", "1000 SILVER");
                assert_eq!(props.token_supply, expected_bal);

                let bal = chain.get_balance(&(&minter.genesis_info().script).into());
                assert_eq!(bal, expected_bal);

                System::current().stop();
                Ok(())
            }),
        );
    })
    .unwrap();
}
<|endoftext|>