<commit_before>#[cfg(not(test))]
use crate::alloc::{self, Layout};
use crate::num::NonZeroUsize;
#[cfg(not(test))]
use crate::slice;
#[cfg(not(test))]
use crate::str;

use super::waitqueue::{
    try_lock_or_false, NotifiedTcs, SpinMutex, SpinMutexGuard, WaitQueue, WaitVariable,
};
use crate::mem;

pub struct RWLock {
    readers: SpinMutex<WaitVariable<Option<NonZeroUsize>>>,
    writer: SpinMutex<WaitVariable<bool>>,
}

// Below is to check at compile time, that RWLock has size of 128 bytes.
#[allow(dead_code)]
unsafe fn rw_lock_size_assert(r: RWLock) {
    mem::transmute::<RWLock, [u8; 128]>(r);
}

impl RWLock {
    pub const fn new() -> RWLock {
        RWLock {
            readers: SpinMutex::new(WaitVariable::new(None)),
            writer: SpinMutex::new(WaitVariable::new(false)),
        }
    }

    #[inline]
    pub unsafe fn read(&self) {
        let mut rguard = self.readers.lock();
        let wguard = self.writer.lock();
        if *wguard.lock_var() || !wguard.queue_empty() {
            // Another thread has or is waiting for the write lock, wait
            drop(wguard);
            WaitQueue::wait(rguard);
            // Another thread has passed the lock to us
        } else {
            // No waiting writers, acquire the read lock
            *rguard.lock_var_mut() =
                NonZeroUsize::new(rguard.lock_var().map_or(0, |n| n.get()) + 1);
        }
    }

    #[inline]
    pub unsafe fn try_read(&self) -> bool {
        let mut rguard = try_lock_or_false!(self.readers);
        let wguard = try_lock_or_false!(self.writer);
        if *wguard.lock_var() || !wguard.queue_empty() {
            // Another thread has or is waiting for the write lock
            false
        } else {
            // No waiting writers, acquire the read lock
            *rguard.lock_var_mut() =
                NonZeroUsize::new(rguard.lock_var().map_or(0, |n| n.get()) + 1);
            true
        }
    }

    #[inline]
    pub unsafe fn write(&self) {
        let rguard = self.readers.lock();
        let mut wguard = self.writer.lock();
        if *wguard.lock_var() || rguard.lock_var().is_some() {
            // Another thread has the lock, wait
            drop(rguard);
            WaitQueue::wait(wguard);
            // Another thread has passed the lock to us
        } else {
            // We are just now obtaining the lock
            *wguard.lock_var_mut() = true;
        }
    }

    #[inline]
    pub unsafe fn try_write(&self) -> bool {
        let rguard = try_lock_or_false!(self.readers);
        let mut wguard = try_lock_or_false!(self.writer);
        if *wguard.lock_var() || rguard.lock_var().is_some() {
            // Another thread has the lock
            false
        } else {
            // We are just now obtaining the lock
            *wguard.lock_var_mut() = true;
            true
        }
    }

    #[inline]
    unsafe fn __read_unlock(
        &self,
        mut rguard: SpinMutexGuard<WaitVariable<Option<NonZeroUsize>>>,
        wguard: SpinMutexGuard<WaitVariable<bool>>,
    ) {
        *rguard.lock_var_mut() = NonZeroUsize::new(rguard.lock_var().unwrap().get() - 1);
        if rguard.lock_var().is_some() {
            // There are other active readers
        } else {
            if let Ok(mut wguard) = WaitQueue::notify_one(wguard) {
                // A writer was waiting, pass the lock
                *wguard.lock_var_mut() = true;
            } else {
                // No writers were waiting, the lock is released
                assert!(rguard.queue_empty());
            }
        }
    }

    #[inline]
    pub unsafe fn read_unlock(&self) {
        let rguard = self.readers.lock();
        let wguard = self.writer.lock();
        self.__read_unlock(rguard, wguard);
    }

    #[inline]
    unsafe fn __write_unlock(
        &self,
        rguard: SpinMutexGuard<WaitVariable<Option<NonZeroUsize>>>,
        wguard: SpinMutexGuard<WaitVariable<bool>>,
    ) {
        if let Err(mut wguard) = WaitQueue::notify_one(wguard) {
            // No writers waiting, release the write lock
            *wguard.lock_var_mut() = false;
            if let Ok(mut rguard) = WaitQueue::notify_all(rguard) {
                // One or more readers were waiting, pass the lock to them
                if let NotifiedTcs::All { count } = rguard.notified_tcs() {
                    *rguard.lock_var_mut() = Some(count)
                } else {
                    unreachable!() // called notify_all
                }
            } else {
                // No readers waiting, the lock is released
            }
        } else {
            // There was a thread waiting for write, just pass the lock
        }
    }

    #[inline]
    pub unsafe fn write_unlock(&self) {
        let rguard = self.readers.lock();
        let wguard = self.writer.lock();
        self.__write_unlock(rguard, wguard);
    }

    // only used by __rust_rwlock_unlock below
    #[inline]
    #[cfg_attr(test, allow(dead_code))]
    unsafe fn unlock(&self) {
        let rguard = self.readers.lock();
        let wguard = self.writer.lock();
        if *wguard.lock_var() == true {
            self.__write_unlock(rguard, wguard);
        } else {
            self.__read_unlock(rguard, wguard);
        }
    }

    #[inline]
    pub unsafe fn destroy(&self) {}
}

#[cfg(not(test))]
const EINVAL: i32 = 22;

// used by libunwind port
#[cfg(not(test))]
#[no_mangle]
pub unsafe extern "C" fn __rust_rwlock_rdlock(p: *mut RWLock) -> i32 {
    if p.is_null() {
        return EINVAL;
    }
    (*p).read();
    return 0;
}

#[cfg(not(test))]
#[no_mangle]
pub unsafe extern "C" fn __rust_rwlock_wrlock(p: *mut RWLock) -> i32 {
    if p.is_null() {
        return EINVAL;
    }
    (*p).write();
    return 0;
}
#[cfg(not(test))]
#[no_mangle]
pub unsafe extern "C" fn __rust_rwlock_unlock(p: *mut RWLock) -> i32 {
    if p.is_null() {
        return EINVAL;
    }
    (*p).unlock();
    return 0;
}

// the following functions are also used by the libunwind port. They're
// included here to make sure parallel codegen and LTO don't mess things up.
#[cfg(not(test))]
#[no_mangle]
pub unsafe extern "C" fn __rust_print_err(m: *mut u8, s: i32) {
    if s < 0 {
        return;
    }
    let buf = slice::from_raw_parts(m as *const u8, s as _);
    if let Ok(s) = str::from_utf8(&buf[..buf.iter().position(|&b| b == 0).unwrap_or(buf.len())]) {
        eprint!("{}", s);
    }
}

#[cfg(not(test))]
#[no_mangle]
// NB. used by both libunwind and libpanic_abort
pub unsafe extern "C" fn __rust_abort() {
    crate::sys::abort_internal();
}

#[cfg(not(test))]
#[no_mangle]
pub unsafe extern "C" fn __rust_c_alloc(size: usize, align: usize) -> *mut u8 {
    alloc::alloc(Layout::from_size_align_unchecked(size, align))
}

#[cfg(not(test))]
#[no_mangle]
pub unsafe extern "C" fn __rust_c_dealloc(ptr: *mut u8, size: usize, align: usize) {
    alloc::dealloc(ptr, Layout::from_size_align_unchecked(size, align))
}

#[cfg(test)]
mod tests {
    use super::*;
    use core::array::FixedSizeArray;
    use crate::mem::{self, MaybeUninit};

    // Verify that the bytes of initialized RWLock are the same as in
    // libunwind. If they change, `src/UnwindRustSgx.h` in libunwind needs to
    // be changed too.
    #[test]
    fn test_c_rwlock_initializer() {
        const RWLOCK_INIT: &[u8] = &[
            0x1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x3, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x3, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
        ];

        #[inline(never)]
        fn zero_stack() {
            test::black_box(MaybeUninit::<[RWLock; 16]>::zeroed());
        }

        #[inline(never)]
        unsafe fn rwlock_new(init: &mut MaybeUninit<RWLock>) {
            init.set(RWLock::new());
        }

        unsafe {
            // try hard to make sure that the padding/unused bytes in RWLock
            // get initialized as 0. If the assertion below fails, that might
            // just be an issue with the test code and not with the value of
            // RWLOCK_INIT.
            zero_stack();
            let mut init = MaybeUninit::<RWLock>::zeroed();
            rwlock_new(&mut init);
            assert_eq!(
                mem::transmute::<_, [u8; 128]>(init.assume_init()).as_slice(),
                RWLOCK_INIT
            )
        };
    }
}
<commit_msg>Fix missed fn rename in #59284<commit_after>#[cfg(not(test))]
use crate::alloc::{self, Layout};
use crate::num::NonZeroUsize;
#[cfg(not(test))]
use crate::slice;
#[cfg(not(test))]
use crate::str;

use super::waitqueue::{
    try_lock_or_false, NotifiedTcs, SpinMutex, SpinMutexGuard, WaitQueue, WaitVariable,
};
use crate::mem;

pub struct RWLock {
    readers: SpinMutex<WaitVariable<Option<NonZeroUsize>>>,
    writer: SpinMutex<WaitVariable<bool>>,
}

// Below is to check at compile time, that RWLock has size of 128 bytes.
#[allow(dead_code)]
unsafe fn rw_lock_size_assert(r: RWLock) {
    mem::transmute::<RWLock, [u8; 128]>(r);
}

impl RWLock {
    pub const fn new() -> RWLock {
        RWLock {
            readers: SpinMutex::new(WaitVariable::new(None)),
            writer: SpinMutex::new(WaitVariable::new(false)),
        }
    }

    #[inline]
    pub unsafe fn read(&self) {
        let mut rguard = self.readers.lock();
        let wguard = self.writer.lock();
        if *wguard.lock_var() || !wguard.queue_empty() {
            // Another thread has or is waiting for the write lock, wait
            drop(wguard);
            WaitQueue::wait(rguard);
            // Another thread has passed the lock to us
        } else {
            // No waiting writers, acquire the read lock
            *rguard.lock_var_mut() =
                NonZeroUsize::new(rguard.lock_var().map_or(0, |n| n.get()) + 1);
        }
    }

    #[inline]
    pub unsafe fn try_read(&self) -> bool {
        let mut rguard = try_lock_or_false!(self.readers);
        let wguard = try_lock_or_false!(self.writer);
        if *wguard.lock_var() || !wguard.queue_empty() {
            // Another thread has or is waiting for the write lock
            false
        } else {
            // No waiting writers, acquire the read lock
            *rguard.lock_var_mut() =
                NonZeroUsize::new(rguard.lock_var().map_or(0, |n| n.get()) + 1);
            true
        }
    }

    #[inline]
    pub unsafe fn write(&self) {
        let rguard = self.readers.lock();
        let mut wguard = self.writer.lock();
        if *wguard.lock_var() || rguard.lock_var().is_some() {
            // Another thread has the lock, wait
            drop(rguard);
            WaitQueue::wait(wguard);
            // Another thread has passed the lock to us
        } else {
            // We are just now obtaining the lock
            *wguard.lock_var_mut() = true;
        }
    }

    #[inline]
    pub unsafe fn try_write(&self) -> bool {
        let rguard = try_lock_or_false!(self.readers);
        let mut wguard = try_lock_or_false!(self.writer);
        if *wguard.lock_var() || rguard.lock_var().is_some() {
            // Another thread has the lock
            false
        } else {
            // We are just now obtaining the lock
            *wguard.lock_var_mut() = true;
            true
        }
    }

    #[inline]
    unsafe fn __read_unlock(
        &self,
        mut rguard: SpinMutexGuard<WaitVariable<Option<NonZeroUsize>>>,
        wguard: SpinMutexGuard<WaitVariable<bool>>,
    ) {
        *rguard.lock_var_mut() = NonZeroUsize::new(rguard.lock_var().unwrap().get() - 1);
        if rguard.lock_var().is_some() {
            // There are other active readers
        } else {
            if let Ok(mut wguard) = WaitQueue::notify_one(wguard) {
                // A writer was waiting, pass the lock
                *wguard.lock_var_mut() = true;
            } else {
                // No writers were waiting, the lock is released
                assert!(rguard.queue_empty());
            }
        }
    }

    #[inline]
    pub unsafe fn read_unlock(&self) {
        let rguard = self.readers.lock();
        let wguard = self.writer.lock();
        self.__read_unlock(rguard, wguard);
    }

    #[inline]
    unsafe fn __write_unlock(
        &self,
        rguard: SpinMutexGuard<WaitVariable<Option<NonZeroUsize>>>,
        wguard: SpinMutexGuard<WaitVariable<bool>>,
    ) {
        if let Err(mut wguard) = WaitQueue::notify_one(wguard) {
            // No writers waiting, release the write lock
            *wguard.lock_var_mut() = false;
            if let Ok(mut rguard) = WaitQueue::notify_all(rguard) {
                // One or more readers were waiting, pass the lock to them
                if let NotifiedTcs::All { count } = rguard.notified_tcs() {
                    *rguard.lock_var_mut() = Some(count)
                } else {
                    unreachable!() // called notify_all
                }
            } else {
                // No readers waiting, the lock is released
            }
        } else {
            // There was a thread waiting for write, just pass the lock
        }
    }

    #[inline]
    pub unsafe fn write_unlock(&self) {
        let rguard = self.readers.lock();
        let wguard = self.writer.lock();
        self.__write_unlock(rguard, wguard);
    }

    // only used by __rust_rwlock_unlock below
    #[inline]
    #[cfg_attr(test, allow(dead_code))]
    unsafe fn unlock(&self) {
        let rguard = self.readers.lock();
        let wguard = self.writer.lock();
        if *wguard.lock_var() == true {
            self.__write_unlock(rguard, wguard);
        } else {
            self.__read_unlock(rguard, wguard);
        }
    }

    #[inline]
    pub unsafe fn destroy(&self) {}
}

#[cfg(not(test))]
const EINVAL: i32 = 22;

// used by libunwind port
#[cfg(not(test))]
#[no_mangle]
pub unsafe extern "C" fn __rust_rwlock_rdlock(p: *mut RWLock) -> i32 {
    if p.is_null() {
        return EINVAL;
    }
    (*p).read();
    return 0;
}

#[cfg(not(test))]
#[no_mangle]
pub unsafe extern "C" fn __rust_rwlock_wrlock(p: *mut RWLock) -> i32 {
    if p.is_null() {
        return EINVAL;
    }
    (*p).write();
    return 0;
}
#[cfg(not(test))]
#[no_mangle]
pub unsafe extern "C" fn __rust_rwlock_unlock(p: *mut RWLock) -> i32 {
    if p.is_null() {
        return EINVAL;
    }
    (*p).unlock();
    return 0;
}

// the following functions are also used by the libunwind port. They're
// included here to make sure parallel codegen and LTO don't mess things up.
#[cfg(not(test))]
#[no_mangle]
pub unsafe extern "C" fn __rust_print_err(m: *mut u8, s: i32) {
    if s < 0 {
        return;
    }
    let buf = slice::from_raw_parts(m as *const u8, s as _);
    if let Ok(s) = str::from_utf8(&buf[..buf.iter().position(|&b| b == 0).unwrap_or(buf.len())]) {
        eprint!("{}", s);
    }
}

#[cfg(not(test))]
#[no_mangle]
// NB. used by both libunwind and libpanic_abort
pub unsafe extern "C" fn __rust_abort() {
    crate::sys::abort_internal();
}

#[cfg(not(test))]
#[no_mangle]
pub unsafe extern "C" fn __rust_c_alloc(size: usize, align: usize) -> *mut u8 {
    alloc::alloc(Layout::from_size_align_unchecked(size, align))
}

#[cfg(not(test))]
#[no_mangle]
pub unsafe extern "C" fn __rust_c_dealloc(ptr: *mut u8, size: usize, align: usize) {
    alloc::dealloc(ptr, Layout::from_size_align_unchecked(size, align))
}

#[cfg(test)]
mod tests {
    use super::*;
    use core::array::FixedSizeArray;
    use crate::mem::{self, MaybeUninit};

    // Verify that the bytes of initialized RWLock are the same as in
    // libunwind. If they change, `src/UnwindRustSgx.h` in libunwind needs to
    // be changed too.
    #[test]
    fn test_c_rwlock_initializer() {
        const RWLOCK_INIT: &[u8] = &[
            0x1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x3, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x1, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x3, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
            0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0, 0x0,
        ];

        #[inline(never)]
        fn zero_stack() {
            test::black_box(MaybeUninit::<[RWLock; 16]>::zeroed());
        }

        #[inline(never)]
        unsafe fn rwlock_new(init: &mut MaybeUninit<RWLock>) {
            init.write(RWLock::new());
        }

        unsafe {
            // try hard to make sure that the padding/unused bytes in RWLock
            // get initialized as 0. If the assertion below fails, that might
            // just be an issue with the test code and not with the value of
            // RWLOCK_INIT.
            zero_stack();
            let mut init = MaybeUninit::<RWLock>::zeroed();
            rwlock_new(&mut init);
            assert_eq!(
                mem::transmute::<_, [u8; 128]>(init.assume_init()).as_slice(),
                RWLOCK_INIT
            )
        };
    }
}
<|endoftext|>


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

extern crate gluon;
extern crate gluon_base as base;
extern crate gluon_check as check;
extern crate gluon_parser as parser;

use std::fs;

use criterion::{Bencher, Criterion};

use gluon::{compiler_pipeline::*, new_vm, Compiler};

fn typecheck_prelude(b: &mut Bencher) {
    let vm = new_vm();
    let mut compiler = Compiler::new();
    let MacroValue { expr } = {
        let text = fs::read_to_string("std/prelude.glu").unwrap();
        text.expand_macro(&mut compiler, &vm, "std.prelude", &text)
            .unwrap_or_else(|(_, err)| panic!("{}", err))
    };
    b.iter(|| {
        let result = MacroValue { expr: expr.clone() }.typecheck(&mut compiler, &vm, "<top>", "");
        if let Err(ref err) = result {
            println!("{}", err);
            assert!(false);
        }
        result
    })
}

fn clone_prelude(b: &mut Bencher) {
    let vm = new_vm();
    let mut compiler = Compiler::new();
    let TypecheckValue { expr, .. } = {
        let text = fs::read_to_string("std/prelude.glu").unwrap();
        text.typecheck(&mut compiler, &vm, "std.prelude", &text)
            .unwrap_or_else(|err| panic!("{}", err))
    };
    b.iter(|| expr.clone())
}

fn typecheck_24(b: &mut Bencher) {
    let vm = new_vm();
    let mut compiler = Compiler::new();
    let MacroValue { expr } = {
        let text = fs::read_to_string("examples/24.glu").unwrap();
        text.expand_macro(&mut compiler, &vm, "examples.24", &text)
            .unwrap_or_else(|(_, err)| panic!("{}", err))
    };
    b.iter(|| {
        let result = MacroValue { expr: expr.clone() }.typecheck(&mut compiler, &vm, "<top>", "");
        if let Err(ref err) = result {
            println!("{}", err);
            assert!(false);
        }
        result
    })
}

fn typecheck_file(b: &mut Bencher, file: &str) {
    let vm = new_vm();
    let mut compiler = Compiler::new();
    let reparsed = {
        let text = fs::read_to_string(file).unwrap();
        text.reparse_infix(&mut compiler, &vm, &base::filename_to_module(file), &text)
            .unwrap_or_else(|(_, err)| panic!("{}", err))
    };
    let InfixReparsed {
        metadata,
        metadata_map,
        expr,
    } = &reparsed;
    b.iter_with_setup(
        || InfixReparsed {
            metadata: metadata.clone(),
            metadata_map: metadata_map.clone(),
            expr: expr.clone(),
        },
        |input| {
            let result = input.typecheck(&mut compiler, &vm, "<top>", "");
            if let Err(ref err) = result {
                println!("{}", err);
                assert!(false);
            }
            result
        },
    )
}

fn clone_benchmark(c: &mut Criterion) {
    let _ = env_logger::try_init();

    c.bench_function("clone prelude", clone_prelude);
}

fn typecheck_benchmark(c: &mut Criterion) {
    let _ = env_logger::try_init();

    c.bench_function("std/prelude", typecheck_prelude);
    c.bench_function("examples/24", typecheck_24);

    c.bench_function("typecheck_only_24", move |b| {
        typecheck_file(b, "examples/24.glu")
    });
}

criterion_group!(
    name = check;
    config = Criterion::default().sample_size(20).configure_from_args();
    targets = typecheck_benchmark, clone_benchmark
);
criterion_main!(check);
<commit_msg>chore: Add a benchmark for compiling the lisp example fully<commit_after>#[macro_use]
extern crate criterion;
extern crate env_logger;

extern crate gluon;
extern crate gluon_base as base;
extern crate gluon_check as check;
extern crate gluon_parser as parser;

use std::fs;

use criterion::{Bencher, Criterion};

use gluon::{compiler_pipeline::*, new_vm, Compiler};

fn typecheck_prelude(b: &mut Bencher) {
    let vm = new_vm();
    let mut compiler = Compiler::new();
    let MacroValue { expr } = {
        let text = fs::read_to_string("std/prelude.glu").unwrap();
        text.expand_macro(&mut compiler, &vm, "std.prelude", &text)
            .unwrap_or_else(|(_, err)| panic!("{}", err))
    };
    b.iter(|| {
        let result = MacroValue { expr: expr.clone() }.typecheck(&mut compiler, &vm, "<top>", "");
        if let Err(ref err) = result {
            println!("{}", err);
            assert!(false);
        }
        result
    })
}

fn clone_prelude(b: &mut Bencher) {
    let vm = new_vm();
    let mut compiler = Compiler::new();
    let TypecheckValue { expr, .. } = {
        let text = fs::read_to_string("std/prelude.glu").unwrap();
        text.typecheck(&mut compiler, &vm, "std.prelude", &text)
            .unwrap_or_else(|err| panic!("{}", err))
    };
    b.iter(|| expr.clone())
}

fn typecheck_24(b: &mut Bencher) {
    let vm = new_vm();
    let mut compiler = Compiler::new();
    let MacroValue { expr } = {
        let text = fs::read_to_string("examples/24.glu").unwrap();
        text.expand_macro(&mut compiler, &vm, "examples.24", &text)
            .unwrap_or_else(|(_, err)| panic!("{}", err))
    };
    b.iter(|| {
        let result = MacroValue { expr: expr.clone() }.typecheck(&mut compiler, &vm, "<top>", "");
        if let Err(ref err) = result {
            println!("{}", err);
            assert!(false);
        }
        result
    })
}

fn typecheck_file(b: &mut Bencher, file: &str) {
    let vm = new_vm();
    let mut compiler = Compiler::new();
    let reparsed = {
        let text = fs::read_to_string(file).unwrap();
        text.reparse_infix(&mut compiler, &vm, &base::filename_to_module(file), &text)
            .unwrap_or_else(|(_, err)| panic!("{}", err))
    };
    let InfixReparsed {
        metadata,
        metadata_map,
        expr,
    } = &reparsed;
    b.iter_with_setup(
        || InfixReparsed {
            metadata: metadata.clone(),
            metadata_map: metadata_map.clone(),
            expr: expr.clone(),
        },
        |input| {
            let result = input.typecheck(&mut compiler, &vm, "<top>", "");
            if let Err(ref err) = result {
                println!("{}", err);
                assert!(false);
            }
            result
        },
    )
}

fn clone_benchmark(c: &mut Criterion) {
    let _ = env_logger::try_init();

    c.bench_function("clone prelude", clone_prelude);
}

fn typecheck_benchmark(c: &mut Criterion) {
    let _ = env_logger::try_init();

    c.bench_function("std/prelude", typecheck_prelude);
    c.bench_function("examples/24", typecheck_24);

    c.bench_function("typecheck_only_24", move |b| {
        typecheck_file(b, "examples/24.glu")
    });

    c.bench_function("compile_lisp", move |b| {
        b.iter(|| {
            let vm = new_vm();
            let mut compiler = Compiler::new();
            compiler.load_file(&vm, "examples/lisp/lisp.glu").unwrap()
        })
    });
}

criterion_group!(
    name = check;
    config = Criterion::default().sample_size(20).configure_from_args();
    targets = typecheck_benchmark, clone_benchmark
);
criterion_main!(check);
<|endoftext|>


<commit_before>use super::super::{AuditLogEntryId, UserId};
use serde::de::{self, Deserialize, Deserializer, MapAccess, Visitor};
use std::fmt;
use std::collections::HashMap;

trait FromNum {
    fn from_num(num: i32) -> Result<Self, String> where Self: Sized;
}

/// Determines to what entity an action was used on.
#[derive(Debug)]
pub enum Target {
    Guild,
    Channel,
    User,
    Role,
    Invite,
    Webhook,
    Emoji,
}

impl FromNum for Target {
    fn from_num(num: i32) -> Result<Self, String> {
        use self::Target::*;

        Ok(match num {
            10 => Guild,
            20 => Channel,
            30 => User,
            40 => Role,
            50 => Invite,
            60 => Webhook,
            70 => Emoji,
            unex => return Err(format!("Unexpected target number: {}", unex)),
        })
    }
}

/// Determines the action that was done on a target.
#[derive(Debug)]
pub enum Action {
    GuildUpdate,
    Channel(ActionChannel),
    ChannelOverwrite(ActionChannelOverwrite),
    Member(ActionMember),
    Role(ActionRole),
    Invite(ActionInvite),
    Webhook(ActionWebhook),
    Emoji(ActionEmoji),
}

impl FromNum for Action {
    fn from_num(num: i32) -> Result<Self, String> {
        Ok(match num {
            1 => Action::GuildUpdate,
            num if num >= 10 && num <= 13 => Action::Channel(ActionChannel::from_num(num)?),
            num if num >= 13 && num <= 15 => Action::ChannelOverwrite(ActionChannelOverwrite::from_num(num)?),
            num if num >= 20 && num <= 25 => Action::Member(ActionMember::from_num(num)?),
            num if num >= 30 && num <= 32 => Action::Role(ActionRole::from_num(num)?),
            num if num >= 40 && num <= 42 => Action::Invite(ActionInvite::from_num(num)?),
            num if num >= 50 && num <= 52 => Action::Webhook(ActionWebhook::from_num(num)?),
            num if num >= 60 && num <= 62 => Action::Emoji(ActionEmoji::from_num(num)?),
            _ => return Err(format!("Unexpected action number: {}", num)),
        })
    }
}

#[derive(Debug)]
pub enum ActionChannel {
    Create,
    Update,
    Delete,
}

impl FromNum for ActionChannel {
    fn from_num(num: i32) -> Result<Self, String> {
        use self::ActionChannel::*;

        Ok(match num {
            10 => Create,
            11 => Update,
            12 => Delete,
            unex => return Err(format!("Unexpected action channel num: {}", unex)),
        })
    }
}

#[derive(Debug)]
pub enum ActionChannelOverwrite {
    Create,
    Update,
    Delete,
}

impl FromNum for ActionChannelOverwrite {
    fn from_num(num: i32) -> Result<Self, String> {
        use self::ActionChannelOverwrite::*;

        Ok(match num {
            13 => Create,
            14 => Update,
            15 => Delete,
            unex => return Err(format!("Unexpected action channel overwrite num: {}", unex)),
        })
    }
}

#[derive(Debug)]
pub enum ActionMember {
    Kick,
    Prune,
    BanAdd,
    BanRemove,
    Update,
    RoleUpdate
}

impl FromNum for ActionMember {
    fn from_num(num: i32) -> Result<Self, String> {
        use self::ActionMember::*;

        Ok(match num {
            20 => Kick,
            21 => Prune,
            22 => BanAdd,
            23 => BanRemove,
            24 => Update,
            25 => RoleUpdate,
            unex => return Err(format!("Unexpected action member num: {}", unex)),
        })
    }
}

#[derive(Debug)]
pub enum ActionRole {
    Create,
    Update,
    Delete,
}

impl FromNum for ActionRole {
    fn from_num(num: i32) -> Result<Self, String> {
        use self::ActionRole::*;

        Ok(match num {
            30 => Create,
            31 => Update,
            32 => Delete,
            unex => return Err(format!("Unexpected action role num: {}", unex)),
        })
    }
}

#[derive(Debug)]
pub enum ActionInvite {
    Create,
    Update,
    Delete,
}

impl FromNum for ActionInvite {
    fn from_num(num: i32) -> Result<Self, String> {
        use self::ActionInvite::*;

        Ok(match num {
            40 => Create,
            41 => Update,
            42 => Delete,
            unex => return Err(format!("Unexpected action invite num: {}", unex)),
        })
    }
}

#[derive(Debug)]
pub enum ActionWebhook {
    Create,
    Update,
    Delete,
}

impl FromNum for ActionWebhook {
    fn from_num(num: i32) -> Result<Self, String> {
        use self::ActionWebhook::*;

        Ok(match num {
            50 => Create,
            51 => Update,
            52 => Delete,
            unex => return Err(format!("Unexpected action webhook num: {}", unex)),
        })
    }
}

#[derive(Debug)]
pub enum ActionEmoji {
    Create,
    Delete,
    Update,
}

impl FromNum for ActionEmoji {
    fn from_num(num: i32) -> Result<Self, String> {
        use self::ActionEmoji::*;

        Ok(match num {
            60 => Create,
            61 => Update,
            62 => Delete,
            unex => return Err(format!("Unexpected action emoji num: {}", unex)),
        })
    }
}

#[derive(Debug, Deserialize)]
pub struct Change {
    #[serde(rename = "key")]
    pub name: String,
    #[serde(rename = "old_value")]
    pub old: String,
    #[serde(rename = "new_value")]
    pub new: String,
}

#[derive(Debug)]
pub struct AuditLogs {
    pub entries: HashMap<AuditLogEntryId, AuditLogEntry>,
}

#[derive(Debug, Deserialize)]
pub struct AuditLogEntry {
    /// Determines to what entity an [`action`] was used on.
    ///
    /// [`action`]: #structfield.action
    #[serde(deserialize_with = "deserialize_target", rename = "target_type")]
    pub target: Target,
    /// Determines what action was done on a [`target`]
    ///
    /// [`target`]: #structfield.target
    #[serde(deserialize_with = "deserialize_action", rename = "action_type")]
    pub action: Action,
    /// What was the reasoning by doing an action on a target? If there was one.
    pub reason: Option<String>,
    /// The user that did this action on a target.
    pub user_id: UserId,
    /// What changes were made.
    pub changes: Vec<Change>,
    /// The id of this entry.
    pub id: AuditLogEntryId,
}

fn deserialize_target<'de, D: Deserializer<'de>>(de: D) -> Result<Target, D::Error> {
    struct TargetVisitor;

    impl<'de> Visitor<'de> for TargetVisitor {
        type Value = Target;

        fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
            formatter.write_str("an integer between 0 to 70")
        }

        fn visit_i32<E: de::Error>(self, value: i32) -> Result<Target, E> {
            Target::from_num(value).map_err(|e| E::custom(e))
        }
    }

    de.deserialize_i32(TargetVisitor)
}

fn deserialize_action<'de, D: Deserializer<'de>>(de: D) -> Result<Action, D::Error> {
    struct ActionVisitor;

    impl<'de> Visitor<'de> for ActionVisitor {
        type Value = Action;

        fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
            formatter.write_str("an integer between 1 to 62")
        }

        fn visit_i32<E: de::Error>(self, value: i32) -> Result<Action, E> {
            Action::from_num(value).map_err(|e| E::custom(e))
        }
    }

    de.deserialize_i32(ActionVisitor)
}

impl<'de> Deserialize<'de> for AuditLogs {
    fn deserialize<D: Deserializer<'de>>(de: D) -> Result<Self, D::Error> {
        #[derive(Deserialize)]
        #[serde(field_identifier)]
        enum Field {
            #[serde(rename = "audit_log_entries")]
            AuditLogEntries,
        }

        struct EntriesVisitor;

        impl<'de> Visitor<'de> for EntriesVisitor {
            type Value = AuditLogs;

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_str("audit log entries")
            }

            fn visit_map<V: MapAccess<'de>>(self, mut map: V) -> Result<AuditLogs, V::Error> {
                let audit_log_entries = loop {
                    if let Some(key) = map.next_key()? {
                        match key {
                            Field::AuditLogEntries => {
                                break map.next_value::<Vec<AuditLogEntry>>()?;
                            },
                        }
                    }
                };

                Ok(AuditLogs {
                    entries: audit_log_entries
                        .into_iter()
                        .map(|entry| (entry.id, entry))
                        .collect(),
                })
            }
        }

        const FIELD: &'static [&'static str] = &["audit_log_entries"];
        de.deserialize_struct("AuditLogs", FIELD, EntriesVisitor)
    }
}
<commit_msg>Reuse `num`<commit_after>use super::super::{AuditLogEntryId, UserId};
use serde::de::{self, Deserialize, Deserializer, MapAccess, Visitor};
use std::fmt;
use std::collections::HashMap;

trait FromNum {
    fn from_num(num: i32) -> Result<Self, String> where Self: Sized;
}

/// Determines to what entity an action was used on.
#[derive(Debug)]
pub enum Target {
    Guild,
    Channel,
    User,
    Role,
    Invite,
    Webhook,
    Emoji,
}

impl FromNum for Target {
    fn from_num(num: i32) -> Result<Self, String> {
        use self::Target::*;

        Ok(match num {
            10 => Guild,
            20 => Channel,
            30 => User,
            40 => Role,
            50 => Invite,
            60 => Webhook,
            70 => Emoji,
            _ => return Err(format!("Unexpected target number: {}", num)),
        })
    }
}

/// Determines the action that was done on a target.
#[derive(Debug)]
pub enum Action {
    GuildUpdate,
    Channel(ActionChannel),
    ChannelOverwrite(ActionChannelOverwrite),
    Member(ActionMember),
    Role(ActionRole),
    Invite(ActionInvite),
    Webhook(ActionWebhook),
    Emoji(ActionEmoji),
}

impl FromNum for Action {
    fn from_num(num: i32) -> Result<Self, String> {
        Ok(match num {
            1 => Action::GuildUpdate,
            num if num >= 10 && num <= 13 => Action::Channel(ActionChannel::from_num(num)?),
            num if num >= 13 && num <= 15 => Action::ChannelOverwrite(ActionChannelOverwrite::from_num(num)?),
            num if num >= 20 && num <= 25 => Action::Member(ActionMember::from_num(num)?),
            num if num >= 30 && num <= 32 => Action::Role(ActionRole::from_num(num)?),
            num if num >= 40 && num <= 42 => Action::Invite(ActionInvite::from_num(num)?),
            num if num >= 50 && num <= 52 => Action::Webhook(ActionWebhook::from_num(num)?),
            num if num >= 60 && num <= 62 => Action::Emoji(ActionEmoji::from_num(num)?),
            _ => return Err(format!("Unexpected action number: {}", num)),
        })
    }
}

#[derive(Debug)]
pub enum ActionChannel {
    Create,
    Update,
    Delete,
}

impl FromNum for ActionChannel {
    fn from_num(num: i32) -> Result<Self, String> {
        use self::ActionChannel::*;

        Ok(match num {
            10 => Create,
            11 => Update,
            12 => Delete,
            _ => return Err(format!("Unexpected action channel num: {}", num)),
        })
    }
}

#[derive(Debug)]
pub enum ActionChannelOverwrite {
    Create,
    Update,
    Delete,
}

impl FromNum for ActionChannelOverwrite {
    fn from_num(num: i32) -> Result<Self, String> {
        use self::ActionChannelOverwrite::*;

        Ok(match num {
            13 => Create,
            14 => Update,
            15 => Delete,
            _ => return Err(format!("Unexpected action channel overwrite num: {}", num)),
        })
    }
}

#[derive(Debug)]
pub enum ActionMember {
    Kick,
    Prune,
    BanAdd,
    BanRemove,
    Update,
    RoleUpdate
}

impl FromNum for ActionMember {
    fn from_num(num: i32) -> Result<Self, String> {
        use self::ActionMember::*;

        Ok(match num {
            20 => Kick,
            21 => Prune,
            22 => BanAdd,
            23 => BanRemove,
            24 => Update,
            25 => RoleUpdate,
            _ => return Err(format!("Unexpected action member num: {}", num)),
        })
    }
}

#[derive(Debug)]
pub enum ActionRole {
    Create,
    Update,
    Delete,
}

impl FromNum for ActionRole {
    fn from_num(num: i32) -> Result<Self, String> {
        use self::ActionRole::*;

        Ok(match num {
            30 => Create,
            31 => Update,
            32 => Delete,
            _ => return Err(format!("Unexpected action role num: {}", num)),
        })
    }
}

#[derive(Debug)]
pub enum ActionInvite {
    Create,
    Update,
    Delete,
}

impl FromNum for ActionInvite {
    fn from_num(num: i32) -> Result<Self, String> {
        use self::ActionInvite::*;

        Ok(match num {
            40 => Create,
            41 => Update,
            42 => Delete,
            _ => return Err(format!("Unexpected action invite num: {}", num)),
        })
    }
}

#[derive(Debug)]
pub enum ActionWebhook {
    Create,
    Update,
    Delete,
}

impl FromNum for ActionWebhook {
    fn from_num(num: i32) -> Result<Self, String> {
        use self::ActionWebhook::*;

        Ok(match num {
            50 => Create,
            51 => Update,
            52 => Delete,
            _ => return Err(format!("Unexpected action webhook num: {}", num)),
        })
    }
}

#[derive(Debug)]
pub enum ActionEmoji {
    Create,
    Delete,
    Update,
}

impl FromNum for ActionEmoji {
    fn from_num(num: i32) -> Result<Self, String> {
        use self::ActionEmoji::*;

        Ok(match num {
            60 => Create,
            61 => Update,
            62 => Delete,
            _ => return Err(format!("Unexpected action emoji num: {}", num)),
        })
    }
}

#[derive(Debug, Deserialize)]
pub struct Change {
    #[serde(rename = "key")]
    pub name: String,
    #[serde(rename = "old_value")]
    pub old: String,
    #[serde(rename = "new_value")]
    pub new: String,
}

#[derive(Debug)]
pub struct AuditLogs {
    pub entries: HashMap<AuditLogEntryId, AuditLogEntry>,
}

#[derive(Debug, Deserialize)]
pub struct AuditLogEntry {
    /// Determines to what entity an [`action`] was used on.
    ///
    /// [`action`]: #structfield.action
    #[serde(deserialize_with = "deserialize_target", rename = "target_type")]
    pub target: Target,
    /// Determines what action was done on a [`target`]
    ///
    /// [`target`]: #structfield.target
    #[serde(deserialize_with = "deserialize_action", rename = "action_type")]
    pub action: Action,
    /// What was the reasoning by doing an action on a target? If there was one.
    pub reason: Option<String>,
    /// The user that did this action on a target.
    pub user_id: UserId,
    /// What changes were made.
    pub changes: Vec<Change>,
    /// The id of this entry.
    pub id: AuditLogEntryId,
}

fn deserialize_target<'de, D: Deserializer<'de>>(de: D) -> Result<Target, D::Error> {
    struct TargetVisitor;

    impl<'de> Visitor<'de> for TargetVisitor {
        type Value = Target;

        fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
            formatter.write_str("an integer between 0 to 70")
        }

        fn visit_i32<E: de::Error>(self, value: i32) -> Result<Target, E> {
            Target::from_num(value).map_err(|e| E::custom(e))
        }
    }

    de.deserialize_i32(TargetVisitor)
}

fn deserialize_action<'de, D: Deserializer<'de>>(de: D) -> Result<Action, D::Error> {
    struct ActionVisitor;

    impl<'de> Visitor<'de> for ActionVisitor {
        type Value = Action;

        fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
            formatter.write_str("an integer between 1 to 62")
        }

        fn visit_i32<E: de::Error>(self, value: i32) -> Result<Action, E> {
            Action::from_num(value).map_err(|e| E::custom(e))
        }
    }

    de.deserialize_i32(ActionVisitor)
}

impl<'de> Deserialize<'de> for AuditLogs {
    fn deserialize<D: Deserializer<'de>>(de: D) -> Result<Self, D::Error> {
        #[derive(Deserialize)]
        #[serde(field_identifier)]
        enum Field {
            #[serde(rename = "audit_log_entries")]
            AuditLogEntries,
        }

        struct EntriesVisitor;

        impl<'de> Visitor<'de> for EntriesVisitor {
            type Value = AuditLogs;

            fn expecting(&self, formatter: &mut fmt::Formatter) -> fmt::Result {
                formatter.write_str("audit log entries")
            }

            fn visit_map<V: MapAccess<'de>>(self, mut map: V) -> Result<AuditLogs, V::Error> {
                let audit_log_entries = loop {
                    if let Some(key) = map.next_key()? {
                        match key {
                            Field::AuditLogEntries => {
                                break map.next_value::<Vec<AuditLogEntry>>()?;
                            },
                        }
                    }
                };

                Ok(AuditLogs {
                    entries: audit_log_entries
                        .into_iter()
                        .map(|entry| (entry.id, entry))
                        .collect(),
                })
            }
        }

        const FIELD: &'static [&'static str] = &["audit_log_entries"];
        de.deserialize_struct("AuditLogs", FIELD, EntriesVisitor)
    }
}
<|endoftext|>


<commit_before>//! Intel ABI implementation.

use ir;
use isa::{RegClass, RegUnit, TargetIsa};
use regalloc::AllocatableSet;
use settings as shared_settings;
use super::registers::{GPR, FPR, RU};
use abi::{ArgAction, ValueConversion, ArgAssigner, legalize_args};
use ir::{AbiParam, ArgumentPurpose, ArgumentLoc, ArgumentExtension, CallConv, InstBuilder};
use ir::immediates::Imm64;
use stack_layout::layout_stack;
use std::i32;
use cursor::{Cursor, EncCursor, CursorPosition};
use result;

/// Argument registers for x86-64
static ARG_GPRS: [RU; 6] = [RU::rdi, RU::rsi, RU::rdx, RU::rcx, RU::r8, RU::r9];

/// Return value registers.
static RET_GPRS: [RU; 3] = [RU::rax, RU::rdx, RU::rcx];

struct Args {
    pointer_bytes: u32,
    pointer_bits: u16,
    pointer_type: ir::Type,
    gpr: &'static [RU],
    gpr_used: usize,
    fpr_limit: usize,
    fpr_used: usize,
    offset: u32,
    call_conv: CallConv,
}

impl Args {
    fn new(bits: u16, gpr: &'static [RU], fpr_limit: usize, call_conv: CallConv) -> Args {
        Args {
            pointer_bytes: u32::from(bits) / 8,
            pointer_bits: bits,
            pointer_type: ir::Type::int(bits).unwrap(),
            gpr,
            gpr_used: 0,
            fpr_limit,
            fpr_used: 0,
            offset: 0,
            call_conv: call_conv,
        }
    }
}

impl ArgAssigner for Args {
    fn assign(&mut self, arg: &AbiParam) -> ArgAction {
        let ty = arg.value_type;

        // Check for a legal type.
        // We don't support SIMD yet, so break all vectors down.
        if ty.is_vector() {
            return ValueConversion::VectorSplit.into();
        }

        // Large integers and booleans are broken down to fit in a register.
        if !ty.is_float() && ty.bits() > self.pointer_bits {
            return ValueConversion::IntSplit.into();
        }

        // Small integers are extended to the size of a pointer register.
        if ty.is_int() && ty.bits() < self.pointer_bits {
            match arg.extension {
                ArgumentExtension::None => {}
                ArgumentExtension::Uext => return ValueConversion::Uext(self.pointer_type).into(),
                ArgumentExtension::Sext => return ValueConversion::Sext(self.pointer_type).into(),
            }
        }

        // Handle special-purpose arguments.
        if ty.is_int() && self.call_conv == CallConv::SpiderWASM {
            match arg.purpose {
                // This is SpiderMonkey's `WasmTlsReg`.
                ArgumentPurpose::VMContext => {
                    return ArgumentLoc::Reg(if self.pointer_bits == 64 {
                        RU::r14
                    } else {
                        RU::rsi
                    } as RegUnit).into()
                }
                // This is SpiderMonkey's `WasmTableCallSigReg`.
                ArgumentPurpose::SignatureId => return ArgumentLoc::Reg(RU::rbx as RegUnit).into(),
                _ => {}
            }
        }

        // Try to use a GPR.
        if !ty.is_float() && self.gpr_used < self.gpr.len() {
            let reg = self.gpr[self.gpr_used] as RegUnit;
            self.gpr_used += 1;
            return ArgumentLoc::Reg(reg).into();
        }

        // Try to use an FPR.
        if ty.is_float() && self.fpr_used < self.fpr_limit {
            let reg = FPR.unit(self.fpr_used);
            self.fpr_used += 1;
            return ArgumentLoc::Reg(reg).into();
        }

        // Assign a stack location.
        let loc = ArgumentLoc::Stack(self.offset as i32);
        self.offset += self.pointer_bytes;
        assert!(self.offset <= i32::MAX as u32);
        loc.into()
    }
}

/// Legalize `sig`.
pub fn legalize_signature(sig: &mut ir::Signature, flags: &shared_settings::Flags, _current: bool) {
    let bits;
    let mut args;

    if flags.is_64bit() {
        bits = 64;
        args = Args::new(bits, &ARG_GPRS, 8, sig.call_conv);
    } else {
        bits = 32;
        args = Args::new(bits, &[], 0, sig.call_conv);
    }

    legalize_args(&mut sig.params, &mut args);

    let mut rets = Args::new(bits, &RET_GPRS, 2, sig.call_conv);
    legalize_args(&mut sig.returns, &mut rets);
}

/// Get register class for a type appearing in a legalized signature.
pub fn regclass_for_abi_type(ty: ir::Type) -> RegClass {
    if ty.is_int() || ty.is_bool() {
        GPR
    } else {
        FPR
    }
}

/// Get the set of allocatable registers for `func`.
pub fn allocatable_registers(
    _func: &ir::Function,
    flags: &shared_settings::Flags,
) -> AllocatableSet {
    let mut regs = AllocatableSet::new();
    regs.take(GPR, RU::rsp as RegUnit);
    regs.take(GPR, RU::rbp as RegUnit);

    // 32-bit arch only has 8 registers.
    if !flags.is_64bit() {
        for i in 8..16 {
            regs.take(GPR, GPR.unit(i));
            regs.take(FPR, FPR.unit(i));
        }
    }

    regs
}

/// Get the set of callee-saved registers.
pub fn callee_saved_registers(flags: &shared_settings::Flags) -> &'static [RU] {
    if flags.is_64bit() {
        &[RU::rbx, RU::r12, RU::r13, RU::r14, RU::r15]
    } else {
        &[RU::rbx, RU::rsi, RU::rdi]
    }
}

/// Insert a System V-compatible prologue and epilogue.
pub fn prologue_epilogue(func: &mut ir::Function, isa: &TargetIsa) -> result::CtonResult {
    let word_size = if isa.flags().is_64bit() { 8 } else { 4 };
    let csr_type = if isa.flags().is_64bit() {
        ir::types::I64
    } else {
        ir::types::I32
    };
    let csrs = callee_saved_registers(isa.flags());

    // The reserved stack area is composed of:
    //   return address + frame pointer + all callee-saved registers
    //
    // Pushing the return address is an implicit function of the `call`
    // instruction. Each of the others we will then push explicitly. Then we
    // will adjust the stack pointer to make room for the rest of the required
    // space for this frame.
    let csr_stack_size = ((csrs.len() + 2) * word_size as usize) as i32;
    func.create_stack_slot(ir::StackSlotData {
        kind: ir::StackSlotKind::IncomingArg,
        size: csr_stack_size as u32,
        offset: -csr_stack_size,
    });

    let total_stack_size = layout_stack(&mut func.stack_slots, word_size)? as i32;
    let local_stack_size = (total_stack_size - csr_stack_size) as i64;

    // Add CSRs to function signature
    let fp_arg = ir::AbiParam::special_reg(
        csr_type,
        ir::ArgumentPurpose::FramePointer,
        RU::rbp as RegUnit,
    );
    func.signature.params.push(fp_arg);
    func.signature.returns.push(fp_arg);

    for csr in csrs.iter() {
        let csr_arg =
            ir::AbiParam::special_reg(csr_type, ir::ArgumentPurpose::CalleeSaved, *csr as RegUnit);
        func.signature.params.push(csr_arg);
        func.signature.returns.push(csr_arg);
    }

    // Set up the cursor and insert the prologue
    let entry_ebb = func.layout.entry_block().expect("missing entry block");
    let mut pos = EncCursor::new(func, isa).at_first_insertion_point(entry_ebb);
    insert_prologue(&mut pos, local_stack_size, csr_type, csrs);

    // Reset the cursor and insert the epilogue
    let mut pos = pos.at_position(CursorPosition::Nowhere);
    insert_epilogues(&mut pos, local_stack_size, csr_type, csrs);

    Ok(())
}

/// Insert the prologue for a given function.
fn insert_prologue(
    pos: &mut EncCursor,
    stack_size: i64,
    csr_type: ir::types::Type,
    csrs: &'static [RU],
) {
    // Append param to entry EBB
    let ebb = pos.current_ebb().expect("missing ebb under cursor");
    let fp = pos.func.dfg.append_ebb_param(ebb, csr_type);
    pos.func.locations[fp] = ir::ValueLoc::Reg(RU::rbp as RegUnit);

    pos.ins().x86_push(fp);
    pos.ins().copy_special(
        RU::rsp as RegUnit,
        RU::rbp as RegUnit,
    );

    for reg in csrs.iter() {
        // Append param to entry EBB
        let csr_arg = pos.func.dfg.append_ebb_param(ebb, csr_type);

        // Assign it a location
        pos.func.locations[csr_arg] = ir::ValueLoc::Reg(*reg as RegUnit);

        // Remember it so we can push it momentarily
        pos.ins().x86_push(csr_arg);
    }

    if stack_size > 0 {
        pos.ins().adjust_sp_imm(Imm64::new(-stack_size));
    }
}

/// Find all `return` instructions and insert epilogues before them.
fn insert_epilogues(
    pos: &mut EncCursor,
    stack_size: i64,
    csr_type: ir::types::Type,
    csrs: &'static [RU],
) {
    while let Some(ebb) = pos.next_ebb() {
        pos.goto_last_inst(ebb);
        if let Some(inst) = pos.current_inst() {
            if pos.func.dfg[inst].opcode().is_return() {
                insert_epilogue(inst, stack_size, pos, csr_type, csrs);
            }
        }
    }
}

/// Insert an epilogue given a specific `return` instruction.
fn insert_epilogue(
    inst: ir::Inst,
    stack_size: i64,
    pos: &mut EncCursor,
    csr_type: ir::types::Type,
    csrs: &'static [RU],
) {
    if stack_size > 0 {
        pos.ins().adjust_sp_imm(Imm64::new(stack_size));
    }

    // Pop all the callee-saved registers, stepping backward each time to
    // preserve the correct order.
    let fp_ret = pos.ins().x86_pop(csr_type);
    pos.prev_inst();

    pos.func.locations[fp_ret] = ir::ValueLoc::Reg(RU::rbp as RegUnit);
    pos.func.dfg.append_inst_arg(inst, fp_ret);

    for reg in csrs.iter() {
        let csr_ret = pos.ins().x86_pop(csr_type);
        pos.prev_inst();

        pos.func.locations[csr_ret] = ir::ValueLoc::Reg(*reg as RegUnit);
        pos.func.dfg.append_inst_arg(inst, csr_ret);
    }
}
<commit_msg>Add separate spiderwasm prologue/epilogue to intel's abi.rs<commit_after>//! Intel ABI implementation.

use ir;
use isa::{RegClass, RegUnit, TargetIsa};
use regalloc::AllocatableSet;
use settings as shared_settings;
use super::registers::{GPR, FPR, RU};
use abi::{ArgAction, ValueConversion, ArgAssigner, legalize_args};
use ir::{AbiParam, ArgumentPurpose, ArgumentLoc, ArgumentExtension, CallConv, InstBuilder};
use ir::stackslot::{StackSize, StackOffset};
use ir::immediates::Imm64;
use stack_layout::layout_stack;
use std::i32;
use cursor::{Cursor, EncCursor, CursorPosition};
use result;

/// Argument registers for x86-64
static ARG_GPRS: [RU; 6] = [RU::rdi, RU::rsi, RU::rdx, RU::rcx, RU::r8, RU::r9];

/// Return value registers.
static RET_GPRS: [RU; 3] = [RU::rax, RU::rdx, RU::rcx];

struct Args {
    pointer_bytes: u32,
    pointer_bits: u16,
    pointer_type: ir::Type,
    gpr: &'static [RU],
    gpr_used: usize,
    fpr_limit: usize,
    fpr_used: usize,
    offset: u32,
    call_conv: CallConv,
}

impl Args {
    fn new(bits: u16, gpr: &'static [RU], fpr_limit: usize, call_conv: CallConv) -> Args {
        Args {
            pointer_bytes: u32::from(bits) / 8,
            pointer_bits: bits,
            pointer_type: ir::Type::int(bits).unwrap(),
            gpr,
            gpr_used: 0,
            fpr_limit,
            fpr_used: 0,
            offset: 0,
            call_conv: call_conv,
        }
    }
}

impl ArgAssigner for Args {
    fn assign(&mut self, arg: &AbiParam) -> ArgAction {
        let ty = arg.value_type;

        // Check for a legal type.
        // We don't support SIMD yet, so break all vectors down.
        if ty.is_vector() {
            return ValueConversion::VectorSplit.into();
        }

        // Large integers and booleans are broken down to fit in a register.
        if !ty.is_float() && ty.bits() > self.pointer_bits {
            return ValueConversion::IntSplit.into();
        }

        // Small integers are extended to the size of a pointer register.
        if ty.is_int() && ty.bits() < self.pointer_bits {
            match arg.extension {
                ArgumentExtension::None => {}
                ArgumentExtension::Uext => return ValueConversion::Uext(self.pointer_type).into(),
                ArgumentExtension::Sext => return ValueConversion::Sext(self.pointer_type).into(),
            }
        }

        // Handle special-purpose arguments.
        if ty.is_int() && self.call_conv == CallConv::SpiderWASM {
            match arg.purpose {
                // This is SpiderMonkey's `WasmTlsReg`.
                ArgumentPurpose::VMContext => {
                    return ArgumentLoc::Reg(if self.pointer_bits == 64 {
                        RU::r14
                    } else {
                        RU::rsi
                    } as RegUnit).into()
                }
                // This is SpiderMonkey's `WasmTableCallSigReg`.
                ArgumentPurpose::SignatureId => return ArgumentLoc::Reg(RU::rbx as RegUnit).into(),
                _ => {}
            }
        }

        // Try to use a GPR.
        if !ty.is_float() && self.gpr_used < self.gpr.len() {
            let reg = self.gpr[self.gpr_used] as RegUnit;
            self.gpr_used += 1;
            return ArgumentLoc::Reg(reg).into();
        }

        // Try to use an FPR.
        if ty.is_float() && self.fpr_used < self.fpr_limit {
            let reg = FPR.unit(self.fpr_used);
            self.fpr_used += 1;
            return ArgumentLoc::Reg(reg).into();
        }

        // Assign a stack location.
        let loc = ArgumentLoc::Stack(self.offset as i32);
        self.offset += self.pointer_bytes;
        assert!(self.offset <= i32::MAX as u32);
        loc.into()
    }
}

/// Legalize `sig`.
pub fn legalize_signature(sig: &mut ir::Signature, flags: &shared_settings::Flags, _current: bool) {
    let bits;
    let mut args;

    if flags.is_64bit() {
        bits = 64;
        args = Args::new(bits, &ARG_GPRS, 8, sig.call_conv);
    } else {
        bits = 32;
        args = Args::new(bits, &[], 0, sig.call_conv);
    }

    legalize_args(&mut sig.params, &mut args);

    let mut rets = Args::new(bits, &RET_GPRS, 2, sig.call_conv);
    legalize_args(&mut sig.returns, &mut rets);
}

/// Get register class for a type appearing in a legalized signature.
pub fn regclass_for_abi_type(ty: ir::Type) -> RegClass {
    if ty.is_int() || ty.is_bool() {
        GPR
    } else {
        FPR
    }
}

/// Get the set of allocatable registers for `func`.
pub fn allocatable_registers(
    _func: &ir::Function,
    flags: &shared_settings::Flags,
) -> AllocatableSet {
    let mut regs = AllocatableSet::new();
    regs.take(GPR, RU::rsp as RegUnit);
    regs.take(GPR, RU::rbp as RegUnit);

    // 32-bit arch only has 8 registers.
    if !flags.is_64bit() {
        for i in 8..16 {
            regs.take(GPR, GPR.unit(i));
            regs.take(FPR, FPR.unit(i));
        }
    }

    regs
}

/// Get the set of callee-saved registers.
pub fn callee_saved_registers(flags: &shared_settings::Flags) -> &'static [RU] {
    if flags.is_64bit() {
        &[RU::rbx, RU::r12, RU::r13, RU::r14, RU::r15]
    } else {
        &[RU::rbx, RU::rsi, RU::rdi]
    }
}

pub fn prologue_epilogue(func: &mut ir::Function, isa: &TargetIsa) -> result::CtonResult {
    match func.signature.call_conv {
        ir::CallConv::Native => native_prologue_epilogue(func, isa),
        ir::CallConv::SpiderWASM => spiderwasm_prologue_epilogue(func, isa),
    }
}

pub fn spiderwasm_prologue_epilogue(
    func: &mut ir::Function,
    isa: &TargetIsa,
) -> result::CtonResult {
    let word_size = if isa.flags().is_64bit() { 8 } else { 4 };
    let bytes = StackSize::from(isa.flags().spiderwasm_prologue_words()) * word_size;

    let mut ss = ir::StackSlotData::new(ir::StackSlotKind::IncomingArg, bytes);
    ss.offset = -(bytes as StackOffset);
    func.stack_slots.push(ss);

    layout_stack(&mut func.stack_slots, word_size)?;
    Ok(())
}

/// Insert a System V-compatible prologue and epilogue.
pub fn native_prologue_epilogue(func: &mut ir::Function, isa: &TargetIsa) -> result::CtonResult {
    let word_size = if isa.flags().is_64bit() { 8 } else { 4 };
    let csr_type = if isa.flags().is_64bit() {
        ir::types::I64
    } else {
        ir::types::I32
    };
    let csrs = callee_saved_registers(isa.flags());

    // The reserved stack area is composed of:
    //   return address + frame pointer + all callee-saved registers
    //
    // Pushing the return address is an implicit function of the `call`
    // instruction. Each of the others we will then push explicitly. Then we
    // will adjust the stack pointer to make room for the rest of the required
    // space for this frame.
    let csr_stack_size = ((csrs.len() + 2) * word_size as usize) as i32;
    func.create_stack_slot(ir::StackSlotData {
        kind: ir::StackSlotKind::IncomingArg,
        size: csr_stack_size as u32,
        offset: -csr_stack_size,
    });

    let total_stack_size = layout_stack(&mut func.stack_slots, word_size)? as i32;
    let local_stack_size = (total_stack_size - csr_stack_size) as i64;

    // Add CSRs to function signature
    let fp_arg = ir::AbiParam::special_reg(
        csr_type,
        ir::ArgumentPurpose::FramePointer,
        RU::rbp as RegUnit,
    );
    func.signature.params.push(fp_arg);
    func.signature.returns.push(fp_arg);

    for csr in csrs.iter() {
        let csr_arg =
            ir::AbiParam::special_reg(csr_type, ir::ArgumentPurpose::CalleeSaved, *csr as RegUnit);
        func.signature.params.push(csr_arg);
        func.signature.returns.push(csr_arg);
    }

    // Set up the cursor and insert the prologue
    let entry_ebb = func.layout.entry_block().expect("missing entry block");
    let mut pos = EncCursor::new(func, isa).at_first_insertion_point(entry_ebb);
    insert_native_prologue(&mut pos, local_stack_size, csr_type, csrs);

    // Reset the cursor and insert the epilogue
    let mut pos = pos.at_position(CursorPosition::Nowhere);
    insert_native_epilogues(&mut pos, local_stack_size, csr_type, csrs);

    Ok(())
}

/// Insert the prologue for a given function.
fn insert_native_prologue(
    pos: &mut EncCursor,
    stack_size: i64,
    csr_type: ir::types::Type,
    csrs: &'static [RU],
) {
    // Append param to entry EBB
    let ebb = pos.current_ebb().expect("missing ebb under cursor");
    let fp = pos.func.dfg.append_ebb_param(ebb, csr_type);
    pos.func.locations[fp] = ir::ValueLoc::Reg(RU::rbp as RegUnit);

    pos.ins().x86_push(fp);
    pos.ins().copy_special(
        RU::rsp as RegUnit,
        RU::rbp as RegUnit,
    );

    for reg in csrs.iter() {
        // Append param to entry EBB
        let csr_arg = pos.func.dfg.append_ebb_param(ebb, csr_type);

        // Assign it a location
        pos.func.locations[csr_arg] = ir::ValueLoc::Reg(*reg as RegUnit);

        // Remember it so we can push it momentarily
        pos.ins().x86_push(csr_arg);
    }

    if stack_size > 0 {
        pos.ins().adjust_sp_imm(Imm64::new(-stack_size));
    }
}

/// Find all `return` instructions and insert epilogues before them.
fn insert_native_epilogues(
    pos: &mut EncCursor,
    stack_size: i64,
    csr_type: ir::types::Type,
    csrs: &'static [RU],
) {
    while let Some(ebb) = pos.next_ebb() {
        pos.goto_last_inst(ebb);
        if let Some(inst) = pos.current_inst() {
            if pos.func.dfg[inst].opcode().is_return() {
                insert_native_epilogue(inst, stack_size, pos, csr_type, csrs);
            }
        }
    }
}

/// Insert an epilogue given a specific `return` instruction.
fn insert_native_epilogue(
    inst: ir::Inst,
    stack_size: i64,
    pos: &mut EncCursor,
    csr_type: ir::types::Type,
    csrs: &'static [RU],
) {
    if stack_size > 0 {
        pos.ins().adjust_sp_imm(Imm64::new(stack_size));
    }

    // Pop all the callee-saved registers, stepping backward each time to
    // preserve the correct order.
    let fp_ret = pos.ins().x86_pop(csr_type);
    pos.prev_inst();

    pos.func.locations[fp_ret] = ir::ValueLoc::Reg(RU::rbp as RegUnit);
    pos.func.dfg.append_inst_arg(inst, fp_ret);

    for reg in csrs.iter() {
        let csr_ret = pos.ins().x86_pop(csr_type);
        pos.prev_inst();

        pos.func.locations[csr_ret] = ir::ValueLoc::Reg(*reg as RegUnit);
        pos.func.dfg.append_inst_arg(inst, csr_ret);
    }
}
<|endoftext|>


<commit_before>use std::vec::Vec;

use byteorder::{ByteOrder, BigEndian, WriteBytesExt};

use raftserver::{Result, other};
use std::mem;

pub const MIN_KEY: &'static [u8] = &[];
pub const MAX_KEY: &'static [u8] = &[0xFF];

// local is in (0x01, 0x02);
pub const LOCAL_PREFIX: u8 = 0x01;
pub const LOCAL_MIN_KEY: &'static [u8] = &[LOCAL_PREFIX];
pub const LOCAL_MAX_KEY: &'static [u8] = &[LOCAL_PREFIX + 1];
pub const META1_PREFIX: u8 = 0x02;
pub const META2_PREFIX: u8 = 0x03;
pub const META1_PREFIX_KEY: &'static [u8] = &[META1_PREFIX];
pub const META2_PREFIX_KEY: &'static [u8] = &[META2_PREFIX];
pub const META_MIN_KEY: &'static [u8] = &[META1_PREFIX];
pub const META_MAX_KEY: &'static [u8] = &[META2_PREFIX + 1];
pub const META1_MAX_KEY: &'static [u8] = &[META1_PREFIX, MAX_KEY[0]];
pub const META2_MAX_KEY: &'static [u8] = &[META2_PREFIX, MAX_KEY[0]];

pub const DATA_PREFIX: u8 = b'z';
pub const DATA_PREFIX_KEY: &'static [u8] = &[DATA_PREFIX];

// Following keys are all local keys, so the first byte must be 0x01.
pub const STORE_IDENT_KEY: &'static [u8] = &[LOCAL_PREFIX, 0x01];
// We save two types region data in DB, for raft and other meta data.
// When the store starts, we should iterate all region meta data to
// construct peer, no need to travel large raft data, so we separate them
// with different prefixes.
pub const REGION_RAFT_PREFIX: u8 = 0x02;
pub const REGION_RAFT_PREFIX_KEY: &'static [u8] = &[LOCAL_PREFIX, REGION_RAFT_PREFIX];
pub const REGION_META_PREFIX: u8 = 0x03;
pub const REGION_META_PREFIX_KEY: &'static [u8] = &[LOCAL_PREFIX, REGION_META_PREFIX];
pub const REGION_META_MIN_KEY: &'static [u8] = &[LOCAL_PREFIX, REGION_META_PREFIX];
pub const REGION_META_MAX_KEY: &'static [u8] = &[LOCAL_PREFIX, REGION_META_PREFIX + 1];

// Following are the suffix after the local prefix.
// For region id
pub const RAFT_LOG_SUFFIX: u8 = 0x01;
pub const RAFT_HARD_STATE_SUFFIX: u8 = 0x02;
pub const RAFT_APPLIED_INDEX_SUFFIX: u8 = 0x03;
pub const RAFT_LAST_INDEX_SUFFIX: u8 = 0x04;
pub const RAFT_TRUNCATED_STATE_SUFFIX: u8 = 0x05;

// For region meta
pub const REGION_INFO_SUFFIX: u8 = 0x01;
pub const REGION_TOMBSTONE_SUFFIX: u8 = 0x02;

pub fn store_ident_key() -> Vec<u8> {
    STORE_IDENT_KEY.to_vec()
}

fn make_region_id_key(region_id: u64, suffix: u8, extra_cap: usize) -> Vec<u8> {
    let mut key = Vec::with_capacity(REGION_RAFT_PREFIX_KEY.len() + mem::size_of::<u64>() +
                                     mem::size_of::<u8>() +
                                     extra_cap);
    key.extend_from_slice(REGION_RAFT_PREFIX_KEY);
    // no need check error here, can't panic;
    key.write_u64::<BigEndian>(region_id).unwrap();
    key.push(suffix);
    key
}

pub fn region_raft_prefix(region_id: u64) -> Vec<u8> {
    let mut key = Vec::with_capacity(REGION_RAFT_PREFIX_KEY.len() + mem::size_of::<u64>());
    key.extend_from_slice(REGION_RAFT_PREFIX_KEY);
    // no need check error here, can't panic;
    key.write_u64::<BigEndian>(region_id).unwrap();
    key
}

pub fn raft_log_key(region_id: u64, log_index: u64) -> Vec<u8> {
    let mut key = make_region_id_key(region_id, RAFT_LOG_SUFFIX, mem::size_of::<u64>());
    // no need check error here, can't panic;
    key.write_u64::<BigEndian>(log_index).unwrap();
    key
}

pub fn raft_log_prefix(region_id: u64) -> Vec<u8> {
    make_region_id_key(region_id, RAFT_LOG_SUFFIX, 0)
}

pub fn raft_hard_state_key(region_id: u64) -> Vec<u8> {
    make_region_id_key(region_id, RAFT_HARD_STATE_SUFFIX, 0)
}

pub fn raft_applied_index_key(region_id: u64) -> Vec<u8> {
    make_region_id_key(region_id, RAFT_APPLIED_INDEX_SUFFIX, 0)
}

pub fn raft_last_index_key(region_id: u64) -> Vec<u8> {
    make_region_id_key(region_id, RAFT_LAST_INDEX_SUFFIX, 0)
}

pub fn raft_truncated_state_key(region_id: u64) -> Vec<u8> {
    make_region_id_key(region_id, RAFT_TRUNCATED_STATE_SUFFIX, 0)
}

fn make_region_meta_key(region_id: u64, suffix: u8) -> Vec<u8> {
    let mut key = Vec::with_capacity(REGION_META_PREFIX_KEY.len() + mem::size_of::<u64>() +
                                     mem::size_of::<u8>());
    key.extend_from_slice(REGION_META_PREFIX_KEY);
    // no need to check error here, can't panic;
    key.write_u64::<BigEndian>(region_id).unwrap();
    key.push(suffix);
    key
}

// Decode region meta key, return the region key and meta suffix type.
pub fn decode_region_meta_key(key: &[u8]) -> Result<(u64, u8)> {
    if REGION_META_PREFIX_KEY.len() + mem::size_of::<u64>() + mem::size_of::<u8>() != key.len() {
        return Err(other(format!("invalid region meta key length for key {:?}", key)));
    }

    if !key.starts_with(REGION_META_PREFIX_KEY) {
        return Err(other(format!("invalid region meta prefix for key {:?}", key)));
    }

    let region_id =
        BigEndian::read_u64(&key[REGION_META_PREFIX_KEY.len()..REGION_META_PREFIX_KEY.len() +
                                                               mem::size_of::<u64>()]);

    Ok((region_id, key[key.len() - 1]))
}

pub fn region_meta_prefix(region_id: u64) -> Vec<u8> {
    let mut key = Vec::with_capacity(REGION_META_PREFIX_KEY.len() + mem::size_of::<u64>());
    key.extend_from_slice(REGION_META_PREFIX_KEY);
    key.write_u64::<BigEndian>(region_id).unwrap();
    key
}

pub fn region_info_key(region_id: u64) -> Vec<u8> {
    make_region_meta_key(region_id, REGION_INFO_SUFFIX)
}

pub fn region_tombstone_key(region_id: u64) -> Vec<u8> {
    make_region_meta_key(region_id, REGION_TOMBSTONE_SUFFIX)
}

// Returns a region route meta (meta1, meta2) indexing key for the
// given key.
// For data key, it returns a meta2 key, e.g, "zabc" -> \0x03"zabc"
// For meta2 key, it returns a meta1 key, e.g, \0x03\"zabc" -> \0x02"zabc"
// For meta1 key, it returns a MIN_KEY, e.g, \x02\"zabc" -> ""
pub fn region_route_meta_key(key: &[u8]) -> Vec<u8> {
    if key.len() == 0 {
        return MIN_KEY.to_vec();
    }

    match key[0] {
        META1_PREFIX => MIN_KEY.to_vec(),
        META2_PREFIX => vec![META1_PREFIX_KEY, &key[1..]].concat(),
        _ => vec![META2_PREFIX_KEY, key].concat(),
    }
}

pub fn validate_region_route_meta_key(key: &[u8]) -> Result<()> {
    if key == MIN_KEY {
        return Ok(());
    }

    // TODO: Maybe no necessary to check this, remove later?
    if key.len() < META1_PREFIX_KEY.len() {
        return Err(other(format!("{:?} is too short", key)));
    }

    let prefix = key[0];
    if prefix != META2_PREFIX && prefix != META1_PREFIX {
        return Err(other(format!("{:?} is not a meta key", key)));
    }

    // TODO: check data prefix later?
    if MAX_KEY < &key[META1_PREFIX_KEY.len()..] {
        return Err(other(format!("{:?} is > {:?}", key, MAX_KEY)));
    }

    Ok(())
}

pub fn validate_data_key(key: &[u8]) -> Result<()> {
    if !key.starts_with(DATA_PREFIX_KEY) {
        return Err(other(format!("invalid data key {:?}, must start with {}",
                                 key,
                                 DATA_PREFIX)));
    }

    Ok(())
}

pub fn data_key(key: &[u8]) -> Vec<u8> {
    let mut v = Vec::with_capacity(DATA_PREFIX_KEY.len() + key.len());
    v.extend_from_slice(DATA_PREFIX_KEY);
    v.extend_from_slice(key);
    v
}

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

    #[test]
    fn test_region_id_key() {
        let region_ids = vec![0, 1, 1024, ::std::u64::MAX];
        for region_id in region_ids {
            let prefix = region_raft_prefix(region_id);

            assert!(raft_log_prefix(region_id).starts_with(&prefix));
            assert!(raft_log_key(region_id, 1).starts_with(&prefix));
            assert!(raft_hard_state_key(region_id).starts_with(&prefix));
            assert!(raft_applied_index_key(region_id).starts_with(&prefix));
            assert!(raft_last_index_key(region_id).starts_with(&prefix));
            assert!(raft_truncated_state_key(region_id).starts_with(&prefix));
        }

        // test sort.
        let tbls = vec![(1, 0, Ordering::Greater), (1, 1, Ordering::Equal), (1, 2, Ordering::Less)];
        for (lid, rid, order) in tbls {
            let lhs = region_raft_prefix(lid);
            let rhs = region_raft_prefix(rid);
            assert_eq!(lhs.partial_cmp(&rhs), Some(order));
        }
    }

    #[test]
    fn test_raft_log_sort() {
        let tbls = vec![(1, 1, 1, 2, Ordering::Less),
                        (2, 1, 1, 2, Ordering::Greater),
                        (1, 1, 1, 1, Ordering::Equal)];

        for (lid, l_log_id, rid, r_log_id, order) in tbls {
            let lhs = raft_log_key(lid, l_log_id);
            let rhs = raft_log_key(rid, r_log_id);
            assert_eq!(lhs.partial_cmp(&rhs), Some(order));
        }
    }

    #[test]
    fn test_region_meta_key() {
        let ids: Vec<u64> = vec![1, 1024, u64::max_value()];
        for id in ids {
            let prefix = region_meta_prefix(id);
            let info_key = region_info_key(id);
            assert!(info_key.starts_with(&prefix));

            assert_eq!(decode_region_meta_key(&info_key).unwrap(),
                       (id, REGION_INFO_SUFFIX));
        }

        // test sort.
        let tbls: Vec<(u64, u64, Ordering)> = vec![
        (1, 2, Ordering::Less),
        (1, 1, Ordering::Equal),
        (2, 1, Ordering::Greater),
        ];

        for (lkey, rkey, order) in tbls {
            let lhs = region_info_key(lkey);
            let rhs = region_info_key(rkey);
            assert_eq!(lhs.partial_cmp(&rhs), Some(order));
        }
    }

    fn route_key(prefix: u8, key: &[u8]) -> Vec<u8> {
        let mut v = vec![];
        v.push(prefix);
        v.extend_from_slice(key);
        v
    }

    fn meta1_key(key: &[u8]) -> Vec<u8> {
        route_key(META1_PREFIX, key)
    }

    fn meta2_key(key: &[u8]) -> Vec<u8> {
        route_key(META2_PREFIX, key)
    }

    #[test]
    fn test_region_route_meta_key() {
        let dkey = data_key(b"abc");
        let tbls = vec![
            (vec![], vec![]),
            (meta1_key(&dkey), vec![]),
            (meta2_key(&dkey), meta1_key(&dkey)),
            (dkey.clone(), meta2_key(&dkey)),
        ];

        for (lkey, rkey) in tbls {
            assert_eq!(region_route_meta_key(&lkey), rkey);
        }

        let tbls = vec![
            (vec![], true),
            (vec![0xFF], false),
            (vec![DATA_PREFIX, 1], false),
            (vec![META1_PREFIX - 1, 1], false),
            (vec![META2_PREFIX + 1, 1], false),
            (vec![META1_PREFIX, 1], true),
            (vec![META2_PREFIX, 1], true),
            (vec![META1_PREFIX, 0xFF, 1], false),
            (vec![META2_PREFIX, 0xFF, 1], false),
        ];

        for (key, ok) in tbls {
            assert_eq!(validate_region_route_meta_key(&key).is_ok(), ok);
        }
    }

    #[test]
    fn test_data_key() {
        validate_data_key(&data_key(b"abc")).unwrap();
        validate_data_key(b"abc").unwrap_err();
    }
}
<commit_msg>raftserver: Add comment for tombstone key.<commit_after>use std::vec::Vec;

use byteorder::{ByteOrder, BigEndian, WriteBytesExt};

use raftserver::{Result, other};
use std::mem;

pub const MIN_KEY: &'static [u8] = &[];
pub const MAX_KEY: &'static [u8] = &[0xFF];

// local is in (0x01, 0x02);
pub const LOCAL_PREFIX: u8 = 0x01;
pub const LOCAL_MIN_KEY: &'static [u8] = &[LOCAL_PREFIX];
pub const LOCAL_MAX_KEY: &'static [u8] = &[LOCAL_PREFIX + 1];
pub const META1_PREFIX: u8 = 0x02;
pub const META2_PREFIX: u8 = 0x03;
pub const META1_PREFIX_KEY: &'static [u8] = &[META1_PREFIX];
pub const META2_PREFIX_KEY: &'static [u8] = &[META2_PREFIX];
pub const META_MIN_KEY: &'static [u8] = &[META1_PREFIX];
pub const META_MAX_KEY: &'static [u8] = &[META2_PREFIX + 1];
pub const META1_MAX_KEY: &'static [u8] = &[META1_PREFIX, MAX_KEY[0]];
pub const META2_MAX_KEY: &'static [u8] = &[META2_PREFIX, MAX_KEY[0]];

pub const DATA_PREFIX: u8 = b'z';
pub const DATA_PREFIX_KEY: &'static [u8] = &[DATA_PREFIX];

// Following keys are all local keys, so the first byte must be 0x01.
pub const STORE_IDENT_KEY: &'static [u8] = &[LOCAL_PREFIX, 0x01];
// We save two types region data in DB, for raft and other meta data.
// When the store starts, we should iterate all region meta data to
// construct peer, no need to travel large raft data, so we separate them
// with different prefixes.
pub const REGION_RAFT_PREFIX: u8 = 0x02;
pub const REGION_RAFT_PREFIX_KEY: &'static [u8] = &[LOCAL_PREFIX, REGION_RAFT_PREFIX];
pub const REGION_META_PREFIX: u8 = 0x03;
pub const REGION_META_PREFIX_KEY: &'static [u8] = &[LOCAL_PREFIX, REGION_META_PREFIX];
pub const REGION_META_MIN_KEY: &'static [u8] = &[LOCAL_PREFIX, REGION_META_PREFIX];
pub const REGION_META_MAX_KEY: &'static [u8] = &[LOCAL_PREFIX, REGION_META_PREFIX + 1];

// Following are the suffix after the local prefix.
// For region id
pub const RAFT_LOG_SUFFIX: u8 = 0x01;
pub const RAFT_HARD_STATE_SUFFIX: u8 = 0x02;
pub const RAFT_APPLIED_INDEX_SUFFIX: u8 = 0x03;
pub const RAFT_LAST_INDEX_SUFFIX: u8 = 0x04;
pub const RAFT_TRUNCATED_STATE_SUFFIX: u8 = 0x05;

// For region meta
pub const REGION_INFO_SUFFIX: u8 = 0x01;
pub const REGION_TOMBSTONE_SUFFIX: u8 = 0x02;

pub fn store_ident_key() -> Vec<u8> {
    STORE_IDENT_KEY.to_vec()
}

fn make_region_id_key(region_id: u64, suffix: u8, extra_cap: usize) -> Vec<u8> {
    let mut key = Vec::with_capacity(REGION_RAFT_PREFIX_KEY.len() + mem::size_of::<u64>() +
                                     mem::size_of::<u8>() +
                                     extra_cap);
    key.extend_from_slice(REGION_RAFT_PREFIX_KEY);
    // no need check error here, can't panic;
    key.write_u64::<BigEndian>(region_id).unwrap();
    key.push(suffix);
    key
}

pub fn region_raft_prefix(region_id: u64) -> Vec<u8> {
    let mut key = Vec::with_capacity(REGION_RAFT_PREFIX_KEY.len() + mem::size_of::<u64>());
    key.extend_from_slice(REGION_RAFT_PREFIX_KEY);
    // no need check error here, can't panic;
    key.write_u64::<BigEndian>(region_id).unwrap();
    key
}

pub fn raft_log_key(region_id: u64, log_index: u64) -> Vec<u8> {
    let mut key = make_region_id_key(region_id, RAFT_LOG_SUFFIX, mem::size_of::<u64>());
    // no need check error here, can't panic;
    key.write_u64::<BigEndian>(log_index).unwrap();
    key
}

pub fn raft_log_prefix(region_id: u64) -> Vec<u8> {
    make_region_id_key(region_id, RAFT_LOG_SUFFIX, 0)
}

pub fn raft_hard_state_key(region_id: u64) -> Vec<u8> {
    make_region_id_key(region_id, RAFT_HARD_STATE_SUFFIX, 0)
}

pub fn raft_applied_index_key(region_id: u64) -> Vec<u8> {
    make_region_id_key(region_id, RAFT_APPLIED_INDEX_SUFFIX, 0)
}

pub fn raft_last_index_key(region_id: u64) -> Vec<u8> {
    make_region_id_key(region_id, RAFT_LAST_INDEX_SUFFIX, 0)
}

pub fn raft_truncated_state_key(region_id: u64) -> Vec<u8> {
    make_region_id_key(region_id, RAFT_TRUNCATED_STATE_SUFFIX, 0)
}

fn make_region_meta_key(region_id: u64, suffix: u8) -> Vec<u8> {
    let mut key = Vec::with_capacity(REGION_META_PREFIX_KEY.len() + mem::size_of::<u64>() +
                                     mem::size_of::<u8>());
    key.extend_from_slice(REGION_META_PREFIX_KEY);
    // no need to check error here, can't panic;
    key.write_u64::<BigEndian>(region_id).unwrap();
    key.push(suffix);
    key
}

// Decode region meta key, return the region key and meta suffix type.
pub fn decode_region_meta_key(key: &[u8]) -> Result<(u64, u8)> {
    if REGION_META_PREFIX_KEY.len() + mem::size_of::<u64>() + mem::size_of::<u8>() != key.len() {
        return Err(other(format!("invalid region meta key length for key {:?}", key)));
    }

    if !key.starts_with(REGION_META_PREFIX_KEY) {
        return Err(other(format!("invalid region meta prefix for key {:?}", key)));
    }

    let region_id =
        BigEndian::read_u64(&key[REGION_META_PREFIX_KEY.len()..REGION_META_PREFIX_KEY.len() +
                                                               mem::size_of::<u64>()]);

    Ok((region_id, key[key.len() - 1]))
}

pub fn region_meta_prefix(region_id: u64) -> Vec<u8> {
    let mut key = Vec::with_capacity(REGION_META_PREFIX_KEY.len() + mem::size_of::<u64>());
    key.extend_from_slice(REGION_META_PREFIX_KEY);
    key.write_u64::<BigEndian>(region_id).unwrap();
    key
}

pub fn region_info_key(region_id: u64) -> Vec<u8> {
    make_region_meta_key(region_id, REGION_INFO_SUFFIX)
}

// When a peer is destroyed, we would record current region max peer id as
// the tombstone value, any newer peer for this region with a peer id <= tombstone
// value is not allowed to create in this store.
pub fn region_tombstone_key(region_id: u64) -> Vec<u8> {
    make_region_meta_key(region_id, REGION_TOMBSTONE_SUFFIX)
}

// Returns a region route meta (meta1, meta2) indexing key for the
// given key.
// For data key, it returns a meta2 key, e.g, "zabc" -> \0x03"zabc"
// For meta2 key, it returns a meta1 key, e.g, \0x03\"zabc" -> \0x02"zabc"
// For meta1 key, it returns a MIN_KEY, e.g, \x02\"zabc" -> ""
pub fn region_route_meta_key(key: &[u8]) -> Vec<u8> {
    if key.len() == 0 {
        return MIN_KEY.to_vec();
    }

    match key[0] {
        META1_PREFIX => MIN_KEY.to_vec(),
        META2_PREFIX => vec![META1_PREFIX_KEY, &key[1..]].concat(),
        _ => vec![META2_PREFIX_KEY, key].concat(),
    }
}

pub fn validate_region_route_meta_key(key: &[u8]) -> Result<()> {
    if key == MIN_KEY {
        return Ok(());
    }

    // TODO: Maybe no necessary to check this, remove later?
    if key.len() < META1_PREFIX_KEY.len() {
        return Err(other(format!("{:?} is too short", key)));
    }

    let prefix = key[0];
    if prefix != META2_PREFIX && prefix != META1_PREFIX {
        return Err(other(format!("{:?} is not a meta key", key)));
    }

    // TODO: check data prefix later?
    if MAX_KEY < &key[META1_PREFIX_KEY.len()..] {
        return Err(other(format!("{:?} is > {:?}", key, MAX_KEY)));
    }

    Ok(())
}

pub fn validate_data_key(key: &[u8]) -> Result<()> {
    if !key.starts_with(DATA_PREFIX_KEY) {
        return Err(other(format!("invalid data key {:?}, must start with {}",
                                 key,
                                 DATA_PREFIX)));
    }

    Ok(())
}

pub fn data_key(key: &[u8]) -> Vec<u8> {
    let mut v = Vec::with_capacity(DATA_PREFIX_KEY.len() + key.len());
    v.extend_from_slice(DATA_PREFIX_KEY);
    v.extend_from_slice(key);
    v
}

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

    #[test]
    fn test_region_id_key() {
        let region_ids = vec![0, 1, 1024, ::std::u64::MAX];
        for region_id in region_ids {
            let prefix = region_raft_prefix(region_id);

            assert!(raft_log_prefix(region_id).starts_with(&prefix));
            assert!(raft_log_key(region_id, 1).starts_with(&prefix));
            assert!(raft_hard_state_key(region_id).starts_with(&prefix));
            assert!(raft_applied_index_key(region_id).starts_with(&prefix));
            assert!(raft_last_index_key(region_id).starts_with(&prefix));
            assert!(raft_truncated_state_key(region_id).starts_with(&prefix));
        }

        // test sort.
        let tbls = vec![(1, 0, Ordering::Greater), (1, 1, Ordering::Equal), (1, 2, Ordering::Less)];
        for (lid, rid, order) in tbls {
            let lhs = region_raft_prefix(lid);
            let rhs = region_raft_prefix(rid);
            assert_eq!(lhs.partial_cmp(&rhs), Some(order));
        }
    }

    #[test]
    fn test_raft_log_sort() {
        let tbls = vec![(1, 1, 1, 2, Ordering::Less),
                        (2, 1, 1, 2, Ordering::Greater),
                        (1, 1, 1, 1, Ordering::Equal)];

        for (lid, l_log_id, rid, r_log_id, order) in tbls {
            let lhs = raft_log_key(lid, l_log_id);
            let rhs = raft_log_key(rid, r_log_id);
            assert_eq!(lhs.partial_cmp(&rhs), Some(order));
        }
    }

    #[test]
    fn test_region_meta_key() {
        let ids: Vec<u64> = vec![1, 1024, u64::max_value()];
        for id in ids {
            let prefix = region_meta_prefix(id);
            let info_key = region_info_key(id);
            assert!(info_key.starts_with(&prefix));

            assert_eq!(decode_region_meta_key(&info_key).unwrap(),
                       (id, REGION_INFO_SUFFIX));
        }

        // test sort.
        let tbls: Vec<(u64, u64, Ordering)> = vec![
        (1, 2, Ordering::Less),
        (1, 1, Ordering::Equal),
        (2, 1, Ordering::Greater),
        ];

        for (lkey, rkey, order) in tbls {
            let lhs = region_info_key(lkey);
            let rhs = region_info_key(rkey);
            assert_eq!(lhs.partial_cmp(&rhs), Some(order));
        }
    }

    fn route_key(prefix: u8, key: &[u8]) -> Vec<u8> {
        let mut v = vec![];
        v.push(prefix);
        v.extend_from_slice(key);
        v
    }

    fn meta1_key(key: &[u8]) -> Vec<u8> {
        route_key(META1_PREFIX, key)
    }

    fn meta2_key(key: &[u8]) -> Vec<u8> {
        route_key(META2_PREFIX, key)
    }

    #[test]
    fn test_region_route_meta_key() {
        let dkey = data_key(b"abc");
        let tbls = vec![
            (vec![], vec![]),
            (meta1_key(&dkey), vec![]),
            (meta2_key(&dkey), meta1_key(&dkey)),
            (dkey.clone(), meta2_key(&dkey)),
        ];

        for (lkey, rkey) in tbls {
            assert_eq!(region_route_meta_key(&lkey), rkey);
        }

        let tbls = vec![
            (vec![], true),
            (vec![0xFF], false),
            (vec![DATA_PREFIX, 1], false),
            (vec![META1_PREFIX - 1, 1], false),
            (vec![META2_PREFIX + 1, 1], false),
            (vec![META1_PREFIX, 1], true),
            (vec![META2_PREFIX, 1], true),
            (vec![META1_PREFIX, 0xFF, 1], false),
            (vec![META2_PREFIX, 0xFF, 1], false),
        ];

        for (key, ok) in tbls {
            assert_eq!(validate_region_route_meta_key(&key).is_ok(), ok);
        }
    }

    #[test]
    fn test_data_key() {
        validate_data_key(&data_key(b"abc")).unwrap();
        validate_data_key(b"abc").unwrap_err();
    }
}
<|endoftext|>


<commit_before>//! A simplified implementation of the `bytes` crate, with different features, less safety.
//!
//! #Examples
//!
//! ```
//! use bytes::rc::Bytes;
//!
//! let bytes = vec![0u8; 1024];
//! let mut shared1 = Bytes::from(bytes);
//! let mut shared2 = shared1.extract_to(100);
//! let mut shared3 = shared1.extract_to(100);
//! let mut shared4 = shared2.extract_to(60);
//!
//! assert_eq!(shared1.len(), 824);
//! assert_eq!(shared2.len(), 40);
//! assert_eq!(shared3.len(), 100);
//! assert_eq!(shared4.len(), 60);
//!
//! for byte in shared1.iter_mut() { *byte = 1u8; }
//! for byte in shared2.iter_mut() { *byte = 2u8; }
//! for byte in shared3.iter_mut() { *byte = 3u8; }
//! for byte in shared4.iter_mut() { *byte = 4u8; }
//!
//! drop(shared1);
//! drop(shared2);
//! drop(shared3);
//!
//! if let Ok(bytes) = shared4.try_recover() {
//!     assert_eq!(bytes[200..1024].to_vec(), [1u8;824].to_vec());
//!     assert_eq!(bytes[60..100].to_vec(),   [2u8;40].to_vec());
//!     assert_eq!(bytes[100..200].to_vec(),  [3u8;100].to_vec());
//!     assert_eq!(bytes[0..60].to_vec(),     [4u8;60].to_vec());
//! }
//! else {
//!     panic!("unrecoverable!");
//! }
//! ```
#![forbid(missing_docs)]

/// An `Rc`-backed mutable byte slice backed by a common allocation.
pub mod rc {

    use std::ops::{Deref, DerefMut};
    use std::rc::Rc;

    /// A thread-local byte buffer backed by a shared allocation.
    pub struct Bytes<B: DerefMut<Target=[u8]>> {
        /// Pointer to the start of this slice (not the allocation).
        ptr: *mut u8,
        /// Length of this slice.
        len: usize,
        /// Shared access to underlying resources.
        ///
        /// Importantly, this is unavailable for as long as the struct exists, which may
        /// prevent shared access to ptr[0 .. len]. I'm not sure I understand Rust's rules
        /// enough to make a strong statement about this.
        sequestered: Rc<B>,
    }

    impl<B: DerefMut<Target=[u8]>> Bytes<B> {

        /// Create a new instance from a byte allocation.
        pub fn from(bytes: B) -> Bytes<B> {

            let mut rc = Rc::new(bytes);

            Bytes {
                ptr: Rc::get_mut(&mut rc).unwrap().as_mut_ptr(),
                len: rc.len(),
                sequestered: rc,
            }
        }

        /// Extracts [0, index) into a new `Bytes` which is returned, updating `self`.
        ///
        /// #Safety
        ///
        /// This method uses an `unsafe` region to advance the pointer by `index`. It first
        /// tests `index` against `self.len`, which should ensure that the offset is in-bounds.
        pub fn extract_to(&mut self, index: usize) -> Bytes<B> {

            assert!(index <= self.len);

            let result = Bytes {
                ptr: self.ptr,
                len: index,
                sequestered: self.sequestered.clone(),
            };

            unsafe { self.ptr = self.ptr.offset(index as isize); }
            self.len -= index;

            result
        }

        /// Recover the underlying storage.
        ///
        /// This method either results in the underlying storage if it is uniquely held, or the
        /// input `Bytes` if it is not uniquely held.
        pub fn try_recover(self) -> Result<B, Bytes<B>> {
            match Rc::try_unwrap(self.sequestered) {
                Ok(bytes) => Ok(bytes),
                Err(rc) => Err(Bytes {
                    ptr: self.ptr,
                    len: self.len,
                    sequestered: rc,
                }),
            }
        }
    }

    impl<B: DerefMut<Target=[u8]>+Clone> Deref for Bytes<B> {
        type Target = [u8];
        fn deref(&self) -> &[u8] {
            unsafe { ::std::slice::from_raw_parts(self.ptr, self.len) }
        }
    }

    impl<B: DerefMut<Target=[u8]>+Clone> DerefMut for Bytes<B> {
        fn deref_mut(&mut self) -> &mut [u8] {
            unsafe { ::std::slice::from_raw_parts_mut(self.ptr, self.len) }
        }
    }
}

/// An `Arc`-backed mutable byte slice backed by a common allocation.
pub mod arc {

    use std::ops::{Deref, DerefMut};
    use std::sync::Arc;

    /// A thread-safe byte buffer backed by a shared allocation.
    pub struct Bytes<B: DerefMut<Target=[u8]>> {
        /// Pointer to the start of this slice (not the allocation).
        ptr: *mut u8,
        /// Length of this slice.
        len: usize,
        /// Shared access to underlying resources.
        ///
        /// Importantly, this is unavailable for as long as the struct exists, which may
        /// prevent shared access to ptr[0 .. len]. I'm not sure I understand Rust's rules
        /// enough to make a strong statement about this.
        sequestered: Arc<B>,
    }

    impl<B: DerefMut<Target=[u8]>> Bytes<B> {

        /// Create a new instance from a byte allocation.
        pub fn from(bytes: B) -> Bytes<B> {

            let mut arc = Arc::new(bytes);

            Bytes {
                ptr: Arc::get_mut(&mut arc).unwrap().as_mut_ptr(),
                len: arc.len(),
                sequestered: arc,
            }
        }

        /// Extracts [0, index) into a new `Bytes` which is returned, updating `self`.
        ///
        /// #Safety
        ///
        /// This method uses an `unsafe` region to advance the pointer by `index`. It first
        /// tests `index` against `self.len`, which should ensure that the offset is in-bounds.
        pub fn extract_to(&mut self, index: usize) -> Bytes<B> {

            assert!(index <= self.len);

            let result = Bytes {
                ptr: self.ptr,
                len: index,
                sequestered: self.sequestered.clone(),
            };

            unsafe { self.ptr = self.ptr.offset(index as isize); }
            self.len -= index;

            result
        }

        /// Recover the underlying storage.
        ///
        /// This method either results in the underlying storage if it is uniquely held, or the
        /// input `Bytes` if it is not uniquely held.
        pub fn try_recover(self) -> Result<B, Bytes<B>> {
            match Arc::try_unwrap(self.sequestered) {
                Ok(bytes) => Ok(bytes),
                Err(arc) => Err(Bytes {
                    ptr: self.ptr,
                    len: self.len,
                    sequestered: arc,
                }),
            }
        }
    }

    impl<B: DerefMut<Target=[u8]>+Clone> Deref for Bytes<B> {
        type Target = [u8];
        fn deref(&self) -> &[u8] {
            unsafe { ::std::slice::from_raw_parts(self.ptr, self.len) }
        }
    }

    impl<B: DerefMut<Target=[u8]>+Clone> DerefMut for Bytes<B> {
        fn deref_mut(&mut self) -> &mut [u8] {
            unsafe { ::std::slice::from_raw_parts_mut(self.ptr, self.len) }
        }
    }
}<commit_msg>remove clone requirement<commit_after>//! A simplified implementation of the `bytes` crate, with different features, less safety.
//!
//! #Examples
//!
//! ```
//! use bytes::rc::Bytes;
//!
//! let bytes = vec![0u8; 1024];
//! let mut shared1 = Bytes::from(bytes);
//! let mut shared2 = shared1.extract_to(100);
//! let mut shared3 = shared1.extract_to(100);
//! let mut shared4 = shared2.extract_to(60);
//!
//! assert_eq!(shared1.len(), 824);
//! assert_eq!(shared2.len(), 40);
//! assert_eq!(shared3.len(), 100);
//! assert_eq!(shared4.len(), 60);
//!
//! for byte in shared1.iter_mut() { *byte = 1u8; }
//! for byte in shared2.iter_mut() { *byte = 2u8; }
//! for byte in shared3.iter_mut() { *byte = 3u8; }
//! for byte in shared4.iter_mut() { *byte = 4u8; }
//!
//! drop(shared1);
//! drop(shared2);
//! drop(shared3);
//!
//! if let Ok(bytes) = shared4.try_recover() {
//!     assert_eq!(bytes[200..1024].to_vec(), [1u8;824].to_vec());
//!     assert_eq!(bytes[60..100].to_vec(),   [2u8;40].to_vec());
//!     assert_eq!(bytes[100..200].to_vec(),  [3u8;100].to_vec());
//!     assert_eq!(bytes[0..60].to_vec(),     [4u8;60].to_vec());
//! }
//! else {
//!     panic!("unrecoverable!");
//! }
//! ```
#![forbid(missing_docs)]

/// An `Rc`-backed mutable byte slice backed by a common allocation.
pub mod rc {

    use std::ops::{Deref, DerefMut};
    use std::rc::Rc;

    /// A thread-local byte buffer backed by a shared allocation.
    pub struct Bytes<B: DerefMut<Target=[u8]>> {
        /// Pointer to the start of this slice (not the allocation).
        ptr: *mut u8,
        /// Length of this slice.
        len: usize,
        /// Shared access to underlying resources.
        ///
        /// Importantly, this is unavailable for as long as the struct exists, which may
        /// prevent shared access to ptr[0 .. len]. I'm not sure I understand Rust's rules
        /// enough to make a strong statement about this.
        sequestered: Rc<B>,
    }

    impl<B: DerefMut<Target=[u8]>> Bytes<B> {

        /// Create a new instance from a byte allocation.
        pub fn from(bytes: B) -> Bytes<B> {

            let mut rc = Rc::new(bytes);

            Bytes {
                ptr: Rc::get_mut(&mut rc).unwrap().as_mut_ptr(),
                len: rc.len(),
                sequestered: rc,
            }
        }

        /// Extracts [0, index) into a new `Bytes` which is returned, updating `self`.
        ///
        /// #Safety
        ///
        /// This method uses an `unsafe` region to advance the pointer by `index`. It first
        /// tests `index` against `self.len`, which should ensure that the offset is in-bounds.
        pub fn extract_to(&mut self, index: usize) -> Bytes<B> {

            assert!(index <= self.len);

            let result = Bytes {
                ptr: self.ptr,
                len: index,
                sequestered: self.sequestered.clone(),
            };

            unsafe { self.ptr = self.ptr.offset(index as isize); }
            self.len -= index;

            result
        }

        /// Recover the underlying storage.
        ///
        /// This method either results in the underlying storage if it is uniquely held, or the
        /// input `Bytes` if it is not uniquely held.
        pub fn try_recover(self) -> Result<B, Bytes<B>> {
            match Rc::try_unwrap(self.sequestered) {
                Ok(bytes) => Ok(bytes),
                Err(rc) => Err(Bytes {
                    ptr: self.ptr,
                    len: self.len,
                    sequestered: rc,
                }),
            }
        }
    }

    impl<B: DerefMut<Target=[u8]>> Deref for Bytes<B> {
        type Target = [u8];
        fn deref(&self) -> &[u8] {
            unsafe { ::std::slice::from_raw_parts(self.ptr, self.len) }
        }
    }

    impl<B: DerefMut<Target=[u8]>> DerefMut for Bytes<B> {
        fn deref_mut(&mut self) -> &mut [u8] {
            unsafe { ::std::slice::from_raw_parts_mut(self.ptr, self.len) }
        }
    }
}

/// An `Arc`-backed mutable byte slice backed by a common allocation.
pub mod arc {

    use std::ops::{Deref, DerefMut};
    use std::sync::Arc;

    /// A thread-safe byte buffer backed by a shared allocation.
    pub struct Bytes<B: DerefMut<Target=[u8]>> {
        /// Pointer to the start of this slice (not the allocation).
        ptr: *mut u8,
        /// Length of this slice.
        len: usize,
        /// Shared access to underlying resources.
        ///
        /// Importantly, this is unavailable for as long as the struct exists, which may
        /// prevent shared access to ptr[0 .. len]. I'm not sure I understand Rust's rules
        /// enough to make a strong statement about this.
        sequestered: Arc<B>,
    }

    impl<B: DerefMut<Target=[u8]>> Bytes<B> {

        /// Create a new instance from a byte allocation.
        pub fn from(bytes: B) -> Bytes<B> {

            let mut arc = Arc::new(bytes);

            Bytes {
                ptr: Arc::get_mut(&mut arc).unwrap().as_mut_ptr(),
                len: arc.len(),
                sequestered: arc,
            }
        }

        /// Extracts [0, index) into a new `Bytes` which is returned, updating `self`.
        ///
        /// #Safety
        ///
        /// This method uses an `unsafe` region to advance the pointer by `index`. It first
        /// tests `index` against `self.len`, which should ensure that the offset is in-bounds.
        pub fn extract_to(&mut self, index: usize) -> Bytes<B> {

            assert!(index <= self.len);

            let result = Bytes {
                ptr: self.ptr,
                len: index,
                sequestered: self.sequestered.clone(),
            };

            unsafe { self.ptr = self.ptr.offset(index as isize); }
            self.len -= index;

            result
        }

        /// Recover the underlying storage.
        ///
        /// This method either results in the underlying storage if it is uniquely held, or the
        /// input `Bytes` if it is not uniquely held.
        pub fn try_recover(self) -> Result<B, Bytes<B>> {
            match Arc::try_unwrap(self.sequestered) {
                Ok(bytes) => Ok(bytes),
                Err(arc) => Err(Bytes {
                    ptr: self.ptr,
                    len: self.len,
                    sequestered: arc,
                }),
            }
        }
    }

    impl<B: DerefMut<Target=[u8]>> Deref for Bytes<B> {
        type Target = [u8];
        fn deref(&self) -> &[u8] {
            unsafe { ::std::slice::from_raw_parts(self.ptr, self.len) }
        }
    }

    impl<B: DerefMut<Target=[u8]>> DerefMut for Bytes<B> {
        fn deref_mut(&mut self) -> &mut [u8] {
            unsafe { ::std::slice::from_raw_parts_mut(self.ptr, self.len) }
        }
    }
}<|endoftext|>