<commit_before>//! A Rust interface for the functionality of the Objective-C runtime.
//!
//! For more information on foreign functions, see Apple's documentation:
//! https://developer.apple.com/library/mac/documentation/Cocoa/Reference/ObjCRuntimeRef/index.html

use std::ffi::{CString, self};
use std::mem;
use std::ptr;
use std::str;
use libc::{c_char, c_int, c_uint, c_void, ptrdiff_t, size_t};
use malloc_buf::{MallocBuffer, MallocString};

use {encode, Encode};

unsafe fn from_c_str<'a>(ptr: *const c_char) -> &'a str {
    let bytes = ffi::c_str_to_bytes(&ptr);
    let s = str::from_utf8(bytes).unwrap();
    mem::transmute(s)
}

/// A type that represents a method selector.
#[repr(C)]
pub struct Sel {
    ptr: *const c_void,
}

/// A type that represents an instance variable.
#[allow(missing_copy_implementations)]
pub enum Ivar { }

/// A type that represents a method in a class definition.
#[allow(missing_copy_implementations)]
pub enum Method { }

/// A type that represents an Objective-C class.
#[allow(missing_copy_implementations)]
pub enum Class { }

/// A type that represents an instance of a class.
#[allow(missing_copy_implementations)]
pub enum Object { }

/// A pointer to the start of a method implementation.
pub type Imp = extern fn(*mut Object, Sel, ...) -> *mut Object;

#[allow(improper_ctypes)]
#[link(name = "Foundation", kind = "framework")]
extern {
    pub fn sel_registerName(name: *const c_char) -> Sel;
    pub fn sel_getName(sel: Sel) -> *const c_char;

    pub fn class_getName(cls: *const Class) -> *const c_char;
    pub fn class_getInstanceSize(cls: *const Class) -> size_t;
    pub fn class_getInstanceMethod(cls: *const Class, sel: Sel) -> *const Method;
    pub fn class_getInstanceVariable(cls: *const Class, name: *const c_char) -> *const Ivar;
    pub fn class_copyMethodList(cls: *const Class, outCount: *mut c_uint) -> *mut *const Method;
    pub fn class_copyIvarList(cls: *const Class, outCount: *mut c_uint) -> *mut *const Ivar;
    pub fn class_addMethod(cls: *mut Class, name: Sel, imp: Imp, types: *const c_char) -> bool;
    pub fn class_addIvar(cls: *mut Class, name: *const c_char, size: size_t, alignment: u8, types: *const c_char) -> bool;

    pub fn objc_allocateClassPair(superclass: *const Class, name: *const c_char, extraBytes: size_t) -> *mut Class;
    pub fn objc_disposeClassPair(cls: *mut Class);
    pub fn objc_registerClassPair(cls: *mut Class);

    pub fn object_getClass(obj: *const Object) -> *const Class;

    pub fn objc_getClassList(buffer: *mut *const Class, bufferLen: c_int) -> c_int;
    pub fn objc_copyClassList(outCount: *mut c_uint) -> *mut *const Class;
    pub fn objc_getClass(name: *const c_char) -> *const Class;

    pub fn ivar_getName(ivar: *const Ivar) -> *const c_char;
    pub fn ivar_getOffset(ivar: *const Ivar) -> ptrdiff_t;
    pub fn ivar_getTypeEncoding(ivar: *const Ivar) -> *const c_char;

    pub fn objc_msgSend(obj: *mut Object, op: Sel, ...) -> *mut Object;

    pub fn method_getName(method: *const Method) -> Sel;
    pub fn method_getImplementation(method: *const Method) -> Imp;
    pub fn method_getTypeEncoding(method: *const Method) -> *const c_char;
    pub fn method_copyReturnType(method: *const Method) -> *mut c_char;
    pub fn method_copyArgumentType(method: *const Method, index: c_uint) -> *mut c_char;
    pub fn method_getNumberOfArguments(method: *const Method) -> c_uint;
    pub fn method_setImplementation(method: *mut Method, imp: Imp) -> Imp;
    pub fn method_exchangeImplementations(m1: *mut Method, m2: *mut Method);
}

impl Sel {
    /// Registers a method with the Objective-C runtime system,
    /// maps the method name to a selector, and returns the selector value.
    pub fn register(name: &str) -> Sel {
        let name = CString::from_slice(name.as_bytes());
        unsafe {
            sel_registerName(name.as_ptr())
        }
    }

    /// Returns the name of the method specified by self.
    pub fn name(&self) -> &str {
        unsafe {
            let name = sel_getName(*self);
            from_c_str(name)
        }
    }
}

impl PartialEq for Sel {
    fn eq(&self, other: &Sel) -> bool {
        self.ptr == other.ptr
    }
}

impl Eq for Sel { }

impl Copy for Sel { }

impl Clone for Sel {
    fn clone(&self) -> Sel { *self }
}

impl Ivar {
    /// Returns the name of self.
    pub fn name(&self) -> &str {
        unsafe {
            let name = ivar_getName(self);
            from_c_str(name)
        }
    }

    /// Returns the offset of self.
    pub fn offset(&self) -> isize {
        let offset = unsafe {
            ivar_getOffset(self)
        };
        offset as isize
    }

    /// Returns the type string of self.
    pub fn type_encoding(&self) -> &str {
        unsafe {
            let encoding = ivar_getTypeEncoding(self);
            from_c_str(encoding)
        }
    }
}

impl Method {
    /// Returns the name of self.
    pub fn name(&self) -> Sel {
        unsafe {
            method_getName(self)
        }
    }

    /// Returns a string describing self's parameter and return types.
    pub fn type_encoding(&self) -> &str {
        unsafe {
            let encoding = method_getTypeEncoding(self);
            from_c_str(encoding)
        }
    }

    /// Returns a string describing self's return type.
    pub fn return_type(&self) -> MallocString {
        unsafe {
            let encoding = method_copyReturnType(self);
            MallocString::new(encoding).unwrap()
        }
    }

    /// Returns a string describing a single parameter type of self, or
    /// `None` if self has no parameter at the given index.
    pub fn argument_type(&self, index: usize) -> Option<MallocString> {
        unsafe {
            let encoding = method_copyArgumentType(self, index as c_uint);
            if encoding.is_null() {
                None
            } else {
                Some(MallocString::new(encoding).unwrap())
            }
        }
    }

    /// Returns the number of arguments accepted by self.
    pub fn arguments_count(&self) -> usize {
        unsafe {
            method_getNumberOfArguments(self) as usize
        }
    }

    /// Returns the implementation of self.
    pub fn implementation(&self) -> Imp {
        unsafe {
            method_getImplementation(self)
        }
    }

    /// Sets the implementation of self.
    /// Unsafe because the caller must ensure the implementation has the
    /// correct self, return, and argument types for the method.
    pub unsafe fn set_implementation(&mut self, imp: Imp) -> Imp {
        method_setImplementation(self, imp)
    }

    /// Exchanges the implementations of self with another `Method`.
    /// Unsafe because the caller must ensure the two methods have the same
    /// self, return, and argument types.
    pub unsafe fn exchange_implementation(&mut self, other: &mut Method) {
        method_exchangeImplementations(self, other);
    }
}

impl Class {
    /// Returns the class definition of a specified class, or `None` if the
    /// class is not registered with the Objective-C runtime.
    pub fn get(name: &str) -> Option<&'static Class> {
        let name = CString::from_slice(name.as_bytes());
        unsafe {
            objc_getClass(name.as_ptr()).as_ref()
        }
    }

    /// Obtains the list of registered class definitions.
    pub fn classes() -> MallocBuffer<&'static Class> {
        unsafe {
            let mut count: c_uint = 0;
            let classes = objc_copyClassList(&mut count);
            MallocBuffer::new(classes as *mut _, count as usize).unwrap()
        }
    }

    /// Returns the total number of registered classes.
    pub fn classes_count() -> usize {
        unsafe {
            objc_getClassList(ptr::null_mut(), 0) as usize
        }
    }

    /// Returns the name of self.
    pub fn name(&self) -> &str {
        unsafe {
            let name = class_getName(self);
            from_c_str(name)
        }
    }

    /// Returns the size of instances of self.
    pub fn instance_size(&self) -> usize {
        unsafe {
            class_getInstanceSize(self) as usize
        }
    }

    /// Returns a specified instance method for self, or `None` if self and
    /// its superclasses do not contain an instance method with the
    /// specified selector.
    pub fn instance_method(&self, sel: Sel) -> Option<&Method> {
        unsafe {
            class_getInstanceMethod(self, sel).as_ref()
        }
    }

    /// Returns the ivar for a specified instance variable of self, or `None`
    /// if self has no ivar with the given name.
    pub fn instance_variable(&self, name: &str) -> Option<&Ivar> {
        let name = CString::from_slice(name.as_bytes());
        unsafe {
            class_getInstanceVariable(self, name.as_ptr()).as_ref()
        }
    }

    /// Describes the instance methods implemented by self.
    pub fn instance_methods(&self) -> MallocBuffer<&Method> {
        unsafe {
            let mut count: c_uint = 0;
            let methods = class_copyMethodList(self, &mut count);
            MallocBuffer::new(methods as *mut _, count as usize).unwrap()
        }

    }

    /// Describes the instance variables declared by self.
    pub fn instance_variables(&self) -> MallocBuffer<&Ivar> {
        unsafe {
            let mut count: c_uint = 0;
            let ivars = class_copyIvarList(self, &mut count);
            MallocBuffer::new(ivars as *mut _, count as usize).unwrap()
        }
    }
}

impl PartialEq for Class {
    fn eq(&self, other: &Class) -> bool {
        self as *const Class == other as *const Class
    }
}

impl Eq for Class { }

impl Object {
    /// Returns the class of self.
    pub fn class(&self) -> &'static Class {
        unsafe {
            &*object_getClass(self)
        }
    }

    /// Returns a reference to the ivar of self with the given name.
    /// Panics if self has no ivar with the given name.
    /// Unsafe because the caller must ensure that the ivar is actually
    /// of type `T`.
    pub unsafe fn get_ivar<T>(&self, name: &str) -> &T where T: Encode {
        let cls = self.class();
        let ptr = match cls.instance_variable(name) {
            Some(ivar) => {
                assert!(ivar.type_encoding() == encode::<T>());
                let offset = ivar.offset();
                let self_ptr = self as *const Object;
                (self_ptr as *const u8).offset(offset) as *const T
            }
            None => panic!("Ivar {} not found on class {}", name, cls.name()),
        };
        &*ptr
    }

    /// Returns a mutable reference to the ivar of self with the given name.
    /// Panics if self has no ivar with the given name.
    /// Unsafe because the caller must ensure that the ivar is actually
    /// of type `T`.
    pub unsafe fn get_mut_ivar<T>(&mut self, name: &str) -> &mut T
            where T: Encode {
        let cls = self.class();
        let ptr = match cls.instance_variable(name) {
            Some(ivar) => {
                assert!(ivar.type_encoding() == encode::<T>());
                let offset = ivar.offset();
                let self_ptr = self as *mut Object;
                (self_ptr as *mut u8).offset(offset) as *mut T
            }
            None => panic!("Ivar {} not found on class {}", name, cls.name()),
        };
        &mut *ptr
    }

    /// Sets the value of the ivar of self with the given name.
    /// Panics if self has no ivar with the given name.
    /// Unsafe because the caller must ensure that the ivar is actually
    /// of type `T`.
    pub unsafe fn set_ivar<T>(&mut self, name: &str, value: T)
            where T: Encode {
        *self.get_mut_ivar::<T>(name) = value;
    }
}

#[cfg(test)]
mod tests {
    use std::mem;
    use super::{Class, Object, Sel};

    #[test]
    fn test_ivar() {
        let cls = Class::get("NSObject").unwrap();
        let ivar = cls.instance_variable("isa").unwrap();
        assert!(ivar.name() == "isa");
        assert!(ivar.type_encoding() == "#");
        assert!(ivar.offset() == 0);

        let ivars = cls.instance_variables();
        assert!(ivars.as_slice().len() > 0);
    }

    #[test]
    fn test_method() {
        let cls = Class::get("NSObject").unwrap();
        let sel = Sel::register("description");
        let method = cls.instance_method(sel).unwrap();
        assert!(method.name().name() == "description");
        assert!(method.type_encoding() != "");
        assert!(method.arguments_count() == 2);
        assert!(method.return_type().as_slice() == "@");
        assert!(method.argument_type(1).unwrap().as_slice() == ":");

        let methods = cls.instance_methods();
        assert!(methods.as_slice().len() > 0);
    }

    #[test]
    fn test_class() {
        let cls = Class::get("NSObject").unwrap();
        assert!(cls.name() == "NSObject");
        assert!(cls.instance_size() == mem::size_of::<*const Class>());
    }

    #[test]
    fn test_classes() {
        let classes_count = Class::classes_count();
        assert!(classes_count > 0);

        let classes = Class::classes();
        assert!(classes.as_slice().len() == classes_count);
    }

    #[test]
    fn test_object() {
        let cls = Class::get("NSObject").unwrap();
        let obj = unsafe {
            let obj: *mut Object = msg_send![cls, alloc];
            let obj: *mut Object = msg_send![obj, init];
            &*obj
        };
        assert!(obj.class() == cls);
        unsafe {
            let _: () = msg_send![obj, release];
        }
    }
}
<commit_msg>Implemented Show for some runtime types.<commit_after>//! A Rust interface for the functionality of the Objective-C runtime.
//!
//! For more information on foreign functions, see Apple's documentation:
//! https://developer.apple.com/library/mac/documentation/Cocoa/Reference/ObjCRuntimeRef/index.html

use std::ffi::{CString, self};
use std::fmt;
use std::mem;
use std::ptr;
use std::str;
use libc::{c_char, c_int, c_uint, c_void, ptrdiff_t, size_t};
use malloc_buf::{MallocBuffer, MallocString};

use {encode, Encode};

unsafe fn from_c_str<'a>(ptr: *const c_char) -> &'a str {
    let bytes = ffi::c_str_to_bytes(&ptr);
    let s = str::from_utf8(bytes).unwrap();
    mem::transmute(s)
}

/// A type that represents a method selector.
#[repr(C)]
pub struct Sel {
    ptr: *const c_void,
}

/// A type that represents an instance variable.
#[allow(missing_copy_implementations)]
pub enum Ivar { }

/// A type that represents a method in a class definition.
#[allow(missing_copy_implementations)]
pub enum Method { }

/// A type that represents an Objective-C class.
#[allow(missing_copy_implementations)]
pub enum Class { }

/// A type that represents an instance of a class.
#[allow(missing_copy_implementations)]
pub enum Object { }

/// A pointer to the start of a method implementation.
pub type Imp = extern fn(*mut Object, Sel, ...) -> *mut Object;

#[allow(improper_ctypes)]
#[link(name = "Foundation", kind = "framework")]
extern {
    pub fn sel_registerName(name: *const c_char) -> Sel;
    pub fn sel_getName(sel: Sel) -> *const c_char;

    pub fn class_getName(cls: *const Class) -> *const c_char;
    pub fn class_getInstanceSize(cls: *const Class) -> size_t;
    pub fn class_getInstanceMethod(cls: *const Class, sel: Sel) -> *const Method;
    pub fn class_getInstanceVariable(cls: *const Class, name: *const c_char) -> *const Ivar;
    pub fn class_copyMethodList(cls: *const Class, outCount: *mut c_uint) -> *mut *const Method;
    pub fn class_copyIvarList(cls: *const Class, outCount: *mut c_uint) -> *mut *const Ivar;
    pub fn class_addMethod(cls: *mut Class, name: Sel, imp: Imp, types: *const c_char) -> bool;
    pub fn class_addIvar(cls: *mut Class, name: *const c_char, size: size_t, alignment: u8, types: *const c_char) -> bool;

    pub fn objc_allocateClassPair(superclass: *const Class, name: *const c_char, extraBytes: size_t) -> *mut Class;
    pub fn objc_disposeClassPair(cls: *mut Class);
    pub fn objc_registerClassPair(cls: *mut Class);

    pub fn object_getClass(obj: *const Object) -> *const Class;

    pub fn objc_getClassList(buffer: *mut *const Class, bufferLen: c_int) -> c_int;
    pub fn objc_copyClassList(outCount: *mut c_uint) -> *mut *const Class;
    pub fn objc_getClass(name: *const c_char) -> *const Class;

    pub fn ivar_getName(ivar: *const Ivar) -> *const c_char;
    pub fn ivar_getOffset(ivar: *const Ivar) -> ptrdiff_t;
    pub fn ivar_getTypeEncoding(ivar: *const Ivar) -> *const c_char;

    pub fn objc_msgSend(obj: *mut Object, op: Sel, ...) -> *mut Object;

    pub fn method_getName(method: *const Method) -> Sel;
    pub fn method_getImplementation(method: *const Method) -> Imp;
    pub fn method_getTypeEncoding(method: *const Method) -> *const c_char;
    pub fn method_copyReturnType(method: *const Method) -> *mut c_char;
    pub fn method_copyArgumentType(method: *const Method, index: c_uint) -> *mut c_char;
    pub fn method_getNumberOfArguments(method: *const Method) -> c_uint;
    pub fn method_setImplementation(method: *mut Method, imp: Imp) -> Imp;
    pub fn method_exchangeImplementations(m1: *mut Method, m2: *mut Method);
}

impl Sel {
    /// Registers a method with the Objective-C runtime system,
    /// maps the method name to a selector, and returns the selector value.
    pub fn register(name: &str) -> Sel {
        let name = CString::from_slice(name.as_bytes());
        unsafe {
            sel_registerName(name.as_ptr())
        }
    }

    /// Returns the name of the method specified by self.
    pub fn name(&self) -> &str {
        unsafe {
            let name = sel_getName(*self);
            from_c_str(name)
        }
    }
}

impl PartialEq for Sel {
    fn eq(&self, other: &Sel) -> bool {
        self.ptr == other.ptr
    }
}

impl Eq for Sel { }

impl Copy for Sel { }

impl Clone for Sel {
    fn clone(&self) -> Sel { *self }
}

impl fmt::Show for Sel {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.name())
    }
}

impl Ivar {
    /// Returns the name of self.
    pub fn name(&self) -> &str {
        unsafe {
            let name = ivar_getName(self);
            from_c_str(name)
        }
    }

    /// Returns the offset of self.
    pub fn offset(&self) -> isize {
        let offset = unsafe {
            ivar_getOffset(self)
        };
        offset as isize
    }

    /// Returns the type string of self.
    pub fn type_encoding(&self) -> &str {
        unsafe {
            let encoding = ivar_getTypeEncoding(self);
            from_c_str(encoding)
        }
    }
}

impl Method {
    /// Returns the name of self.
    pub fn name(&self) -> Sel {
        unsafe {
            method_getName(self)
        }
    }

    /// Returns a string describing self's parameter and return types.
    pub fn type_encoding(&self) -> &str {
        unsafe {
            let encoding = method_getTypeEncoding(self);
            from_c_str(encoding)
        }
    }

    /// Returns a string describing self's return type.
    pub fn return_type(&self) -> MallocString {
        unsafe {
            let encoding = method_copyReturnType(self);
            MallocString::new(encoding).unwrap()
        }
    }

    /// Returns a string describing a single parameter type of self, or
    /// `None` if self has no parameter at the given index.
    pub fn argument_type(&self, index: usize) -> Option<MallocString> {
        unsafe {
            let encoding = method_copyArgumentType(self, index as c_uint);
            if encoding.is_null() {
                None
            } else {
                Some(MallocString::new(encoding).unwrap())
            }
        }
    }

    /// Returns the number of arguments accepted by self.
    pub fn arguments_count(&self) -> usize {
        unsafe {
            method_getNumberOfArguments(self) as usize
        }
    }

    /// Returns the implementation of self.
    pub fn implementation(&self) -> Imp {
        unsafe {
            method_getImplementation(self)
        }
    }

    /// Sets the implementation of self.
    /// Unsafe because the caller must ensure the implementation has the
    /// correct self, return, and argument types for the method.
    pub unsafe fn set_implementation(&mut self, imp: Imp) -> Imp {
        method_setImplementation(self, imp)
    }

    /// Exchanges the implementations of self with another `Method`.
    /// Unsafe because the caller must ensure the two methods have the same
    /// self, return, and argument types.
    pub unsafe fn exchange_implementation(&mut self, other: &mut Method) {
        method_exchangeImplementations(self, other);
    }
}

impl Class {
    /// Returns the class definition of a specified class, or `None` if the
    /// class is not registered with the Objective-C runtime.
    pub fn get(name: &str) -> Option<&'static Class> {
        let name = CString::from_slice(name.as_bytes());
        unsafe {
            objc_getClass(name.as_ptr()).as_ref()
        }
    }

    /// Obtains the list of registered class definitions.
    pub fn classes() -> MallocBuffer<&'static Class> {
        unsafe {
            let mut count: c_uint = 0;
            let classes = objc_copyClassList(&mut count);
            MallocBuffer::new(classes as *mut _, count as usize).unwrap()
        }
    }

    /// Returns the total number of registered classes.
    pub fn classes_count() -> usize {
        unsafe {
            objc_getClassList(ptr::null_mut(), 0) as usize
        }
    }

    /// Returns the name of self.
    pub fn name(&self) -> &str {
        unsafe {
            let name = class_getName(self);
            from_c_str(name)
        }
    }

    /// Returns the size of instances of self.
    pub fn instance_size(&self) -> usize {
        unsafe {
            class_getInstanceSize(self) as usize
        }
    }

    /// Returns a specified instance method for self, or `None` if self and
    /// its superclasses do not contain an instance method with the
    /// specified selector.
    pub fn instance_method(&self, sel: Sel) -> Option<&Method> {
        unsafe {
            class_getInstanceMethod(self, sel).as_ref()
        }
    }

    /// Returns the ivar for a specified instance variable of self, or `None`
    /// if self has no ivar with the given name.
    pub fn instance_variable(&self, name: &str) -> Option<&Ivar> {
        let name = CString::from_slice(name.as_bytes());
        unsafe {
            class_getInstanceVariable(self, name.as_ptr()).as_ref()
        }
    }

    /// Describes the instance methods implemented by self.
    pub fn instance_methods(&self) -> MallocBuffer<&Method> {
        unsafe {
            let mut count: c_uint = 0;
            let methods = class_copyMethodList(self, &mut count);
            MallocBuffer::new(methods as *mut _, count as usize).unwrap()
        }

    }

    /// Describes the instance variables declared by self.
    pub fn instance_variables(&self) -> MallocBuffer<&Ivar> {
        unsafe {
            let mut count: c_uint = 0;
            let ivars = class_copyIvarList(self, &mut count);
            MallocBuffer::new(ivars as *mut _, count as usize).unwrap()
        }
    }
}

impl PartialEq for Class {
    fn eq(&self, other: &Class) -> bool {
        self as *const Class == other as *const Class
    }
}

impl Eq for Class { }

impl fmt::Show for Class {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}", self.name())
    }
}

impl Object {
    /// Returns the class of self.
    pub fn class(&self) -> &'static Class {
        unsafe {
            &*object_getClass(self)
        }
    }

    /// Returns a reference to the ivar of self with the given name.
    /// Panics if self has no ivar with the given name.
    /// Unsafe because the caller must ensure that the ivar is actually
    /// of type `T`.
    pub unsafe fn get_ivar<T>(&self, name: &str) -> &T where T: Encode {
        let cls = self.class();
        let ptr = match cls.instance_variable(name) {
            Some(ivar) => {
                assert!(ivar.type_encoding() == encode::<T>());
                let offset = ivar.offset();
                let self_ptr = self as *const Object;
                (self_ptr as *const u8).offset(offset) as *const T
            }
            None => panic!("Ivar {} not found on class {}", name, cls.name()),
        };
        &*ptr
    }

    /// Returns a mutable reference to the ivar of self with the given name.
    /// Panics if self has no ivar with the given name.
    /// Unsafe because the caller must ensure that the ivar is actually
    /// of type `T`.
    pub unsafe fn get_mut_ivar<T>(&mut self, name: &str) -> &mut T
            where T: Encode {
        let cls = self.class();
        let ptr = match cls.instance_variable(name) {
            Some(ivar) => {
                assert!(ivar.type_encoding() == encode::<T>());
                let offset = ivar.offset();
                let self_ptr = self as *mut Object;
                (self_ptr as *mut u8).offset(offset) as *mut T
            }
            None => panic!("Ivar {} not found on class {}", name, cls.name()),
        };
        &mut *ptr
    }

    /// Sets the value of the ivar of self with the given name.
    /// Panics if self has no ivar with the given name.
    /// Unsafe because the caller must ensure that the ivar is actually
    /// of type `T`.
    pub unsafe fn set_ivar<T>(&mut self, name: &str, value: T)
            where T: Encode {
        *self.get_mut_ivar::<T>(name) = value;
    }
}

impl fmt::Show for Object {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "<{:?}: {:?}>", self.class(), self as *const Object)
    }
}

#[cfg(test)]
mod tests {
    use std::mem;
    use super::{Class, Object, Sel};

    #[test]
    fn test_ivar() {
        let cls = Class::get("NSObject").unwrap();
        let ivar = cls.instance_variable("isa").unwrap();
        assert!(ivar.name() == "isa");
        assert!(ivar.type_encoding() == "#");
        assert!(ivar.offset() == 0);

        let ivars = cls.instance_variables();
        assert!(ivars.as_slice().len() > 0);
    }

    #[test]
    fn test_method() {
        let cls = Class::get("NSObject").unwrap();
        let sel = Sel::register("description");
        let method = cls.instance_method(sel).unwrap();
        assert!(method.name().name() == "description");
        assert!(method.type_encoding() != "");
        assert!(method.arguments_count() == 2);
        assert!(method.return_type().as_slice() == "@");
        assert!(method.argument_type(1).unwrap().as_slice() == ":");

        let methods = cls.instance_methods();
        assert!(methods.as_slice().len() > 0);
    }

    #[test]
    fn test_class() {
        let cls = Class::get("NSObject").unwrap();
        assert!(cls.name() == "NSObject");
        assert!(cls.instance_size() == mem::size_of::<*const Class>());
    }

    #[test]
    fn test_classes() {
        let classes_count = Class::classes_count();
        assert!(classes_count > 0);

        let classes = Class::classes();
        assert!(classes.as_slice().len() == classes_count);
    }

    #[test]
    fn test_object() {
        let cls = Class::get("NSObject").unwrap();
        let obj = unsafe {
            let obj: *mut Object = msg_send![cls, alloc];
            let obj: *mut Object = msg_send![obj, init];
            &*obj
        };
        assert!(obj.class() == cls);
        unsafe {
            let _: () = msg_send![obj, release];
        }
    }
}
<|endoftext|>


<commit_before>use bus::{Addressable, OAM_START, OAM_END, VIDEO_RAM_START, VIDEO_RAM_END};
use enum_primitive::FromPrimitive;

pub const SCREEN_WIDTH: usize = 160;
pub const SCREEN_HEIGHT: usize = 144;

const ADDR_LCDC: u16 = 0xFF40;
const ADDR_STAT: u16 = 0xFF41;
const ADDR_SCY: u16  = 0xFF42;
const ADDR_SCX: u16  = 0xFF43;
const ADDR_LY: u16   = 0xFF44;
const ADDR_LYC: u16  = 0xFF45;
const ADDR_BGP: u16  = 0xFF47;
const ADDR_OBP0: u16 = 0xFF48;
const ADDR_OBP1: u16 = 0xFF49;
const ADDR_WY: u16   = 0xFF4A;
const ADDR_WX: u16   = 0xFF4B;

bitflags! {
    struct Lcdc: u8 {
        const LCDC_ENABLED            = 0b1000_0000;
        const LCDC_WIN_TILE_9C        = 0b0100_0000;
        const LCDC_WIN_DISPLAY        = 0b0010_0000;
        const LCDC_UNSIGNED_TILE_DATA = 0b0001_0000;
        const LCDC_BG_TILE_9C         = 0b0000_1000;
        const LCDC_16PX_SPRITE        = 0b0000_0100;
        const LCDC_SPRITE_DISPLAY     = 0b0000_0010;
        const LCDC_BG_DISPLAY         = 0b0000_0001;
    }
}

// STAT is mostly flags, but the lower 2 bits are the current mode
// Mode is will stored separately so the Stat struct is only treated as bitflags.
bitflags! {
    struct Stat: u8 {
        const STAT_COINCIDENCE_INT   = 0b0100_0000;
        const STAT_OAM_INT           = 0b0010_0000;
        const STAT_VBLANK_INT        = 0b0001_0000;
        const STAT_HBLANK_INT        = 0b0000_1000;
        const STAT_COINCIDENCE_EQUAL = 0b0000_0100;
    }
}

enum_from_primitive! {
    #[derive(Copy, Clone, PartialEq)]
    pub enum Mode {
        VBlank   = 0,
        HBlank   = 1,
        Oam      = 2,
        Transfer = 3
    }
}

enum_from_primitive! {
    #[derive(Copy, Clone, PartialEq)]
    enum Shade {
        White       = 0,
        LightGray   = 1,
        DarkGray    = 2,
        Black       = 3
    }
}

const CYCLES_PER_OAM_READ: usize = 80;
const CYCLES_PER_TRANSFER: usize = 172;
const CYCLES_PER_HBLANK: usize = 204;
const CYCLES_PER_LINE: usize = 456;

#[derive(Copy, Clone)]
struct Palette(u8);

impl Palette {
    fn rgb(&self, shade: Shade) -> u32 {
        match shade {
            Shade::White        => 0x9CBD0F,
            Shade::LightGray    => 0x8CAD0F,
            Shade::DarkGray     => 0x306230,
            Shade::Black        => 0x0F380F
        }
    }
}

#[derive(Default)]
pub struct StepResult {
    pub int_vblank: bool,
    pub int_stat: bool
}

pub struct Lcd {
    vram: [u8; (VIDEO_RAM_END - VIDEO_RAM_START) as usize + 1],
    oam: [u8; (OAM_END - OAM_START) as usize + 1],
    lcdc: Lcdc,
    stat: Stat,
    scy: u8,
    scx: u8,
    ly: u8,
    lyc: u8,
    bgp: Palette,
    obp0: Palette,
    obp1: Palette,
    wy: u8,
    wx: u8,
    mode: Mode,
    mode_cycles: usize,
    draw_pending: bool
}

impl Lcd {
    pub fn new() -> Self {
        Self {
            vram: [0; (VIDEO_RAM_END - VIDEO_RAM_START) as usize + 1],
            oam: [0; (OAM_END - OAM_START) as usize + 1],
            lcdc: Lcdc::from_bits(0x91).unwrap(),
            stat: Stat::empty(),
            scy: 0,
            scx: 0,
            ly: 0,
            lyc: 0,
            bgp: Palette(0xFC),
            obp0: Palette(0xFF),
            obp1: Palette(0xFF),
            wy: 0,
            wx: 0,
            mode: Mode::Oam,
            mode_cycles: 0,
            draw_pending: false
        }
    }

    pub fn step(&mut self, cycles: usize, screen_buffer: &mut [u32]) -> StepResult {
        let mut result = StepResult::default();

        if self.lcdc.contains(Lcdc::LCDC_ENABLED) {
            let previous_mode = self.mode;
            self.mode_cycles = self.mode_cycles + cycles;

            match self.mode {
                Mode::Oam => {
                    // Process starts with the LCD controller reading information from OAM.
                    // This lasts 77-83 cycles
                    if self.mode_cycles >= CYCLES_PER_OAM_READ {
                        // Move onto transfer stage
                        self.mode = Mode::Transfer;
                        self.mode_cycles = 0;
                    }
                },
                Mode::Transfer => {
                    // Data is being actively read by the LCD driver from OAM & VRAM.
                    // This lasts 169-175 cycles
                    if self.mode_cycles >= CYCLES_PER_TRANSFER {
                        // Once transfer is complete, enter HBLANK
                        self.mode = Mode::HBlank;
                        self.mode_cycles = 0;
                        self.draw_pending = true;
                    }
                },
                Mode::HBlank => {
                    // LCD has reached the end of a line.
                    // HBlank phase lasts 201-207 cycles
                    if self.mode_cycles >= CYCLES_PER_HBLANK {
                        // LCD is ready to begin the next line
                        self.ly = self.ly + 1;
                        self.mode_cycles = 0;

                        // If the LCD has line 144, then enter VBLANK. Else, enter OAM read.
                        if self.ly == 144 {
                            self.mode = Mode::VBlank;
                            self.mode_cycles = 0;
                            result.int_vblank = true;

                            if self.lcdc.contains(Lcdc::LCDC_SPRITE_DISPLAY) {
                                self.draw_sprites(screen_buffer);
                            }
                        } else {
                            self.mode = Mode::Oam;
                        };

                        self.check_coincidence(&mut result);
                    } else {
                        if self.ly < 144 && self.draw_pending {
                            if self.lcdc.contains(Lcdc::LCDC_BG_DISPLAY) {
                                self.draw_background_line(screen_buffer);
                            }

                            self.draw_pending = false;
                        }
                    }
                },
                Mode::VBlank => {
                    // LCD is writing to VBlank lines
                    if self.mode_cycles >= CYCLES_PER_LINE {
                        self.mode_cycles = 0;
                        self.ly = self.ly + 1;

                        if self.ly > 153 {
                            // Finished with VBlank. Enter OAM with the first line.
                            self.ly = 0;
                            self.mode = Mode::Oam;
                        }

                        self.check_coincidence(&mut result);
                    }
                }
            }

            // Check if a new mode has been entered during this step.
            // If so, an interrupt will be raised if the respective flag is set in the STAT register
            if self.mode != previous_mode {
                let flag = match self.mode {
                    Mode::VBlank   => Some(Stat::STAT_VBLANK_INT),
                    Mode::HBlank   => Some(Stat::STAT_HBLANK_INT),
                    Mode::Oam      => Some(Stat::STAT_OAM_INT),
                    Mode::Transfer => None
                };

                if flag.is_some() && self.stat.contains(flag.unwrap()) {
                    result.int_stat = true;
                }
            }
        }

        result
    }

    fn draw_background_line(&self, screen_buffer: &mut [u32]) {
        let map_y = self.scy.wrapping_add(self.ly);
        let map_tile_y = (map_y / 8) as u16;
        let tile_row = map_y & 7;

        let map_addr_base = if self.lcdc.contains(Lcdc::LCDC_BG_TILE_9C) {
            0x9C00
        } else {
            0x9800
        } as u16;

        for screen_x in 0..SCREEN_WIDTH {
            let map_x = self.scx.wrapping_add(screen_x as u8);
            let map_tile_x = (map_x / 8) as u16;
            let tile_column = map_x & 7;

            let tile_idx_addr = map_addr_base + (map_tile_y * 32) + map_tile_x;
            let tile_idx = self.vram[(tile_idx_addr - VIDEO_RAM_START) as usize];
            let tile_addr = self.bg_tile_address(tile_idx);
            let shade = self.tile_pixel_shade(tile_addr, tile_row, tile_column);

            screen_buffer[(self.ly as usize * SCREEN_WIDTH) + screen_x] = self.bgp.rgb(shade);
        }
    }

    fn draw_sprites(&self, screen_buffer: &mut [u32]) {
        // Sprites are 40 blocks in OAM. Each block is 32 bits
        for idx in 0..40 {
            let y = self.oam[idx * 4] as usize;
            let x = self.oam[(idx * 4) + 1] as usize;
            let pattern =  self.oam[(idx * 4) + 2];
            let flags = self.oam[(idx * 4) + 3];

            if x >= 8 && y >= 16 {
                let tile_addr = self.sprite_tile_address(pattern);
                let screen_x = x.wrapping_sub(8) as usize;
                let screen_y = y.wrapping_sub(16) as usize;
                let palette = if flags & 0b10000 == 0b10000 { self.obp1 } else { self.obp0 };

                for row in 0..8 as usize {
                    for column in 0..8 as usize {
                        let idx = ((screen_y + row) * SCREEN_WIDTH) + screen_x + column;
                        let shade = self.tile_pixel_shade(tile_addr, row as u8, column as u8);

                        if idx < screen_buffer.len() && shade != Shade::White {
                            screen_buffer[idx] = palette.rgb(shade);
                        }
                    }
                }
            }
        }
    }

    fn tile_pixel_shade(&self, address: u16, row: u8, column: u8) -> Shade {
        let upper_byte = self.vram[((address + (row * 2) as u16) + 1 - VIDEO_RAM_START) as usize];
        let lower_byte = self.vram[((address + (row * 2) as u16) - VIDEO_RAM_START) as usize];

        let shift = 7 - column;
        let mask = 1 << shift;
        let upper_bit = (upper_byte & mask) >> shift;
        let lower_bit = (lower_byte & mask) >> shift;

        Shade::from_u8((upper_bit << 1) | lower_bit).unwrap()
    }

    fn check_coincidence(&mut self, result: &mut StepResult) {
        if self.stat.contains(Stat::STAT_COINCIDENCE_INT) {
            if self.ly == self.lyc {
                self.stat.insert(Stat::STAT_COINCIDENCE_EQUAL);
                result.int_stat = true;
            } else {
                self.stat.remove(Stat::STAT_COINCIDENCE_EQUAL);
            }
        } else {
            self.stat.remove(Stat::STAT_COINCIDENCE_EQUAL);
        }
    }

    fn bg_tile_address(&self, tile_index: u8) -> u16 {
        let tile_addr_base = if self.lcdc.contains(Lcdc::LCDC_UNSIGNED_TILE_DATA) {
            0x8000
        } else {
            0x9000
        } as u16;

        if self.lcdc.contains(Lcdc::LCDC_UNSIGNED_TILE_DATA) {
            (tile_addr_base + ((tile_index as u16) * 16))
        } else {
            let signed_index = (tile_index as i8) as i16;
            ((tile_addr_base as i16) + ((signed_index) * 16)) as u16
        }
    }

    fn sprite_tile_address(&self, tile_index: u8) -> u16 {
        0x8000 + ((tile_index as u16) * 16)
    }
}

impl Addressable for Lcd {
    fn read(&self, addr: u16) -> u8 {
        match addr {
            VIDEO_RAM_START..=VIDEO_RAM_END => self.vram[(addr - VIDEO_RAM_START) as usize],
            OAM_START..=OAM_END => self.oam[(addr - OAM_START) as usize],
            ADDR_LCDC => self.lcdc.bits,
            ADDR_STAT => 0b1000_0000 | ((self.stat.bits & 0b1111_1100) | (self.mode as u8)),
            ADDR_SCY => self.scy,
            ADDR_SCX => self.scx,
            ADDR_LY => self.ly,
            ADDR_LYC => self.lyc,
            ADDR_BGP => self.bgp.0,
            ADDR_OBP0 => self.obp0.0,
            ADDR_OBP1 => self.obp1.0,
            ADDR_WY => self.wy,
            ADDR_WX => self.wx,
            _ => { println!("LCD IO read unimplemented ({:#X})", addr); 0 }
        }
    }

    fn write(&mut self, addr: u16, val: u8) {
        match addr {
            VIDEO_RAM_START..=VIDEO_RAM_END => self.vram[(addr - VIDEO_RAM_START) as usize] = val,
            OAM_START..=OAM_END => self.oam[(addr - OAM_START) as usize] = val,
            ADDR_LCDC => self.lcdc.bits = val,
            ADDR_STAT => {
                self.stat.bits = val & 0b1111_1100;
                self.mode = Mode::from_u8(val & 0b11).unwrap();
            },
            ADDR_SCY => self.scy = val,
            ADDR_SCX => self.scx = val,
            ADDR_LY => (), // read-only
            ADDR_LYC => self.lyc = val,
            ADDR_BGP => self.bgp = Palette(val),
            ADDR_OBP0 => self.obp0 = Palette(val),
            ADDR_OBP1 => self.obp1 = Palette(val),
            ADDR_WY => self.wy = val,
            ADDR_WX => self.wx = val,
            _ => println!("LCD IO write unimplemented {:#X} -> {:#X}", val, addr)
        }
    }
}<commit_msg>LCD fixes<commit_after>use bus::{Addressable, OAM_START, OAM_END, VIDEO_RAM_START, VIDEO_RAM_END};
use enum_primitive::FromPrimitive;

pub const SCREEN_WIDTH: usize = 160;
pub const SCREEN_HEIGHT: usize = 144;

const ADDR_LCDC: u16 = 0xFF40;
const ADDR_STAT: u16 = 0xFF41;
const ADDR_SCY: u16  = 0xFF42;
const ADDR_SCX: u16  = 0xFF43;
const ADDR_LY: u16   = 0xFF44;
const ADDR_LYC: u16  = 0xFF45;
const ADDR_BGP: u16  = 0xFF47;
const ADDR_OBP0: u16 = 0xFF48;
const ADDR_OBP1: u16 = 0xFF49;
const ADDR_WY: u16   = 0xFF4A;
const ADDR_WX: u16   = 0xFF4B;

bitflags! {
    struct Lcdc: u8 {
        const LCDC_ENABLED            = 0b1000_0000;
        const LCDC_WIN_TILE_9C        = 0b0100_0000;
        const LCDC_WIN_DISPLAY        = 0b0010_0000;
        const LCDC_UNSIGNED_TILE_DATA = 0b0001_0000;
        const LCDC_BG_TILE_9C         = 0b0000_1000;
        const LCDC_16PX_SPRITE        = 0b0000_0100;
        const LCDC_SPRITE_DISPLAY     = 0b0000_0010;
        const LCDC_BG_DISPLAY         = 0b0000_0001;
    }
}

// STAT is mostly flags, but the lower 2 bits are the current mode
// Mode is will stored separately so the Stat struct is only treated as bitflags.
bitflags! {
    struct Stat: u8 {
        const STAT_COINCIDENCE_INT   = 0b0100_0000;
        const STAT_OAM_INT           = 0b0010_0000;
        const STAT_VBLANK_INT        = 0b0001_0000;
        const STAT_HBLANK_INT        = 0b0000_1000;
        const STAT_COINCIDENCE_EQUAL = 0b0000_0100;
    }
}

enum_from_primitive! {
    #[derive(Copy, Clone, PartialEq)]
    pub enum Mode {
        HBlank   = 0,
        VBlank   = 1,
        Oam      = 2,
        Transfer = 3
    }
}

enum_from_primitive! {
    #[derive(Copy, Clone, PartialEq)]
    enum Shade {
        White       = 0,
        LightGray   = 1,
        DarkGray    = 2,
        Black       = 3
    }
}

const CYCLES_PER_OAM_READ: usize = 80;
const CYCLES_PER_TRANSFER: usize = 172;
const CYCLES_PER_HBLANK: usize = 204;
const CYCLES_PER_LINE: usize = 456;

#[derive(Copy, Clone)]
struct Palette(u8);

impl Palette {
    fn rgb(&self, shade: Shade) -> u32 {
        match shade {
            Shade::White        => 0x9CBD0F,
            Shade::LightGray    => 0x8CAD0F,
            Shade::DarkGray     => 0x306230,
            Shade::Black        => 0x0F380F
        }
    }
}

#[derive(Default)]
pub struct StepResult {
    pub int_vblank: bool,
    pub int_stat: bool
}

pub struct Lcd {
    vram: [u8; (VIDEO_RAM_END - VIDEO_RAM_START) as usize + 1],
    oam: [u8; (OAM_END - OAM_START) as usize + 1],
    lcdc: Lcdc,
    stat: Stat,
    scy: u8,
    scx: u8,
    ly: u8,
    lyc: u8,
    bgp: Palette,
    obp0: Palette,
    obp1: Palette,
    wy: u8,
    wx: u8,
    mode: Mode,
    mode_cycles: usize,
    draw_pending: bool
}

impl Lcd {
    pub fn new() -> Self {
        Self {
            vram: [0; (VIDEO_RAM_END - VIDEO_RAM_START) as usize + 1],
            oam: [0; (OAM_END - OAM_START) as usize + 1],
            lcdc: Lcdc::from_bits(0x91).unwrap(),
            stat: Stat::empty(),
            scy: 0,
            scx: 0,
            ly: 0,
            lyc: 0,
            bgp: Palette(0xFC),
            obp0: Palette(0xFF),
            obp1: Palette(0xFF),
            wy: 0,
            wx: 0,
            mode: Mode::Oam,
            mode_cycles: 0,
            draw_pending: false
        }
    }

    pub fn step(&mut self, cycles: usize, screen_buffer: &mut [u32]) -> StepResult {
        let mut result = StepResult::default();

        if self.lcdc.contains(Lcdc::LCDC_ENABLED) {
            let previous_mode = self.mode;
            self.mode_cycles += cycles;

            match self.mode {
                Mode::Oam => {
                    // Process starts with the LCD controller reading information from OAM.
                    // This lasts 77-83 cycles
                    if self.mode_cycles >= CYCLES_PER_OAM_READ {
                        // Move onto transfer stage
                        self.mode = Mode::Transfer;
                        self.mode_cycles -= CYCLES_PER_OAM_READ;
                    }
                },
                Mode::Transfer => {
                    // Data is being actively read by the LCD driver from OAM & VRAM.
                    // This lasts 169-175 cycles
                    if self.mode_cycles >= CYCLES_PER_TRANSFER {
                        // Once transfer is complete, enter HBLANK
                        self.mode = Mode::HBlank;
                        self.mode_cycles -= CYCLES_PER_TRANSFER;
                        self.draw_pending = true;
                    }
                },
                Mode::HBlank => {
                    // LCD has reached the end of a line.
                    // HBlank phase lasts 201-207 cycles
                    if self.mode_cycles >= CYCLES_PER_HBLANK {
                        // LCD is ready to begin the next line
                        self.ly += 1;
                        self.mode_cycles -= CYCLES_PER_HBLANK;

                        // If the LCD has line 144, then enter VBLANK. Else, enter OAM read.
                        if self.ly == 144 {
                            self.mode = Mode::VBlank;
                            result.int_vblank = true;

                            if self.lcdc.contains(Lcdc::LCDC_SPRITE_DISPLAY) {
                                self.draw_sprites(screen_buffer);
                            }
                        } else {
                            self.mode = Mode::Oam;

                            self.check_coincidence(&mut result);
                        };
                    } else {
                        if self.ly < 144 && self.draw_pending {
                            if self.lcdc.contains(Lcdc::LCDC_BG_DISPLAY) {
                                self.draw_background_line(screen_buffer);
                            }

                            self.draw_pending = false;
                        }
                    }
                },
                Mode::VBlank => {
                    // LCD is writing to VBlank lines
                    if self.mode_cycles >= CYCLES_PER_LINE {
                        self.mode_cycles -= CYCLES_PER_LINE;
                        self.ly = self.ly + 1;

                        if self.ly > 153 {
                            // Finished with VBlank. Enter OAM with the first line.
                            self.ly = 0;
                            self.mode = Mode::Oam;

                            self.check_coincidence(&mut result);
                        }
                    }
                }
            }

            // Check if a new mode has been entered during this step.
            // If so, an interrupt will be raised if the respective flag is set in the STAT register
            if self.mode != previous_mode {
                let raise = match self.mode {
                    Mode::HBlank   => self.stat.contains(Stat::STAT_HBLANK_INT) || self.stat.contains(Stat::STAT_OAM_INT),
                    Mode::VBlank   => self.stat.contains(Stat::STAT_VBLANK_INT),
                    Mode::Oam      => self.stat.contains(Stat::STAT_OAM_INT),
                    Mode::Transfer => false
                };

                if raise {
                    result.int_stat = true;
                }
            }
        }

        result
    }

    fn draw_background_line(&self, screen_buffer: &mut [u32]) {
        let map_y = self.scy.wrapping_add(self.ly);
        let map_tile_y = (map_y / 8) as u16;
        let tile_row = map_y & 7;

        let map_addr_base = if self.lcdc.contains(Lcdc::LCDC_BG_TILE_9C) {
            0x9C00
        } else {
            0x9800
        } as u16;

        for screen_x in 0..SCREEN_WIDTH {
            let map_x = self.scx.wrapping_add(screen_x as u8);
            let map_tile_x = (map_x / 8) as u16;
            let tile_column = map_x & 7;

            let tile_idx_addr = map_addr_base + (map_tile_y * 32) + map_tile_x;
            let tile_idx = self.vram[(tile_idx_addr - VIDEO_RAM_START) as usize];
            let tile_addr = self.bg_tile_address(tile_idx);
            let shade = self.tile_pixel_shade(tile_addr, tile_row, tile_column);

            screen_buffer[(self.ly as usize * SCREEN_WIDTH) + screen_x] = self.bgp.rgb(shade);
        }
    }

    fn draw_sprites(&self, screen_buffer: &mut [u32]) {
        // Sprites are 40 blocks in OAM. Each block is 32 bits
        for idx in 0..40 {
            let y = self.oam[idx * 4] as usize;
            let x = self.oam[(idx * 4) + 1] as usize;
            let pattern =  self.oam[(idx * 4) + 2];
            let flags = self.oam[(idx * 4) + 3];

            if x >= 8 && y >= 16 {
                let tile_addr = self.sprite_tile_address(pattern);
                let screen_x = x.wrapping_sub(8) as usize;
                let screen_y = y.wrapping_sub(16) as usize;
                let palette = if flags & 0b10000 == 0b10000 { self.obp1 } else { self.obp0 };

                for row in 0..8 as usize {
                    for column in 0..8 as usize {
                        let idx = ((screen_y + row) * SCREEN_WIDTH) + screen_x + column;
                        let shade = self.tile_pixel_shade(tile_addr, row as u8, column as u8);

                        if idx < screen_buffer.len() && shade != Shade::White {
                            screen_buffer[idx] = palette.rgb(shade);
                        }
                    }
                }
            }
        }
    }

    fn tile_pixel_shade(&self, address: u16, row: u8, column: u8) -> Shade {
        let upper_byte = self.vram[((address + (row * 2) as u16) + 1 - VIDEO_RAM_START) as usize];
        let lower_byte = self.vram[((address + (row * 2) as u16) - VIDEO_RAM_START) as usize];

        let shift = 7 - column;
        let mask = 1 << shift;
        let upper_bit = (upper_byte & mask) >> shift;
        let lower_bit = (lower_byte & mask) >> shift;

        Shade::from_u8((upper_bit << 1) | lower_bit).unwrap()
    }

    fn check_coincidence(&mut self, result: &mut StepResult) {
        if self.stat.contains(Stat::STAT_COINCIDENCE_INT) {
            if self.ly == self.lyc {
                self.stat.insert(Stat::STAT_COINCIDENCE_EQUAL);
                result.int_stat = true;
            } else {
                self.stat.remove(Stat::STAT_COINCIDENCE_EQUAL);
            }
        } else {
            self.stat.remove(Stat::STAT_COINCIDENCE_EQUAL);
        }
    }

    fn bg_tile_address(&self, tile_index: u8) -> u16 {
        let tile_addr_base = if self.lcdc.contains(Lcdc::LCDC_UNSIGNED_TILE_DATA) {
            0x8000
        } else {
            0x9000
        } as u16;

        if self.lcdc.contains(Lcdc::LCDC_UNSIGNED_TILE_DATA) {
            (tile_addr_base + ((tile_index as u16) * 16))
        } else {
            let signed_index = (tile_index as i8) as i16;
            ((tile_addr_base as i16) + ((signed_index) * 16)) as u16
        }
    }

    fn sprite_tile_address(&self, tile_index: u8) -> u16 {
        0x8000 + ((tile_index as u16) * 16)
    }
}

impl Addressable for Lcd {
    fn read(&self, addr: u16) -> u8 {
        match addr {
            VIDEO_RAM_START..=VIDEO_RAM_END => self.vram[(addr - VIDEO_RAM_START) as usize],
            OAM_START..=OAM_END => self.oam[(addr - OAM_START) as usize],
            ADDR_LCDC => self.lcdc.bits,
            ADDR_STAT => 0b1000_0000 | ((self.stat.bits & 0b0111_1100) | (self.mode as u8)),
            ADDR_SCY => self.scy,
            ADDR_SCX => self.scx,
            ADDR_LY => self.ly,
            ADDR_LYC => self.lyc,
            ADDR_BGP => self.bgp.0,
            ADDR_OBP0 => self.obp0.0,
            ADDR_OBP1 => self.obp1.0,
            ADDR_WY => self.wy,
            ADDR_WX => self.wx,
            _ => { println!("LCD IO read unimplemented ({:#X})", addr); 0 }
        }
    }

    fn write(&mut self, addr: u16, val: u8) {
        match addr {
            VIDEO_RAM_START..=VIDEO_RAM_END => self.vram[(addr - VIDEO_RAM_START) as usize] = val,
            OAM_START..=OAM_END => self.oam[(addr - OAM_START) as usize] = val,
            ADDR_LCDC => {
                self.lcdc = Lcdc::from_bits(val).unwrap();

                // Reset internal counters if display disabled
                // LCDC mode goes back to 0 (HBlank)
                if !self.lcdc.contains(Lcdc::LCDC_ENABLED) {
                    self.mode_cycles = 0;
                    self.ly = 0;
                    self.mode = Mode::HBlank;

                }
            },
            ADDR_STAT => self.stat.bits = val & 0b1111_1100,
            ADDR_SCY => self.scy = val,
            ADDR_SCX => self.scx = val,
            ADDR_LY => (), // read-only
            ADDR_LYC => self.lyc = val,
            ADDR_BGP => self.bgp = Palette(val),
            ADDR_OBP0 => self.obp0 = Palette(val),
            ADDR_OBP1 => self.obp1 = Palette(val),
            ADDR_WY => self.wy = val,
            ADDR_WX => self.wx = val,
            _ => println!("LCD IO write unimplemented {:#X} -> {:#X}", val, addr)
        }
    }
}<|endoftext|>


<commit_before>#![allow(unsafe_code)]

use std::{
    borrow::{Borrow, BorrowMut},
    convert::TryFrom,
    hash::{Hash, Hasher},
    mem::MaybeUninit,
    sync::atomic::{AtomicPtr, AtomicUsize, Ordering},
};

use crate::{
    atomic_shim::AtomicU64,
    debug_delay,
    dll::{DoublyLinkedList, Node},
    fastlock::FastLock,
    FastSet8, Guard, PageId,
};

#[cfg(any(test, feature = "lock_free_delays"))]
const MAX_QUEUE_ITEMS: usize = 4;

#[cfg(not(any(test, feature = "lock_free_delays")))]
const MAX_QUEUE_ITEMS: usize = 64;

#[cfg(any(test, feature = "lock_free_delays"))]
const N_SHARDS: usize = 2;

#[cfg(not(any(test, feature = "lock_free_delays")))]
const N_SHARDS: usize = 256;

struct AccessBlock {
    len: AtomicUsize,
    block: [AtomicU64; MAX_QUEUE_ITEMS],
    next: AtomicPtr<AccessBlock>,
}

impl Default for AccessBlock {
    fn default() -> AccessBlock {
        AccessBlock {
            len: AtomicUsize::new(0),
            block: unsafe { MaybeUninit::zeroed().assume_init() },
            next: AtomicPtr::default(),
        }
    }
}

struct AccessQueue {
    writing: AtomicPtr<AccessBlock>,
    full_list: AtomicPtr<AccessBlock>,
}

impl Default for AccessQueue {
    fn default() -> AccessQueue {
        AccessQueue {
            writing: AtomicPtr::new(Box::into_raw(Box::new(
                AccessBlock::default(),
            ))),
            full_list: AtomicPtr::default(),
        }
    }
}

impl AccessQueue {
    fn push(&self, item: CacheAccess) -> bool {
        let mut filled = false;
        let item_u64: u64 = item.into();
        assert_ne!(item_u64, 0);
        loop {
            debug_delay();
            let head = self.writing.load(Ordering::Acquire);
            let block = unsafe { &*head };

            debug_delay();
            let offset = block.len.fetch_add(1, Ordering::Release);

            if offset < MAX_QUEUE_ITEMS {
                debug_delay();
                unsafe {
                    block
                        .block
                        .get_unchecked(offset)
                        .store(item_u64, Ordering::Release);
                }
                return filled;
            } else {
                // install new writer
                let new = Box::into_raw(Box::new(AccessBlock::default()));
                debug_delay();
                let res = self.writing.compare_exchange_weak(
                    head,
                    new,
                    Ordering::Release,
                    Ordering::Relaxed,
                );
                if res.is_err() {
                    // we lost the CAS, free the new item that was
                    // never published to other threads
                    unsafe {
                        drop(Box::from_raw(new));
                    }
                    continue;
                }

                // push the now-full item to the full list for future
                // consumption
                let mut ret;
                let mut full_list_ptr = self.full_list.load(Ordering::Acquire);
                while {
                    // we loop because maybe other threads are pushing stuff too
                    block.next.store(full_list_ptr, Ordering::Release);
                    debug_delay();
                    ret = self.full_list.compare_exchange_weak(
                        full_list_ptr,
                        head,
                        Ordering::Release,
                        Ordering::Relaxed,
                    );
                    ret.is_err()
                } {
                    full_list_ptr = ret.unwrap_err();
                }
                filled = true;
            }
        }
    }

    fn take<'a>(&self, guard: &'a Guard) -> CacheAccessIter<'a> {
        debug_delay();
        let ptr = self.full_list.swap(std::ptr::null_mut(), Ordering::AcqRel);

        CacheAccessIter { guard, current_offset: 0, current_block: ptr }
    }
}

impl Drop for AccessQueue {
    fn drop(&mut self) {
        debug_delay();
        let writing = self.writing.load(Ordering::Acquire);
        unsafe {
            Box::from_raw(writing);
        }
        debug_delay();
        let mut head = self.full_list.load(Ordering::Acquire);
        while !head.is_null() {
            unsafe {
                debug_delay();
                let next =
                    (*head).next.swap(std::ptr::null_mut(), Ordering::Release);
                Box::from_raw(head);
                head = next;
            }
        }
    }
}

struct CacheAccessIter<'a> {
    guard: &'a Guard,
    current_offset: usize,
    current_block: *mut AccessBlock,
}

impl<'a> Iterator for CacheAccessIter<'a> {
    type Item = CacheAccess;

    fn next(&mut self) -> Option<CacheAccess> {
        while !self.current_block.is_null() {
            let current_block = unsafe { &*self.current_block };

            debug_delay();
            if self.current_offset >= MAX_QUEUE_ITEMS {
                let to_drop_ptr = self.current_block;
                debug_delay();
                self.current_block = current_block.next.load(Ordering::Acquire);
                self.current_offset = 0;
                debug_delay();
                let to_drop = unsafe { Box::from_raw(to_drop_ptr) };
                self.guard.defer(|| to_drop);
                continue;
            }

            let mut next = 0;
            while next == 0 {
                // we spin here because there's a race between bumping
                // the offset and setting the value to something other
                // than 0 (and 0 is an invalid value)
                debug_delay();
                next = current_block.block[self.current_offset]
                    .load(Ordering::Acquire);
            }
            self.current_offset += 1;
            return Some(CacheAccess::from(next));
        }

        None
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct CacheAccess {
    // safe because MAX_PID_BITS / N_SHARDS < u32::MAX
    //                     2**37 /   2**8   < 2**32
    pub pid: u32,
    pub sz: u8,
}

impl From<CacheAccess> for u64 {
    fn from(ca: CacheAccess) -> u64 {
        (u64::from(ca.pid) << 8) | u64::from(ca.sz)
    }
}

#[allow(clippy::fallible_impl_from)]
impl From<u64> for CacheAccess {
    fn from(u: u64) -> CacheAccess {
        let sz = usize::try_from((u << 56) >> 56).unwrap();
        assert_ne!(sz, 0);
        let pid = u >> 8;
        assert!(pid < u64::from(std::u32::MAX));
        CacheAccess {
            pid: u32::try_from(pid).unwrap(),
            sz: u8::try_from(sz).unwrap(),
        }
    }
}

impl CacheAccess {
    fn size(&self) -> usize {
        1 << usize::from(self.sz)
    }

    fn new(pid: PageId, sz: usize) -> CacheAccess {
        let rounded_up_power_of_2 =
            u8::try_from(sz.next_power_of_two().trailing_zeros()).unwrap();

        CacheAccess {
            pid: u32::try_from(pid).expect("expected caller to shift pid down"),
            sz: rounded_up_power_of_2,
        }
    }
}

/// A simple LRU cache.
pub struct Lru {
    shards: Vec<(AccessQueue, FastLock<Shard>)>,
}

impl Lru {
    /// Instantiates a new `Lru` cache.
    pub(crate) fn new(cache_capacity: usize) -> Self {
        assert!(
            cache_capacity >= N_SHARDS,
            "Please configure the cache \
             capacity to be at least 256 bytes"
        );
        let shard_capacity = cache_capacity / N_SHARDS;

        let mut shards = Vec::with_capacity(N_SHARDS);
        shards.resize_with(N_SHARDS, || {
            (AccessQueue::default(), FastLock::new(Shard::new(shard_capacity)))
        });

        Self { shards }
    }

    /// Called when an item is accessed. Returns a Vec of items to be
    /// evicted. Uses flat-combining to avoid blocking on what can
    /// be an asynchronous operation.
    ///
    /// layout:
    ///   items:   1 2 3 4 5 6 7 8 9 10
    ///   shards:  1 0 1 0 1 0 1 0 1 0
    ///   shard 0:   2   4   6   8   10
    ///   shard 1: 1   3   5   7   9
    pub(crate) fn accessed(
        &self,
        id: PageId,
        item_size: usize,
        guard: &Guard,
    ) -> Vec<PageId> {
        const SHARD_BITS: usize = N_SHARDS.trailing_zeros() as usize;

        let mut ret = vec![];
        let shards = N_SHARDS as u64;
        let (shard_idx, shifted_pid) = (id % shards, id >> SHARD_BITS);
        let (access_queue, shard_mu) = &self.shards[safe_usize(shard_idx)];

        let cache_access = CacheAccess::new(shifted_pid, item_size);
        let filled = access_queue.push(cache_access);

        if filled {
            // only try to acquire this if the access queue has filled
            // an entire segment
            if let Some(mut shard) = shard_mu.try_lock() {
                let accesses = access_queue.take(guard);
                for item in accesses {
                    let to_evict = shard.accessed(item);
                    // map shard internal offsets to global items ids
                    for pos in to_evict {
                        let item =
                            (PageId::from(pos) << SHARD_BITS) + shard_idx;
                        ret.push(item);
                    }
                }
            }
        }
        ret
    }
}

#[derive(Eq)]
struct Entry(*mut Node);

unsafe impl Send for Entry {}

impl Ord for Entry {
    fn cmp(&self, other: &Entry) -> std::cmp::Ordering {
        let left_pid: u32 = *self.borrow();
        let right_pid: u32 = *other.borrow();
        left_pid.cmp(&right_pid)
    }
}

impl PartialOrd<Entry> for Entry {
    fn partial_cmp(&self, other: &Entry) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl PartialEq for Entry {
    fn eq(&self, other: &Entry) -> bool {
        unsafe { (*self.0).pid == (*other.0).pid }
    }
}

impl BorrowMut<CacheAccess> for Entry {
    fn borrow_mut(&mut self) -> &mut CacheAccess {
        unsafe { &mut *self.0 }
    }
}

impl Borrow<CacheAccess> for Entry {
    fn borrow(&self) -> &CacheAccess {
        unsafe { &*self.0 }
    }
}

impl Borrow<u32> for Entry {
    fn borrow(&self) -> &u32 {
        unsafe { &(*self.0).pid }
    }
}

// we only hash on pid, since we will change
// sz sometimes and we access the item by pid
impl Hash for Entry {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        unsafe { (*self.0).pid.hash(hasher) }
    }
}

struct Shard {
    dll: DoublyLinkedList,
    entries: FastSet8<Entry>,
    capacity: usize,
    size: usize,
}

impl Shard {
    fn new(capacity: usize) -> Self {
        assert!(capacity > 0, "shard capacity must be non-zero");

        Self {
            dll: DoublyLinkedList::default(),
            entries: FastSet8::default(),
            capacity,
            size: 0,
        }
    }

    /// `PageId`s in the shard list are indexes of the entries.
    fn accessed(&mut self, cache_access: CacheAccess) -> Vec<u32> {
        if let Some(entry) = self.entries.get(&cache_access.pid) {
            let old_sz_po2 = unsafe { (*entry.0).swap_sz(cache_access.sz) };
            let old_size = 1 << usize::from(old_sz_po2);

            self.size -= old_size;
            self.dll.promote(entry.0);
        } else {
            let ptr = self.dll.push_head(cache_access);
            self.entries.insert(Entry(ptr));
        };

        self.size += cache_access.size();

        let mut to_evict = vec![];

        while self.size > self.capacity {
            if self.dll.len() == 1 {
                // don't evict what we just added
                break;
            }

            let min_pid = self.dll.pop_tail().unwrap();

            self.entries.remove(&min_pid.pid);

            to_evict.push(min_pid.pid);

            self.size -= min_pid.size();
        }

        to_evict
    }
}

#[inline]
fn safe_usize(value: PageId) -> usize {
    usize::try_from(value).unwrap()
}

#[test]
fn lru_smoke_test() {
    use crate::pin;

    let lru = Lru::new(2);
    for i in 0..1000 {
        let guard = pin();
        lru.accessed(i, 16, &guard);
    }
}

#[test]
fn lru_access_test() {
    use crate::pin;

    let ci = CacheAccess::new(6, 20667);
    assert_eq!(ci.size(), 32 * 1024);

    let lru = Lru::new(4096);

    let guard = pin();

    assert_eq!(lru.accessed(0, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(2, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(4, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(6, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(8, 20667, &guard), vec![0, 2, 4]);
    assert_eq!(lru.accessed(10, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(12, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(14, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(16, 20667, &guard), vec![6, 8, 10, 12]);
    assert_eq!(lru.accessed(18, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(20, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(22, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(24, 20667, &guard), vec![14, 16, 18, 20]);
}
<commit_msg>Avoid looping in AccessQueue::push<commit_after>#![allow(unsafe_code)]

use std::{
    borrow::{Borrow, BorrowMut},
    convert::TryFrom,
    hash::{Hash, Hasher},
    mem::MaybeUninit,
    sync::atomic::{AtomicPtr, AtomicUsize, Ordering},
};

use crate::{
    atomic_shim::AtomicU64,
    debug_delay,
    dll::{DoublyLinkedList, Node},
    fastlock::FastLock,
    FastSet8, Guard, PageId,
};

#[cfg(any(test, feature = "lock_free_delays"))]
const MAX_QUEUE_ITEMS: usize = 4;

#[cfg(not(any(test, feature = "lock_free_delays")))]
const MAX_QUEUE_ITEMS: usize = 64;

#[cfg(any(test, feature = "lock_free_delays"))]
const N_SHARDS: usize = 2;

#[cfg(not(any(test, feature = "lock_free_delays")))]
const N_SHARDS: usize = 256;

struct AccessBlock {
    len: AtomicUsize,
    block: [AtomicU64; MAX_QUEUE_ITEMS],
    next: AtomicPtr<AccessBlock>,
}

impl Default for AccessBlock {
    fn default() -> AccessBlock {
        AccessBlock {
            len: AtomicUsize::new(0),
            block: unsafe { MaybeUninit::zeroed().assume_init() },
            next: AtomicPtr::default(),
        }
    }
}

struct AccessQueue {
    writing: AtomicPtr<AccessBlock>,
    full_list: AtomicPtr<AccessBlock>,
}

impl AccessBlock {
    fn new(item: CacheAccess) -> AccessBlock {
        let mut ret = AccessBlock {
            len: AtomicUsize::new(1),
            block: unsafe { MaybeUninit::zeroed().assume_init() },
            next: AtomicPtr::default(),
        };
        ret.block[0] = AtomicU64::from(u64::from(item));
        ret
    }
}

impl Default for AccessQueue {
    fn default() -> AccessQueue {
        AccessQueue {
            writing: AtomicPtr::new(Box::into_raw(Box::new(
                AccessBlock::default(),
            ))),
            full_list: AtomicPtr::default(),
        }
    }
}

impl AccessQueue {
    fn push(&self, item: CacheAccess) -> bool {
        loop {
            debug_delay();
            let head = self.writing.load(Ordering::Relaxed);
            let block = unsafe { &*head };

            debug_delay();
            let offset = block.len.fetch_add(1, Ordering::Acquire);

            if offset < MAX_QUEUE_ITEMS {
                let item_u64: u64 = item.into();
                assert_ne!(item_u64, 0);
                debug_delay();
                unsafe {
                    block
                        .block
                        .get_unchecked(offset)
                        .store(item_u64, Ordering::Release);
                }
                return false;
            } else {
                // install new writer
                let new = Box::into_raw(Box::new(AccessBlock::new(item)));
                debug_delay();
                let res = self.writing.compare_exchange_weak(
                    head,
                    new,
                    Ordering::Release,
                    Ordering::Relaxed,
                );
                if res.is_err() {
                    // we lost the CAS, free the new item that was
                    // never published to other threads
                    unsafe {
                        drop(Box::from_raw(new));
                    }
                    continue;
                }

                // push the now-full item to the full list for future
                // consumption
                let mut ret;
                let mut full_list_ptr = self.full_list.load(Ordering::Acquire);
                while {
                    // we loop because maybe other threads are pushing stuff too
                    block.next.store(full_list_ptr, Ordering::Release);
                    debug_delay();
                    ret = self.full_list.compare_exchange_weak(
                        full_list_ptr,
                        head,
                        Ordering::Release,
                        Ordering::Relaxed,
                    );
                    ret.is_err()
                } {
                    full_list_ptr = ret.unwrap_err();
                }
                return true;
            }
        }
    }

    fn take<'a>(&self, guard: &'a Guard) -> CacheAccessIter<'a> {
        debug_delay();
        let ptr = self.full_list.swap(std::ptr::null_mut(), Ordering::AcqRel);

        CacheAccessIter { guard, current_offset: 0, current_block: ptr }
    }
}

impl Drop for AccessQueue {
    fn drop(&mut self) {
        debug_delay();
        let writing = self.writing.load(Ordering::Acquire);
        unsafe {
            Box::from_raw(writing);
        }
        debug_delay();
        let mut head = self.full_list.load(Ordering::Acquire);
        while !head.is_null() {
            unsafe {
                debug_delay();
                let next =
                    (*head).next.swap(std::ptr::null_mut(), Ordering::Release);
                Box::from_raw(head);
                head = next;
            }
        }
    }
}

struct CacheAccessIter<'a> {
    guard: &'a Guard,
    current_offset: usize,
    current_block: *mut AccessBlock,
}

impl<'a> Iterator for CacheAccessIter<'a> {
    type Item = CacheAccess;

    fn next(&mut self) -> Option<CacheAccess> {
        while !self.current_block.is_null() {
            let current_block = unsafe { &*self.current_block };

            debug_delay();
            if self.current_offset >= MAX_QUEUE_ITEMS {
                let to_drop_ptr = self.current_block;
                debug_delay();
                self.current_block = current_block.next.load(Ordering::Acquire);
                self.current_offset = 0;
                debug_delay();
                let to_drop = unsafe { Box::from_raw(to_drop_ptr) };
                self.guard.defer(|| to_drop);
                continue;
            }

            let mut next = 0;
            while next == 0 {
                // we spin here because there's a race between bumping
                // the offset and setting the value to something other
                // than 0 (and 0 is an invalid value)
                debug_delay();
                next = current_block.block[self.current_offset]
                    .load(Ordering::Acquire);
            }
            self.current_offset += 1;
            return Some(CacheAccess::from(next));
        }

        None
    }
}

#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub(crate) struct CacheAccess {
    // safe because MAX_PID_BITS / N_SHARDS < u32::MAX
    //                     2**37 /   2**8   < 2**32
    pub pid: u32,
    pub sz: u8,
}

impl From<CacheAccess> for u64 {
    fn from(ca: CacheAccess) -> u64 {
        (u64::from(ca.pid) << 8) | u64::from(ca.sz)
    }
}

#[allow(clippy::fallible_impl_from)]
impl From<u64> for CacheAccess {
    fn from(u: u64) -> CacheAccess {
        let sz = usize::try_from((u << 56) >> 56).unwrap();
        assert_ne!(sz, 0);
        let pid = u >> 8;
        assert!(pid < u64::from(std::u32::MAX));
        CacheAccess {
            pid: u32::try_from(pid).unwrap(),
            sz: u8::try_from(sz).unwrap(),
        }
    }
}

impl CacheAccess {
    fn size(&self) -> usize {
        1 << usize::from(self.sz)
    }

    fn new(pid: PageId, sz: usize) -> CacheAccess {
        let rounded_up_power_of_2 =
            u8::try_from(sz.next_power_of_two().trailing_zeros()).unwrap();

        CacheAccess {
            pid: u32::try_from(pid).expect("expected caller to shift pid down"),
            sz: rounded_up_power_of_2,
        }
    }
}

/// A simple LRU cache.
pub struct Lru {
    shards: Vec<(AccessQueue, FastLock<Shard>)>,
}

impl Lru {
    /// Instantiates a new `Lru` cache.
    pub(crate) fn new(cache_capacity: usize) -> Self {
        assert!(
            cache_capacity >= N_SHARDS,
            "Please configure the cache \
             capacity to be at least 256 bytes"
        );
        let shard_capacity = cache_capacity / N_SHARDS;

        let mut shards = Vec::with_capacity(N_SHARDS);
        shards.resize_with(N_SHARDS, || {
            (AccessQueue::default(), FastLock::new(Shard::new(shard_capacity)))
        });

        Self { shards }
    }

    /// Called when an item is accessed. Returns a Vec of items to be
    /// evicted. Uses flat-combining to avoid blocking on what can
    /// be an asynchronous operation.
    ///
    /// layout:
    ///   items:   1 2 3 4 5 6 7 8 9 10
    ///   shards:  1 0 1 0 1 0 1 0 1 0
    ///   shard 0:   2   4   6   8   10
    ///   shard 1: 1   3   5   7   9
    pub(crate) fn accessed(
        &self,
        id: PageId,
        item_size: usize,
        guard: &Guard,
    ) -> Vec<PageId> {
        const SHARD_BITS: usize = N_SHARDS.trailing_zeros() as usize;

        let mut ret = vec![];
        let shards = N_SHARDS as u64;
        let (shard_idx, shifted_pid) = (id % shards, id >> SHARD_BITS);
        let (access_queue, shard_mu) = &self.shards[safe_usize(shard_idx)];

        let cache_access = CacheAccess::new(shifted_pid, item_size);
        let filled = access_queue.push(cache_access);

        if filled {
            // only try to acquire this if the access queue has filled
            // an entire segment
            if let Some(mut shard) = shard_mu.try_lock() {
                let accesses = access_queue.take(guard);
                for item in accesses {
                    let to_evict = shard.accessed(item);
                    // map shard internal offsets to global items ids
                    for pos in to_evict {
                        let item =
                            (PageId::from(pos) << SHARD_BITS) + shard_idx;
                        ret.push(item);
                    }
                }
            }
        }
        ret
    }
}

#[derive(Eq)]
struct Entry(*mut Node);

unsafe impl Send for Entry {}

impl Ord for Entry {
    fn cmp(&self, other: &Entry) -> std::cmp::Ordering {
        let left_pid: u32 = *self.borrow();
        let right_pid: u32 = *other.borrow();
        left_pid.cmp(&right_pid)
    }
}

impl PartialOrd<Entry> for Entry {
    fn partial_cmp(&self, other: &Entry) -> Option<std::cmp::Ordering> {
        Some(self.cmp(other))
    }
}

impl PartialEq for Entry {
    fn eq(&self, other: &Entry) -> bool {
        unsafe { (*self.0).pid == (*other.0).pid }
    }
}

impl BorrowMut<CacheAccess> for Entry {
    fn borrow_mut(&mut self) -> &mut CacheAccess {
        unsafe { &mut *self.0 }
    }
}

impl Borrow<CacheAccess> for Entry {
    fn borrow(&self) -> &CacheAccess {
        unsafe { &*self.0 }
    }
}

impl Borrow<u32> for Entry {
    fn borrow(&self) -> &u32 {
        unsafe { &(*self.0).pid }
    }
}

// we only hash on pid, since we will change
// sz sometimes and we access the item by pid
impl Hash for Entry {
    fn hash<H: Hasher>(&self, hasher: &mut H) {
        unsafe { (*self.0).pid.hash(hasher) }
    }
}

struct Shard {
    dll: DoublyLinkedList,
    entries: FastSet8<Entry>,
    capacity: usize,
    size: usize,
}

impl Shard {
    fn new(capacity: usize) -> Self {
        assert!(capacity > 0, "shard capacity must be non-zero");

        Self {
            dll: DoublyLinkedList::default(),
            entries: FastSet8::default(),
            capacity,
            size: 0,
        }
    }

    /// `PageId`s in the shard list are indexes of the entries.
    fn accessed(&mut self, cache_access: CacheAccess) -> Vec<u32> {
        if let Some(entry) = self.entries.get(&cache_access.pid) {
            let old_sz_po2 = unsafe { (*entry.0).swap_sz(cache_access.sz) };
            let old_size = 1 << usize::from(old_sz_po2);

            self.size -= old_size;
            self.dll.promote(entry.0);
        } else {
            let ptr = self.dll.push_head(cache_access);
            self.entries.insert(Entry(ptr));
        };

        self.size += cache_access.size();

        let mut to_evict = vec![];

        while self.size > self.capacity {
            if self.dll.len() == 1 {
                // don't evict what we just added
                break;
            }

            let min_pid = self.dll.pop_tail().unwrap();

            self.entries.remove(&min_pid.pid);

            to_evict.push(min_pid.pid);

            self.size -= min_pid.size();
        }

        to_evict
    }
}

#[inline]
fn safe_usize(value: PageId) -> usize {
    usize::try_from(value).unwrap()
}

#[test]
fn lru_smoke_test() {
    use crate::pin;

    let lru = Lru::new(2);
    for i in 0..1000 {
        let guard = pin();
        lru.accessed(i, 16, &guard);
    }
}

#[test]
fn lru_access_test() {
    use crate::pin;

    let ci = CacheAccess::new(6, 20667);
    assert_eq!(ci.size(), 32 * 1024);

    let lru = Lru::new(4096);

    let guard = pin();

    assert_eq!(lru.accessed(0, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(2, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(4, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(6, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(8, 20667, &guard), vec![0, 2, 4]);
    assert_eq!(lru.accessed(10, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(12, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(14, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(16, 20667, &guard), vec![6, 8, 10, 12]);
    assert_eq!(lru.accessed(18, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(20, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(22, 20667, &guard), vec![]);
    assert_eq!(lru.accessed(24, 20667, &guard), vec![14, 16, 18, 20]);
}
<|endoftext|>