<commit_before>//! Defines structs and methods that describe a bindings.

use std::collections::HashSet;
use std::fmt;
use std::ops::{Index, IndexMut};
use std::hash::Hash;

pub trait RBind {
    type SSARef: fmt::Debug + Clone;

    fn is_argument(&self) -> bool;
    fn is_local(&self) -> bool;
    fn is_modified(&self) -> bool;
    fn is_preserved(&self) -> bool;
    fn is_return(&self) -> bool;

    fn mark_argument(&mut self);
    fn mark_modified(&mut self);
    fn mark_preserved(&mut self);
    fn mark_return(&mut self);

    fn name(&self) -> String;

    fn is_register(&self) -> bool;
    fn mark_register(&mut self, String);

    fn is_memory(&self) -> bool {
        !self.is_register() && !self.is_unknown()
    }
    fn mark_memory(&mut self);

    fn is_fn_local(&self) -> bool;
    fn mark_fn_local(&mut self, usize, i64);
    fn local_info(&self) -> LocalInfo;

    fn is_stack(&self) -> bool;
    fn mark_stack(&mut self);

    fn is_global(&self) -> bool;
    fn mark_global(&mut self, u64);
    fn global_offset(&self) -> u64;

    fn is_unknown(&self) -> bool;

    fn add_refs(&mut self, Vec<Self::SSARef>);
    fn refs(&self) -> ::std::slice::Iter<Self::SSARef>;
}

/// Trait that describes variable bindings for a function.
pub trait RBindings {
    type BTy: RBind;
    type Idx: Clone + fmt::Debug + Eq + Ord + Hash + From<usize>;

    fn new() -> Self;
    fn insert(&mut self, Self::BTy) -> Self::Idx;

    fn binding(&self, &Self::Idx) -> Option<&Self::BTy>;
    fn binding_mut(&mut self, &Self::Idx) -> Option<&mut Self::BTy>;

    fn bindings(&self) -> RBinds<Self::BTy>;
    fn bindings_mut(&mut self) -> RBindsMut<Self::BTy>;
}

#[derive(Debug)]
pub struct RBinds<'a, T: 'a + RBind> {
    binds: ::std::slice::Iter<'a, T>,
}

impl<'a, T: 'a + RBind> Iterator for RBinds<'a, T> {
    type Item = &'a T;
    fn next(&mut self) -> Option<&'a T> {
        self.binds.next()
    }
}

#[derive(Debug)]
pub struct RBindsMut<'a, T: 'a + RBind> {
    binds: ::std::slice::IterMut<'a, T>,
}

#[derive(Clone, Debug)]
pub struct RadecoBindings<T: RBind> {
    binds: Vec<T>,
}

impl<BTy: RBind> Index<usize> for RadecoBindings<BTy> {
    type Output = BTy;
    fn index(&self, index: usize) -> &Self::Output {
        self.binding(&index).unwrap()
    }
}

impl<BTy: RBind> IndexMut<usize> for RadecoBindings<BTy> {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        self.binding_mut(&index).unwrap()
    }
}

impl<BTy: RBind> RBindings for RadecoBindings<BTy> {
    type BTy = BTy;
    type Idx = usize;

    fn new() -> Self {
        RadecoBindings { binds: Vec::new() }
    }

    fn insert(&mut self, bind: BTy) -> usize {
        self.binds.push(bind);
        self.binds.len() - 1
    }

    fn binding(&self, idx: &usize) -> Option<&BTy> {
        self.binds.get(*idx)
    }

    fn binding_mut(&mut self, idx: &usize) -> Option<&mut BTy> {
        self.binds.get_mut(*idx)
    }

    fn bindings(&self) -> RBinds<BTy> {
        RBinds { binds: self.binds.iter() }
    }

    fn bindings_mut(&mut self) -> RBindsMut<BTy> {
        RBindsMut { binds: self.binds.iter_mut() }
    }
}

#[derive(Clone, Debug)]
pub struct Binding<T: fmt::Debug + Clone> {
    // Optional name used to represent this binding.
    pub named: Option<String>,
    // Type of binding.
    vloc: VarLocation,
    dty: VarDataType,
    // Node indices in the function ssa that corresponds to this binding.
    ssa_refs: Vec<T>,
    set: HashSet<VarSet>,
}

impl<'a, T: 'a + RBind> Iterator for RBindsMut<'a, T> {
    type Item = &'a mut T;
    fn next(&mut self) -> Option<&'a mut T> {
        self.binds.next()
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum VarLocation {
    Register {
        name: String,
    },
    Memory(MemoryRegion),
    Unknown,
}

#[derive(Clone, Copy, Debug, PartialEq)]
pub enum MemoryRegion {
    // Base, offset for function local.
    FunctionLocal {
        base: usize,
        offset: i64,
    },
    Stack,
    Global {
        offset: u64,
    },
    Heap,
    Unknown,
}

// Represents the set to which a variable belongs to in the call summary
// information.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
enum VarSet {
    Argument,
    Local,
    Modified,
    Preserved,
    Returned,
}

#[derive(Clone, Copy, Debug)]
pub enum VarDataType {
    Unknown(u64),
    Int,
    IntRef,
    Char,
    CharRef,
    Float,
    FloatRef,
    Double,
    DoubleRef,
}

#[derive(Clone, Debug)]
pub struct LocalInfo {
    pub base: usize,
    pub offset: i64,
}

impl<T: Clone + fmt::Debug> Default for Binding<T> {
    fn default() -> Binding<T> {
        Binding {
            named: None,
            vloc: VarLocation::Unknown,
            dty: VarDataType::Unknown(64),
            ssa_refs: Vec::new(),
            set: HashSet::new(),
        }
    }
}

impl<T: Clone + fmt::Debug> RBind for Binding<T> {
    type SSARef = T;

    fn is_argument(&self) -> bool {
        self.set.contains(&VarSet::Argument)
    }

    fn is_local(&self) -> bool {
        self.set.contains(&VarSet::Local)
    }

    fn is_modified(&self) -> bool {
        self.set.contains(&VarSet::Modified)
    }

    fn is_preserved(&self) -> bool {
        self.set.contains(&VarSet::Preserved)
    }

    fn is_return(&self) -> bool {
        self.set.contains(&VarSet::Returned)
    }

    fn mark_argument(&mut self) {
        self.set.insert(VarSet::Argument);
    }

    fn mark_modified(&mut self) {
        self.set.insert(VarSet::Modified);
    }

    fn mark_preserved(&mut self) {
        self.set.insert(VarSet::Preserved);
    }

    fn mark_return(&mut self) {
        self.set.insert(VarSet::Returned);
    }

    fn name(&self) -> String {
        self.named.clone().unwrap_or_else(String::new)
    }

    fn is_register(&self) -> bool {
        match self.vloc {
            VarLocation::Register { .. } => true,
            _ => false,
        }
    }

    fn mark_register(&mut self, name: String) {
        self.vloc = VarLocation::Register { name: name };
    }

    fn mark_memory(&mut self) {
        self.vloc = VarLocation::Memory(MemoryRegion::Unknown);
    }

    fn is_fn_local(&self) -> bool {
        match self.vloc {
            VarLocation::Memory(MemoryRegion::FunctionLocal { .. }) => true,
            _ => false,
        }
    }

    fn mark_fn_local(&mut self, base: usize, offset: i64) {
        self.vloc = VarLocation::Memory(MemoryRegion::FunctionLocal {
            base: base,
            offset: offset,
        });
    }

    fn local_info(&self) -> LocalInfo {
        if let VarLocation::Memory(MemoryRegion::FunctionLocal { base, offset }) = self.vloc {
            LocalInfo {
                base: base,
                offset: offset,
            }
        } else {
            panic!()
        }
    }

    fn is_stack(&self) -> bool {
        if let VarLocation::Memory(MemoryRegion::Stack) = self.vloc {
            true
        } else {
            false
        }
    }

    fn mark_stack(&mut self) {
        self.vloc = VarLocation::Memory(MemoryRegion::Stack);
    }

    fn is_global(&self) -> bool {
        match self.vloc {
            VarLocation::Memory(MemoryRegion::Global { offset }) => true,
            _ => false,
        }
    }

    fn mark_global(&mut self, offset: u64) {
        self.vloc = VarLocation::Memory(MemoryRegion::Global { offset: offset });
    }

    fn global_offset(&self) -> u64 {
        if let VarLocation::Memory(MemoryRegion::Global { offset }) = self.vloc {
            offset
        } else {
            panic!()
        }
    }

    fn is_unknown(&self) -> bool {
        match self.vloc {
            VarLocation::Unknown => true,
            _ => false,
        }
    }

    fn add_refs(&mut self, refs: Vec<T>) {
        self.ssa_refs.extend(refs);
    }

    fn refs(&self) -> ::std::slice::Iter<T> {
        self.ssa_refs.iter()
    }
}
<commit_msg>Add API for binding's name<commit_after>//! Defines structs and methods that describe a bindings.

use std::collections::HashSet;
use std::fmt;
use std::ops::{Index, IndexMut};
use std::hash::Hash;

pub trait RBind {
    type SSARef: fmt::Debug + Clone;

    fn is_argument(&self) -> bool;
    fn is_local(&self) -> bool;
    fn is_modified(&self) -> bool;
    fn is_preserved(&self) -> bool;
    fn is_return(&self) -> bool;

    fn mark_argument(&mut self);
    fn mark_modified(&mut self);
    fn mark_preserved(&mut self);
    fn mark_return(&mut self);

    fn name(&self) -> String;
    fn set_name(&mut self, String);

    fn is_register(&self) -> bool;
    fn mark_register(&mut self, String);

    fn is_memory(&self) -> bool {
        !self.is_register() && !self.is_unknown()
    }
    fn mark_memory(&mut self);

    fn is_fn_local(&self) -> bool;
    fn mark_fn_local(&mut self, usize, i64);
    fn local_info(&self) -> LocalInfo;

    fn is_stack(&self) -> bool;
    fn mark_stack(&mut self);

    fn is_global(&self) -> bool;
    fn mark_global(&mut self, u64);
    fn global_offset(&self) -> u64;

    fn is_unknown(&self) -> bool;

    fn add_refs(&mut self, Vec<Self::SSARef>);
    fn refs(&self) -> ::std::slice::Iter<Self::SSARef>;
}

/// Trait that describes variable bindings for a function.
pub trait RBindings {
    type BTy: RBind;
    type Idx: Clone + fmt::Debug + Eq + Ord + Hash + From<usize>;

    fn new() -> Self;
    fn insert(&mut self, Self::BTy) -> Self::Idx;

    fn binding(&self, &Self::Idx) -> Option<&Self::BTy>;
    fn binding_mut(&mut self, &Self::Idx) -> Option<&mut Self::BTy>;

    fn bindings(&self) -> RBinds<Self::BTy>;
    fn bindings_mut(&mut self) -> RBindsMut<Self::BTy>;
}

#[derive(Debug)]
pub struct RBinds<'a, T: 'a + RBind> {
    binds: ::std::slice::Iter<'a, T>,
}

impl<'a, T: 'a + RBind> Iterator for RBinds<'a, T> {
    type Item = &'a T;
    fn next(&mut self) -> Option<&'a T> {
        self.binds.next()
    }
}

#[derive(Debug)]
pub struct RBindsMut<'a, T: 'a + RBind> {
    binds: ::std::slice::IterMut<'a, T>,
}

#[derive(Clone, Debug)]
pub struct RadecoBindings<T: RBind> {
    binds: Vec<T>,
}

impl<BTy: RBind> Index<usize> for RadecoBindings<BTy> {
    type Output = BTy;
    fn index(&self, index: usize) -> &Self::Output {
        self.binding(&index).unwrap()
    }
}

impl<BTy: RBind> IndexMut<usize> for RadecoBindings<BTy> {
    fn index_mut(&mut self, index: usize) -> &mut Self::Output {
        self.binding_mut(&index).unwrap()
    }
}

impl<BTy: RBind> RBindings for RadecoBindings<BTy> {
    type BTy = BTy;
    type Idx = usize;

    fn new() -> Self {
        RadecoBindings { binds: Vec::new() }
    }

    fn insert(&mut self, bind: BTy) -> usize {
        self.binds.push(bind);
        self.binds.len() - 1
    }

    fn binding(&self, idx: &usize) -> Option<&BTy> {
        self.binds.get(*idx)
    }

    fn binding_mut(&mut self, idx: &usize) -> Option<&mut BTy> {
        self.binds.get_mut(*idx)
    }

    fn bindings(&self) -> RBinds<BTy> {
        RBinds { binds: self.binds.iter() }
    }

    fn bindings_mut(&mut self) -> RBindsMut<BTy> {
        RBindsMut { binds: self.binds.iter_mut() }
    }
}

#[derive(Clone, Debug)]
pub struct Binding<T: fmt::Debug + Clone> {
    // Optional name used to represent this binding.
    pub named: Option<String>,
    // Type of binding.
    vloc: VarLocation,
    dty: VarDataType,
    // Node indices in the function ssa that corresponds to this binding.
    ssa_refs: Vec<T>,
    set: HashSet<VarSet>,
}

impl<'a, T: 'a + RBind> Iterator for RBindsMut<'a, T> {
    type Item = &'a mut T;
    fn next(&mut self) -> Option<&'a mut T> {
        self.binds.next()
    }
}

#[derive(Clone, Debug, PartialEq)]
pub enum VarLocation {
    Register {
        name: String,
    },
    Memory(MemoryRegion),
    Unknown,
}

#[derive(Clone, Copy, Debug, PartialEq)]
pub enum MemoryRegion {
    // Base, offset for function local.
    FunctionLocal {
        base: usize,
        offset: i64,
    },
    Stack,
    Global {
        offset: u64,
    },
    Heap,
    Unknown,
}

// Represents the set to which a variable belongs to in the call summary
// information.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)]
enum VarSet {
    Argument,
    Local,
    Modified,
    Preserved,
    Returned,
}

#[derive(Clone, Copy, Debug)]
pub enum VarDataType {
    Unknown(u64),
    Int,
    IntRef,
    Char,
    CharRef,
    Float,
    FloatRef,
    Double,
    DoubleRef,
}

#[derive(Clone, Debug)]
pub struct LocalInfo {
    pub base: usize,
    pub offset: i64,
}

impl<T: Clone + fmt::Debug> Default for Binding<T> {
    fn default() -> Binding<T> {
        Binding {
            named: None,
            vloc: VarLocation::Unknown,
            dty: VarDataType::Unknown(64),
            ssa_refs: Vec::new(),
            set: HashSet::new(),
        }
    }
}

impl<T: Clone + fmt::Debug> RBind for Binding<T> {
    type SSARef = T;

    fn is_argument(&self) -> bool {
        self.set.contains(&VarSet::Argument)
    }

    fn is_local(&self) -> bool {
        self.set.contains(&VarSet::Local)
    }

    fn is_modified(&self) -> bool {
        self.set.contains(&VarSet::Modified)
    }

    fn is_preserved(&self) -> bool {
        self.set.contains(&VarSet::Preserved)
    }

    fn is_return(&self) -> bool {
        self.set.contains(&VarSet::Returned)
    }

    fn mark_argument(&mut self) {
        self.set.insert(VarSet::Argument);
    }

    fn mark_modified(&mut self) {
        self.set.insert(VarSet::Modified);
    }

    fn mark_preserved(&mut self) {
        self.set.insert(VarSet::Preserved);
    }

    fn mark_return(&mut self) {
        self.set.insert(VarSet::Returned);
    }

    fn name(&self) -> String {
        self.named.clone().unwrap_or_else(String::new)
    }

    fn set_name(&mut self, name: String) {
        self.named = Some(name);
    }

    fn is_register(&self) -> bool {
        match self.vloc {
            VarLocation::Register { .. } => true,
            _ => false,
        }
    }

    fn mark_register(&mut self, name: String) {
        self.vloc = VarLocation::Register { name: name };
    }

    fn mark_memory(&mut self) {
        self.vloc = VarLocation::Memory(MemoryRegion::Unknown);
    }

    fn is_fn_local(&self) -> bool {
        match self.vloc {
            VarLocation::Memory(MemoryRegion::FunctionLocal { .. }) => true,
            _ => false,
        }
    }

    fn mark_fn_local(&mut self, base: usize, offset: i64) {
        self.vloc = VarLocation::Memory(MemoryRegion::FunctionLocal {
            base: base,
            offset: offset,
        });
    }

    fn local_info(&self) -> LocalInfo {
        if let VarLocation::Memory(MemoryRegion::FunctionLocal { base, offset }) = self.vloc {
            LocalInfo {
                base: base,
                offset: offset,
            }
        } else {
            panic!()
        }
    }

    fn is_stack(&self) -> bool {
        if let VarLocation::Memory(MemoryRegion::Stack) = self.vloc {
            true
        } else {
            false
        }
    }

    fn mark_stack(&mut self) {
        self.vloc = VarLocation::Memory(MemoryRegion::Stack);
    }

    fn is_global(&self) -> bool {
        match self.vloc {
            VarLocation::Memory(MemoryRegion::Global { offset }) => true,
            _ => false,
        }
    }

    fn mark_global(&mut self, offset: u64) {
        self.vloc = VarLocation::Memory(MemoryRegion::Global { offset: offset });
    }

    fn global_offset(&self) -> u64 {
        if let VarLocation::Memory(MemoryRegion::Global { offset }) = self.vloc {
            offset
        } else {
            panic!()
        }
    }

    fn is_unknown(&self) -> bool {
        match self.vloc {
            VarLocation::Unknown => true,
            _ => false,
        }
    }

    fn add_refs(&mut self, refs: Vec<T>) {
        self.ssa_refs.extend(refs);
    }

    fn refs(&self) -> ::std::slice::Iter<T> {
        self.ssa_refs.iter()
    }
}
<|endoftext|>


<commit_before>use core::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, Ordering};

/// This must be used by the kernel to ensure that context switches are done atomically
/// Compare and exchange this to true when beginning a context switch on any CPU
/// The Context::switch_to function will set it back to false, allowing other CPU's to switch
/// This must be done, as no locks can be held on the stack during switch
pub static CONTEXT_SWITCH_LOCK: AtomicBool = ATOMIC_BOOL_INIT;

#[derive(Clone, Debug)]
pub struct Context {
    /// Page table pointer
    cr3: usize,
    /// RFLAGS register
    rflags: usize,
    /// RBX register
    rbx: usize,
    /// R12 register
    r12: usize,
    /// R13 register
    r13: usize,
    /// R14 register
    r14: usize,
    /// R15 register
    r15: usize,
    /// Base pointer
    rbp: usize,
    /// Stack pointer
    rsp: usize
}

impl Context {
    pub fn new() -> Context {
        Context {
            cr3: 0,
            rflags: 0,
            rbx: 0,
            r12: 0,
            r13: 0,
            r14: 0,
            r15: 0,
            rbp: 0,
            rsp: 0
        }
    }

    pub fn get_page_table(&self) -> usize {
        self.cr3
    }

    pub fn set_page_table(&mut self, address: usize) {
        self.cr3 = address;
    }

    pub fn set_stack(&mut self, address: usize) {
        self.rsp = address;
    }

    /// Switch to the next context by restoring its stack and registers
    #[cold]
    #[inline(never)]
    #[naked]
    pub unsafe fn switch_to(&mut self, next: &mut Context) {
/*
        asm!("fxsave [$0]" : : "r"(self.fx) : "memory" : "intel", "volatile");
        self.loadable = true;
        if next.loadable {
            asm!("fxrstor [$0]" : : "r"(next.fx) : "memory" : "intel", "volatile");
        }else{
            asm!("fninit" : : : "memory" : "intel", "volatile");
        }
*/

        asm!("mov $0, cr3" : "=r"(self.cr3) : : "memory" : "intel", "volatile");
        if next.cr3 != self.cr3 {
            asm!("mov cr3, $0" : : "r"(next.cr3) : "memory" : "intel", "volatile");
        }

        asm!("pushfq ; pop $0" : "=r"(self.rflags) : : "memory" : "intel", "volatile");
        asm!("push $0 ; popfq" : : "r"(next.rflags) : "memory" : "intel", "volatile");

        asm!("mov $0, rbx" : "=r"(self.rbx) : : "memory" : "intel", "volatile");
        asm!("mov rbx, $0" : : "r"(next.rbx) : "memory" : "intel", "volatile");

        asm!("mov $0, r12" : "=r"(self.r12) : : "memory" : "intel", "volatile");
        asm!("mov r12, $0" : : "r"(next.r12) : "memory" : "intel", "volatile");

        asm!("mov $0, r13" : "=r"(self.r13) : : "memory" : "intel", "volatile");
        asm!("mov r13, $0" : : "r"(next.r13) : "memory" : "intel", "volatile");

        asm!("mov $0, r14" : "=r"(self.r14) : : "memory" : "intel", "volatile");
        asm!("mov r14, $0" : : "r"(next.r14) : "memory" : "intel", "volatile");

        asm!("mov $0, r15" : "=r"(self.r15) : : "memory" : "intel", "volatile");
        asm!("mov r15, $0" : : "r"(next.r15) : "memory" : "intel", "volatile");

        asm!("mov $0, rsp" : "=r"(self.rsp) : : "memory" : "intel", "volatile");
        asm!("mov rsp, $0" : : "r"(next.rsp) : "memory" : "intel", "volatile");

        asm!("mov $0, rbp" : "=r"(self.rbp) : : "memory" : "intel", "volatile");
        asm!("mov rbp, $0" : : "r"(next.rbp) : "memory" : "intel", "volatile");

        asm!("call context_switch_unlock" : : : "memory" : "intel", "volatile");
    }
}

/// Unset global lock, set inside of kernel
#[no_mangle]
extern fn context_switch_unlock(){
    CONTEXT_SWITCH_LOCK.store(false, Ordering::SeqCst);
}
<commit_msg>Fix context_switch_unlock being optimized out<commit_after>use core::sync::atomic::{AtomicBool, ATOMIC_BOOL_INIT, Ordering};

/// This must be used by the kernel to ensure that context switches are done atomically
/// Compare and exchange this to true when beginning a context switch on any CPU
/// The Context::switch_to function will set it back to false, allowing other CPU's to switch
/// This must be done, as no locks can be held on the stack during switch
pub static CONTEXT_SWITCH_LOCK: AtomicBool = ATOMIC_BOOL_INIT;

#[derive(Clone, Debug)]
pub struct Context {
    /// Page table pointer
    cr3: usize,
    /// RFLAGS register
    rflags: usize,
    /// RBX register
    rbx: usize,
    /// R12 register
    r12: usize,
    /// R13 register
    r13: usize,
    /// R14 register
    r14: usize,
    /// R15 register
    r15: usize,
    /// Base pointer
    rbp: usize,
    /// Stack pointer
    rsp: usize
}

impl Context {
    pub fn new() -> Context {
        Context {
            cr3: 0,
            rflags: 0,
            rbx: 0,
            r12: 0,
            r13: 0,
            r14: 0,
            r15: 0,
            rbp: 0,
            rsp: 0
        }
    }

    pub fn get_page_table(&self) -> usize {
        self.cr3
    }

    pub fn set_page_table(&mut self, address: usize) {
        self.cr3 = address;
    }

    pub fn set_stack(&mut self, address: usize) {
        self.rsp = address;
    }

    /// Switch to the next context by restoring its stack and registers
    #[cold]
    #[inline(never)]
    #[naked]
    pub unsafe fn switch_to(&mut self, next: &mut Context) {
/*
        asm!("fxsave [$0]" : : "r"(self.fx) : "memory" : "intel", "volatile");
        self.loadable = true;
        if next.loadable {
            asm!("fxrstor [$0]" : : "r"(next.fx) : "memory" : "intel", "volatile");
        }else{
            asm!("fninit" : : : "memory" : "intel", "volatile");
        }
*/

        asm!("mov $0, cr3" : "=r"(self.cr3) : : "memory" : "intel", "volatile");
        if next.cr3 != self.cr3 {
            asm!("mov cr3, $0" : : "r"(next.cr3) : "memory" : "intel", "volatile");
        }

        asm!("pushfq ; pop $0" : "=r"(self.rflags) : : "memory" : "intel", "volatile");
        asm!("push $0 ; popfq" : : "r"(next.rflags) : "memory" : "intel", "volatile");

        asm!("mov $0, rbx" : "=r"(self.rbx) : : "memory" : "intel", "volatile");
        asm!("mov rbx, $0" : : "r"(next.rbx) : "memory" : "intel", "volatile");

        asm!("mov $0, r12" : "=r"(self.r12) : : "memory" : "intel", "volatile");
        asm!("mov r12, $0" : : "r"(next.r12) : "memory" : "intel", "volatile");

        asm!("mov $0, r13" : "=r"(self.r13) : : "memory" : "intel", "volatile");
        asm!("mov r13, $0" : : "r"(next.r13) : "memory" : "intel", "volatile");

        asm!("mov $0, r14" : "=r"(self.r14) : : "memory" : "intel", "volatile");
        asm!("mov r14, $0" : : "r"(next.r14) : "memory" : "intel", "volatile");

        asm!("mov $0, r15" : "=r"(self.r15) : : "memory" : "intel", "volatile");
        asm!("mov r15, $0" : : "r"(next.r15) : "memory" : "intel", "volatile");

        asm!("mov $0, rsp" : "=r"(self.rsp) : : "memory" : "intel", "volatile");
        asm!("mov rsp, $0" : : "r"(next.rsp) : "memory" : "intel", "volatile");

        asm!("mov $0, rbp" : "=r"(self.rbp) : : "memory" : "intel", "volatile");
        asm!("mov rbp, $0" : : "r"(next.rbp) : "memory" : "intel", "volatile");

        asm!("call context_switch_unlock" : : : "memory" : "intel", "volatile");
    }
}

/// Unset global lock, set inside of kernel
#[no_mangle]
pub extern fn context_switch_unlock(){
    CONTEXT_SWITCH_LOCK.store(false, Ordering::SeqCst);
}
<|endoftext|>


<commit_before>use std::rc::Rc;

use geometry::{Intersectable, Shape, Transformable};
use intersection::Intersection;
use material::Material;
use math::EPSILON;
use math::{Point3, Transform, Vector3};
use ray::Ray;
use texture::TextureCoord;

#[cfg(feature = "stats")]
pub mod stats {
    use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};

    static TRIANGLE_INTERSECTION_TEST_COUNT: AtomicUsize = ATOMIC_USIZE_INIT;
    static TRIANGLE_INTERSECTION_HIT_COUNT: AtomicUsize = ATOMIC_USIZE_INIT;

    pub fn record_triangle_intersection() {
        TRIANGLE_INTERSECTION_TEST_COUNT.fetch_add(1, Ordering::SeqCst);
    }

    pub fn record_triangle_hit() {
        TRIANGLE_INTERSECTION_HIT_COUNT.fetch_add(1, Ordering::SeqCst);
    }

    pub fn number_of_triangle_intersections() -> usize {
        TRIANGLE_INTERSECTION_TEST_COUNT.load(Ordering::SeqCst)
    }

    pub fn number_of_triangle_hits() -> usize {
        TRIANGLE_INTERSECTION_HIT_COUNT.load(Ordering::SeqCst)
    }
}

#[cfg(feature = "stats")]
fn record_triangle_intersection() {
    stats::record_triangle_intersection();
}

#[cfg(feature = "stats")]
fn record_triangle_hit() {
    stats::record_triangle_hit();
}

#[derive(Debug)]
pub enum Normal {
    Face(Vector3),
    Vertex(Vector3, Vector3, Vector3),
}

#[derive(Debug)]
pub struct Triangle {
    // A, B, C
    pub vertices: [Point3; 3],
    pub ab: Vector3, // B - A
    pub ac: Vector3, // C - A
    pub normal: Normal,
    // Texture coords for A, b, C
    texture_coords: Option<[TextureCoord; 3]>,
    material: Rc<Material>,
}

impl Triangle {
    pub fn new(
        a: Point3,
        b: Point3,
        c: Point3,
        normal: Normal,
        texture_coords: Option<[TextureCoord; 3]>,
        material: Rc<Material>,
    ) -> Self {
        let ab = b - a;
        let ac = c - a;

        Self {
            vertices: [a, b, c],
            ab,
            ac,
            normal,
            texture_coords,
            material,
        }
    }
}

impl Shape for Triangle {
    fn material(&self) -> &Material {
        &self.material
    }
}

impl Intersectable for Triangle {
    fn intersect(&self, ray: Ray, cull: bool) -> Option<Intersection> {
        #[cfg(feature = "stats")]
        record_triangle_intersection();

        let pvec = ray.direction.cross(&self.ac);
        let det = self.ab.dot(&pvec);

        if cull && det < EPSILON {
            return None;
        }

        if det.abs() < EPSILON {
            return None;
        }

        let inv_det = 1.0 / det;

        let tvec = ray.origin - self.vertices[0];
        let u = tvec.dot(&pvec) * inv_det;

        if u < 0.0 || u > 1.0 {
            return None;
        }

        let qvec = tvec.cross(&self.ab);
        let v = ray.direction.dot(&qvec) * inv_det;

        if v < 0.0 || u + v > 1.0 {
            return None;
        }

        let t = self.ac.dot(&qvec) * inv_det;
        if t > EPSILON {
            let intersection_point = (ray.origin + ray.direction * t).as_point();
            let texture_coord = if let Some([ta, tb, tc]) = self.texture_coords {
                let w = 1.0 - v - u;

                Some(TextureCoord::new(
                    w * (ta.x as f32) + u * (tb.x as f32) + v * (tc.x as f32),
                    w * (ta.y as f32) + u * (tb.y as f32) + v * (tc.y as f32),
                ))
            } else {
                None
            };

            let intersection = Intersection::new(
                t,
                self,
                intersection_point,
                ray,
                self.normal_at_intersection(u, v),
                false,
                texture_coord,
            );

            #[cfg(feature = "stats")]
            record_triangle_hit();
            return Some(intersection);
        }

        None
    }
}

impl Triangle {
    fn normal_at_intersection(&self, u: f32, v: f32) -> Vector3 {
        match self.normal {
            Normal::Face(normal) => normal,
            Normal::Vertex(n0, n1, n2) => (n0 * (1.0 - u - v) + n1 * u + n2 * v).normalize(),
        }
    }
}

impl Transformable for Triangle {
    fn transform(&mut self, transform: &Transform) {
        let matrix = transform.matrix;
        let normal_matrix = transform.matrix;

        // TODO: Consider doing this as a 4x4 matrix calculation instead
        self.vertices[0] = self.vertices[0] * matrix;
        self.vertices[1] = self.vertices[1] * matrix;
        self.vertices[2] = self.vertices[2] * matrix;
        self.ab = self.vertices[1] - self.vertices[0];
        self.ac = self.vertices[2] - self.vertices[0];
        self.normal = match self.normal {
            Normal::Face(normal) => Normal::Face((normal * normal_matrix).normalize()),
            Normal::Vertex(n0, n1, n2) => Normal::Vertex(
                (n0 * normal_matrix).normalize(),
                (n1 * normal_matrix).normalize(),
                (n2 * normal_matrix).normalize(),
            ),
        };
    }
}
<commit_msg>[Triangle] Adjust epsilon<commit_after>use std::rc::Rc;

use geometry::{Intersectable, Shape, Transformable};
use intersection::Intersection;
use material::Material;
use math::EPSILON;
use math::{Point3, Transform, Vector3};
use ray::Ray;
use texture::TextureCoord;

#[cfg(feature = "stats")]
pub mod stats {
    use std::sync::atomic::{AtomicUsize, Ordering, ATOMIC_USIZE_INIT};

    static TRIANGLE_INTERSECTION_TEST_COUNT: AtomicUsize = ATOMIC_USIZE_INIT;
    static TRIANGLE_INTERSECTION_HIT_COUNT: AtomicUsize = ATOMIC_USIZE_INIT;

    pub fn record_triangle_intersection() {
        TRIANGLE_INTERSECTION_TEST_COUNT.fetch_add(1, Ordering::SeqCst);
    }

    pub fn record_triangle_hit() {
        TRIANGLE_INTERSECTION_HIT_COUNT.fetch_add(1, Ordering::SeqCst);
    }

    pub fn number_of_triangle_intersections() -> usize {
        TRIANGLE_INTERSECTION_TEST_COUNT.load(Ordering::SeqCst)
    }

    pub fn number_of_triangle_hits() -> usize {
        TRIANGLE_INTERSECTION_HIT_COUNT.load(Ordering::SeqCst)
    }
}

#[cfg(feature = "stats")]
fn record_triangle_intersection() {
    stats::record_triangle_intersection();
}

#[cfg(feature = "stats")]
fn record_triangle_hit() {
    stats::record_triangle_hit();
}

#[derive(Debug)]
pub enum Normal {
    Face(Vector3),
    Vertex(Vector3, Vector3, Vector3),
}

#[derive(Debug)]
pub struct Triangle {
    // A, B, C
    pub vertices: [Point3; 3],
    pub ab: Vector3, // B - A
    pub ac: Vector3, // C - A
    pub normal: Normal,
    // Texture coords for A, b, C
    texture_coords: Option<[TextureCoord; 3]>,
    material: Rc<Material>,
}

impl Triangle {
    pub fn new(
        a: Point3,
        b: Point3,
        c: Point3,
        normal: Normal,
        texture_coords: Option<[TextureCoord; 3]>,
        material: Rc<Material>,
    ) -> Self {
        let ab = b - a;
        let ac = c - a;

        Self {
            vertices: [a, b, c],
            ab,
            ac,
            normal,
            texture_coords,
            material,
        }
    }
}

impl Shape for Triangle {
    fn material(&self) -> &Material {
        &self.material
    }
}

impl Intersectable for Triangle {
    fn intersect(&self, ray: Ray, cull: bool) -> Option<Intersection> {
        #[cfg(feature = "stats")]
        record_triangle_intersection();

        let pvec = ray.direction.cross(&self.ac);
        let det = self.ab.dot(&pvec);

        if cull && det < 1e-7 {
            return None;
        }

        if !cull && det.abs() < 1e-7 {
            return None;
        }

        let inv_det = 1.0 / det;

        let tvec = ray.origin - self.vertices[0];
        let u = tvec.dot(&pvec) * inv_det;

        if u < 0.0 || u > 1.0 {
            return None;
        }

        let qvec = tvec.cross(&self.ab);
        let v = ray.direction.dot(&qvec) * inv_det;

        if v < 0.0 || u + v > 1.0 {
            return None;
        }

        let t = self.ac.dot(&qvec) * inv_det;
        if t > EPSILON {
            let intersection_point = (ray.origin + ray.direction * t).as_point();
            let texture_coord = if let Some([ta, tb, tc]) = self.texture_coords {
                let w = 1.0 - v - u;

                Some(TextureCoord::new(
                    w * (ta.x as f32) + u * (tb.x as f32) + v * (tc.x as f32),
                    w * (ta.y as f32) + u * (tb.y as f32) + v * (tc.y as f32),
                ))
            } else {
                None
            };

            let intersection = Intersection::new(
                t,
                self,
                intersection_point,
                ray,
                self.normal_at_intersection(u, v),
                false,
                texture_coord,
            );

            #[cfg(feature = "stats")]
            record_triangle_hit();
            return Some(intersection);
        }

        None
    }
}

impl Triangle {
    fn normal_at_intersection(&self, u: f32, v: f32) -> Vector3 {
        match self.normal {
            Normal::Face(normal) => normal,
            Normal::Vertex(n0, n1, n2) => (n0 * (1.0 - u - v) + n1 * u + n2 * v).normalize(),
        }
    }
}

impl Transformable for Triangle {
    fn transform(&mut self, transform: &Transform) {
        let matrix = transform.matrix;
        let normal_matrix = transform.matrix;

        // TODO: Consider doing this as a 4x4 matrix calculation instead
        self.vertices[0] = self.vertices[0] * matrix;
        self.vertices[1] = self.vertices[1] * matrix;
        self.vertices[2] = self.vertices[2] * matrix;
        self.ab = self.vertices[1] - self.vertices[0];
        self.ac = self.vertices[2] - self.vertices[0];
        self.normal = match self.normal {
            Normal::Face(normal) => Normal::Face((normal * normal_matrix).normalize()),
            Normal::Vertex(n0, n1, n2) => Normal::Vertex(
                (n0 * normal_matrix).normalize(),
                (n1 * normal_matrix).normalize(),
                (n2 * normal_matrix).normalize(),
            ),
        };
    }
}
<|endoftext|>


<commit_before>use internal::ColorComponent;
use current::{ Get, Set };
use context::{ Transform, GetTransform, SetTransform };
use context::{ ViewTransform, GetViewTransform, SetViewTransform };
use vecmath::{
    get_scale,
    hsv,
    identity,
    margin_rectangle,
    multiply,
    orient,
    relative_rectangle,
    relative_source_rectangle,
    rotate_radians,
    scale,
    shear,
    translate,
    Matrix2d,
    Scalar,
    Vec2d,
};
use radians::Radians;
use Color;
use Rect;
use SrcRect;

/// Implemented by contexts that contains color.
pub trait RelativeColor: Get<Color> + Set<Color> + Clone {
    /// Multiplies with red, green, blue and alpha values.
    #[inline(always)]
    fn mul_rgba(
        &self,
        r: ColorComponent,
        g: ColorComponent,
        b: ColorComponent,
        a: ColorComponent
    ) -> Self {
        let Color(col) = self.get();
        self.clone().set(Color([col[0] * r, col[1] * g, col[2] * b, col[3] * a]))
    }

    /// Mixes the current color with white.
    ///
    /// 0 is black and 1 is white.
    #[inline(always)]
    fn tint(&self, f: ColorComponent) -> Self {
        self.mul_rgba(f, f, f, 1.0)
    }

    /// Mixes the current color with black.
    ///
    /// 0 is white and 1 is black.
    #[inline(always)]
    fn shade(&self, f: ColorComponent) -> Self {
        let f = 1.0 - f;
        self.mul_rgba(f, f, f, 1.0)
    }

    /// Rotates hue by degrees.
    #[inline(always)]
    fn hue_deg(&self, angle: ColorComponent) -> Self {
        let pi: ColorComponent = Radians::_180();
        self.hue_rad(angle * pi / 180.0)
    }

    /// Rotates hue by radians.
    #[inline(always)]
    fn hue_rad(&self, angle: ColorComponent) -> Self {
        let Color(val) = self.get();
        self.clone().set(Color(hsv(val, angle, 1.0, 1.0)))
    }
}

impl<T: Get<Color> + Set<Color> + Clone> RelativeColor for T {}

/// Should be implemented by contexts that have rectangle information.
pub trait RelativeRectangle: Get<Rect> + Set<Rect> + Clone {
    /// Shrinks the current rectangle equally by all sides.
    #[inline(always)]
    fn margin(&self, m: Scalar) -> Self {
        let Rect(val) = self.get();
        self.clone().set(Rect(margin_rectangle(val, m)))
    }

    /// Expands the current rectangle equally by all sides.
    #[inline(always)]
    fn expand(&self, m: Scalar) -> Self {
        self.margin(-m)
    }

    /// Moves to a relative rectangle using the current rectangle as tile.
    #[inline(always)]
    fn rel(&self, x: Scalar, y: Scalar) -> Self {
        let Rect(val) = self.get();
        self.clone().set(Rect(relative_rectangle(val, [x, y])))
    }
}

impl<T: Get<Rect> + Set<Rect> + Clone> RelativeRectangle for T {}

/// Should be implemented by contexts that
/// have source rectangle information.
pub trait RelativeSourceRectangle: Get<SrcRect> + Set<SrcRect> + Clone {
    /// Adds a source rectangle.
    #[inline(always)]
    fn src_rect(&self, x: i32, y: i32, w: i32, h: i32) -> Self {
        self.clone().set(SrcRect([x, y, w, h]))
    }

    /// Moves to a relative source rectangle using
    /// the current source rectangle as tile.
    #[inline(always)]
    fn src_rel(&self, x: i32, y: i32) -> Self {
        let SrcRect(val) = self.get();
        self.clone().set(SrcRect(
            relative_source_rectangle(val, x, y)
        ))
    }

    /// Flips the source rectangle horizontally.
    #[inline(always)]
    fn src_flip_h(&self) -> Self {
        let SrcRect(source_rect) = self.get();
        self.clone().set(SrcRect([
            source_rect[0] + source_rect[2],
            source_rect[1],
            -source_rect[2],
            source_rect[3]
        ]))
    }

    /// Flips the source rectangle vertically.
    #[inline(always)]
    fn src_flip_v(&self) -> Self {
        let SrcRect(source_rect) = self.get();
        self.clone().set(SrcRect([
            source_rect[0],
            source_rect[1] + source_rect[3],
            source_rect[2],
            -source_rect[3]
        ]))
    }

    /// Flips the source rectangle horizontally and vertically.
    #[inline(always)]
    fn src_flip_hv(&self) -> Self {
        let SrcRect(source_rect) = self.get();
        self.clone().set(SrcRect([
            source_rect[0] + source_rect[2],
            source_rect[1] + source_rect[3],
            -source_rect[2],
            -source_rect[3]
        ]))
    }
}

impl<T: Get<SrcRect>
      + Set<SrcRect>
      + Clone
> RelativeSourceRectangle for T {}

/// Implemented by contexts that can transform.
pub trait RelativeTransform: GetTransform + SetTransform + Clone {
    /// Appends transform to the current one.
    #[inline(always)]
    fn append_transform(&self, transform: Matrix2d) -> Self {
        let mut res = self.clone();
        let Transform(mat) = self.get_transform();
        res.set_transform(Transform(multiply(mat, transform)));
        res
    }

    /// Prepends transform to the current one.
    #[inline(always)]
    fn prepend_transform(&self, transform: Matrix2d) -> Self {
        let mut res = self.clone();
        let Transform(mat) = self.get_transform();
        res.set_transform(Transform(multiply(transform, mat)));
        res
    }

    /// Translate x an y in local coordinates.
    #[inline(always)]
    fn trans(&self, x: Scalar, y: Scalar) -> Self {
        let trans = translate([x, y]);
        let mut res = self.clone();
        let Transform(mat) = self.get_transform();
        res.set_transform(Transform(multiply(mat, trans)));
        res
    }

    /// Rotates degrees in local coordinates.
    #[inline(always)]
    fn rot_deg(&self, angle: Scalar) -> Self {
        let pi: Scalar = Radians::_180();
        self.rot_rad(angle * pi / 180.0)
    }

    /// Rotate radians in local coordinates.
    #[inline(always)]
    fn rot_rad(&self, angle: Scalar) -> Self {
        let rot = rotate_radians(angle);
        let mut res = self.clone();
        let Transform(mat) = self.get_transform();
        res.set_transform(Transform(multiply(mat, rot)));
        res
    }

    /// Orients x axis to look at point locally.
    ///
    /// Leaves x axis unchanged if the point to
    /// look at is the origin.
    #[inline(always)]
    fn orient(&self, x: Scalar, y: Scalar) -> Self {
        let orient = orient(x, y);
        let mut res = self.clone();
        let Transform(mat) = self.get_transform();
        res.set_transform(Transform(multiply(mat, orient)));
        res
    }

    /// Scales in local coordinates.
    #[inline(always)]
    fn scale(&self, sx: Scalar, sy: Scalar) -> Self {
        let scale = scale(sx, sy);
        let mut res = self.clone();
        let Transform(mat) = self.get_transform();
        res.set_transform(Transform(multiply(mat, scale)));
        res
    }

    /// Scales in both directions in local coordinates.
    #[inline(always)]
    fn zoom(&self, s: Scalar) -> Self {
        self.scale(s, s)
    }

    /// Flips vertically in local coordinates.
    #[inline(always)]
    fn flip_v(&self) -> Self {
        self.scale(1.0, -1.0)
    }

    /// Flips horizontally in local coordinates.
    #[inline(always)]
    fn flip_h(&self) -> Self {
        self.scale(-1.0, 1.0)
    }

    /// Flips horizontally and vertically in local coordinates.
    #[inline(always)]
    fn flip_hv(&self) -> Self {
        self.scale(-1.0, -1.0)
    }

    /// Shears in local coordinates.
    #[inline(always)]
    fn shear(&self, v: Vec2d) -> Self {
        let shear = shear(v);
        let mut res = self.clone();
        let Transform(mat) = self.get_transform();
        res.set_transform(Transform(multiply(mat, shear)));
        res
    }
}

impl<T: GetTransform + SetTransform + Clone> RelativeTransform for T {}

/// Should be implemented by contexts that
/// draws something relative to view.
pub trait RelativeViewTransform:
    GetViewTransform + SetViewTransform
  + GetTransform + SetTransform
  + Clone
{
    /// Moves the current transform to the view coordinate system.
    ///
    /// This is usually [0.0, 0.0] in the upper left corner
    /// with the x axis pointing to the right
    /// and the y axis pointing down.
    #[inline(always)]
    fn view(&self) -> Self {
        let mut res = self.clone();
        let ViewTransform(mat) = self.get_view_transform();
        res.set_transform(Transform(mat));
        res
    }

    /// Moves the current transform to the default coordinate system.
    ///
    /// This is usually [0.0, 0.0] in the center
    /// with the x axis pointing to the right
    /// and the y axis pointing up.
    #[inline(always)]
    fn reset(&self) -> Self {
        let mut res = self.clone();
        res.set_transform(Transform(identity()));
        res
    }

    /// Stores the current transform as new view.
    #[inline(always)]
    fn store_view(&self) -> Self {
        let mut res = self.clone();
        let Transform(mat) = self.get_transform();
        res.set_view_transform(ViewTransform(mat));
        res
    }

    /// Computes the current view size.
    #[inline(always)]
    fn get_view_size(&self) -> (Scalar, Scalar) {
        let ViewTransform(mat) = self.get_view_transform();
        let scale = get_scale(mat);
        (2.0 / scale[0], 2.0 / scale[1])
    }
}

impl<
    T: GetViewTransform
     + GetTransform
     + SetViewTransform
     + SetTransform
     + Clone
> RelativeViewTransform for T {}

<commit_msg>Using Get<Transform> instead of GetTransform<commit_after>use internal::ColorComponent;
use current::{ Get, Set };
use context::{ Transform, GetTransform, SetTransform };
use context::{ ViewTransform, GetViewTransform, SetViewTransform };
use vecmath::{
    get_scale,
    hsv,
    identity,
    margin_rectangle,
    multiply,
    orient,
    relative_rectangle,
    relative_source_rectangle,
    rotate_radians,
    scale,
    shear,
    translate,
    Matrix2d,
    Scalar,
    Vec2d,
};
use radians::Radians;
use Color;
use Rect;
use SrcRect;

/// Implemented by contexts that contains color.
pub trait RelativeColor: Get<Color> + Set<Color> + Clone {
    /// Multiplies with red, green, blue and alpha values.
    #[inline(always)]
    fn mul_rgba(
        &self,
        r: ColorComponent,
        g: ColorComponent,
        b: ColorComponent,
        a: ColorComponent
    ) -> Self {
        let Color(col) = self.get();
        self.clone().set(Color([col[0] * r, col[1] * g, col[2] * b, col[3] * a]))
    }

    /// Mixes the current color with white.
    ///
    /// 0 is black and 1 is white.
    #[inline(always)]
    fn tint(&self, f: ColorComponent) -> Self {
        self.mul_rgba(f, f, f, 1.0)
    }

    /// Mixes the current color with black.
    ///
    /// 0 is white and 1 is black.
    #[inline(always)]
    fn shade(&self, f: ColorComponent) -> Self {
        let f = 1.0 - f;
        self.mul_rgba(f, f, f, 1.0)
    }

    /// Rotates hue by degrees.
    #[inline(always)]
    fn hue_deg(&self, angle: ColorComponent) -> Self {
        let pi: ColorComponent = Radians::_180();
        self.hue_rad(angle * pi / 180.0)
    }

    /// Rotates hue by radians.
    #[inline(always)]
    fn hue_rad(&self, angle: ColorComponent) -> Self {
        let Color(val) = self.get();
        self.clone().set(Color(hsv(val, angle, 1.0, 1.0)))
    }
}

impl<T: Get<Color> + Set<Color> + Clone> RelativeColor for T {}

/// Should be implemented by contexts that have rectangle information.
pub trait RelativeRectangle: Get<Rect> + Set<Rect> + Clone {
    /// Shrinks the current rectangle equally by all sides.
    #[inline(always)]
    fn margin(&self, m: Scalar) -> Self {
        let Rect(val) = self.get();
        self.clone().set(Rect(margin_rectangle(val, m)))
    }

    /// Expands the current rectangle equally by all sides.
    #[inline(always)]
    fn expand(&self, m: Scalar) -> Self {
        self.margin(-m)
    }

    /// Moves to a relative rectangle using the current rectangle as tile.
    #[inline(always)]
    fn rel(&self, x: Scalar, y: Scalar) -> Self {
        let Rect(val) = self.get();
        self.clone().set(Rect(relative_rectangle(val, [x, y])))
    }
}

impl<T: Get<Rect> + Set<Rect> + Clone> RelativeRectangle for T {}

/// Should be implemented by contexts that
/// have source rectangle information.
pub trait RelativeSourceRectangle: Get<SrcRect> + Set<SrcRect> + Clone {
    /// Adds a source rectangle.
    #[inline(always)]
    fn src_rect(&self, x: i32, y: i32, w: i32, h: i32) -> Self {
        self.clone().set(SrcRect([x, y, w, h]))
    }

    /// Moves to a relative source rectangle using
    /// the current source rectangle as tile.
    #[inline(always)]
    fn src_rel(&self, x: i32, y: i32) -> Self {
        let SrcRect(val) = self.get();
        self.clone().set(SrcRect(
            relative_source_rectangle(val, x, y)
        ))
    }

    /// Flips the source rectangle horizontally.
    #[inline(always)]
    fn src_flip_h(&self) -> Self {
        let SrcRect(source_rect) = self.get();
        self.clone().set(SrcRect([
            source_rect[0] + source_rect[2],
            source_rect[1],
            -source_rect[2],
            source_rect[3]
        ]))
    }

    /// Flips the source rectangle vertically.
    #[inline(always)]
    fn src_flip_v(&self) -> Self {
        let SrcRect(source_rect) = self.get();
        self.clone().set(SrcRect([
            source_rect[0],
            source_rect[1] + source_rect[3],
            source_rect[2],
            -source_rect[3]
        ]))
    }

    /// Flips the source rectangle horizontally and vertically.
    #[inline(always)]
    fn src_flip_hv(&self) -> Self {
        let SrcRect(source_rect) = self.get();
        self.clone().set(SrcRect([
            source_rect[0] + source_rect[2],
            source_rect[1] + source_rect[3],
            -source_rect[2],
            -source_rect[3]
        ]))
    }
}

impl<T: Get<SrcRect>
      + Set<SrcRect>
      + Clone
> RelativeSourceRectangle for T {}

/// Implemented by contexts that can transform.
pub trait RelativeTransform: Get<Transform> + Set<Transform> + Clone {
    /// Appends transform to the current one.
    #[inline(always)]
    fn append_transform(&self, transform: Matrix2d) -> Self {
        let Transform(mat) = self.get_transform();
        self.clone().set(Transform(multiply(mat, transform)))
    }

    /// Prepends transform to the current one.
    #[inline(always)]
    fn prepend_transform(&self, transform: Matrix2d) -> Self {
        let Transform(mat) = self.get_transform();
        self.clone().set(Transform(multiply(transform, mat)))
    }

    /// Translate x an y in local coordinates.
    #[inline(always)]
    fn trans(&self, x: Scalar, y: Scalar) -> Self {
        let trans = translate([x, y]);
        let Transform(mat) = self.get_transform();
        self.clone().set(Transform(multiply(mat, trans)))
    }

    /// Rotates degrees in local coordinates.
    #[inline(always)]
    fn rot_deg(&self, angle: Scalar) -> Self {
        let pi: Scalar = Radians::_180();
        self.rot_rad(angle * pi / 180.0)
    }

    /// Rotate radians in local coordinates.
    #[inline(always)]
    fn rot_rad(&self, angle: Scalar) -> Self {
        let rot = rotate_radians(angle);
        let Transform(mat) = self.get_transform();
        self.clone().set(Transform(multiply(mat, rot)))
    }

    /// Orients x axis to look at point locally.
    ///
    /// Leaves x axis unchanged if the point to
    /// look at is the origin.
    #[inline(always)]
    fn orient(&self, x: Scalar, y: Scalar) -> Self {
        let orient = orient(x, y);
        let Transform(mat) = self.get_transform();
        self.clone().set(Transform(multiply(mat, orient)))
    }

    /// Scales in local coordinates.
    #[inline(always)]
    fn scale(&self, sx: Scalar, sy: Scalar) -> Self {
        let scale = scale(sx, sy);
        let Transform(mat) = self.get_transform();
        self.clone().set(Transform(multiply(mat, scale)))
    }

    /// Scales in both directions in local coordinates.
    #[inline(always)]
    fn zoom(&self, s: Scalar) -> Self {
        self.scale(s, s)
    }

    /// Flips vertically in local coordinates.
    #[inline(always)]
    fn flip_v(&self) -> Self {
        self.scale(1.0, -1.0)
    }

    /// Flips horizontally in local coordinates.
    #[inline(always)]
    fn flip_h(&self) -> Self {
        self.scale(-1.0, 1.0)
    }

    /// Flips horizontally and vertically in local coordinates.
    #[inline(always)]
    fn flip_hv(&self) -> Self {
        self.scale(-1.0, -1.0)
    }

    /// Shears in local coordinates.
    #[inline(always)]
    fn shear(&self, v: Vec2d) -> Self {
        let shear = shear(v);
        let Transform(mat) = self.get_transform();
        self.clone().set(Transform(multiply(mat, shear)))
    }
}

impl<T: Get<Transform> + Set<Transform> + Clone> RelativeTransform for T {}

/// Should be implemented by contexts that
/// draws something relative to view.
pub trait RelativeViewTransform:
    GetViewTransform + SetViewTransform
  + GetTransform + SetTransform
  + Clone
{
    /// Moves the current transform to the view coordinate system.
    ///
    /// This is usually [0.0, 0.0] in the upper left corner
    /// with the x axis pointing to the right
    /// and the y axis pointing down.
    #[inline(always)]
    fn view(&self) -> Self {
        let mut res = self.clone();
        let ViewTransform(mat) = self.get_view_transform();
        res.set_transform(Transform(mat));
        res
    }

    /// Moves the current transform to the default coordinate system.
    ///
    /// This is usually [0.0, 0.0] in the center
    /// with the x axis pointing to the right
    /// and the y axis pointing up.
    #[inline(always)]
    fn reset(&self) -> Self {
        let mut res = self.clone();
        res.set_transform(Transform(identity()));
        res
    }

    /// Stores the current transform as new view.
    #[inline(always)]
    fn store_view(&self) -> Self {
        let mut res = self.clone();
        let Transform(mat) = self.get_transform();
        res.set_view_transform(ViewTransform(mat));
        res
    }

    /// Computes the current view size.
    #[inline(always)]
    fn get_view_size(&self) -> (Scalar, Scalar) {
        let ViewTransform(mat) = self.get_view_transform();
        let scale = get_scale(mat);
        (2.0 / scale[0], 2.0 / scale[1])
    }
}

impl<
    T: GetViewTransform
     + GetTransform
     + SetViewTransform
     + SetTransform
     + Clone
> RelativeViewTransform for T {}

<|endoftext|>


<commit_before>//! Defines an instance of some piece of geometry in the scene, instances
//! can re-use loaded geometry but apply different transformations and materials
//! to them

use std::sync::Arc;

use geometry::{Intersection, Boundable, BBox, BoundableGeom, Receiver, Emitter,
               SampleableGeom};
use material::Material;
use linalg::{Transform, Point, Ray};
use film::Colorf;

/// Defines an instance of some geometry with its own transform and material
pub enum Instance {
    Emitter(Emitter),
    Receiver(Receiver),
}

impl Instance {
    /// Create an instance of the geometry in the scene that will only receive light.
    pub fn receiver(geom: Arc<BoundableGeom + Send + Sync>, material: Arc<Material + Send + Sync>,
               transform: Transform, tag: &str) -> Instance {
        Instance::Receiver(Receiver::new(geom, material, transform, tag))
    }
    /// Create an instance of the geometry in the scene that will emit and receive light
    pub fn area_light(geom: Arc<SampleableGeom + Send + Sync>, material: Arc<Material + Send + Sync>,
               emission: Colorf, transform: Transform, tag: &str) -> Instance {
        Instance::Emitter(Emitter::area(geom, material, emission, transform, tag))
    }
    pub fn point_light(pos: Point, emission: Colorf, tag: &str) ->  Instance {
        Instance::Emitter(Emitter::point(pos, emission, tag))
    }
    /// Test the ray for intersection against this insance of geometry.
    /// returns Some(Intersection) if an intersection was found and None if not.
    /// If an intersection is found `ray.max_t` will be set accordingly
    pub fn intersect(&self, ray: &mut Ray) -> Option<Intersection> {
        let hit = match self {
            &Instance::Emitter(ref e) => e.intersect(ray),
            &Instance::Receiver(ref r) => r.intersect(ray),
        };
        match hit {
            Some((dg, mat)) => Some(Intersection::new(dg, self, mat)),
            None => None,
        }
    }
}

impl Boundable for Instance {
    fn bounds(&self) -> BBox {
        match self {
            &Instance::Emitter(ref e) => e.bounds(),
            &Instance::Receiver(ref r) => r.bounds(),
        }
    }
}

<commit_msg>Add method to get tag of an instance<commit_after>//! Defines an instance of some piece of geometry in the scene, instances
//! can re-use loaded geometry but apply different transformations and materials
//! to them

use std::sync::Arc;

use geometry::{Intersection, Boundable, BBox, BoundableGeom, Receiver, Emitter,
               SampleableGeom};
use material::Material;
use linalg::{Transform, Point, Ray};
use film::Colorf;

/// Defines an instance of some geometry with its own transform and material
pub enum Instance {
    Emitter(Emitter),
    Receiver(Receiver),
}

impl Instance {
    /// Create an instance of the geometry in the scene that will only receive light.
    pub fn receiver(geom: Arc<BoundableGeom + Send + Sync>, material: Arc<Material + Send + Sync>,
               transform: Transform, tag: &str) -> Instance {
        Instance::Receiver(Receiver::new(geom, material, transform, tag))
    }
    /// Create an instance of the geometry in the scene that will emit and receive light
    pub fn area_light(geom: Arc<SampleableGeom + Send + Sync>, material: Arc<Material + Send + Sync>,
               emission: Colorf, transform: Transform, tag: &str) -> Instance {
        Instance::Emitter(Emitter::area(geom, material, emission, transform, tag))
    }
    pub fn point_light(pos: Point, emission: Colorf, tag: &str) ->  Instance {
        Instance::Emitter(Emitter::point(pos, emission, tag))
    }
    /// Test the ray for intersection against this insance of geometry.
    /// returns Some(Intersection) if an intersection was found and None if not.
    /// If an intersection is found `ray.max_t` will be set accordingly
    pub fn intersect(&self, ray: &mut Ray) -> Option<Intersection> {
        let hit = match self {
            &Instance::Emitter(ref e) => e.intersect(ray),
            &Instance::Receiver(ref r) => r.intersect(ray),
        };
        match hit {
            Some((dg, mat)) => Some(Intersection::new(dg, self, mat)),
            None => None,
        }
    }
    /// Get the tag for this instance
    pub fn tag(&self) -> &str {
        match self {
            &Instance::Emitter(ref e) => &e.tag[..],
            &Instance::Receiver(ref r) => &r.tag[..],
        }
    }
}

impl Boundable for Instance {
    fn bounds(&self) -> BBox {
        match self {
            &Instance::Emitter(ref e) => e.bounds(),
            &Instance::Receiver(ref r) => r.bounds(),
        }
    }
}

<|endoftext|>


<commit_before><commit_msg>Remove useless file<commit_after><|endoftext|>


<commit_before>//!

use super::{
    context::GraphicsContext,
    draw::{DrawParam, DrawUniforms, Std140DrawUniforms},
    gpu::arc::ArcBuffer,
    Image, WgpuContext,
};
use crevice::std140::AsStd140;
use std::collections::BTreeMap;

/// Array of instances for fast rendering of many meshes.
///
/// Traditionally known as a "batch".
#[derive(Debug)]
pub struct InstanceArray {
    pub(crate) buffer: ArcBuffer,
    pub(crate) indices: ArcBuffer,
    pub(crate) image: Image,
    pub(crate) ordered: bool,
    dirty: bool,
    capacity: u32,
    uniforms: Vec<Std140DrawUniforms>,
    params: Vec<DrawParam>,
}

impl InstanceArray {
    /// Creates a new [InstanceArray] capable of storing up to n-`capacity` instances (this can be changed and is resized automatically when needed).
    ///
    /// If `image` is `None`, a 1x1 white image will be used which can be used to draw solid rectangles.
    pub fn new(
        gfx: &GraphicsContext,
        image: impl Into<Option<Image>>,
        capacity: u32,
        ordered: bool,
    ) -> Self {
        InstanceArray::new_wgpu(
            &gfx.wgpu,
            image.into().unwrap_or_else(|| gfx.white_image.clone()),
            capacity,
            ordered,
        )
    }

    fn new_wgpu(wgpu: &WgpuContext, image: Image, capacity: u32, ordered: bool) -> Self {
        assert!(capacity > 0);

        let buffer = ArcBuffer::new(wgpu.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: DrawUniforms::std140_size_static() as u64 * capacity as u64,
            usage: wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_DST
                | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        }));

        let indices = ArcBuffer::new(wgpu.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: if ordered {
                std::mem::size_of::<u32>() as u64 * capacity as u64
            } else {
                4 // min for layout
            },
            usage: wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_SRC
                | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        }));

        let uniforms = Vec::with_capacity(capacity as usize);
        let params = Vec::with_capacity(capacity as usize);

        InstanceArray {
            buffer,
            indices,
            image,
            ordered,
            dirty: false,
            capacity,
            uniforms,
            params,
        }
    }

    /// Resets all the instance data to a set of `DrawParam`.
    pub fn set(&mut self, instances: impl IntoIterator<Item = DrawParam>) {
        self.dirty = true;
        (self.uniforms, self.params) = instances
            .into_iter()
            .map(|x| {
                (
                    DrawUniforms::from_param(
                        &x,
                        [self.image.width() as f32, self.image.height() as f32].into(),
                    )
                    .as_std140(),
                    x,
                )
            })
            .unzip();
    }

    /// Pushes a new instance onto the end.
    pub fn push(&mut self, instance: DrawParam) {
        self.dirty = true;
        self.uniforms.push(
            DrawUniforms::from_param(
                &instance,
                [self.image.width() as f32, self.image.height() as f32].into(),
            )
            .as_std140(),
        );
        self.params.push(instance);
    }

    /// Updates an existing instance at a given index, if it is valid.
    pub fn update(&mut self, index: u32, instance: DrawParam) {
        if let Some((uniform, param)) = self
            .uniforms
            .get_mut(index as usize)
            .and_then(|x| Some((x, self.params.get_mut(index as usize)?)))
        {
            self.dirty = true;
            *uniform = DrawUniforms::from_param(
                &instance,
                [self.image.width() as f32, self.image.height() as f32].into(),
            )
            .as_std140();
            *param = instance;
        }
    }

    /// Clears all instance data.
    pub fn clear(&mut self) {
        // don't need to set dirty here
        self.uniforms.clear();
        self.params.clear();
    }

    /// Returns whether the instance data has been changed without [`InstanceArray::flush()`] being called.
    #[inline]
    pub fn is_dirty(&self) -> bool {
        self.dirty
    }

    /// Returns an immutable slice of all the instance data in this [`InstanceArray`].
    #[inline]
    pub fn instances(&self) -> &[DrawParam] {
        &self.params
    }

    /// Uploads the instance data to the GPU.
    ///
    /// You do not usually need to call this yourself as it is done automatically during drawing.
    pub fn flush(&mut self, gfx: &GraphicsContext) {
        self.flush_wgpu(&gfx.wgpu)
    }

    #[allow(unsafe_code)]
    pub(crate) fn flush_wgpu(&mut self, wgpu: &WgpuContext) {
        if !self.dirty {
            return;
        } else {
            self.dirty = false;
        }

        let len = self.uniforms.len() as u32;
        if len > self.capacity {
            let resized = InstanceArray::new_wgpu(wgpu, self.image.clone(), len, self.ordered);
            self.buffer = resized.buffer;
            self.indices = resized.indices;
            self.capacity = len;
        }

        wgpu.queue.write_buffer(&self.buffer, 0, unsafe {
            std::slice::from_raw_parts(
                self.uniforms.as_ptr() as *const u8,
                self.uniforms.len() * DrawUniforms::std140_size_static(),
            )
        });

        if self.ordered {
            let mut layers = BTreeMap::<_, Vec<_>>::new();
            for (i, param) in self.params.iter().enumerate() {
                layers.entry(param.z).or_default().push(i);
            }
            let indices = layers.into_values().flatten().collect::<Vec<_>>();
            wgpu.queue.write_buffer(&self.indices, 0, unsafe {
                std::slice::from_raw_parts(
                    indices.as_ptr() as *const u8,
                    indices.len() * std::mem::size_of::<u32>(),
                )
            });
        }
    }

    /// Changes the capacity of this `InstanceArray` while preserving instances.
    ///
    /// If `new_capacity` is less than the `len`, the instances will be truncated.
    pub fn resize(&mut self, gfx: &GraphicsContext, new_capacity: u32) {
        let resized = InstanceArray::new(gfx, self.image.clone(), new_capacity, self.ordered);
        self.buffer = resized.buffer;
        self.indices = resized.indices;
        self.capacity = new_capacity;
        self.dirty = true;
        self.uniforms.truncate(new_capacity as usize);
        self.params.truncate(new_capacity as usize);
    }

    /// Returns this `InstanceArray`'s associated `image`.
    #[inline]
    pub fn image(&self) -> Image {
        self.image.clone()
    }

    /// Returns the number of instances this [InstanceArray] is capable of holding.
    /// This number was specified when creating the [InstanceArray], or if the [InstanceArray]
    /// was automatically resized, the greatest length of instances.
    #[inline]
    pub fn capacity(&self) -> usize {
        self.capacity as usize
    }
}
<commit_msg>improve instancearray::set performance<commit_after>//!

use super::{
    context::GraphicsContext,
    draw::{DrawParam, DrawUniforms, Std140DrawUniforms},
    gpu::arc::ArcBuffer,
    Image, WgpuContext,
};
use crevice::std140::AsStd140;
use std::collections::BTreeMap;

/// Array of instances for fast rendering of many meshes.
///
/// Traditionally known as a "batch".
#[derive(Debug)]
pub struct InstanceArray {
    pub(crate) buffer: ArcBuffer,
    pub(crate) indices: ArcBuffer,
    pub(crate) image: Image,
    pub(crate) ordered: bool,
    dirty: bool,
    capacity: u32,
    uniforms: Vec<Std140DrawUniforms>,
    params: Vec<DrawParam>,
}

impl InstanceArray {
    /// Creates a new [InstanceArray] capable of storing up to n-`capacity` instances (this can be changed and is resized automatically when needed).
    ///
    /// If `image` is `None`, a 1x1 white image will be used which can be used to draw solid rectangles.
    pub fn new(
        gfx: &GraphicsContext,
        image: impl Into<Option<Image>>,
        capacity: u32,
        ordered: bool,
    ) -> Self {
        InstanceArray::new_wgpu(
            &gfx.wgpu,
            image.into().unwrap_or_else(|| gfx.white_image.clone()),
            capacity,
            ordered,
        )
    }

    fn new_wgpu(wgpu: &WgpuContext, image: Image, capacity: u32, ordered: bool) -> Self {
        assert!(capacity > 0);

        let buffer = ArcBuffer::new(wgpu.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: DrawUniforms::std140_size_static() as u64 * capacity as u64,
            usage: wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_DST
                | wgpu::BufferUsages::COPY_SRC,
            mapped_at_creation: false,
        }));

        let indices = ArcBuffer::new(wgpu.device.create_buffer(&wgpu::BufferDescriptor {
            label: None,
            size: if ordered {
                std::mem::size_of::<u32>() as u64 * capacity as u64
            } else {
                4 // min for layout
            },
            usage: wgpu::BufferUsages::STORAGE
                | wgpu::BufferUsages::COPY_SRC
                | wgpu::BufferUsages::COPY_DST,
            mapped_at_creation: false,
        }));

        let uniforms = Vec::with_capacity(capacity as usize);
        let params = Vec::with_capacity(capacity as usize);

        InstanceArray {
            buffer,
            indices,
            image,
            ordered,
            dirty: false,
            capacity,
            uniforms,
            params,
        }
    }

    /// Resets all the instance data to a set of `DrawParam`.
    pub fn set(&mut self, instances: impl IntoIterator<Item = DrawParam>) {
        self.dirty = true;
        self.params.clear();
        self.params.extend(instances);
        self.uniforms.clear();
        self.uniforms.extend(self.params.iter().map(|x| {
            DrawUniforms::from_param(
                &x,
                [self.image.width() as f32, self.image.height() as f32].into(),
            )
            .as_std140()
        }));
    }

    /// Pushes a new instance onto the end.
    pub fn push(&mut self, instance: DrawParam) {
        self.dirty = true;
        self.uniforms.push(
            DrawUniforms::from_param(
                &instance,
                [self.image.width() as f32, self.image.height() as f32].into(),
            )
            .as_std140(),
        );
        self.params.push(instance);
    }

    /// Updates an existing instance at a given index, if it is valid.
    pub fn update(&mut self, index: u32, instance: DrawParam) {
        if let Some((uniform, param)) = self
            .uniforms
            .get_mut(index as usize)
            .and_then(|x| Some((x, self.params.get_mut(index as usize)?)))
        {
            self.dirty = true;
            *uniform = DrawUniforms::from_param(
                &instance,
                [self.image.width() as f32, self.image.height() as f32].into(),
            )
            .as_std140();
            *param = instance;
        }
    }

    /// Clears all instance data.
    pub fn clear(&mut self) {
        // don't need to set dirty here
        self.uniforms.clear();
        self.params.clear();
    }

    /// Returns whether the instance data has been changed without [`InstanceArray::flush()`] being called.
    #[inline]
    pub fn is_dirty(&self) -> bool {
        self.dirty
    }

    /// Returns an immutable slice of all the instance data in this [`InstanceArray`].
    #[inline]
    pub fn instances(&self) -> &[DrawParam] {
        &self.params
    }

    /// Uploads the instance data to the GPU.
    ///
    /// You do not usually need to call this yourself as it is done automatically during drawing.
    pub fn flush(&mut self, gfx: &GraphicsContext) {
        self.flush_wgpu(&gfx.wgpu)
    }

    #[allow(unsafe_code)]
    pub(crate) fn flush_wgpu(&mut self, wgpu: &WgpuContext) {
        if !self.dirty {
            return;
        } else {
            self.dirty = false;
        }

        let len = self.uniforms.len() as u32;
        if len > self.capacity {
            let resized = InstanceArray::new_wgpu(wgpu, self.image.clone(), len, self.ordered);
            self.buffer = resized.buffer;
            self.indices = resized.indices;
            self.capacity = len;
        }

        wgpu.queue.write_buffer(&self.buffer, 0, unsafe {
            std::slice::from_raw_parts(
                self.uniforms.as_ptr() as *const u8,
                self.uniforms.len() * DrawUniforms::std140_size_static(),
            )
        });

        if self.ordered {
            let mut layers = BTreeMap::<_, Vec<_>>::new();
            for (i, param) in self.params.iter().enumerate() {
                layers.entry(param.z).or_default().push(i);
            }
            let indices = layers.into_values().flatten().collect::<Vec<_>>();
            wgpu.queue.write_buffer(&self.indices, 0, unsafe {
                std::slice::from_raw_parts(
                    indices.as_ptr() as *const u8,
                    indices.len() * std::mem::size_of::<u32>(),
                )
            });
        }
    }

    /// Changes the capacity of this `InstanceArray` while preserving instances.
    ///
    /// If `new_capacity` is less than the `len`, the instances will be truncated.
    pub fn resize(&mut self, gfx: &GraphicsContext, new_capacity: u32) {
        let resized = InstanceArray::new(gfx, self.image.clone(), new_capacity, self.ordered);
        self.buffer = resized.buffer;
        self.indices = resized.indices;
        self.capacity = new_capacity;
        self.dirty = true;
        self.uniforms.truncate(new_capacity as usize);
        self.params.truncate(new_capacity as usize);
        self.uniforms
            .reserve((new_capacity as usize).min(self.uniforms.len()) - self.uniforms.len());
        self.params
            .reserve((new_capacity as usize).min(self.params.len()) - self.params.len());
    }

    /// Returns this `InstanceArray`'s associated `image`.
    #[inline]
    pub fn image(&self) -> Image {
        self.image.clone()
    }

    /// Returns the number of instances this [InstanceArray] is capable of holding.
    /// This number was specified when creating the [InstanceArray], or if the [InstanceArray]
    /// was automatically resized, the greatest length of instances.
    #[inline]
    pub fn capacity(&self) -> usize {
        self.capacity as usize
    }
}
<|endoftext|>


<commit_before>// Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use std::os::unix::fs::FileExt;
use std::os::unix::io::AsRawFd;
use std::thread;
use std::time::Duration;

use utils::epoll::{ControlOperation, Epoll, EpollEvent, EventSet};
use utils::eventfd::EventFd;
use utils::kernel_version::{min_kernel_version_for_io_uring, KernelVersion};
use utils::skip_if_io_uring_unsupported;
use utils::tempfile::TempFile;
use vm_memory::{Bytes, MmapRegion, VolatileMemory};

mod test_utils;
use io_uring::operation::{OpCode, Operation};
use io_uring::restriction::Restriction;
use io_uring::{Error, IoUring, SQueueError};

use crate::test_utils::drive_submission_and_completion;

const NUM_ENTRIES: u32 = 128;

#[test]
fn test_ring_new() {
    skip_if_io_uring_unsupported!();

    // Invalid entries count: 0.
    assert!(matches!(
        IoUring::new(0, vec![], vec![], None),
        Err(Error::Setup(err)) if err.kind() == std::io::ErrorKind::InvalidInput
    ));
    // Try to register too many files.
    let dummy_file = TempFile::new().unwrap().into_file();
    assert!(matches!(
        IoUring::new(10, vec![&dummy_file; 40000usize], vec![], None), // Max is 32768.
        Err(Error::RegisterFileLimitExceeded)
    ));
}

#[test]
fn test_eventfd() {
    skip_if_io_uring_unsupported!();
    // Test that events get delivered.
    let eventfd = EventFd::new(0).unwrap();

    let file = TempFile::new().unwrap().into_file();
    let mut ring =
        IoUring::new(NUM_ENTRIES, vec![&file], vec![], Some(eventfd.as_raw_fd())).unwrap();
    let user_data: u8 = 71;
    let buf = [0; 4];
    let epoll = Epoll::new().unwrap();
    let mut ready_events = vec![EpollEvent::default(); 1];

    epoll
        .ctl(
            ControlOperation::Add,
            eventfd.as_raw_fd(),
            EpollEvent::new(EventSet::IN, 0),
        )
        .unwrap();

    unsafe {
        ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data))
            .unwrap()
    };
    ring.submit().unwrap();

    assert_eq!(epoll.wait(500, &mut ready_events[..]).unwrap(), 1);
    assert_eq!(ready_events[0].event_set(), EventSet::IN);
}

#[test]
fn test_restrictions() {
    skip_if_io_uring_unsupported!();

    // Check that only the allowlisted opcodes are permitted.
    {
        let file = TempFile::new().unwrap().into_file();
        let mut ring = IoUring::new(
            NUM_ENTRIES,
            vec![&file],
            vec![
                Restriction::RequireFixedFds,
                Restriction::AllowOpCode(OpCode::Read),
            ],
            None,
        )
        .unwrap();
        let buf = [0; 4];

        // Read operations are allowed.
        unsafe {
            ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, 71))
                .unwrap()
        };
        assert_eq!(ring.submit_and_wait_all().unwrap(), 1);
        assert!(unsafe { ring.pop::<u8>().unwrap().unwrap().result().is_ok() });

        // Other operations are not allowed.
        unsafe {
            ring.push(Operation::write(0, buf.as_ptr() as usize, 4, 0, 71))
                .unwrap()
        };
        assert_eq!(ring.submit_and_wait_all().unwrap(), 1);
        assert!(
            matches!(unsafe{ring.pop::<u8>().unwrap().unwrap().result()},
            Err(err) if err.kind() == std::io::ErrorKind::Other)
        );
    }
}

#[test]
fn test_ring_push() {
    skip_if_io_uring_unsupported!();

    // Forgot to register file.
    {
        let buf = [0; 4];
        let mut ring = IoUring::new(NUM_ENTRIES, vec![], vec![], None).unwrap();

        assert!(matches!(
            unsafe { ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, 71)) },
            Err((Error::NoRegisteredFds, 71))
        ));
        assert_eq!(ring.pending_sqes().unwrap(), 0);
    }

    // Now register file.
    {
        let file = TempFile::new().unwrap().into_file();
        let mut ring = IoUring::new(NUM_ENTRIES, vec![&file], vec![], None).unwrap();
        let user_data: u8 = 71;
        let buf = [0; 4];

        // Invalid fd.
        assert!(matches!(
            unsafe { ring.push(Operation::read(1, buf.as_ptr() as usize, 4, 0, user_data)) },
            Err((Error::InvalidFixedFd(1), 71))
        ));
        assert_eq!(ring.pending_sqes().unwrap(), 0);
        assert_eq!(ring.num_ops(), 0);

        // Valid fd.
        assert!(
            unsafe { ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data)) }
                .is_ok()
        );

        assert_eq!(ring.pending_sqes().unwrap(), 1);
        assert_eq!(ring.num_ops(), 1);

        // Full Queue.
        for _ in 1..(NUM_ENTRIES) {
            assert!(unsafe {
                ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data))
            }
            .is_ok());
        }

        assert_eq!(ring.pending_sqes().unwrap(), NUM_ENTRIES);
        assert_eq!(ring.num_ops(), NUM_ENTRIES);

        assert!(matches!(
            unsafe { ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data)) },
            Err((Error::SQueue(SQueueError::FullQueue), 71))
        ));

        assert_eq!(ring.pending_sqes().unwrap(), NUM_ENTRIES);
        assert_eq!(ring.num_ops(), NUM_ENTRIES);

        // We didn't get to submit so pop() should return None.
        assert!(unsafe { ring.pop::<u8>().unwrap().is_none() });
        assert_eq!(ring.num_ops(), NUM_ENTRIES);

        // Full Ring.
        assert!(ring.submit().is_ok());
        // Wait for the io_uring ops to reach the CQ
        thread::sleep(Duration::from_millis(150));
        for _ in 0..NUM_ENTRIES {
            assert!(unsafe {
                ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data))
            }
            .is_ok());
        }
        assert!(ring.submit().is_ok());
        // Wait for the io_uring ops to reach the CQ
        thread::sleep(Duration::from_millis(150));
        assert_eq!(ring.num_ops(), NUM_ENTRIES * 2);
        // The CQ should be full now
        assert!(matches!(
            unsafe { ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data)) },
            Err((Error::FullCQueue, 71))
        ));

        // Check if there are NUM_ENTRIES * 2 cqes
        let mut num_cqes = 0;
        while let Ok(Some(_entry)) = unsafe { ring.pop::<u8>() } {
            num_cqes += 1;
        }
        assert_eq!(num_cqes, NUM_ENTRIES * 2);
        assert_eq!(ring.num_ops(), 0);
    }
}

#[test]
fn test_ring_submit() {
    skip_if_io_uring_unsupported!();

    {
        let file = TempFile::new().unwrap().into_file();
        let mut ring = IoUring::new(NUM_ENTRIES, vec![&file], vec![], None).unwrap();
        let user_data: u8 = 71;
        let buf = [0; 4];

        // Return 0 if we didn't push any sqes.
        assert_eq!(ring.submit().unwrap(), 0);
        assert_eq!(ring.num_ops(), 0);

        // Now push an sqe.
        unsafe {
            ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data))
                .unwrap();
        }
        assert_eq!(ring.num_ops(), 1);
        assert_eq!(ring.submit().unwrap(), 1);
        // Now push & submit some more.
        unsafe {
            ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data))
                .unwrap();
            ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data))
                .unwrap();
        }
        assert_eq!(ring.num_ops(), 3);
        assert_eq!(ring.submit().unwrap(), 2);
    }
}

#[test]
fn test_submit_and_wait_all() {
    skip_if_io_uring_unsupported!();

    let file = TempFile::new().unwrap().into_file();
    let mut ring = IoUring::new(NUM_ENTRIES, vec![&file], vec![], None).unwrap();
    let user_data: u8 = 71;
    let buf = [0; 4];

    // Return 0 if we didn't push any sqes.
    assert_eq!(ring.submit_and_wait_all().unwrap(), 0);

    // Now push an sqe.
    unsafe { ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data)) }.unwrap();
    assert_eq!(ring.pending_sqes().unwrap(), 1);
    assert_eq!(ring.num_ops(), 1);

    // A correct waiting period yields the completed entries.
    assert_eq!(ring.submit_and_wait_all().unwrap(), 1);
    assert_eq!(
        unsafe { ring.pop::<u8>().unwrap().unwrap().user_data() },
        user_data
    );
    assert_eq!(ring.pending_sqes().unwrap(), 0);
    assert_eq!(ring.num_ops(), 0);

    // Now push, submit & wait for some more entries.
    unsafe {
        ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, 72))
            .unwrap();
        ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, 73))
            .unwrap();
        ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, 74))
            .unwrap();
        ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, 75))
            .unwrap();
    }
    assert_eq!(ring.pending_sqes().unwrap(), 4);
    assert_eq!(ring.num_ops(), 4);
    assert_eq!(ring.submit_and_wait_all().unwrap(), 4);
    assert_eq!(ring.pending_sqes().unwrap(), 0);
    assert_eq!(ring.num_ops(), 4);

    unsafe {
        assert!(ring.pop::<u8>().unwrap().is_some());
        assert!(ring.pop::<u8>().unwrap().is_some());
        assert!(ring.pop::<u8>().unwrap().is_some());
        assert!(ring.pop::<u8>().unwrap().is_some());
        assert!(ring.pop::<u8>().unwrap().is_none());
    }
    assert_eq!(ring.num_ops(), 0);
}

#[test]
fn test_write() {
    skip_if_io_uring_unsupported!();

    // Test that writing the sorted values 1-100 into a file works correctly.

    const NUM_BYTES: usize = 100;
    // Setup.
    let file = TempFile::new().unwrap().into_file();
    let mut ring = IoUring::new(NUM_ENTRIES, vec![&file], vec![], None).unwrap();

    // Create & init a memory mapping for storing the write buffers.
    let mem_region: MmapRegion = MmapRegion::build(
        None,
        NUM_BYTES as usize,
        libc::PROT_READ | libc::PROT_WRITE,
        libc::MAP_ANONYMOUS | libc::MAP_PRIVATE,
    )
    .unwrap();

    let expected_result: Vec<u8> = (0..(NUM_BYTES as u8)).collect();

    for i in 0..NUM_BYTES {
        mem_region
            .as_volatile_slice()
            .write_obj(i as u8, i as usize)
            .unwrap();
    }

    // Init the file with all zeros.
    file.write_all_at(&[0; NUM_BYTES], 0).unwrap();

    // Perform the IO.
    drive_submission_and_completion(&mut ring, &mem_region, OpCode::Write, NUM_BYTES);

    // Verify the result.
    let mut buf = [0u8; NUM_BYTES];
    file.read_exact_at(&mut buf, 0).unwrap();
    assert_eq!(buf, &expected_result[..]);
}

#[test]
fn test_read() {
    skip_if_io_uring_unsupported!();

    // Test that reading the sorted values 1-100 from a file works correctly.

    const NUM_BYTES: usize = 100;
    // Setup.
    let file = TempFile::new().unwrap().into_file();
    let mut ring = IoUring::new(NUM_ENTRIES, vec![&file], vec![], None).unwrap();

    // Create & init a memory mapping for storing the read buffers.
    let mem_region: MmapRegion = MmapRegion::build(
        None,
        NUM_BYTES as usize,
        libc::PROT_READ | libc::PROT_WRITE,
        libc::MAP_ANONYMOUS | libc::MAP_PRIVATE,
    )
    .unwrap();

    // Init the file with 1-100.
    let init_contents: Vec<u8> = (0..(NUM_BYTES as u8)).collect();
    file.write_all_at(&init_contents, 0).unwrap();

    // Perform the IO.
    drive_submission_and_completion(&mut ring, &mem_region, OpCode::Read, NUM_BYTES);

    let mut buf = [0; NUM_BYTES];
    mem_region
        .as_volatile_slice()
        .read_slice(&mut buf, 0)
        .unwrap();
    // Verify the result.
    assert_eq!(buf, &init_contents[..]);
}
<commit_msg>Fix `test_restrictions` io_uring test<commit_after>// Copyright 2021 Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

use std::os::unix::fs::FileExt;
use std::os::unix::io::AsRawFd;
use std::thread;
use std::time::Duration;

use utils::epoll::{ControlOperation, Epoll, EpollEvent, EventSet};
use utils::eventfd::EventFd;
use utils::kernel_version::{min_kernel_version_for_io_uring, KernelVersion};
use utils::skip_if_io_uring_unsupported;
use utils::tempfile::TempFile;
use vm_memory::{Bytes, MmapRegion, VolatileMemory};

mod test_utils;
use io_uring::operation::{OpCode, Operation};
use io_uring::restriction::Restriction;
use io_uring::{Error, IoUring, SQueueError};

use crate::test_utils::drive_submission_and_completion;

const NUM_ENTRIES: u32 = 128;

#[test]
fn test_ring_new() {
    skip_if_io_uring_unsupported!();

    // Invalid entries count: 0.
    assert!(matches!(
        IoUring::new(0, vec![], vec![], None),
        Err(Error::Setup(err)) if err.kind() == std::io::ErrorKind::InvalidInput
    ));
    // Try to register too many files.
    let dummy_file = TempFile::new().unwrap().into_file();
    assert!(matches!(
        IoUring::new(10, vec![&dummy_file; 40000usize], vec![], None), // Max is 32768.
        Err(Error::RegisterFileLimitExceeded)
    ));
}

#[test]
fn test_eventfd() {
    skip_if_io_uring_unsupported!();
    // Test that events get delivered.
    let eventfd = EventFd::new(0).unwrap();

    let file = TempFile::new().unwrap().into_file();
    let mut ring =
        IoUring::new(NUM_ENTRIES, vec![&file], vec![], Some(eventfd.as_raw_fd())).unwrap();
    let user_data: u8 = 71;
    let buf = [0; 4];
    let epoll = Epoll::new().unwrap();
    let mut ready_events = vec![EpollEvent::default(); 1];

    epoll
        .ctl(
            ControlOperation::Add,
            eventfd.as_raw_fd(),
            EpollEvent::new(EventSet::IN, 0),
        )
        .unwrap();

    unsafe {
        ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data))
            .unwrap()
    };
    ring.submit().unwrap();

    assert_eq!(epoll.wait(500, &mut ready_events[..]).unwrap(), 1);
    assert_eq!(ready_events[0].event_set(), EventSet::IN);
}

#[test]
fn test_restrictions() {
    skip_if_io_uring_unsupported!();

    // Check that only the allowlisted opcodes are permitted.
    {
        let file = TempFile::new().unwrap().into_file();
        let mut ring = IoUring::new(
            NUM_ENTRIES,
            vec![&file],
            vec![
                Restriction::RequireFixedFds,
                Restriction::AllowOpCode(OpCode::Read),
            ],
            None,
        )
        .unwrap();
        let buf = [0; 4];

        // Read operations are allowed.
        unsafe {
            ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, 71))
                .unwrap()
        };
        assert_eq!(ring.submit_and_wait_all().unwrap(), 1);
        assert!(unsafe { ring.pop::<u8>().unwrap().unwrap().result().is_ok() });

        // Other operations are not allowed.
        unsafe {
            ring.push(Operation::write(0, buf.as_ptr() as usize, 4, 0, 71))
                .unwrap()
        };
        assert_eq!(ring.submit_and_wait_all().unwrap(), 1);
        assert!(unsafe { ring.pop::<u8>().unwrap().unwrap().result().is_err() });
    }
}

#[test]
fn test_ring_push() {
    skip_if_io_uring_unsupported!();

    // Forgot to register file.
    {
        let buf = [0; 4];
        let mut ring = IoUring::new(NUM_ENTRIES, vec![], vec![], None).unwrap();

        assert!(matches!(
            unsafe { ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, 71)) },
            Err((Error::NoRegisteredFds, 71))
        ));
        assert_eq!(ring.pending_sqes().unwrap(), 0);
    }

    // Now register file.
    {
        let file = TempFile::new().unwrap().into_file();
        let mut ring = IoUring::new(NUM_ENTRIES, vec![&file], vec![], None).unwrap();
        let user_data: u8 = 71;
        let buf = [0; 4];

        // Invalid fd.
        assert!(matches!(
            unsafe { ring.push(Operation::read(1, buf.as_ptr() as usize, 4, 0, user_data)) },
            Err((Error::InvalidFixedFd(1), 71))
        ));
        assert_eq!(ring.pending_sqes().unwrap(), 0);
        assert_eq!(ring.num_ops(), 0);

        // Valid fd.
        assert!(
            unsafe { ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data)) }
                .is_ok()
        );

        assert_eq!(ring.pending_sqes().unwrap(), 1);
        assert_eq!(ring.num_ops(), 1);

        // Full Queue.
        for _ in 1..(NUM_ENTRIES) {
            assert!(unsafe {
                ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data))
            }
            .is_ok());
        }

        assert_eq!(ring.pending_sqes().unwrap(), NUM_ENTRIES);
        assert_eq!(ring.num_ops(), NUM_ENTRIES);

        assert!(matches!(
            unsafe { ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data)) },
            Err((Error::SQueue(SQueueError::FullQueue), 71))
        ));

        assert_eq!(ring.pending_sqes().unwrap(), NUM_ENTRIES);
        assert_eq!(ring.num_ops(), NUM_ENTRIES);

        // We didn't get to submit so pop() should return None.
        assert!(unsafe { ring.pop::<u8>().unwrap().is_none() });
        assert_eq!(ring.num_ops(), NUM_ENTRIES);

        // Full Ring.
        assert!(ring.submit().is_ok());
        // Wait for the io_uring ops to reach the CQ
        thread::sleep(Duration::from_millis(150));
        for _ in 0..NUM_ENTRIES {
            assert!(unsafe {
                ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data))
            }
            .is_ok());
        }
        assert!(ring.submit().is_ok());
        // Wait for the io_uring ops to reach the CQ
        thread::sleep(Duration::from_millis(150));
        assert_eq!(ring.num_ops(), NUM_ENTRIES * 2);
        // The CQ should be full now
        assert!(matches!(
            unsafe { ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data)) },
            Err((Error::FullCQueue, 71))
        ));

        // Check if there are NUM_ENTRIES * 2 cqes
        let mut num_cqes = 0;
        while let Ok(Some(_entry)) = unsafe { ring.pop::<u8>() } {
            num_cqes += 1;
        }
        assert_eq!(num_cqes, NUM_ENTRIES * 2);
        assert_eq!(ring.num_ops(), 0);
    }
}

#[test]
fn test_ring_submit() {
    skip_if_io_uring_unsupported!();

    {
        let file = TempFile::new().unwrap().into_file();
        let mut ring = IoUring::new(NUM_ENTRIES, vec![&file], vec![], None).unwrap();
        let user_data: u8 = 71;
        let buf = [0; 4];

        // Return 0 if we didn't push any sqes.
        assert_eq!(ring.submit().unwrap(), 0);
        assert_eq!(ring.num_ops(), 0);

        // Now push an sqe.
        unsafe {
            ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data))
                .unwrap();
        }
        assert_eq!(ring.num_ops(), 1);
        assert_eq!(ring.submit().unwrap(), 1);
        // Now push & submit some more.
        unsafe {
            ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data))
                .unwrap();
            ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data))
                .unwrap();
        }
        assert_eq!(ring.num_ops(), 3);
        assert_eq!(ring.submit().unwrap(), 2);
    }
}

#[test]
fn test_submit_and_wait_all() {
    skip_if_io_uring_unsupported!();

    let file = TempFile::new().unwrap().into_file();
    let mut ring = IoUring::new(NUM_ENTRIES, vec![&file], vec![], None).unwrap();
    let user_data: u8 = 71;
    let buf = [0; 4];

    // Return 0 if we didn't push any sqes.
    assert_eq!(ring.submit_and_wait_all().unwrap(), 0);

    // Now push an sqe.
    unsafe { ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, user_data)) }.unwrap();
    assert_eq!(ring.pending_sqes().unwrap(), 1);
    assert_eq!(ring.num_ops(), 1);

    // A correct waiting period yields the completed entries.
    assert_eq!(ring.submit_and_wait_all().unwrap(), 1);
    assert_eq!(
        unsafe { ring.pop::<u8>().unwrap().unwrap().user_data() },
        user_data
    );
    assert_eq!(ring.pending_sqes().unwrap(), 0);
    assert_eq!(ring.num_ops(), 0);

    // Now push, submit & wait for some more entries.
    unsafe {
        ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, 72))
            .unwrap();
        ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, 73))
            .unwrap();
        ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, 74))
            .unwrap();
        ring.push(Operation::read(0, buf.as_ptr() as usize, 4, 0, 75))
            .unwrap();
    }
    assert_eq!(ring.pending_sqes().unwrap(), 4);
    assert_eq!(ring.num_ops(), 4);
    assert_eq!(ring.submit_and_wait_all().unwrap(), 4);
    assert_eq!(ring.pending_sqes().unwrap(), 0);
    assert_eq!(ring.num_ops(), 4);

    unsafe {
        assert!(ring.pop::<u8>().unwrap().is_some());
        assert!(ring.pop::<u8>().unwrap().is_some());
        assert!(ring.pop::<u8>().unwrap().is_some());
        assert!(ring.pop::<u8>().unwrap().is_some());
        assert!(ring.pop::<u8>().unwrap().is_none());
    }
    assert_eq!(ring.num_ops(), 0);
}

#[test]
fn test_write() {
    skip_if_io_uring_unsupported!();

    // Test that writing the sorted values 1-100 into a file works correctly.

    const NUM_BYTES: usize = 100;
    // Setup.
    let file = TempFile::new().unwrap().into_file();
    let mut ring = IoUring::new(NUM_ENTRIES, vec![&file], vec![], None).unwrap();

    // Create & init a memory mapping for storing the write buffers.
    let mem_region: MmapRegion = MmapRegion::build(
        None,
        NUM_BYTES as usize,
        libc::PROT_READ | libc::PROT_WRITE,
        libc::MAP_ANONYMOUS | libc::MAP_PRIVATE,
    )
    .unwrap();

    let expected_result: Vec<u8> = (0..(NUM_BYTES as u8)).collect();

    for i in 0..NUM_BYTES {
        mem_region
            .as_volatile_slice()
            .write_obj(i as u8, i as usize)
            .unwrap();
    }

    // Init the file with all zeros.
    file.write_all_at(&[0; NUM_BYTES], 0).unwrap();

    // Perform the IO.
    drive_submission_and_completion(&mut ring, &mem_region, OpCode::Write, NUM_BYTES);

    // Verify the result.
    let mut buf = [0u8; NUM_BYTES];
    file.read_exact_at(&mut buf, 0).unwrap();
    assert_eq!(buf, &expected_result[..]);
}

#[test]
fn test_read() {
    skip_if_io_uring_unsupported!();

    // Test that reading the sorted values 1-100 from a file works correctly.

    const NUM_BYTES: usize = 100;
    // Setup.
    let file = TempFile::new().unwrap().into_file();
    let mut ring = IoUring::new(NUM_ENTRIES, vec![&file], vec![], None).unwrap();

    // Create & init a memory mapping for storing the read buffers.
    let mem_region: MmapRegion = MmapRegion::build(
        None,
        NUM_BYTES as usize,
        libc::PROT_READ | libc::PROT_WRITE,
        libc::MAP_ANONYMOUS | libc::MAP_PRIVATE,
    )
    .unwrap();

    // Init the file with 1-100.
    let init_contents: Vec<u8> = (0..(NUM_BYTES as u8)).collect();
    file.write_all_at(&init_contents, 0).unwrap();

    // Perform the IO.
    drive_submission_and_completion(&mut ring, &mem_region, OpCode::Read, NUM_BYTES);

    let mut buf = [0; NUM_BYTES];
    mem_region
        .as_volatile_slice()
        .read_slice(&mut buf, 0)
        .unwrap();
    // Verify the result.
    assert_eq!(buf, &init_contents[..]);
}
<|endoftext|>


<commit_before>use std::sync::Mutex;
use std::sync::mpsc::{channel, Receiver, Sender, SendError};

lazy_static! {
    static ref INTERRUPT_TX: Mutex<Option<Sender<()>>> = Mutex::new(None);
}

/// On Unix platforms, spawn a thread and use the signal crate
/// to relay signals back to the main thread.
#[cfg(unix)]
pub fn install() -> Receiver<()> {
    use std::thread;
    use nix::sys::signal::{SIGTERM, SIGINT};
    use signal::trap::Trap;

    let trap = Trap::trap(&[SIGTERM, SIGINT]);
    let rx = create_channel();

    thread::spawn(move || {
        for _ in trap {
            let result = send_interrupt();
            if result.is_err() {
                break;
            }
        }
    });

    rx
}

/// On Windows, use SetConsoleCtrlHandler() to send an interrupt
/// SetConsoleCtrlHandler runs in it's own thread, so it's safe.
#[cfg(windows)]
pub fn install() -> Receiver<()> {
    use kernel32::SetConsoleCtrlHandler;
    use winapi::{BOOL, DWORD};

    pub unsafe extern "system" fn ctrl_handler(_: DWORD) -> BOOL {
        send_interrupt();
    }

    let rx = create_channel();
    SetConsoleCtrlHandler(Some(ctrl_handler), TRUE);

    rx
}

fn create_channel() -> Receiver<()> {
    let mut guard = INTERRUPT_TX.lock().unwrap();
    if (*guard).is_some() {
        panic!("interrupt_handler::install() already called!");
    }

    let (tx, rx) = channel();
    (*guard) = Some(tx);

    rx
}

fn send_interrupt() -> Result<(), SendError<()>> {
    if let Some(ref mut tx) = *INTERRUPT_TX.lock().unwrap() {
        tx.send(())
    } else {
        Err(SendError(()))
    }
}
<commit_msg>Fix build errors on Windows<commit_after>use std::sync::Mutex;
use std::sync::mpsc::{channel, Receiver, Sender, SendError};

lazy_static! {
    static ref INTERRUPT_TX: Mutex<Option<Sender<()>>> = Mutex::new(None);
}

/// On Unix platforms, spawn a thread and use the signal crate
/// to relay signals back to the main thread.
#[cfg(unix)]
pub fn install() -> Receiver<()> {
    use std::thread;
    use nix::sys::signal::{SIGTERM, SIGINT};
    use signal::trap::Trap;

    let trap = Trap::trap(&[SIGTERM, SIGINT]);
    let rx = create_channel();

    thread::spawn(move || {
        for _ in trap {
            let result = send_interrupt();
            if result.is_err() {
                break;
            }
        }
    });

    rx
}

/// On Windows, use SetConsoleCtrlHandler() to send an interrupt
/// SetConsoleCtrlHandler runs in it's own thread, so it's safe.
#[cfg(windows)]
pub fn install() -> Receiver<()> {
    use kernel32::SetConsoleCtrlHandler;
    use winapi::{BOOL, DWORD, TRUE};

    pub unsafe extern "system" fn ctrl_handler(_: DWORD) -> BOOL {
        send_interrupt();
        TRUE
    }

    let rx = create_channel();
    unsafe {
        SetConsoleCtrlHandler(Some(ctrl_handler), TRUE);
    }

    rx
}

fn create_channel() -> Receiver<()> {
    let mut guard = INTERRUPT_TX.lock().unwrap();
    if (*guard).is_some() {
        panic!("interrupt_handler::install() already called!");
    }

    let (tx, rx) = channel();
    (*guard) = Some(tx);

    rx
}

fn send_interrupt() -> Result<(), SendError<()>> {
    if let Some(ref mut tx) = *INTERRUPT_TX.lock().unwrap() {
        tx.send(())
    } else {
        Err(SendError(()))
    }
}
<|endoftext|>