<commit_before>use std::any::Any;
use std::cell::RefCell;
use std::io;
use std::os::unix::io::RawFd;
use std::rc::Rc;
use std::time::Duration;

use mio::unix::EventedFd;
use mio::{Events, Poll, PollOpt, Ready, Token};

use sources::*;
use Implementation;

use super::DisplayInner;

pub(crate) struct EventLoopInner {
    pub(crate) display: Option<Rc<RefCell<DisplayInner>>>,
    sources_poll: SourcesPoll,
    idles: RefCell<Vec<Rc<RefCell<Option<Box<Implementation<(), ()>>>>>>>,
}

#[derive(Clone)]
pub(crate) struct SourcesPoll {
    poll: Rc<Poll>,
    sources: Rc<RefCell<SourceList>>,
}

impl SourcesPoll {
    fn dispatch(&self, timeout: Option<u32>) -> io::Result<()> {
        let mut evts = Events::with_capacity(32);

        self.poll
            .poll(&mut evts, timeout.map(|s| Duration::from_millis(s as u64)))?;

        for evt in evts {
            let Token(idx) = evt.token();
            self.sources.borrow().send_ready(idx, evt.readiness());
        }

        Ok(())
    }

    pub(crate) fn insert_source<Impl, E>(
        &self,
        fd: RawFd,
        interest: Ready,
        implementation: Impl,
        evt: E,
    ) -> Result<SourceInner<E>, (io::Error, Impl)>
    where
        Impl: Implementation<(), E> + 'static,
        SourceDispatcher<E>: EventDispatcher,
    {
        let implem = Rc::new(RefCell::new(SourceDispatcher {
            fd,
            implem: Box::new(implementation),
            evt,
        }));
        let token = self.sources.borrow_mut().add_source(fd, implem.clone());
        let source = SourceInner {
            dispatcher: implem,
            list: self.sources.clone(),
            poll: self.poll.clone(),
            token,
            fd,
        };
        match self.poll
            .register(&EventedFd(&fd), token, interest, PollOpt::empty())
        {
            Ok(()) => Ok(source),
            Err(e) => {
                let implem = source.remove();
                // convert the implem back
                let implem = unsafe { Box::from_raw(Box::into_raw(implem) as _) };
                Err((e, *implem))
            }
        }
    }
}

impl EventLoopInner {
    pub(crate) fn new() -> EventLoopInner {
        EventLoopInner {
            display: None,
            sources_poll: SourcesPoll {
                poll: Rc::new(Poll::new().unwrap()),
                sources: Rc::new(RefCell::new(SourceList::new())),
            },
            idles: RefCell::new(Vec::new()),
        }
    }

    pub(crate) fn get_poll(&self) -> SourcesPoll {
        self.sources_poll.clone()
    }

    pub(crate) fn dispatch(&self, timeout: Option<u32>) -> io::Result<u32> {
        self.dispatch_idles();

        self.sources_poll.dispatch(timeout)?;

        self.dispatch_idles();

        Ok(0)
    }

    pub(crate) fn flush_clients_if_display(&self) {
        if let Some(ref display) = self.display {
            display.borrow_mut().flush_clients();
        }
    }

    pub(crate) fn add_fd_event_source<Impl>(
        &self,
        fd: RawFd,
        interest: FdInterest,
        implementation: Impl,
    ) -> Result<SourceInner<FdEvent>, (io::Error, Impl)>
    where
        Impl: Implementation<(), FdEvent> + 'static,
    {
        self.sources_poll.insert_source(
            fd,
            interest.into(),
            implementation,
            FdEvent::Ready { fd, mask: interest },
        )
    }

    pub(crate) fn add_timer_event_source<Impl>(
        &self,
        implementation: Impl,
    ) -> Result<SourceInner<TimerEvent>, (io::Error, Impl)>
    where
        Impl: Implementation<(), TimerEvent> + 'static,
    {
        use libc::{timerfd_create, CLOCK_MONOTONIC, TFD_CLOEXEC, TFD_NONBLOCK};
        let fd = unsafe { timerfd_create(CLOCK_MONOTONIC, TFD_CLOEXEC | TFD_NONBLOCK) };
        if fd < 0 {
            return Err((io::Error::last_os_error(), implementation));
        }
        self.sources_poll
            .insert_source(fd, Ready::readable(), implementation, TimerEvent)
    }

    pub(crate) fn add_signal_event_source<Impl>(
        &self,
        signal: ::nix::sys::signal::Signal,
        implementation: Impl,
    ) -> Result<SourceInner<SignalEvent>, (io::Error, Impl)>
    where
        Impl: Implementation<(), SignalEvent> + 'static,
    {
        use nix::sys::signal::SigSet;
        use nix::sys::signalfd::{signalfd, SfdFlags};

        let mut set = SigSet::empty();
        set.add(signal);

        let fd = match signalfd(-1, &set, SfdFlags::SFD_NONBLOCK | SfdFlags::SFD_CLOEXEC) {
            Ok(fd) => fd,
            Err(::nix::Error::Sys(e)) => return Err((e.into(), implementation)),
            Err(_) => unreachable!(),
        };

        self.sources_poll
            .insert_source(fd, Ready::readable(), implementation, SignalEvent(signal))
    }

    pub(crate) fn add_idle_event_source<Impl>(&self, implementation: Impl) -> IdleSourceInner
    where
        Impl: Implementation<(), ()> + 'static,
    {
        let implem = Rc::new(RefCell::new(Some(Box::new(implementation) as Box<_>)));
        self.idles.borrow_mut().push(implem.clone());
        IdleSourceInner { implem }
    }

    fn dispatch_idles(&self) {
        let idles = ::std::mem::replace(&mut *self.idles.borrow_mut(), Vec::new());
        for idle in idles {
            if let Some(ref mut implem) = *idle.borrow_mut() {
                implem.receive((), ());
            }
        }
    }
}

// SourceList

pub(crate) trait EventDispatcher: Any {
    fn ready(&mut self, ready: Ready);
    fn error(&mut self, error: io::Error);
}

struct SourceList {
    sources: Vec<Option<(RawFd, Rc<RefCell<EventDispatcher>>)>>,
}

impl SourceList {
    fn new() -> SourceList {
        SourceList { sources: Vec::new() }
    }

    fn send_ready(&self, idx: usize, ready: Ready) {
        match self.sources.get(idx) {
            Some(Some((_, dispatcher))) => dispatcher.borrow_mut().ready(ready),
            _ => (),
        }
    }

    fn add_source(&mut self, fd: RawFd, source: Rc<RefCell<EventDispatcher>>) -> Token {
        let free_id = self.sources.iter().position(Option::is_none);
        if let Some(id) = free_id {
            self.sources[id] = Some((fd, source));
            Token(id)
        } else {
            self.sources.push(Some((fd, source)));
            Token(self.sources.len() - 1)
        }
    }

    fn del_source(&mut self, source: Rc<RefCell<EventDispatcher>>) {
        for src in &mut self.sources {
            let found = if let Some(ref s) = src {
                Rc::ptr_eq(&source, &s.1)
            } else {
                false
            };
            if found {
                *src = None;
                return;
            }
        }
    }
}

// Sources

pub(crate) struct SourceDispatcher<E> {
    implem: Box<Implementation<(), E>>,
    fd: RawFd,
    evt: E,
}

pub(crate) struct SourceInner<E> {
    dispatcher: Rc<RefCell<SourceDispatcher<E>>>,
    list: Rc<RefCell<SourceList>>,
    poll: Rc<Poll>,
    token: Token,
    fd: RawFd,
}

impl<E: 'static> SourceInner<E>
where
    SourceDispatcher<E>: EventDispatcher,
{
    pub(crate) fn remove(self) -> Box<Implementation<(), E>> {
        let _ = self.poll.deregister(&EventedFd(&self.fd));
        self.list
            .borrow_mut()
            .del_source(self.dispatcher.clone() as Rc<_>);
        let dispatcher = match Rc::try_unwrap(self.dispatcher) {
            Ok(d) => d,
            Err(_) => unreachable!(),
        };
        dispatcher.into_inner().implem
    }
}

// FD event source

impl SourceInner<FdEvent> {
    pub(crate) fn update_mask(&mut self, mask: FdInterest) {
        let _ = self.poll
            .reregister(&EventedFd(&self.fd), self.token, mask.into(), PollOpt::empty());
    }
}

impl EventDispatcher for SourceDispatcher<FdEvent> {
    fn ready(&mut self, ready: Ready) {
        self.implem.receive(
            FdEvent::Ready {
                fd: self.fd,
                mask: ready.into(),
            },
            (),
        );
    }

    fn error(&mut self, error: io::Error) {
        self.implem.receive(FdEvent::Error { fd: self.fd, error }, ());
    }
}

// Timer event source

impl SourceInner<TimerEvent> {
    pub(crate) fn set_delay_ms(&mut self, delay: i32) {
        use libc::{c_long, itimerspec, time_t, timerfd_settime, timespec};

        let spec = itimerspec {
            it_interval: timespec {
                tv_sec: 0,
                tv_nsec: 0,
            },
            it_value: timespec {
                tv_sec: delay as time_t / 1000,
                tv_nsec: (delay as c_long % 1000) * 1000 * 1000,
            },
        };

        unsafe {
            timerfd_settime(self.fd, 0, &spec, ::std::ptr::null_mut());
        }
    }
}

impl EventDispatcher for SourceDispatcher<TimerEvent> {
    fn ready(&mut self, _: Ready) {
        self.implem.receive(self.evt, ());
    }

    fn error(&mut self, _: io::Error) {}
}

// Signal event source

impl EventDispatcher for SourceDispatcher<SignalEvent> {
    fn ready(&mut self, _: Ready) {
        self.implem.receive(self.evt, ());
    }

    fn error(&mut self, _: io::Error) {}
}

// Idle event source

pub(crate) struct IdleSourceInner {
    implem: Rc<RefCell<Option<Box<Implementation<(), ()>>>>>,
}

impl IdleSourceInner {
    pub(crate) fn remove(self) -> Box<Implementation<(), ()>> {
        self.implem.borrow_mut().take().unwrap()
    }
}

impl From<FdInterest> for Ready {
    fn from(mask: FdInterest) -> Ready {
        let mut interest = Ready::empty();
        if mask.contains(FdInterest::READ) {
            interest.insert(Ready::readable())
        }
        if mask.contains(FdInterest::WRITE) {
            interest.insert(Ready::writable())
        }
        interest
    }
}

impl From<Ready> for FdInterest {
    fn from(ready: Ready) -> FdInterest {
        let mut mask = FdInterest::empty();
        if ready.contains(Ready::readable()) {
            mask |= FdInterest::READ
        }
        if ready.contains(Ready::writable()) {
            mask |= FdInterest::WRITE
        }
        mask
    }
}
<commit_msg>server: fix double-borrow in event loop dispatch<commit_after>use std::any::Any;
use std::cell::RefCell;
use std::io;
use std::os::unix::io::RawFd;
use std::rc::Rc;
use std::time::Duration;

use mio::unix::EventedFd;
use mio::{Events, Poll, PollOpt, Ready, Token};

use sources::*;
use Implementation;

use super::DisplayInner;

pub(crate) struct EventLoopInner {
    pub(crate) display: Option<Rc<RefCell<DisplayInner>>>,
    sources_poll: SourcesPoll,
    idles: RefCell<Vec<Rc<RefCell<Option<Box<Implementation<(), ()>>>>>>>,
}

#[derive(Clone)]
pub(crate) struct SourcesPoll {
    poll: Rc<Poll>,
    sources: Rc<RefCell<SourceList>>,
}

impl SourcesPoll {
    fn dispatch(&self, timeout: Option<u32>) -> io::Result<()> {
        let mut evts = Events::with_capacity(32);

        self.poll
            .poll(&mut evts, timeout.map(|s| Duration::from_millis(s as u64)))?;

        for evt in evts {
            let Token(idx) = evt.token();
            let dispatcher = self.sources.borrow().get_dispatcher(idx);
            if let Some(dispatcher) = dispatcher {
                dispatcher.borrow_mut().ready(evt.readiness());
            }
        }

        Ok(())
    }

    pub(crate) fn insert_source<Impl, E>(
        &self,
        fd: RawFd,
        interest: Ready,
        implementation: Impl,
        evt: E,
    ) -> Result<SourceInner<E>, (io::Error, Impl)>
    where
        Impl: Implementation<(), E> + 'static,
        SourceDispatcher<E>: EventDispatcher,
    {
        let implem = Rc::new(RefCell::new(SourceDispatcher {
            fd,
            implem: Box::new(implementation),
            evt,
        }));
        let token = self.sources.borrow_mut().add_source(fd, implem.clone());
        let source = SourceInner {
            dispatcher: implem,
            list: self.sources.clone(),
            poll: self.poll.clone(),
            token,
            fd,
        };
        match self.poll
            .register(&EventedFd(&fd), token, interest, PollOpt::empty())
        {
            Ok(()) => Ok(source),
            Err(e) => {
                let implem = source.remove();
                // convert the implem back
                let implem = unsafe { Box::from_raw(Box::into_raw(implem) as _) };
                Err((e, *implem))
            }
        }
    }
}

impl EventLoopInner {
    pub(crate) fn new() -> EventLoopInner {
        EventLoopInner {
            display: None,
            sources_poll: SourcesPoll {
                poll: Rc::new(Poll::new().unwrap()),
                sources: Rc::new(RefCell::new(SourceList::new())),
            },
            idles: RefCell::new(Vec::new()),
        }
    }

    pub(crate) fn get_poll(&self) -> SourcesPoll {
        self.sources_poll.clone()
    }

    pub(crate) fn dispatch(&self, timeout: Option<u32>) -> io::Result<u32> {
        self.dispatch_idles();

        self.sources_poll.dispatch(timeout)?;

        self.dispatch_idles();

        Ok(0)
    }

    pub(crate) fn flush_clients_if_display(&self) {
        if let Some(ref display) = self.display {
            display.borrow_mut().flush_clients();
        }
    }

    pub(crate) fn add_fd_event_source<Impl>(
        &self,
        fd: RawFd,
        interest: FdInterest,
        implementation: Impl,
    ) -> Result<SourceInner<FdEvent>, (io::Error, Impl)>
    where
        Impl: Implementation<(), FdEvent> + 'static,
    {
        self.sources_poll.insert_source(
            fd,
            interest.into(),
            implementation,
            FdEvent::Ready { fd, mask: interest },
        )
    }

    pub(crate) fn add_timer_event_source<Impl>(
        &self,
        implementation: Impl,
    ) -> Result<SourceInner<TimerEvent>, (io::Error, Impl)>
    where
        Impl: Implementation<(), TimerEvent> + 'static,
    {
        use libc::{timerfd_create, CLOCK_MONOTONIC, TFD_CLOEXEC, TFD_NONBLOCK};
        let fd = unsafe { timerfd_create(CLOCK_MONOTONIC, TFD_CLOEXEC | TFD_NONBLOCK) };
        if fd < 0 {
            return Err((io::Error::last_os_error(), implementation));
        }
        self.sources_poll
            .insert_source(fd, Ready::readable(), implementation, TimerEvent)
    }

    pub(crate) fn add_signal_event_source<Impl>(
        &self,
        signal: ::nix::sys::signal::Signal,
        implementation: Impl,
    ) -> Result<SourceInner<SignalEvent>, (io::Error, Impl)>
    where
        Impl: Implementation<(), SignalEvent> + 'static,
    {
        use nix::sys::signal::SigSet;
        use nix::sys::signalfd::{signalfd, SfdFlags};

        let mut set = SigSet::empty();
        set.add(signal);

        let fd = match signalfd(-1, &set, SfdFlags::SFD_NONBLOCK | SfdFlags::SFD_CLOEXEC) {
            Ok(fd) => fd,
            Err(::nix::Error::Sys(e)) => return Err((e.into(), implementation)),
            Err(_) => unreachable!(),
        };

        self.sources_poll
            .insert_source(fd, Ready::readable(), implementation, SignalEvent(signal))
    }

    pub(crate) fn add_idle_event_source<Impl>(&self, implementation: Impl) -> IdleSourceInner
    where
        Impl: Implementation<(), ()> + 'static,
    {
        let implem = Rc::new(RefCell::new(Some(Box::new(implementation) as Box<_>)));
        self.idles.borrow_mut().push(implem.clone());
        IdleSourceInner { implem }
    }

    fn dispatch_idles(&self) {
        let idles = ::std::mem::replace(&mut *self.idles.borrow_mut(), Vec::new());
        for idle in idles {
            if let Some(ref mut implem) = *idle.borrow_mut() {
                implem.receive((), ());
            }
        }
    }
}

// SourceList

pub(crate) trait EventDispatcher: Any {
    fn ready(&mut self, ready: Ready);
    fn error(&mut self, error: io::Error);
}

struct SourceList {
    sources: Vec<Option<(RawFd, Rc<RefCell<EventDispatcher>>)>>,
}

impl SourceList {
    fn new() -> SourceList {
        SourceList { sources: Vec::new() }
    }

    fn get_dispatcher(&self, idx: usize) -> Option<Rc<RefCell<EventDispatcher>>> {
        match self.sources.get(idx) {
            Some(Some((_, dispatcher))) => Some(dispatcher.clone()),
            _ => None,
        }
    }

    fn add_source(&mut self, fd: RawFd, source: Rc<RefCell<EventDispatcher>>) -> Token {
        let free_id = self.sources.iter().position(Option::is_none);
        if let Some(id) = free_id {
            self.sources[id] = Some((fd, source));
            Token(id)
        } else {
            self.sources.push(Some((fd, source)));
            Token(self.sources.len() - 1)
        }
    }

    fn del_source(&mut self, source: Rc<RefCell<EventDispatcher>>) {
        for src in &mut self.sources {
            let found = if let Some(ref s) = src {
                Rc::ptr_eq(&source, &s.1)
            } else {
                false
            };
            if found {
                *src = None;
                return;
            }
        }
    }
}

// Sources

pub(crate) struct SourceDispatcher<E> {
    implem: Box<Implementation<(), E>>,
    fd: RawFd,
    evt: E,
}

pub(crate) struct SourceInner<E> {
    dispatcher: Rc<RefCell<SourceDispatcher<E>>>,
    list: Rc<RefCell<SourceList>>,
    poll: Rc<Poll>,
    token: Token,
    fd: RawFd,
}

impl<E: 'static> SourceInner<E>
where
    SourceDispatcher<E>: EventDispatcher,
{
    pub(crate) fn remove(self) -> Box<Implementation<(), E>> {
        let _ = self.poll.deregister(&EventedFd(&self.fd));
        self.list
            .borrow_mut()
            .del_source(self.dispatcher.clone() as Rc<_>);
        let dispatcher = match Rc::try_unwrap(self.dispatcher) {
            Ok(d) => d,
            Err(_) => unreachable!(),
        };
        dispatcher.into_inner().implem
    }
}

// FD event source

impl SourceInner<FdEvent> {
    pub(crate) fn update_mask(&mut self, mask: FdInterest) {
        let _ = self.poll
            .reregister(&EventedFd(&self.fd), self.token, mask.into(), PollOpt::empty());
    }
}

impl EventDispatcher for SourceDispatcher<FdEvent> {
    fn ready(&mut self, ready: Ready) {
        self.implem.receive(
            FdEvent::Ready {
                fd: self.fd,
                mask: ready.into(),
            },
            (),
        );
    }

    fn error(&mut self, error: io::Error) {
        self.implem.receive(FdEvent::Error { fd: self.fd, error }, ());
    }
}

// Timer event source

impl SourceInner<TimerEvent> {
    pub(crate) fn set_delay_ms(&mut self, delay: i32) {
        use libc::{c_long, itimerspec, time_t, timerfd_settime, timespec};

        let spec = itimerspec {
            it_interval: timespec {
                tv_sec: 0,
                tv_nsec: 0,
            },
            it_value: timespec {
                tv_sec: delay as time_t / 1000,
                tv_nsec: (delay as c_long % 1000) * 1000 * 1000,
            },
        };

        unsafe {
            timerfd_settime(self.fd, 0, &spec, ::std::ptr::null_mut());
        }
    }
}

impl EventDispatcher for SourceDispatcher<TimerEvent> {
    fn ready(&mut self, _: Ready) {
        self.implem.receive(self.evt, ());
    }

    fn error(&mut self, _: io::Error) {}
}

// Signal event source

impl EventDispatcher for SourceDispatcher<SignalEvent> {
    fn ready(&mut self, _: Ready) {
        self.implem.receive(self.evt, ());
    }

    fn error(&mut self, _: io::Error) {}
}

// Idle event source

pub(crate) struct IdleSourceInner {
    implem: Rc<RefCell<Option<Box<Implementation<(), ()>>>>>,
}

impl IdleSourceInner {
    pub(crate) fn remove(self) -> Box<Implementation<(), ()>> {
        self.implem.borrow_mut().take().unwrap()
    }
}

impl From<FdInterest> for Ready {
    fn from(mask: FdInterest) -> Ready {
        let mut interest = Ready::empty();
        if mask.contains(FdInterest::READ) {
            interest.insert(Ready::readable())
        }
        if mask.contains(FdInterest::WRITE) {
            interest.insert(Ready::writable())
        }
        interest
    }
}

impl From<Ready> for FdInterest {
    fn from(ready: Ready) -> FdInterest {
        let mut mask = FdInterest::empty();
        if ready.contains(Ready::readable()) {
            mask |= FdInterest::READ
        }
        if ready.contains(Ready::writable()) {
            mask |= FdInterest::WRITE
        }
        mask
    }
}
<|endoftext|>


<commit_before>use super::raw;
use std::fmt;
use std::cell::RefCell;

pub struct Document<'d> {
    storage: &'d raw::Storage,
    connections: RefCell<&'d raw::Connections>,
}

impl<'d> Document<'d> {
    pub fn new(storage: &'d raw::Storage, connections: &'d raw::Connections) -> Document<'d> {
        Document {
            storage: storage,
            connections: RefCell::new(connections),
        }
    }

    fn wrap_element(&'d self, node: *mut raw::Element) -> Element<'d> {
        Element {
            document: self,
            node: node,
        }
    }

    fn wrap_attribute(&'d self, node: *mut raw::Attribute) -> Attribute<'d> {
        Attribute {
            document: self,
            node: node,
        }
    }

    pub fn create_element(&'d self, name: &str) -> Element<'d> {
        self.wrap_element(self.storage.create_element(name))
    }
}

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

impl<'d> fmt::Show for Document<'d> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Document {{ {} }}", self as *const Document)
    }
}

pub struct Element<'d> {
    document: &'d Document<'d>,
    node: *mut raw::Element,
}

impl<'d> Element<'d> {
    fn node(&self) -> &raw::Element { unsafe { &*self.node } }

    pub fn document(&self) -> &'d Document<'d> { self.document }
    pub fn name(&self) -> &str { self.node().name() }

    pub fn set_name(&self, name: &str) {
        self.document.storage.element_set_name(self.node, name)
    }

    pub fn parent(&'d self) -> Option<ParentOfChild<'d>> {
        let connections = self.document.connections.borrow();

        connections.element_parent(self.node).map(|n| {
            match n {
                raw::ElementPOC(n) => ElementPOC(self.document.wrap_element(n)),
            }
        })
    }

    pub fn append_child(&self, child: Element) {
        let connections = self.document.connections.borrow_mut();
        connections.append_element_child(self.node, child.node)
    }

    pub fn children(&'d self) -> Vec<ChildOfElement<'d>> {
        let connections = self.document.connections.borrow();
        // This is safe because we copy of the children, and the
        // children are never deallocated.
        unsafe {
            let raw_children = connections.element_children(self.node);
            raw_children.iter().map(|n| match n {
                &raw::ElementCOE(n) => ElementCOE(self.document.wrap_element(n)),
            }).collect()
        }
    }

    pub fn attributes(&self) -> Vec<Attribute<'d>> {
        let connections = self.document.connections.borrow();
        // This is safe because we copy of the children, and the
        // children are never deallocated.
        unsafe {
            connections.attributes(self.node).iter().map(|n| {
                self.document.wrap_attribute(*n)
            }).collect()
        }
    }

    pub fn set_attribute_value(&self, name: &str, value: &str) -> Attribute<'d> {
        let attr = self.document.storage.create_attribute(name, value);
        let connections = self.document.connections.borrow_mut();
        connections.set_attribute(self.node, attr);
        self.document.wrap_attribute(attr)
    }

    pub fn attribute_value(&self, name: &str) -> Option<&'d str> {
        let connections = self.document.connections.borrow();
        connections.attribute(self.node, name).map(|a| {
            let a_r = unsafe { &*a };
            a_r.value()
        })
    }
}

impl<'d> PartialEq for Element<'d> {
    fn eq(&self, other: &Element<'d>) -> bool {
        self.node == other.node
    }
}

impl<'d> fmt::Show for Element<'d> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let node = unsafe { &* self.node };
        write!(f, "Element {{ name: {}  }}", node.name())
    }
}

pub struct Attribute<'d> {
    document: &'d Document<'d>,
    node: *mut raw::Attribute,
}

impl<'d> Attribute<'d> {
    fn node(&self) -> &raw::Attribute { unsafe { &*self.node } }

    pub fn name(&self)  -> &str { self.node().name() }
    pub fn value(&self) -> &str { self.node().value() }

    pub fn parent(&self) -> Option<Element<'d>> {
        let connections = self.document.connections.borrow();
        connections.attribute_parent(self.node).map(|n| {
            self.document.wrap_element(n)
        })
    }
}

#[deriving(PartialEq,Show)]
pub enum ChildOfElement<'d> {
    ElementCOE(Element<'d>)
}

impl<'d> ChildOfElement<'d> {
    pub fn element(self) -> Option<Element<'d>> {
        match self {
            ElementCOE(n) => Some(n),
        }
    }
}

#[deriving(PartialEq,Show)]
pub enum ParentOfChild<'d> {
    ElementPOC(Element<'d>),
}

#[cfg(test)]
mod test {
    use super::super::Package;
    use super::{ElementCOE};
    use super::{ElementPOC};

    #[test]
    fn elements_belong_to_a_document() {
        let package = Package::new();
        let doc = package.as_document();

        let element = doc.create_element("alpha");

        assert_eq!(&doc, element.document());
    }

    #[test]
    fn elements_can_have_element_children() {
        let package = Package::new();
        let doc = package.as_document();

        let alpha = doc.create_element("alpha");
        let beta  = doc.create_element("beta");

        alpha.append_child(beta);

        let children = alpha.children();

        assert_eq!(children[0], ElementCOE(beta));
    }

    #[test]
    fn element_children_are_ordered() {
        let package = Package::new();
        let doc = package.as_document();

        let greek = doc.create_element("greek");
        let alpha = doc.create_element("alpha");
        let omega = doc.create_element("omega");

        greek.append_child(alpha);
        greek.append_child(omega);

        let children = greek.children();

        assert_eq!(children[0], ElementCOE(alpha));
        assert_eq!(children[1], ElementCOE(omega));
    }

    #[test]
    fn element_children_know_their_parent() {
        let package = Package::new();
        let doc = package.as_document();

        let alpha = doc.create_element("alpha");
        let beta  = doc.create_element("beta");

        alpha.append_child(beta);

        assert_eq!(Some(ElementPOC(alpha)), beta.parent());
    }

    #[test]
    fn changing_parent_of_element_removes_element_from_original_parent() {
        let package = Package::new();
        let doc = package.as_document();

        let parent1 = doc.create_element("parent1");
        let parent2 = doc.create_element("parent2");
        let child = doc.create_element("child");

        parent1.append_child(child);
        parent2.append_child(child);

        assert!(parent1.children().is_empty());
        assert_eq!(1, parent2.children().len());
    }

    #[test]
    fn elements_can_be_renamed() {
        let package = Package::new();
        let doc = package.as_document();

        let alpha = doc.create_element("alpha");
        alpha.set_name("beta");
        assert_eq!(alpha.name(), "beta");
    }

    #[test]
    fn elements_have_attributes() {
        let package = Package::new();
        let doc = package.as_document();

        let element = doc.create_element("element");

        element.set_attribute_value("hello", "world");

        assert_eq!(Some("world"), element.attribute_value("hello"));
    }

    #[test]
    fn attributes_know_their_element() {
        let package = Package::new();
        let doc = package.as_document();

        let element = doc.create_element("element");
        let attr = element.set_attribute_value("hello", "world");

        assert_eq!(Some(element), attr.parent());
    }

    #[test]
    fn attributes_can_be_reset() {
        let package = Package::new();
        let doc = package.as_document();

        let element = doc.create_element("element");

        element.set_attribute_value("hello", "world");
        element.set_attribute_value("hello", "galaxy");

        assert_eq!(Some("galaxy"), element.attribute_value("hello"));
    }

    #[test]
    fn attributes_can_be_iterated() {
        let package = Package::new();
        let doc = package.as_document();

        let element = doc.create_element("element");

        element.set_attribute_value("name1", "value1");
        element.set_attribute_value("name2", "value2");

        let mut attrs = element.attributes();
        attrs.sort_by(|a, b| a.name().cmp(b.name()));

        assert_eq!(2, attrs.len());
        assert_eq!("name1",  attrs[0].name());
        assert_eq!("value1", attrs[0].value());
        assert_eq!("name2",  attrs[1].name());
        assert_eq!("value2", attrs[1].value());
    }
}
<commit_msg>Fat DOM attributes belong to a document<commit_after>use super::raw;
use std::fmt;
use std::cell::RefCell;

pub struct Document<'d> {
    storage: &'d raw::Storage,
    connections: RefCell<&'d raw::Connections>,
}

impl<'d> Document<'d> {
    pub fn new(storage: &'d raw::Storage, connections: &'d raw::Connections) -> Document<'d> {
        Document {
            storage: storage,
            connections: RefCell::new(connections),
        }
    }

    fn wrap_element(&'d self, node: *mut raw::Element) -> Element<'d> {
        Element {
            document: self,
            node: node,
        }
    }

    fn wrap_attribute(&'d self, node: *mut raw::Attribute) -> Attribute<'d> {
        Attribute {
            document: self,
            node: node,
        }
    }

    pub fn create_element(&'d self, name: &str) -> Element<'d> {
        self.wrap_element(self.storage.create_element(name))
    }
}

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

impl<'d> fmt::Show for Document<'d> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "Document {{ {} }}", self as *const Document)
    }
}

macro_rules! node(
    ($name:ident, $raw:ty) => (
        pub struct $name<'d> {
            document: &'d Document<'d>,
            node: *mut $raw,
        }

        impl<'d> $name<'d> {
            fn node(&self) -> &$raw { unsafe { &*self.node } }

            pub fn document(&self) -> &'d Document<'d> { self.document }
        }

        impl<'d> PartialEq for $name<'d> {
            fn eq(&self, other: &$name<'d>) -> bool {
                self.node == other.node
            }
        }
    )
)

node!(Element, raw::Element)

impl<'d> Element<'d> {
    pub fn name(&self) -> &str { self.node().name() }

    pub fn set_name(&self, name: &str) {
        self.document.storage.element_set_name(self.node, name)
    }

    pub fn parent(&'d self) -> Option<ParentOfChild<'d>> {
        let connections = self.document.connections.borrow();

        connections.element_parent(self.node).map(|n| {
            match n {
                raw::ElementPOC(n) => ElementPOC(self.document.wrap_element(n)),
            }
        })
    }

    pub fn append_child(&self, child: Element) {
        let connections = self.document.connections.borrow_mut();
        connections.append_element_child(self.node, child.node)
    }

    pub fn children(&'d self) -> Vec<ChildOfElement<'d>> {
        let connections = self.document.connections.borrow();
        // This is safe because we copy of the children, and the
        // children are never deallocated.
        unsafe {
            let raw_children = connections.element_children(self.node);
            raw_children.iter().map(|n| match n {
                &raw::ElementCOE(n) => ElementCOE(self.document.wrap_element(n)),
            }).collect()
        }
    }

    pub fn attributes(&self) -> Vec<Attribute<'d>> {
        let connections = self.document.connections.borrow();
        // This is safe because we copy of the children, and the
        // children are never deallocated.
        unsafe {
            connections.attributes(self.node).iter().map(|n| {
                self.document.wrap_attribute(*n)
            }).collect()
        }
    }

    pub fn set_attribute_value(&self, name: &str, value: &str) -> Attribute<'d> {
        let attr = self.document.storage.create_attribute(name, value);
        let connections = self.document.connections.borrow_mut();
        connections.set_attribute(self.node, attr);
        self.document.wrap_attribute(attr)
    }

    pub fn attribute_value(&self, name: &str) -> Option<&'d str> {
        let connections = self.document.connections.borrow();
        connections.attribute(self.node, name).map(|a| {
            let a_r = unsafe { &*a };
            a_r.value()
        })
    }
}

impl<'d> fmt::Show for Element<'d> {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        let node = unsafe { &* self.node };
        write!(f, "Element {{ name: {}  }}", node.name())
    }
}

node!(Attribute, raw::Attribute)

impl<'d> Attribute<'d> {
    pub fn name(&self)  -> &str { self.node().name() }
    pub fn value(&self) -> &str { self.node().value() }

    pub fn parent(&self) -> Option<Element<'d>> {
        let connections = self.document.connections.borrow();
        connections.attribute_parent(self.node).map(|n| {
            self.document.wrap_element(n)
        })
    }
}

#[deriving(PartialEq,Show)]
pub enum ChildOfElement<'d> {
    ElementCOE(Element<'d>)
}

impl<'d> ChildOfElement<'d> {
    pub fn element(self) -> Option<Element<'d>> {
        match self {
            ElementCOE(n) => Some(n),
        }
    }
}

#[deriving(PartialEq,Show)]
pub enum ParentOfChild<'d> {
    ElementPOC(Element<'d>),
}

#[cfg(test)]
mod test {
    use super::super::Package;
    use super::{ElementCOE};
    use super::{ElementPOC};

    #[test]
    fn elements_belong_to_a_document() {
        let package = Package::new();
        let doc = package.as_document();

        let element = doc.create_element("alpha");

        assert_eq!(&doc, element.document());
    }

    #[test]
    fn elements_can_have_element_children() {
        let package = Package::new();
        let doc = package.as_document();

        let alpha = doc.create_element("alpha");
        let beta  = doc.create_element("beta");

        alpha.append_child(beta);

        let children = alpha.children();

        assert_eq!(children[0], ElementCOE(beta));
    }

    #[test]
    fn element_children_are_ordered() {
        let package = Package::new();
        let doc = package.as_document();

        let greek = doc.create_element("greek");
        let alpha = doc.create_element("alpha");
        let omega = doc.create_element("omega");

        greek.append_child(alpha);
        greek.append_child(omega);

        let children = greek.children();

        assert_eq!(children[0], ElementCOE(alpha));
        assert_eq!(children[1], ElementCOE(omega));
    }

    #[test]
    fn element_children_know_their_parent() {
        let package = Package::new();
        let doc = package.as_document();

        let alpha = doc.create_element("alpha");
        let beta  = doc.create_element("beta");

        alpha.append_child(beta);

        assert_eq!(Some(ElementPOC(alpha)), beta.parent());
    }

    #[test]
    fn changing_parent_of_element_removes_element_from_original_parent() {
        let package = Package::new();
        let doc = package.as_document();

        let parent1 = doc.create_element("parent1");
        let parent2 = doc.create_element("parent2");
        let child = doc.create_element("child");

        parent1.append_child(child);
        parent2.append_child(child);

        assert!(parent1.children().is_empty());
        assert_eq!(1, parent2.children().len());
    }

    #[test]
    fn elements_can_be_renamed() {
        let package = Package::new();
        let doc = package.as_document();

        let alpha = doc.create_element("alpha");
        alpha.set_name("beta");
        assert_eq!(alpha.name(), "beta");
    }

    #[test]
    fn attributes_belong_to_a_document() {
        let package = Package::new();
        let doc = package.as_document();

        let element = doc.create_element("alpha");
        let attr = element.set_attribute_value("hello", "world");

        assert_eq!(&doc, attr.document());
    }

    #[test]
    fn elements_have_attributes() {
        let package = Package::new();
        let doc = package.as_document();

        let element = doc.create_element("element");

        element.set_attribute_value("hello", "world");

        assert_eq!(Some("world"), element.attribute_value("hello"));
    }

    #[test]
    fn attributes_know_their_element() {
        let package = Package::new();
        let doc = package.as_document();

        let element = doc.create_element("element");
        let attr = element.set_attribute_value("hello", "world");

        assert_eq!(Some(element), attr.parent());
    }

    #[test]
    fn attributes_can_be_reset() {
        let package = Package::new();
        let doc = package.as_document();

        let element = doc.create_element("element");

        element.set_attribute_value("hello", "world");
        element.set_attribute_value("hello", "galaxy");

        assert_eq!(Some("galaxy"), element.attribute_value("hello"));
    }

    #[test]
    fn attributes_can_be_iterated() {
        let package = Package::new();
        let doc = package.as_document();

        let element = doc.create_element("element");

        element.set_attribute_value("name1", "value1");
        element.set_attribute_value("name2", "value2");

        let mut attrs = element.attributes();
        attrs.sort_by(|a, b| a.name().cmp(b.name()));

        assert_eq!(2, attrs.len());
        assert_eq!("name1",  attrs[0].name());
        assert_eq!("value1", attrs[0].value());
        assert_eq!("name2",  attrs[1].name());
        assert_eq!("value2", attrs[1].value());
    }
}
<|endoftext|>


<commit_before>use crate::{
    audio, interrupts,
    memory::io::{Interrupt, IoRegisters, Timer},
    scheduler::{EventTag, Scheduler},
    Gba,
};

pub fn started(idx: usize, ioregs: &mut IoRegisters, scheduler: &Scheduler) {
    let timer = &mut ioregs.timers[idx];
    timer.origin = ioregs.time;
    timer.counter = timer.reload;

    if timer.control.count_up_timing() {
        return;
    }

    let event_tag = EventTag::timer(idx);
    let overflow_fn = match idx {
        0 => overflow::<0>,
        1 => overflow::<1>,
        2 => overflow::<2>,
        3 => overflow::<3>,
        _ => panic!("invalid timer"),
    };

    let overflow_cycles = (0x10000u32 - timer.counter as u32) << timer.control.prescaler_shift();
    scheduler.unschedule(event_tag);
    scheduler.schedule(overflow_fn, overflow_cycles, event_tag);
}

pub fn stopped(timer: usize, _timers: &mut [Timer; 4], scheduler: &Scheduler) {
    scheduler.unschedule(EventTag::timer(timer));
}

pub fn reschedule(idx: usize, timers: &mut [Timer; 4], scheduler: &Scheduler) {
    let event_tag = EventTag::timer(idx);
    let timer = &mut timers[idx];
    scheduler.unschedule(event_tag);

    if !timer.control.count_up_timing() {
        let overflow_fn = match idx {
            0 => overflow::<0>,
            1 => overflow::<1>,
            2 => overflow::<2>,
            3 => overflow::<3>,
            _ => panic!("invalid timer"),
        };

        let overflow_cycles =
            (0x10000u32 - timer.counter as u32) << timer.control.prescaler_shift();
        scheduler.schedule(overflow_fn, overflow_cycles, event_tag);
    }
}

pub fn flush(timer: &mut Timer, now: u64) {
    if !timer.control.started() || timer.control.count_up_timing() {
        return;
    }

    // This should NEVER overflow before the overflow function is called by the scheduler, which should reset the timer's counter.
    timer.counter += ((now - timer.origin) >> timer.control.prescaler_shift()) as u16;
    timer.origin = now;
}

pub fn overflow<const TIMER: usize>(gba: &mut Gba) {
    let mut overflows = 1u32;

    {
        let timer = &mut gba.mem.ioregs.timers[TIMER];
        timer.counter = timer.reload;
        timer.origin = gba.mem.ioregs.time;

        let overflow_cycles =
            (0x10000u32 - timer.counter as u32) << timer.control.prescaler_shift();
        gba.scheduler
            .schedule(overflow::<TIMER>, overflow_cycles, EventTag::timer(TIMER));
    }
    after_overflow(TIMER, gba);

    let mut idx = TIMER + 1;
    while overflows > 0 && idx < 4 {
        let timer = &mut gba.mem.ioregs.timers[idx];
        if !timer.control.count_up_timing() || !timer.control.started() {
            break;
        }

        let mut rem = overflows;
        overflows = 0;
        while rem >= (0x10000u32 - timer.counter as u32) {
            overflows += 1;
            rem -= 0x10000u32 - timer.counter as u32;
            timer.counter = timer.reload;
        }
        timer.counter += rem as u16;
        for _ in 0..overflows {
            after_overflow(idx, gba);
        }
        idx += 1;
    }
}

fn after_overflow(idx: usize, gba: &mut Gba) {
    if gba.mem.ioregs.timers[idx].control.irq_enable() {
        interrupts::raise(Interrupt::timer(idx), &mut gba.mem.ioregs, &gba.scheduler);
    }

    if idx == 0 || idx == 1 {
        audio::check_fifo_timer_overflow(idx, gba);
    }
}
<commit_msg>Make timers more accurate by not flushing them as much.<commit_after>use crate::{
    audio, interrupts,
    memory::io::{Interrupt, IoRegisters, Timer},
    scheduler::{EventTag, Scheduler},
    Gba,
};

pub fn started(idx: usize, ioregs: &mut IoRegisters, scheduler: &Scheduler) {
    let timer = &mut ioregs.timers[idx];
    timer.origin = ioregs.time;
    timer.counter = timer.reload;

    if timer.control.count_up_timing() {
        return;
    }

    let event_tag = EventTag::timer(idx);
    let overflow_fn = match idx {
        0 => overflow::<0>,
        1 => overflow::<1>,
        2 => overflow::<2>,
        3 => overflow::<3>,
        _ => panic!("invalid timer"),
    };

    let overflow_cycles = (0x10000u32 - timer.counter as u32) << timer.control.prescaler_shift();
    scheduler.unschedule(event_tag);
    scheduler.schedule(overflow_fn, overflow_cycles, event_tag);
}

pub fn stopped(timer: usize, _timers: &mut [Timer; 4], scheduler: &Scheduler) {
    scheduler.unschedule(EventTag::timer(timer));
}

pub fn reschedule(idx: usize, timers: &mut [Timer; 4], scheduler: &Scheduler) {
    let event_tag = EventTag::timer(idx);
    let timer = &mut timers[idx];
    scheduler.unschedule(event_tag);

    if !timer.control.count_up_timing() {
        let overflow_fn = match idx {
            0 => overflow::<0>,
            1 => overflow::<1>,
            2 => overflow::<2>,
            3 => overflow::<3>,
            _ => panic!("invalid timer"),
        };

        let overflow_cycles =
            (0x10000u32 - timer.counter as u32) << timer.control.prescaler_shift();
        scheduler.schedule(overflow_fn, overflow_cycles, event_tag);
    }
}

pub fn flush(timer: &mut Timer, now: u64) {
    if !timer.control.started() || timer.control.count_up_timing() {
        return;
    }

    // This should NEVER overflow before the overflow function is called by the scheduler, which should reset the timer's counter.
    let delta = now - timer.origin;
    if delta >= (1 << timer.control.prescaler_shift()) {
        timer.counter += ((now - timer.origin) >> timer.control.prescaler_shift()) as u16;
        timer.origin = now;
    }
}

pub fn overflow<const TIMER: usize>(gba: &mut Gba) {
    let mut overflows = 1u32;

    {
        let timer = &mut gba.mem.ioregs.timers[TIMER];
        timer.counter = timer.reload;
        timer.origin = gba.mem.ioregs.time;

        let overflow_cycles =
            (0x10000u32 - timer.counter as u32) << timer.control.prescaler_shift();
        gba.scheduler
            .schedule(overflow::<TIMER>, overflow_cycles, EventTag::timer(TIMER));
    }
    after_overflow(TIMER, gba);

    let mut idx = TIMER + 1;
    while overflows > 0 && idx < 4 {
        let timer = &mut gba.mem.ioregs.timers[idx];
        if !timer.control.count_up_timing() || !timer.control.started() {
            break;
        }

        let mut rem = overflows;
        overflows = 0;
        while rem >= (0x10000u32 - timer.counter as u32) {
            overflows += 1;
            rem -= 0x10000u32 - timer.counter as u32;
            timer.counter = timer.reload;
        }
        timer.counter += rem as u16;
        for _ in 0..overflows {
            after_overflow(idx, gba);
        }
        idx += 1;
    }
}

fn after_overflow(idx: usize, gba: &mut Gba) {
    if gba.mem.ioregs.timers[idx].control.irq_enable() {
        interrupts::raise(Interrupt::timer(idx), &mut gba.mem.ioregs, &gba.scheduler);
    }

    if idx == 0 || idx == 1 {
        audio::check_fifo_timer_overflow(idx, gba);
    }
}
<|endoftext|>


<commit_before>use pyo3;
use pyo3::types::{PyBytes};
use pyo3::prelude::*;
use failure::{Error};

use lancelot::workspace;
use lancelot::arch::{RVA};

#[pymodule]
fn pylancelot(_py: Python, m: &PyModule) -> PyResult<()> {
    #[pyfn(m, "from_bytes")]
    /// from_bytes(filename, buf, /)
    /// --
    ///
    /// Create a workspace from the given bytes.
    ///
    /// Args:
    ///   filename (str): the source filename
    ///   buf (bytes): the bytes containing a PE, shellcode, etc.
    ///
    /// Raises:
    ///   ValueError: if failed to create the workspace. TODO: more specific.
    ///
    fn from_bytes(_py: Python, filename: String, buf: &PyBytes) -> PyResult<PyWorkspace> {
        let ws = match workspace::Workspace::from_bytes(&filename, buf.as_bytes()).load() {
            Err(_) => return Err(pyo3::exceptions::ValueError::py_err("failed to create workspace")),
            Ok(ws) => ws,
        };
        Ok(PyWorkspace{ws})
    }

    #[pyclass]
    pub struct PyWorkspace {
        ws: workspace::Workspace
    }

    #[pyclass]
    pub struct PySection {
        #[pyo3(get)]
        pub addr: i64,
        #[pyo3(get)]
        pub length: u64,
        #[pyo3(get)]
        pub perms: u8,
        #[pyo3(get)]
        pub name: String,
    }

    m.add("PERM_NONE", lancelot::loader::Permissions::empty().bits())?;
    m.add("PERM_R", lancelot::loader::Permissions::R.bits())?;
    m.add("PERM_W", lancelot::loader::Permissions::W.bits())?;
    m.add("PERM_X", lancelot::loader::Permissions::X.bits())?;
    m.add("PERM_RW", lancelot::loader::Permissions::RW.bits())?;
    m.add("PERM_RX", lancelot::loader::Permissions::RX.bits())?;
    m.add("PERM_RWX", lancelot::loader::Permissions::RWX.bits())?;

    #[pymethods]
    impl PyWorkspace {
        #[getter]
        /// filename(self, /)
        /// --
        ///
        /// Fetch the filename.
        /// ```
        pub fn filename(&self) -> PyResult<String> {
            Ok(self.ws.filename.clone())
        }

        #[getter]
        /// loader(self, /)
        /// --
        ///
        /// Fetch the name of the loader used to create the workspace.
        pub fn loader(&self) -> PyResult<String> {
            Ok(self.ws.loader.get_name())
        }

        #[getter]
        /// base_address(self, /)
        /// --
        ///
        /// Fetch the base address to which the module was loaded.
        pub fn base_address(&self) -> PyResult<u64> {
            Ok(self.ws.module.base_address.into())
        }

        #[getter]
        pub fn sections(&self) -> PyResult<Vec<PySection>> {
            Ok(self.ws.module.sections.iter()
                .map(|section| PySection{
                    addr: section.addr.into(),
                    length: section.buf.len() as u64,
                    perms: section.perms.bits(),
                    name: section.name.clone(),
                })
                .collect())
        }

        /// probe(self, rva, length=1, /)
        /// --
        ///
        /// Is the given address mapped?
        pub fn probe(&self, rva: i64, length: Option<usize>) -> PyResult<bool> {
            match length {
                Some(length) => Ok(self.ws.probe(RVA::from(rva), length)),
                None =>  Ok(self.ws.probe(RVA::from(rva), 1)),
            }
        }

        /// read_bytes(self, rva, length, /)
        /// --
        ///
        /// Read bytes from the given memory address.
        pub fn read_bytes<'p>(&self, py: Python<'p>, rva: i64, length: usize) -> PyResult<&'p pyo3::types::PyBytes> {
            match self.ws.read_bytes(RVA::from(rva), length) {
                Ok(buf) => Ok(pyo3::types::PyBytes::new(py, buf)),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        pub fn read_u8(&self, rva: i64) -> PyResult<u8> {
            match self.ws.read_u8(RVA::from(rva)) {
                Ok(b) => Ok(b),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        pub fn read_u16(&self, rva: i64) -> PyResult<u16> {
            match self.ws.read_u16(RVA::from(rva)) {
                Ok(b) => Ok(b),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        pub fn read_u32(&self, rva: i64) -> PyResult<u32> {
            match self.ws.read_u32(RVA::from(rva)) {
                Ok(b) => Ok(b),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        pub fn read_u64(&self, rva: i64) -> PyResult<u64> {
            match self.ws.read_u64(RVA::from(rva)) {
                Ok(b) => Ok(b),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        pub fn read_rva(&self, rva: i64) -> PyResult<i64> {
            match self.ws.read_rva(RVA::from(rva)) {
                Ok(b) => Ok(b.into()),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        pub fn read_va(&self, rva: i64) -> PyResult<u64> {
            match self.ws.read_va(RVA::from(rva)) {
                Ok(b) => Ok(b.into()),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }
    }

    impl PyWorkspace {
        /// translate from a failure::Error returned by a Workspace
        /// while reading an element from memory into a pyo3/Python error
        /// corresponding to out-of-bounds-access.
        fn translate_buffer_error(e: Error) -> PyErr {
            match e.downcast::<lancelot::workspace::WorkspaceError>() {
                // we have to use the documentation for the associated routine,
                // such as `read_u8`, to determine what to inspect here.
                Ok(lancelot::workspace::WorkspaceError::InvalidAddress) => {
                    pyo3::exceptions::LookupError::py_err("invalid address")
                },
                Ok(lancelot::workspace::WorkspaceError::BufferOverrun) => {
                    pyo3::exceptions::LookupError::py_err("buffer overrun")
                }
                Ok(_) => {
                    // this should never be hit
                    pyo3::exceptions::RuntimeError::py_err("unexpected error")
                },
                Err(_) => {
                    // this should never be hit
                    pyo3::exceptions::RuntimeError::py_err("unexpected error")
                },
            }
        }
    }

    #[pyproto]
    impl pyo3::class::basic::PyObjectProtocol for PyWorkspace {
        fn __str__(&self) -> PyResult<String> {
            PyWorkspace::__repr__(self)
        }

        fn __repr__(&self) -> PyResult<String> {
            Ok(format!("PyWorkspace(filename: {} loader: {})",
                self.ws.filename.clone(),
                self.ws.loader.get_name(),
            ))
        }
    }

    #[pyproto]
    impl pyo3::class::basic::PyObjectProtocol for PySection {
        fn __str__(&self) -> PyResult<String> {
            PySection::__repr__(self)
        }

        fn __repr__(&self) -> PyResult<String> {
            Ok(format!("PySection(addr: {:#x} length: {:#x} perms: {:#x} name: {})",
                self.addr,
                self.length,
                self.perms,
                self.name,
            ))
        }
    }

    Ok(())
}

<commit_msg>pylancelot: add doc to read_*<commit_after>use pyo3;
use pyo3::types::{PyBytes};
use pyo3::prelude::*;
use failure::{Error};

use lancelot::workspace;
use lancelot::arch::{RVA};

#[pymodule]
fn pylancelot(_py: Python, m: &PyModule) -> PyResult<()> {
    #[pyfn(m, "from_bytes")]
    /// from_bytes(filename, buf, /)
    /// --
    ///
    /// Create a workspace from the given bytes.
    ///
    /// Args:
    ///   filename (str): the source filename
    ///   buf (bytes): the bytes containing a PE, shellcode, etc.
    ///
    /// Raises:
    ///   ValueError: if failed to create the workspace. TODO: more specific.
    ///
    fn from_bytes(_py: Python, filename: String, buf: &PyBytes) -> PyResult<PyWorkspace> {
        let ws = match workspace::Workspace::from_bytes(&filename, buf.as_bytes()).load() {
            Err(_) => return Err(pyo3::exceptions::ValueError::py_err("failed to create workspace")),
            Ok(ws) => ws,
        };
        Ok(PyWorkspace{ws})
    }

    #[pyclass]
    pub struct PyWorkspace {
        ws: workspace::Workspace
    }

    #[pyclass]
    pub struct PySection {
        #[pyo3(get)]
        pub addr: i64,
        #[pyo3(get)]
        pub length: u64,
        #[pyo3(get)]
        pub perms: u8,
        #[pyo3(get)]
        pub name: String,
    }

    m.add("PERM_NONE", lancelot::loader::Permissions::empty().bits())?;
    m.add("PERM_R", lancelot::loader::Permissions::R.bits())?;
    m.add("PERM_W", lancelot::loader::Permissions::W.bits())?;
    m.add("PERM_X", lancelot::loader::Permissions::X.bits())?;
    m.add("PERM_RW", lancelot::loader::Permissions::RW.bits())?;
    m.add("PERM_RX", lancelot::loader::Permissions::RX.bits())?;
    m.add("PERM_RWX", lancelot::loader::Permissions::RWX.bits())?;

    #[pymethods]
    impl PyWorkspace {
        #[getter]
        /// filename(self, /)
        /// --
        ///
        /// Fetch the filename.
        /// ```
        pub fn filename(&self) -> PyResult<String> {
            Ok(self.ws.filename.clone())
        }

        #[getter]
        /// loader(self, /)
        /// --
        ///
        /// Fetch the name of the loader used to create the workspace.
        pub fn loader(&self) -> PyResult<String> {
            Ok(self.ws.loader.get_name())
        }

        #[getter]
        /// base_address(self, /)
        /// --
        ///
        /// Fetch the base address to which the module was loaded.
        pub fn base_address(&self) -> PyResult<u64> {
            Ok(self.ws.module.base_address.into())
        }

        #[getter]
        pub fn sections(&self) -> PyResult<Vec<PySection>> {
            Ok(self.ws.module.sections.iter()
                .map(|section| PySection{
                    addr: section.addr.into(),
                    length: section.buf.len() as u64,
                    perms: section.perms.bits(),
                    name: section.name.clone(),
                })
                .collect())
        }

        /// probe(self, rva, length=1, /)
        /// --
        ///
        /// Is the given address mapped?
        pub fn probe(&self, rva: i64, length: Option<usize>) -> PyResult<bool> {
            match length {
                Some(length) => Ok(self.ws.probe(RVA::from(rva), length)),
                None =>  Ok(self.ws.probe(RVA::from(rva), 1)),
            }
        }

        /// read_bytes(self, rva, length, /)
        /// --
        ///
        /// Read bytes from the given memory address.
        ///
        /// raises:
        ///   - LookupError: if the address is not mapped, or the length overruns.
        pub fn read_bytes<'p>(&self, py: Python<'p>, rva: i64, length: usize) -> PyResult<&'p pyo3::types::PyBytes> {
            match self.ws.read_bytes(RVA::from(rva), length) {
                Ok(buf) => Ok(pyo3::types::PyBytes::new(py, buf)),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        /// read_u8(self, rva, /)
        /// --
        ///
        pub fn read_u8(&self, rva: i64) -> PyResult<u8> {
            match self.ws.read_u8(RVA::from(rva)) {
                Ok(b) => Ok(b),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        /// read_u16(self, rva, /)
        /// --
        ///
        pub fn read_u16(&self, rva: i64) -> PyResult<u16> {
            match self.ws.read_u16(RVA::from(rva)) {
                Ok(b) => Ok(b),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        /// read_u32(self, rva, /)
        /// --
        ///
        pub fn read_u32(&self, rva: i64) -> PyResult<u32> {
            match self.ws.read_u32(RVA::from(rva)) {
                Ok(b) => Ok(b),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        /// read_u64(self, rva, /)
        /// --
        ///
        pub fn read_u64(&self, rva: i64) -> PyResult<u64> {
            match self.ws.read_u64(RVA::from(rva)) {
                Ok(b) => Ok(b),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        /// read_rva(self, rva, /)
        /// --
        ///
        /// Read the integer relative virtual address at the given address.
        /// The size of the RVA is dependent upon the bitness of the current workspace.
        pub fn read_rva(&self, rva: i64) -> PyResult<i64> {
            match self.ws.read_rva(RVA::from(rva)) {
                Ok(b) => Ok(b.into()),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }

        /// read_va(self, rva, /)
        /// --
        ///
        /// Read the integer virtual address at the given address.
        /// The size of the VA is dependent upon the bitness of the current workspace.
        pub fn read_va(&self, rva: i64) -> PyResult<u64> {
            match self.ws.read_va(RVA::from(rva)) {
                Ok(b) => Ok(b.into()),
                Err(e) => Err(PyWorkspace::translate_buffer_error(e)),
            }
        }
    }

    impl PyWorkspace {
        /// translate from a failure::Error returned by a Workspace
        /// while reading an element from memory into a pyo3/Python error
        /// corresponding to out-of-bounds-access.
        fn translate_buffer_error(e: Error) -> PyErr {
            match e.downcast::<lancelot::workspace::WorkspaceError>() {
                // we have to use the documentation for the associated routine,
                // such as `read_u8`, to determine what to inspect here.
                Ok(lancelot::workspace::WorkspaceError::InvalidAddress) => {
                    pyo3::exceptions::LookupError::py_err("invalid address")
                },
                Ok(lancelot::workspace::WorkspaceError::BufferOverrun) => {
                    pyo3::exceptions::LookupError::py_err("buffer overrun")
                }
                Ok(_) => {
                    // this should never be hit
                    pyo3::exceptions::RuntimeError::py_err("unexpected error")
                },
                Err(_) => {
                    // this should never be hit
                    pyo3::exceptions::RuntimeError::py_err("unexpected error")
                },
            }
        }
    }

    #[pyproto]
    impl pyo3::class::basic::PyObjectProtocol for PyWorkspace {
        fn __str__(&self) -> PyResult<String> {
            PyWorkspace::__repr__(self)
        }

        fn __repr__(&self) -> PyResult<String> {
            Ok(format!("PyWorkspace(filename: {} loader: {})",
                self.ws.filename.clone(),
                self.ws.loader.get_name(),
            ))
        }
    }

    #[pyproto]
    impl pyo3::class::basic::PyObjectProtocol for PySection {
        fn __str__(&self) -> PyResult<String> {
            PySection::__repr__(self)
        }

        fn __repr__(&self) -> PyResult<String> {
            Ok(format!("PySection(addr: {:#x} length: {:#x} perms: {:#x} name: {})",
                self.addr,
                self.length,
                self.perms,
                self.name,
            ))
        }
    }

    Ok(())
}

<|endoftext|>


<commit_before>use crate::gl;
use crate::version::Version;
use crate::version::Api;

use crate::CapabilitiesSource;
use crate::backend::Facade;
use crate::context::Context;
use crate::ContextExt;

use std::{ffi, ptr};
use std::rc::Rc;

use crate::GlObject;
use crate::Handle;

use crate::program::{ProgramCreationError, ShaderType, SpirvEntryPoint};

/// A single, compiled but unlinked, shader.
pub struct Shader {
    context: Rc<Context>,
    id: Handle,
}

impl GlObject for Shader {
    type Id = Handle;

    #[inline]
    fn get_id(&self) -> Handle {
        self.id
    }
}

impl Drop for Shader {
    fn drop(&mut self) {
        let ctxt = self.context.make_current();

        unsafe {
            match self.id {
                Handle::Id(id) => {
                    assert!(ctxt.version >= &Version(Api::Gl, 2, 0) ||
                            ctxt.version >= &Version(Api::GlEs, 2, 0));
                    ctxt.gl.DeleteShader(id);
                },
                Handle::Handle(id) => {
                    assert!(ctxt.extensions.gl_arb_shader_objects);
                    ctxt.gl.DeleteObjectARB(id);
                }
            }
        }
    }
}

/// Builds an individual shader.
pub fn build_shader<F: ?Sized>(facade: &F, shader_type: gl::types::GLenum, source_code: &str)
                       -> Result<Shader, ProgramCreationError> where F: Facade
{
    unsafe {
        let mut ctxt = facade.get_context().make_current();

        if ctxt.capabilities.supported_glsl_versions.is_empty() {
            return Err(ProgramCreationError::CompilationNotSupported);
        }

        if !check_shader_type_compatibility(&mut ctxt, shader_type) {
            return Err(ProgramCreationError::ShaderTypeNotSupported);
        }

        let source_code = ffi::CString::new(source_code.as_bytes()).unwrap();

        let id = if ctxt.version >= &Version(Api::Gl, 2, 0) ||
                    ctxt.version >= &Version(Api::GlEs, 2, 0)
        {
            Handle::Id(ctxt.gl.CreateShader(shader_type))
        } else if ctxt.extensions.gl_arb_shader_objects {
            Handle::Handle(ctxt.gl.CreateShaderObjectARB(shader_type))
        } else {
            unreachable!()
        };

        if id == Handle::Id(0) || id == Handle::Handle(0 as gl::types::GLhandleARB) {
            return Err(ProgramCreationError::ShaderTypeNotSupported);
        }

        match id {
            Handle::Id(id) => {
                assert!(ctxt.version >= &Version(Api::Gl, 2, 0) ||
                        ctxt.version >= &Version(Api::GlEs, 2, 0));
                ctxt.gl.ShaderSource(id, 1, [ source_code.as_ptr() ].as_ptr(), ptr::null());
            },
            Handle::Handle(id) => {
                assert!(ctxt.extensions.gl_arb_shader_objects);
                ctxt.gl.ShaderSourceARB(id, 1, [ source_code.as_ptr() ].as_ptr(), ptr::null());
            }
        }

        // compiling
        {
            ctxt.report_debug_output_errors.set(false);

            match id {
                Handle::Id(id) => {
                    assert!(ctxt.version >= &Version(Api::Gl, 2, 0)||
                            ctxt.version >= &Version(Api::GlEs, 2, 0));
                    ctxt.gl.CompileShader(id);
                },
                Handle::Handle(id) => {
                    assert!(ctxt.extensions.gl_arb_shader_objects);
                    ctxt.gl.CompileShaderARB(id);
                }
            }

            ctxt.report_debug_output_errors.set(true);
        }

        // checking compilation success by reading a flag on the shader
        let compilation_success = {
            let mut compilation_success: gl::types::GLint = 0;
            match id {
                Handle::Id(id) => {
                    assert!(ctxt.version >= &Version(Api::Gl, 2, 0) ||
                            ctxt.version >= &Version(Api::GlEs, 2, 0));
                    ctxt.gl.GetShaderiv(id, gl::COMPILE_STATUS, &mut compilation_success);
                },
                Handle::Handle(id) => {
                    assert!(ctxt.extensions.gl_arb_shader_objects);
                    ctxt.gl.GetObjectParameterivARB(id, gl::OBJECT_COMPILE_STATUS_ARB,
                                                    &mut compilation_success);
                }
            }
            compilation_success
        };

        if compilation_success == 1 {
            Ok(Shader {
                context: facade.get_context().clone(),
                id
            })

        } else {
            // compilation error
            let mut error_log_size: gl::types::GLint = 0;

            match id {
                Handle::Id(id) => {
                    assert!(ctxt.version >= &Version(Api::Gl, 2, 0) ||
                            ctxt.version >= &Version(Api::GlEs, 2, 0));
                    ctxt.gl.GetShaderiv(id, gl::INFO_LOG_LENGTH, &mut error_log_size);
                },
                Handle::Handle(id) => {
                    assert!(ctxt.extensions.gl_arb_shader_objects);
                    ctxt.gl.GetObjectParameterivARB(id, gl::OBJECT_INFO_LOG_LENGTH_ARB,
                                                    &mut error_log_size);
                }
            }

            let mut error_log: Vec<u8> = Vec::with_capacity(error_log_size as usize);

            match id {
                Handle::Id(id) => {
                    assert!(ctxt.version >= &Version(Api::Gl, 2, 0) ||
                            ctxt.version >= &Version(Api::GlEs, 2, 0));
                    ctxt.gl.GetShaderInfoLog(id, error_log_size, &mut error_log_size,
                                             error_log.as_mut_ptr() as *mut gl::types::GLchar);
                },
                Handle::Handle(id) => {
                    assert!(ctxt.extensions.gl_arb_shader_objects);
                    ctxt.gl.GetInfoLogARB(id, error_log_size, &mut error_log_size,
                                          error_log.as_mut_ptr() as *mut gl::types::GLchar);
                }
            }

            error_log.set_len(error_log_size as usize);

            match String::from_utf8(error_log) {
                Ok(msg) => Err(ProgramCreationError::CompilationError(msg, ShaderType::from_opengl_type(shader_type))),
                Err(_) => Err(
                    ProgramCreationError::CompilationError("Could not convert the log \
                                                            message to UTF-8".to_owned(), ShaderType::from_opengl_type(shader_type))
                ),
            }
        }
    }
}

/// Builds an individual shader from a SPIR-V binary.
pub fn build_spirv_shader<F: ?Sized>(facade: &F, shader_type: gl::types::GLenum, spirv: &SpirvEntryPoint)
                       -> Result<Shader, ProgramCreationError> where F: Facade
{
    unsafe {
        let mut ctxt = facade.get_context().make_current();

        if ctxt.capabilities.supported_glsl_versions.is_empty() {
            return Err(ProgramCreationError::CompilationNotSupported);
        }

        if !check_shader_type_compatibility(&mut ctxt, shader_type) {
            return Err(ProgramCreationError::ShaderTypeNotSupported);
        }

        let id = if ctxt.version >= &Version(Api::Gl, 4, 6) ||
            (ctxt.version >= &Version(Api::Gl, 4, 1) && ctxt.extensions.gl_arb_gl_spirv)
        {
            ctxt.gl.CreateShader(shader_type)
        } else {
            return Err(ProgramCreationError::CompilationNotSupported);
        };

        if id == 0 {
            return Err(ProgramCreationError::ShaderTypeNotSupported);
        }

        let binary = &spirv.binary;
        let entry_point = ffi::CString::new(spirv.entry_point.as_bytes()).unwrap();

        const GL_SHADER_BINARY_FORMAT_SPIR_V: gl::types::GLenum = 0x9551;

        ctxt.gl.ShaderBinary(1, &id, GL_SHADER_BINARY_FORMAT_SPIR_V, binary.as_ptr() as _, binary.len() as _);

        // compiling
        {
            ctxt.report_debug_output_errors.set(false);

            ctxt.gl.SpecializeShader(id, entry_point.as_ptr() as _, 0, ptr::null(), ptr::null());

            ctxt.report_debug_output_errors.set(true);
        }

        // checking compilation success by reading a flag on the shader
        let compilation_success = {
            let mut compilation_success: gl::types::GLint = 0;
            ctxt.gl.GetShaderiv(id, gl::COMPILE_STATUS, &mut compilation_success);
            compilation_success
        };

        if compilation_success == 1 {
            Ok(Shader {
                context: facade.get_context().clone(),
                id: Handle::Id(id)
            })
        } else {
            // compilation error
            let mut error_log_size: gl::types::GLint = 0;

            ctxt.gl.GetShaderiv(id, gl::INFO_LOG_LENGTH, &mut error_log_size);

            let mut error_log: Vec<u8> = Vec::with_capacity(error_log_size as usize);

            ctxt.gl.GetShaderInfoLog(id, error_log_size, &mut error_log_size,
                                     error_log.as_mut_ptr() as *mut gl::types::GLchar);

            error_log.set_len(error_log_size as usize);

            match String::from_utf8(error_log) {
                Ok(msg) => Err(ProgramCreationError::CompilationError(msg, ShaderType::from_opengl_type(shader_type))),
                Err(_) => Err(
                    ProgramCreationError::CompilationError("Could not convert the log \
                                                            message to UTF-8".to_owned(), ShaderType::from_opengl_type(shader_type))
                ),
            }
        }
    }
}

pub fn check_shader_type_compatibility<C: ?Sized>(ctxt: &C, shader_type: gl::types::GLenum)
                                          -> bool where C: CapabilitiesSource
{
    match shader_type {
        gl::VERTEX_SHADER | gl::FRAGMENT_SHADER => (),
        gl::GEOMETRY_SHADER => {
            if !(ctxt.get_version() >= &Version(Api::Gl, 3, 2))
                && !(ctxt.get_version() >= &Version(Api::GlEs, 3, 2))
                && !ctxt.get_extensions().gl_arb_geometry_shader4
                && !ctxt.get_extensions().gl_ext_geometry_shader4
                && !ctxt.get_extensions().gl_ext_geometry_shader
                && !ctxt.get_extensions().gl_oes_geometry_shader
            {
                return false;
            }
        },
        gl::TESS_CONTROL_SHADER | gl::TESS_EVALUATION_SHADER => {
            if !(ctxt.get_version() >= &Version(Api::Gl, 4, 0))
                && !(ctxt.get_version() >= &Version(Api::GlEs, 3, 2))
                && !ctxt.get_extensions().gl_arb_tessellation_shader
                && !ctxt.get_extensions().gl_oes_tessellation_shader
            {
                return false;
            }
        },
        gl::COMPUTE_SHADER => {
            if !(ctxt.get_version() >= &Version(Api::Gl, 4, 3))
                && !(ctxt.get_version() >= &Version(Api::GlEs, 3, 1))
                && !ctxt.get_extensions().gl_arb_compute_shader
            {
                return false;
            }
        },
        _ => unreachable!()
    };

    true
}
<commit_msg>Remove redundant mut<commit_after>use crate::gl;
use crate::version::Version;
use crate::version::Api;

use crate::CapabilitiesSource;
use crate::backend::Facade;
use crate::context::Context;
use crate::ContextExt;

use std::{ffi, ptr};
use std::rc::Rc;

use crate::GlObject;
use crate::Handle;

use crate::program::{ProgramCreationError, ShaderType, SpirvEntryPoint};

/// A single, compiled but unlinked, shader.
pub struct Shader {
    context: Rc<Context>,
    id: Handle,
}

impl GlObject for Shader {
    type Id = Handle;

    #[inline]
    fn get_id(&self) -> Handle {
        self.id
    }
}

impl Drop for Shader {
    fn drop(&mut self) {
        let ctxt = self.context.make_current();

        unsafe {
            match self.id {
                Handle::Id(id) => {
                    assert!(ctxt.version >= &Version(Api::Gl, 2, 0) ||
                            ctxt.version >= &Version(Api::GlEs, 2, 0));
                    ctxt.gl.DeleteShader(id);
                },
                Handle::Handle(id) => {
                    assert!(ctxt.extensions.gl_arb_shader_objects);
                    ctxt.gl.DeleteObjectARB(id);
                }
            }
        }
    }
}

/// Builds an individual shader.
pub fn build_shader<F: ?Sized>(facade: &F, shader_type: gl::types::GLenum, source_code: &str)
                       -> Result<Shader, ProgramCreationError> where F: Facade
{
    unsafe {
        let ctxt = facade.get_context().make_current();

        if ctxt.capabilities.supported_glsl_versions.is_empty() {
            return Err(ProgramCreationError::CompilationNotSupported);
        }

        if !check_shader_type_compatibility(&ctxt, shader_type) {
            return Err(ProgramCreationError::ShaderTypeNotSupported);
        }

        let source_code = ffi::CString::new(source_code.as_bytes()).unwrap();

        let id = if ctxt.version >= &Version(Api::Gl, 2, 0) ||
                    ctxt.version >= &Version(Api::GlEs, 2, 0)
        {
            Handle::Id(ctxt.gl.CreateShader(shader_type))
        } else if ctxt.extensions.gl_arb_shader_objects {
            Handle::Handle(ctxt.gl.CreateShaderObjectARB(shader_type))
        } else {
            unreachable!()
        };

        if id == Handle::Id(0) || id == Handle::Handle(0 as gl::types::GLhandleARB) {
            return Err(ProgramCreationError::ShaderTypeNotSupported);
        }

        match id {
            Handle::Id(id) => {
                assert!(ctxt.version >= &Version(Api::Gl, 2, 0) ||
                        ctxt.version >= &Version(Api::GlEs, 2, 0));
                ctxt.gl.ShaderSource(id, 1, [ source_code.as_ptr() ].as_ptr(), ptr::null());
            },
            Handle::Handle(id) => {
                assert!(ctxt.extensions.gl_arb_shader_objects);
                ctxt.gl.ShaderSourceARB(id, 1, [ source_code.as_ptr() ].as_ptr(), ptr::null());
            }
        }

        // compiling
        {
            ctxt.report_debug_output_errors.set(false);

            match id {
                Handle::Id(id) => {
                    assert!(ctxt.version >= &Version(Api::Gl, 2, 0)||
                            ctxt.version >= &Version(Api::GlEs, 2, 0));
                    ctxt.gl.CompileShader(id);
                },
                Handle::Handle(id) => {
                    assert!(ctxt.extensions.gl_arb_shader_objects);
                    ctxt.gl.CompileShaderARB(id);
                }
            }

            ctxt.report_debug_output_errors.set(true);
        }

        // checking compilation success by reading a flag on the shader
        let compilation_success = {
            let mut compilation_success: gl::types::GLint = 0;
            match id {
                Handle::Id(id) => {
                    assert!(ctxt.version >= &Version(Api::Gl, 2, 0) ||
                            ctxt.version >= &Version(Api::GlEs, 2, 0));
                    ctxt.gl.GetShaderiv(id, gl::COMPILE_STATUS, &mut compilation_success);
                },
                Handle::Handle(id) => {
                    assert!(ctxt.extensions.gl_arb_shader_objects);
                    ctxt.gl.GetObjectParameterivARB(id, gl::OBJECT_COMPILE_STATUS_ARB,
                                                    &mut compilation_success);
                }
            }
            compilation_success
        };

        if compilation_success == 1 {
            Ok(Shader {
                context: facade.get_context().clone(),
                id
            })

        } else {
            // compilation error
            let mut error_log_size: gl::types::GLint = 0;

            match id {
                Handle::Id(id) => {
                    assert!(ctxt.version >= &Version(Api::Gl, 2, 0) ||
                            ctxt.version >= &Version(Api::GlEs, 2, 0));
                    ctxt.gl.GetShaderiv(id, gl::INFO_LOG_LENGTH, &mut error_log_size);
                },
                Handle::Handle(id) => {
                    assert!(ctxt.extensions.gl_arb_shader_objects);
                    ctxt.gl.GetObjectParameterivARB(id, gl::OBJECT_INFO_LOG_LENGTH_ARB,
                                                    &mut error_log_size);
                }
            }

            let mut error_log: Vec<u8> = Vec::with_capacity(error_log_size as usize);

            match id {
                Handle::Id(id) => {
                    assert!(ctxt.version >= &Version(Api::Gl, 2, 0) ||
                            ctxt.version >= &Version(Api::GlEs, 2, 0));
                    ctxt.gl.GetShaderInfoLog(id, error_log_size, &mut error_log_size,
                                             error_log.as_mut_ptr() as *mut gl::types::GLchar);
                },
                Handle::Handle(id) => {
                    assert!(ctxt.extensions.gl_arb_shader_objects);
                    ctxt.gl.GetInfoLogARB(id, error_log_size, &mut error_log_size,
                                          error_log.as_mut_ptr() as *mut gl::types::GLchar);
                }
            }

            error_log.set_len(error_log_size as usize);

            match String::from_utf8(error_log) {
                Ok(msg) => Err(ProgramCreationError::CompilationError(msg, ShaderType::from_opengl_type(shader_type))),
                Err(_) => Err(
                    ProgramCreationError::CompilationError("Could not convert the log \
                                                            message to UTF-8".to_owned(), ShaderType::from_opengl_type(shader_type))
                ),
            }
        }
    }
}

/// Builds an individual shader from a SPIR-V binary.
pub fn build_spirv_shader<F: ?Sized>(facade: &F, shader_type: gl::types::GLenum, spirv: &SpirvEntryPoint)
                       -> Result<Shader, ProgramCreationError> where F: Facade
{
    unsafe {
        let ctxt = facade.get_context().make_current();

        if ctxt.capabilities.supported_glsl_versions.is_empty() {
            return Err(ProgramCreationError::CompilationNotSupported);
        }

        if !check_shader_type_compatibility(&ctxt, shader_type) {
            return Err(ProgramCreationError::ShaderTypeNotSupported);
        }

        let id = if ctxt.version >= &Version(Api::Gl, 4, 6) ||
            (ctxt.version >= &Version(Api::Gl, 4, 1) && ctxt.extensions.gl_arb_gl_spirv)
        {
            ctxt.gl.CreateShader(shader_type)
        } else {
            return Err(ProgramCreationError::CompilationNotSupported);
        };

        if id == 0 {
            return Err(ProgramCreationError::ShaderTypeNotSupported);
        }

        let binary = &spirv.binary;
        let entry_point = ffi::CString::new(spirv.entry_point.as_bytes()).unwrap();

        const GL_SHADER_BINARY_FORMAT_SPIR_V: gl::types::GLenum = 0x9551;

        ctxt.gl.ShaderBinary(1, &id, GL_SHADER_BINARY_FORMAT_SPIR_V, binary.as_ptr() as _, binary.len() as _);

        // compiling
        {
            ctxt.report_debug_output_errors.set(false);

            ctxt.gl.SpecializeShader(id, entry_point.as_ptr() as _, 0, ptr::null(), ptr::null());

            ctxt.report_debug_output_errors.set(true);
        }

        // checking compilation success by reading a flag on the shader
        let compilation_success = {
            let mut compilation_success: gl::types::GLint = 0;
            ctxt.gl.GetShaderiv(id, gl::COMPILE_STATUS, &mut compilation_success);
            compilation_success
        };

        if compilation_success == 1 {
            Ok(Shader {
                context: facade.get_context().clone(),
                id: Handle::Id(id)
            })
        } else {
            // compilation error
            let mut error_log_size: gl::types::GLint = 0;

            ctxt.gl.GetShaderiv(id, gl::INFO_LOG_LENGTH, &mut error_log_size);

            let mut error_log: Vec<u8> = Vec::with_capacity(error_log_size as usize);

            ctxt.gl.GetShaderInfoLog(id, error_log_size, &mut error_log_size,
                                     error_log.as_mut_ptr() as *mut gl::types::GLchar);

            error_log.set_len(error_log_size as usize);

            match String::from_utf8(error_log) {
                Ok(msg) => Err(ProgramCreationError::CompilationError(msg, ShaderType::from_opengl_type(shader_type))),
                Err(_) => Err(
                    ProgramCreationError::CompilationError("Could not convert the log \
                                                            message to UTF-8".to_owned(), ShaderType::from_opengl_type(shader_type))
                ),
            }
        }
    }
}

pub fn check_shader_type_compatibility<C: ?Sized>(ctxt: &C, shader_type: gl::types::GLenum)
                                          -> bool where C: CapabilitiesSource
{
    match shader_type {
        gl::VERTEX_SHADER | gl::FRAGMENT_SHADER => (),
        gl::GEOMETRY_SHADER => {
            if !(ctxt.get_version() >= &Version(Api::Gl, 3, 2))
                && !(ctxt.get_version() >= &Version(Api::GlEs, 3, 2))
                && !ctxt.get_extensions().gl_arb_geometry_shader4
                && !ctxt.get_extensions().gl_ext_geometry_shader4
                && !ctxt.get_extensions().gl_ext_geometry_shader
                && !ctxt.get_extensions().gl_oes_geometry_shader
            {
                return false;
            }
        },
        gl::TESS_CONTROL_SHADER | gl::TESS_EVALUATION_SHADER => {
            if !(ctxt.get_version() >= &Version(Api::Gl, 4, 0))
                && !(ctxt.get_version() >= &Version(Api::GlEs, 3, 2))
                && !ctxt.get_extensions().gl_arb_tessellation_shader
                && !ctxt.get_extensions().gl_oes_tessellation_shader
            {
                return false;
            }
        },
        gl::COMPUTE_SHADER => {
            if !(ctxt.get_version() >= &Version(Api::Gl, 4, 3))
                && !(ctxt.get_version() >= &Version(Api::GlEs, 3, 1))
                && !ctxt.get_extensions().gl_arb_compute_shader
            {
                return false;
            }
        },
        _ => unreachable!()
    };

    true
}
<|endoftext|>