<commit_before>// Copyright © 2017 Mozilla Foundation
//
// This program is made available under an ISC-style license.  See the
// accompanying file LICENSE for details.

use ::*;
use context;
use ffi;
use operation;
use std::ffi::CStr;
use std::mem::{self, forget, MaybeUninit};
use std::os::raw::{c_int, c_void};
use std::ptr;
use util::*;

#[derive(Debug)]
pub struct Stream(*mut ffi::pa_stream);

impl Stream {
    pub fn new<'a, CM>(c: &Context, name: &::std::ffi::CStr, ss: &SampleSpec, map: CM) -> Option<Self>
        where CM: Into<Option<&'a ChannelMap>>
    {
        let ptr = unsafe {
            ffi::pa_stream_new(c.raw_mut(),
                               name.as_ptr(),
                               ss as *const _,
                               to_ptr(map.into()))
        };
        if ptr.is_null() {
            None
        } else {
            Some(Stream(ptr))
        }
    }

    #[doc(hidden)]
    pub fn raw_mut(&self) -> &mut ffi::pa_stream {
        unsafe { &mut *self.0 }
    }

    pub fn unref(self) {
        unsafe {
            ffi::pa_stream_unref(self.raw_mut());
        }
    }

    pub fn get_state(&self) -> StreamState {
        StreamState::try_from(unsafe {
            ffi::pa_stream_get_state(self.raw_mut())
        }).expect("pa_stream_get_state returned invalid StreamState")
    }

    pub fn get_context(&self) -> Option<Context> {
        let ptr = unsafe { ffi::pa_stream_get_context(self.raw_mut()) };
        if ptr.is_null() {
            return None;
        }

        let ctx = unsafe { context::from_raw_ptr(ptr) };
        Some(ctx)
    }

    pub fn get_index(&self) -> u32 {
        unsafe { ffi::pa_stream_get_index(self.raw_mut()) }
    }

    pub fn get_device_name<'a>(&'a self) -> Result<&'a CStr> {
        let r = unsafe { ffi::pa_stream_get_device_name(self.raw_mut()) };
        if r.is_null() {
            let err = if let Some(c) = self.get_context() {
                c.errno()
            } else {
                ffi::PA_ERR_UNKNOWN
            };
            return Err(ErrorCode::from_error_code(err));
        }
        Ok(unsafe { CStr::from_ptr(r) })
    }

    pub fn is_suspended(&self) -> Result<bool> {
        let r = unsafe { ffi::pa_stream_is_suspended(self.raw_mut()) };
        error_result!(r != 0, r)
    }

    pub fn is_corked(&self) -> Result<bool> {
        let r = unsafe { ffi::pa_stream_is_corked(self.raw_mut()) };
        error_result!(r != 0, r)
    }

    pub fn connect_playback<'a, D, A, V, S>(&self,
                                            dev: D,
                                            attr: A,
                                            flags: StreamFlags,
                                            volume: V,
                                            sync_stream: S)
                                            -> Result<()>
        where D: Into<Option<&'a CStr>>,
              A: Into<Option<&'a BufferAttr>>,
              V: Into<Option<&'a CVolume>>,
              S: Into<Option<&'a mut Stream>>
    {
        let r = unsafe {
            ffi::pa_stream_connect_playback(self.raw_mut(),
                                            str_to_ptr(dev.into()),
                                            to_ptr(attr.into()),
                                            flags.into(),
                                            to_ptr(volume.into()),
                                            map_to_mut_ptr(sync_stream.into(), |p| p.0))
        };
        error_result!((), r)
    }

    pub fn connect_record<'a, D, A>(&self, dev: D, attr: A, flags: StreamFlags) -> Result<()>
        where D: Into<Option<&'a CStr>>,
              A: Into<Option<&'a BufferAttr>>
    {
        let r = unsafe {
            ffi::pa_stream_connect_record(self.raw_mut(),
                                          str_to_ptr(dev.into()),
                                          to_ptr(attr.into()),
                                          flags.into())
        };
        error_result!((), r)
    }

    pub fn disconnect(&self) -> Result<()> {
        let r = unsafe { ffi::pa_stream_disconnect(self.raw_mut()) };
        error_result!((), r)
    }

    pub fn begin_write(&self, req_bytes: usize) -> Result<(*mut c_void, usize)> {
        let mut data: *mut c_void = ptr::null_mut();
        let mut nbytes = req_bytes;
        let r = unsafe { ffi::pa_stream_begin_write(self.raw_mut(), &mut data, &mut nbytes) };
        error_result!((data, nbytes), r)
    }

    pub fn cancel_write(&self) -> Result<()> {
        let r = unsafe { ffi::pa_stream_cancel_write(self.raw_mut()) };
        error_result!((), r)
    }

    pub fn write(&self, data: *const c_void, nbytes: usize, offset: i64, seek: SeekMode) -> Result<()> {
        let r = unsafe { ffi::pa_stream_write(self.raw_mut(), data, nbytes, None, offset, seek.into()) };
        error_result!((), r)
    }

    pub unsafe fn peek(&self, data: *mut *const c_void, length: *mut usize) -> Result<()> {
        let r = ffi::pa_stream_peek(self.raw_mut(), data, length);
        error_result!((), r)
    }

    pub fn drop(&self) -> Result<()> {
        let r = unsafe { ffi::pa_stream_drop(self.raw_mut()) };
        error_result!((), r)
    }

    pub fn writable_size(&self) -> Result<usize> {
        let r = unsafe { ffi::pa_stream_writable_size(self.raw_mut()) };
        if r == ::std::usize::MAX {
            let err = if let Some(c) = self.get_context() {
                c.errno()
            } else {
                ffi::PA_ERR_UNKNOWN
            };
            return Err(ErrorCode::from_error_code(err));
        }
        Ok(r)
    }

    pub fn readable_size(&self) -> Result<usize> {
        let r = unsafe { ffi::pa_stream_readable_size(self.raw_mut()) };
        if r == ::std::usize::MAX {
            let err = if let Some(c) = self.get_context() {
                c.errno()
            } else {
                ffi::PA_ERR_UNKNOWN
            };
            return Err(ErrorCode::from_error_code(err));
        }
        Ok(r)
    }

    pub fn update_timing_info<CB>(&self, _: CB, userdata: *mut c_void) -> Result<Operation>
        where CB: Fn(&Stream, i32, *mut c_void)
    {
        assert_eq!(mem::size_of::<CB>(), 0);

        // See: A note about `wrapped` functions
        unsafe extern "C" fn wrapped<F>(s: *mut ffi::pa_stream, success: c_int, userdata: *mut c_void)
            where F: Fn(&Stream, i32, *mut c_void)
        {
            let mut stm = stream::from_raw_ptr(s);
            let cb = MaybeUninit::<F>::uninit();
            let result = (*cb.as_ptr())(&mut stm, success, userdata);
            forget(stm);

            result
        }

        let r = unsafe { ffi::pa_stream_update_timing_info(self.raw_mut(), Some(wrapped::<CB>), userdata) };
        if r.is_null() {
            let err = if let Some(c) = self.get_context() {
                c.errno()
            } else {
                ffi::PA_ERR_UNKNOWN
            };
            return Err(ErrorCode::from_error_code(err));
        }
        Ok(unsafe { operation::from_raw_ptr(r) })
    }

    pub fn clear_state_callback(&self) {
        unsafe {
            ffi::pa_stream_set_state_callback(self.raw_mut(), None, ptr::null_mut());
        }
    }

    pub fn set_state_callback<CB>(&self, _: CB, userdata: *mut c_void)
        where CB: Fn(&Stream, *mut c_void)
    {
        assert_eq!(mem::size_of::<CB>(), 0);

        // See: A note about `wrapped` functions
        unsafe extern "C" fn wrapped<F>(s: *mut ffi::pa_stream, userdata: *mut c_void)
            where F: Fn(&Stream, *mut c_void)
        {
            let mut stm = stream::from_raw_ptr(s);
            let cb = MaybeUninit::<F>::uninit();
            let result = (*cb.as_ptr())(&mut stm, userdata);
            forget(stm);

            result
        }

        unsafe {
            ffi::pa_stream_set_state_callback(self.raw_mut(), Some(wrapped::<CB>), userdata);
        }
    }

    pub fn clear_write_callback(&self) {
        unsafe {
            ffi::pa_stream_set_write_callback(self.raw_mut(), None, ptr::null_mut());
        }
    }

    pub fn set_write_callback<CB>(&self, _: CB, userdata: *mut c_void)
        where CB: Fn(&Stream, usize, *mut c_void)
    {
        assert_eq!(mem::size_of::<CB>(), 0);

        // See: A note about `wrapped` functions
        unsafe extern "C" fn wrapped<F>(s: *mut ffi::pa_stream, nbytes: usize, userdata: *mut c_void)
            where F: Fn(&Stream, usize, *mut c_void)
        {
            let mut stm = stream::from_raw_ptr(s);
            let cb = MaybeUninit::<F>::uninit();
            let result = (*cb.as_ptr())(&mut stm, nbytes, userdata);
            forget(stm);

            result
        }

        unsafe {
            ffi::pa_stream_set_write_callback(self.raw_mut(), Some(wrapped::<CB>), userdata);
        }
    }

    pub fn clear_read_callback(&self) {
        unsafe {
            ffi::pa_stream_set_read_callback(self.raw_mut(), None, ptr::null_mut());
        }
    }

    pub fn set_read_callback<CB>(&self, _: CB, userdata: *mut c_void)
        where CB: Fn(&Stream, usize, *mut c_void)
    {
        assert_eq!(mem::size_of::<CB>(), 0);

        // See: A note about `wrapped` functions
        unsafe extern "C" fn wrapped<F>(s: *mut ffi::pa_stream, nbytes: usize, userdata: *mut c_void)
            where F: Fn(&Stream, usize, *mut c_void)
        {
            let mut stm = stream::from_raw_ptr(s);
            let cb = MaybeUninit::<F>::uninit();
            let result = (*cb.as_ptr())(&mut stm, nbytes, userdata);
            forget(stm);

            result
        }

        unsafe {
            ffi::pa_stream_set_read_callback(self.raw_mut(), Some(wrapped::<CB>), userdata);
        }
    }

    pub fn cork<CB>(&self, b: i32, _: CB, userdata: *mut c_void) -> Result<Operation>
        where CB: Fn(&Stream, i32, *mut c_void)
    {
        assert_eq!(mem::size_of::<CB>(), 0);

        // See: A note about `wrapped` functions
        unsafe extern "C" fn wrapped<F>(s: *mut ffi::pa_stream, success: c_int, userdata: *mut c_void)
            where F: Fn(&Stream, i32, *mut c_void)
        {
            let mut stm = stream::from_raw_ptr(s);
            let cb = MaybeUninit::<F>::uninit();
            let result = (*cb.as_ptr())(&mut stm, success, userdata);
            forget(stm);

            result
        }

        let r = unsafe { ffi::pa_stream_cork(self.raw_mut(), b, Some(wrapped::<CB>), userdata) };
        if r.is_null() {
            let err = if let Some(c) = self.get_context() {
                c.errno()
            } else {
                ffi::PA_ERR_UNKNOWN
            };
            return Err(ErrorCode::from_error_code(err));
        }
        Ok(unsafe { operation::from_raw_ptr(r) })
    }

    pub fn get_time(&self) -> Result<(USec)> {
        let mut usec: USec = 0;
        let r = unsafe { ffi::pa_stream_get_time(self.raw_mut(), &mut usec) };
        error_result!(usec, r)
    }

    pub fn get_latency(&self) -> Result<StreamLatency> {
        let mut usec: u64 = 0;
        let mut negative: i32 = 0;
        let r = unsafe { ffi::pa_stream_get_latency(self.raw_mut(), &mut usec, &mut negative) };
        error_result!(
            if negative == 0 {
                StreamLatency::Positive(usec)
            } else {
                StreamLatency::Negative(usec)
            },
            r
        )
    }

    pub fn get_sample_spec(&self) -> &SampleSpec {
        unsafe {
            let ptr = ffi::pa_stream_get_sample_spec(self.raw_mut());
            debug_assert!(!ptr.is_null());
            &*ptr
        }
    }

    pub fn get_channel_map(&self) -> &ChannelMap {
        unsafe {
            let ptr = ffi::pa_stream_get_channel_map(self.raw_mut());
            debug_assert!(!ptr.is_null());
            &*ptr
        }
    }

    pub fn get_buffer_attr(&self) -> &BufferAttr {
        unsafe {
            let ptr = ffi::pa_stream_get_buffer_attr(self.raw_mut());
            debug_assert!(!ptr.is_null());
            &*ptr
        }
    }
}

#[doc(hidden)]
pub unsafe fn from_raw_ptr(ptr: *mut ffi::pa_stream) -> Stream {
    Stream(ptr)
}
<commit_msg>Fix warning<commit_after>// Copyright © 2017 Mozilla Foundation
//
// This program is made available under an ISC-style license.  See the
// accompanying file LICENSE for details.

use ::*;
use context;
use ffi;
use operation;
use std::ffi::CStr;
use std::mem::{self, forget, MaybeUninit};
use std::os::raw::{c_int, c_void};
use std::ptr;
use util::*;

#[derive(Debug)]
pub struct Stream(*mut ffi::pa_stream);

impl Stream {
    pub fn new<'a, CM>(c: &Context, name: &::std::ffi::CStr, ss: &SampleSpec, map: CM) -> Option<Self>
        where CM: Into<Option<&'a ChannelMap>>
    {
        let ptr = unsafe {
            ffi::pa_stream_new(c.raw_mut(),
                               name.as_ptr(),
                               ss as *const _,
                               to_ptr(map.into()))
        };
        if ptr.is_null() {
            None
        } else {
            Some(Stream(ptr))
        }
    }

    #[doc(hidden)]
    pub fn raw_mut(&self) -> &mut ffi::pa_stream {
        unsafe { &mut *self.0 }
    }

    pub fn unref(self) {
        unsafe {
            ffi::pa_stream_unref(self.raw_mut());
        }
    }

    pub fn get_state(&self) -> StreamState {
        StreamState::try_from(unsafe {
            ffi::pa_stream_get_state(self.raw_mut())
        }).expect("pa_stream_get_state returned invalid StreamState")
    }

    pub fn get_context(&self) -> Option<Context> {
        let ptr = unsafe { ffi::pa_stream_get_context(self.raw_mut()) };
        if ptr.is_null() {
            return None;
        }

        let ctx = unsafe { context::from_raw_ptr(ptr) };
        Some(ctx)
    }

    pub fn get_index(&self) -> u32 {
        unsafe { ffi::pa_stream_get_index(self.raw_mut()) }
    }

    pub fn get_device_name<'a>(&'a self) -> Result<&'a CStr> {
        let r = unsafe { ffi::pa_stream_get_device_name(self.raw_mut()) };
        if r.is_null() {
            let err = if let Some(c) = self.get_context() {
                c.errno()
            } else {
                ffi::PA_ERR_UNKNOWN
            };
            return Err(ErrorCode::from_error_code(err));
        }
        Ok(unsafe { CStr::from_ptr(r) })
    }

    pub fn is_suspended(&self) -> Result<bool> {
        let r = unsafe { ffi::pa_stream_is_suspended(self.raw_mut()) };
        error_result!(r != 0, r)
    }

    pub fn is_corked(&self) -> Result<bool> {
        let r = unsafe { ffi::pa_stream_is_corked(self.raw_mut()) };
        error_result!(r != 0, r)
    }

    pub fn connect_playback<'a, D, A, V, S>(&self,
                                            dev: D,
                                            attr: A,
                                            flags: StreamFlags,
                                            volume: V,
                                            sync_stream: S)
                                            -> Result<()>
        where D: Into<Option<&'a CStr>>,
              A: Into<Option<&'a BufferAttr>>,
              V: Into<Option<&'a CVolume>>,
              S: Into<Option<&'a mut Stream>>
    {
        let r = unsafe {
            ffi::pa_stream_connect_playback(self.raw_mut(),
                                            str_to_ptr(dev.into()),
                                            to_ptr(attr.into()),
                                            flags.into(),
                                            to_ptr(volume.into()),
                                            map_to_mut_ptr(sync_stream.into(), |p| p.0))
        };
        error_result!((), r)
    }

    pub fn connect_record<'a, D, A>(&self, dev: D, attr: A, flags: StreamFlags) -> Result<()>
        where D: Into<Option<&'a CStr>>,
              A: Into<Option<&'a BufferAttr>>
    {
        let r = unsafe {
            ffi::pa_stream_connect_record(self.raw_mut(),
                                          str_to_ptr(dev.into()),
                                          to_ptr(attr.into()),
                                          flags.into())
        };
        error_result!((), r)
    }

    pub fn disconnect(&self) -> Result<()> {
        let r = unsafe { ffi::pa_stream_disconnect(self.raw_mut()) };
        error_result!((), r)
    }

    pub fn begin_write(&self, req_bytes: usize) -> Result<(*mut c_void, usize)> {
        let mut data: *mut c_void = ptr::null_mut();
        let mut nbytes = req_bytes;
        let r = unsafe { ffi::pa_stream_begin_write(self.raw_mut(), &mut data, &mut nbytes) };
        error_result!((data, nbytes), r)
    }

    pub fn cancel_write(&self) -> Result<()> {
        let r = unsafe { ffi::pa_stream_cancel_write(self.raw_mut()) };
        error_result!((), r)
    }

    pub fn write(&self, data: *const c_void, nbytes: usize, offset: i64, seek: SeekMode) -> Result<()> {
        let r = unsafe { ffi::pa_stream_write(self.raw_mut(), data, nbytes, None, offset, seek.into()) };
        error_result!((), r)
    }

    pub unsafe fn peek(&self, data: *mut *const c_void, length: *mut usize) -> Result<()> {
        let r = ffi::pa_stream_peek(self.raw_mut(), data, length);
        error_result!((), r)
    }

    pub fn drop(&self) -> Result<()> {
        let r = unsafe { ffi::pa_stream_drop(self.raw_mut()) };
        error_result!((), r)
    }

    pub fn writable_size(&self) -> Result<usize> {
        let r = unsafe { ffi::pa_stream_writable_size(self.raw_mut()) };
        if r == ::std::usize::MAX {
            let err = if let Some(c) = self.get_context() {
                c.errno()
            } else {
                ffi::PA_ERR_UNKNOWN
            };
            return Err(ErrorCode::from_error_code(err));
        }
        Ok(r)
    }

    pub fn readable_size(&self) -> Result<usize> {
        let r = unsafe { ffi::pa_stream_readable_size(self.raw_mut()) };
        if r == ::std::usize::MAX {
            let err = if let Some(c) = self.get_context() {
                c.errno()
            } else {
                ffi::PA_ERR_UNKNOWN
            };
            return Err(ErrorCode::from_error_code(err));
        }
        Ok(r)
    }

    pub fn update_timing_info<CB>(&self, _: CB, userdata: *mut c_void) -> Result<Operation>
        where CB: Fn(&Stream, i32, *mut c_void)
    {
        assert_eq!(mem::size_of::<CB>(), 0);

        // See: A note about `wrapped` functions
        unsafe extern "C" fn wrapped<F>(s: *mut ffi::pa_stream, success: c_int, userdata: *mut c_void)
            where F: Fn(&Stream, i32, *mut c_void)
        {
            let mut stm = stream::from_raw_ptr(s);
            let cb = MaybeUninit::<F>::uninit();
            let result = (*cb.as_ptr())(&mut stm, success, userdata);
            forget(stm);

            result
        }

        let r = unsafe { ffi::pa_stream_update_timing_info(self.raw_mut(), Some(wrapped::<CB>), userdata) };
        if r.is_null() {
            let err = if let Some(c) = self.get_context() {
                c.errno()
            } else {
                ffi::PA_ERR_UNKNOWN
            };
            return Err(ErrorCode::from_error_code(err));
        }
        Ok(unsafe { operation::from_raw_ptr(r) })
    }

    pub fn clear_state_callback(&self) {
        unsafe {
            ffi::pa_stream_set_state_callback(self.raw_mut(), None, ptr::null_mut());
        }
    }

    pub fn set_state_callback<CB>(&self, _: CB, userdata: *mut c_void)
        where CB: Fn(&Stream, *mut c_void)
    {
        assert_eq!(mem::size_of::<CB>(), 0);

        // See: A note about `wrapped` functions
        unsafe extern "C" fn wrapped<F>(s: *mut ffi::pa_stream, userdata: *mut c_void)
            where F: Fn(&Stream, *mut c_void)
        {
            let mut stm = stream::from_raw_ptr(s);
            let cb = MaybeUninit::<F>::uninit();
            let result = (*cb.as_ptr())(&mut stm, userdata);
            forget(stm);

            result
        }

        unsafe {
            ffi::pa_stream_set_state_callback(self.raw_mut(), Some(wrapped::<CB>), userdata);
        }
    }

    pub fn clear_write_callback(&self) {
        unsafe {
            ffi::pa_stream_set_write_callback(self.raw_mut(), None, ptr::null_mut());
        }
    }

    pub fn set_write_callback<CB>(&self, _: CB, userdata: *mut c_void)
        where CB: Fn(&Stream, usize, *mut c_void)
    {
        assert_eq!(mem::size_of::<CB>(), 0);

        // See: A note about `wrapped` functions
        unsafe extern "C" fn wrapped<F>(s: *mut ffi::pa_stream, nbytes: usize, userdata: *mut c_void)
            where F: Fn(&Stream, usize, *mut c_void)
        {
            let mut stm = stream::from_raw_ptr(s);
            let cb = MaybeUninit::<F>::uninit();
            let result = (*cb.as_ptr())(&mut stm, nbytes, userdata);
            forget(stm);

            result
        }

        unsafe {
            ffi::pa_stream_set_write_callback(self.raw_mut(), Some(wrapped::<CB>), userdata);
        }
    }

    pub fn clear_read_callback(&self) {
        unsafe {
            ffi::pa_stream_set_read_callback(self.raw_mut(), None, ptr::null_mut());
        }
    }

    pub fn set_read_callback<CB>(&self, _: CB, userdata: *mut c_void)
        where CB: Fn(&Stream, usize, *mut c_void)
    {
        assert_eq!(mem::size_of::<CB>(), 0);

        // See: A note about `wrapped` functions
        unsafe extern "C" fn wrapped<F>(s: *mut ffi::pa_stream, nbytes: usize, userdata: *mut c_void)
            where F: Fn(&Stream, usize, *mut c_void)
        {
            let mut stm = stream::from_raw_ptr(s);
            let cb = MaybeUninit::<F>::uninit();
            let result = (*cb.as_ptr())(&mut stm, nbytes, userdata);
            forget(stm);

            result
        }

        unsafe {
            ffi::pa_stream_set_read_callback(self.raw_mut(), Some(wrapped::<CB>), userdata);
        }
    }

    pub fn cork<CB>(&self, b: i32, _: CB, userdata: *mut c_void) -> Result<Operation>
        where CB: Fn(&Stream, i32, *mut c_void)
    {
        assert_eq!(mem::size_of::<CB>(), 0);

        // See: A note about `wrapped` functions
        unsafe extern "C" fn wrapped<F>(s: *mut ffi::pa_stream, success: c_int, userdata: *mut c_void)
            where F: Fn(&Stream, i32, *mut c_void)
        {
            let mut stm = stream::from_raw_ptr(s);
            let cb = MaybeUninit::<F>::uninit();
            let result = (*cb.as_ptr())(&mut stm, success, userdata);
            forget(stm);

            result
        }

        let r = unsafe { ffi::pa_stream_cork(self.raw_mut(), b, Some(wrapped::<CB>), userdata) };
        if r.is_null() {
            let err = if let Some(c) = self.get_context() {
                c.errno()
            } else {
                ffi::PA_ERR_UNKNOWN
            };
            return Err(ErrorCode::from_error_code(err));
        }
        Ok(unsafe { operation::from_raw_ptr(r) })
    }

    pub fn get_time(&self) -> Result<USec> {
        let mut usec: USec = 0;
        let r = unsafe { ffi::pa_stream_get_time(self.raw_mut(), &mut usec) };
        error_result!(usec, r)
    }

    pub fn get_latency(&self) -> Result<StreamLatency> {
        let mut usec: u64 = 0;
        let mut negative: i32 = 0;
        let r = unsafe { ffi::pa_stream_get_latency(self.raw_mut(), &mut usec, &mut negative) };
        error_result!(
            if negative == 0 {
                StreamLatency::Positive(usec)
            } else {
                StreamLatency::Negative(usec)
            },
            r
        )
    }

    pub fn get_sample_spec(&self) -> &SampleSpec {
        unsafe {
            let ptr = ffi::pa_stream_get_sample_spec(self.raw_mut());
            debug_assert!(!ptr.is_null());
            &*ptr
        }
    }

    pub fn get_channel_map(&self) -> &ChannelMap {
        unsafe {
            let ptr = ffi::pa_stream_get_channel_map(self.raw_mut());
            debug_assert!(!ptr.is_null());
            &*ptr
        }
    }

    pub fn get_buffer_attr(&self) -> &BufferAttr {
        unsafe {
            let ptr = ffi::pa_stream_get_buffer_attr(self.raw_mut());
            debug_assert!(!ptr.is_null());
            &*ptr
        }
    }
}

#[doc(hidden)]
pub unsafe fn from_raw_ptr(ptr: *mut ffi::pa_stream) -> Stream {
    Stream(ptr)
}
<|endoftext|>


<commit_before>//! Utilities for working with the WTF-8 encoding.
//!
//! WTF-8 is a hack that allows to represent potentially ill-formed UTF-16 byte
//! strings in a way somewhat compatible with UTF-8. Such bugs UTF-16 strings
//! can be found for example in Windows paths for legacy reasons.
//!
//! See the official [WTF-8][wtf8] specification for more information.
//!
//! [wtf8]: https://simonsapin.github.io/wtf-8

/// Decodes the given potentially ill-formed UTF-16 into a WTF-8 byte sequence.
///
/// If the given byte sequence is a well-formed UTF-16 string, then the result
/// is guaranteed to be a valid UTF-8string.
pub fn decode_ill_formed_utf16(units: impl Iterator<Item=u16>) -> Vec<u8> {
    let mut res = Vec::new();

    let mut iter = units.peekable();
    loop {
        let unit = match iter.next() {
            Some(unit) => unit,
            None => return res,
        };

        // Potentially ill-formed UTF-16 to code point conversion based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#decoding-ill-formed-utf-16
        let mut point = unit as u32;
        if is_lead_surrogate(&unit) {
            if let Some(next) = iter.next_if(is_trail_surrogate) {
                let lead = (unit as u32 - 0xD800) << 10;
                let trail = next as u32 - 0xDC00;
                point = 0x10000 + lead + trail;
            }
        }

        // Code point to WTF-8 encoding procedure based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#encoding-wtf-8
        match point {
            0x0000..=0x007F => {
                res.push(point as u8);
            }
            0x0080..=0x07FF => {
                res.push(0xC0 | (point >> 6) as u8);
                res.push(0x80 | (point & 0x3F) as u8);
            }
            0x0800..=0xFFFF => {
                res.push(0xE0 | (point >> 12) as u8);
                res.push(0x80 | ((point >> 6) as u8 & 0x3F));
                res.push(0x80 | (point & 0x3F) as u8);
            }
            0x10000..=0x10FFFF => {
                res.push(0xF0 | (point >> 18) as u8);
                res.push(0x80 | ((point >> 12) as u8 & 0x3F));
                res.push(0x80 | ((point >> 6) as u8 & 0x3F));
                res.push(0x80 | (point & 0x3F) as u8);
            }
            _ => panic!(), // Not possible by construction.
        }
    }
}

/// Encodes the given WTF-8 byte sequence into potentially ill-formed UTF-16.
///
/// If the given byte sequence is a valid UTF-8 string, then the result is
/// guaranteed to be a well-formed UTF-16 string.
pub fn encode_ill_formed_utf16(units: impl Iterator<Item=u8>) -> Vec<u16> {
    let mut res = Vec::new();

    let mut iter = units.peekable();
    loop {
        let byte1 = match iter.next() {
            Some(byte1) => byte1,
            None => return res,
        };

        // TODO: Add support for proper error handling.

        // WTF-8 to a code point decoding procedure based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#decoding-wtf-8
        let mut point = 0;
        match byte1 {
            0x00..=0x7F => {
                point += byte1 as u32;
            }
            0xC2..=0xDF => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x1F) as u32) << 6;
                point += ((byte2 & 0x3F) as u32) << 0;
            }
            0xE0..=0xEF => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                let byte3 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x0F) as u32) << 12;
                point += ((byte2 & 0x3F) as u32) << 6;
                point += ((byte3 & 0x3F) as u32) << 0;
            }
            0xF0..=0xF4 => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                let byte3 = iter.next().unwrap_or_else(|| todo!());
                let byte4 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x07) as u32) << 18;
                point += ((byte2 & 0x3F) as u32) << 12;
                point += ((byte3 & 0x3F) as u32) << 6;
                point += ((byte4 & 0x3F) as u32) << 0;
            }
            _ => todo!(),
        }

        // Code point to potentially ill-formed UTF-16 coversion based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#encoding-ill-formed-utf-16
        if is_supplementary(&point) {
            res.push(((point - 0x10000) >> 10) as u16 + 0xD800);
            res.push(((point - 0x10000) as u16 & 0x3FF) + 0xDC00);
        } else {
            res.push(point as u16);
        }
    }
}

/// Determines whether the given UTF-16 code unit is a lead surrogate.
#[inline]
fn is_lead_surrogate(unit: &u16) -> bool {
    matches!(unit, 0xD800..=0xDBFF)
}

/// Determines whether the given UTF-16 code unit is a trail surrogate.
#[inline]
fn is_trail_surrogate(unit: &u16) -> bool {
    matches!(unit, 0xDC00..=0xDFFF)
}

/// Determines whether the given code point is in the supplementary.
#[inline]
fn is_supplementary(point: &u32) -> bool {
    matches!(point, 0x10000..=0x10FFFF)
}
<commit_msg>Write tests for the UTF-16 to WTF-8 conversion.<commit_after>//! Utilities for working with the WTF-8 encoding.
//!
//! WTF-8 is a hack that allows to represent potentially ill-formed UTF-16 byte
//! strings in a way somewhat compatible with UTF-8. Such bugs UTF-16 strings
//! can be found for example in Windows paths for legacy reasons.
//!
//! See the official [WTF-8][wtf8] specification for more information.
//!
//! [wtf8]: https://simonsapin.github.io/wtf-8

/// Decodes the given potentially ill-formed UTF-16 into a WTF-8 byte sequence.
///
/// If the given byte sequence is a well-formed UTF-16 string, then the result
/// is guaranteed to be a valid UTF-8string.
pub fn decode_ill_formed_utf16(units: impl Iterator<Item=u16>) -> Vec<u8> {
    let mut res = Vec::new();

    let mut iter = units.peekable();
    loop {
        let unit = match iter.next() {
            Some(unit) => unit,
            None => return res,
        };

        // Potentially ill-formed UTF-16 to code point conversion based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#decoding-ill-formed-utf-16
        let mut point = unit as u32;
        if is_lead_surrogate(&unit) {
            if let Some(next) = iter.next_if(is_trail_surrogate) {
                let lead = (unit as u32 - 0xD800) << 10;
                let trail = next as u32 - 0xDC00;
                point = 0x10000 + lead + trail;
            }
        }

        // Code point to WTF-8 encoding procedure based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#encoding-wtf-8
        match point {
            0x0000..=0x007F => {
                res.push(point as u8);
            }
            0x0080..=0x07FF => {
                res.push(0xC0 | (point >> 6) as u8);
                res.push(0x80 | (point & 0x3F) as u8);
            }
            0x0800..=0xFFFF => {
                res.push(0xE0 | (point >> 12) as u8);
                res.push(0x80 | ((point >> 6) as u8 & 0x3F));
                res.push(0x80 | (point & 0x3F) as u8);
            }
            0x10000..=0x10FFFF => {
                res.push(0xF0 | (point >> 18) as u8);
                res.push(0x80 | ((point >> 12) as u8 & 0x3F));
                res.push(0x80 | ((point >> 6) as u8 & 0x3F));
                res.push(0x80 | (point & 0x3F) as u8);
            }
            _ => panic!(), // Not possible by construction.
        }
    }
}

/// Encodes the given WTF-8 byte sequence into potentially ill-formed UTF-16.
///
/// If the given byte sequence is a valid UTF-8 string, then the result is
/// guaranteed to be a well-formed UTF-16 string.
pub fn encode_ill_formed_utf16(units: impl Iterator<Item=u8>) -> Vec<u16> {
    let mut res = Vec::new();

    let mut iter = units.peekable();
    loop {
        let byte1 = match iter.next() {
            Some(byte1) => byte1,
            None => return res,
        };

        // TODO: Add support for proper error handling.

        // WTF-8 to a code point decoding procedure based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#decoding-wtf-8
        let mut point = 0;
        match byte1 {
            0x00..=0x7F => {
                point += byte1 as u32;
            }
            0xC2..=0xDF => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x1F) as u32) << 6;
                point += ((byte2 & 0x3F) as u32) << 0;
            }
            0xE0..=0xEF => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                let byte3 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x0F) as u32) << 12;
                point += ((byte2 & 0x3F) as u32) << 6;
                point += ((byte3 & 0x3F) as u32) << 0;
            }
            0xF0..=0xF4 => {
                let byte2 = iter.next().unwrap_or_else(|| todo!());
                let byte3 = iter.next().unwrap_or_else(|| todo!());
                let byte4 = iter.next().unwrap_or_else(|| todo!());
                point += ((byte1 & 0x07) as u32) << 18;
                point += ((byte2 & 0x3F) as u32) << 12;
                point += ((byte3 & 0x3F) as u32) << 6;
                point += ((byte4 & 0x3F) as u32) << 0;
            }
            _ => todo!(),
        }

        // Code point to potentially ill-formed UTF-16 coversion based on [1].
        //
        // [1]: https://simonsapin.github.io/wtf-8/#encoding-ill-formed-utf-16
        if is_supplementary(&point) {
            res.push(((point - 0x10000) >> 10) as u16 + 0xD800);
            res.push(((point - 0x10000) as u16 & 0x3FF) + 0xDC00);
        } else {
            res.push(point as u16);
        }
    }
}

/// Determines whether the given UTF-16 code unit is a lead surrogate.
#[inline]
fn is_lead_surrogate(unit: &u16) -> bool {
    matches!(unit, 0xD800..=0xDBFF)
}

/// Determines whether the given UTF-16 code unit is a trail surrogate.
#[inline]
fn is_trail_surrogate(unit: &u16) -> bool {
    matches!(unit, 0xDC00..=0xDFFF)
}

/// Determines whether the given code point is in the supplementary.
#[inline]
fn is_supplementary(point: &u32) -> bool {
    matches!(point, 0x10000..=0x10FFFF)
}

#[cfg(test)]
mod tests {

    use crate::wtf8;

    #[test]
    fn decode_empty() {
        assert_eq!(wtf8::decode_ill_formed_utf16(std::iter::empty()), b"");
    }

    #[test]
    fn decode_small_bytes() {
        let bytes_utf16 = vec![0x04, 0x08, 0x16, 0x23, 0x42];
        let bytes_wtf8 = wtf8::decode_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\x04\x08\x16\x23\x42");
    }

    #[test]
    fn decode_bmp_bytes() {
        let bytes_utf16 = vec![0xABCD, 0xBEEF, 0xFFFF];
        let bytes_wtf8 = wtf8::decode_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\xEA\xAF\x8D\xEB\xBB\xAF\xEF\xBF\xBF");
    }

    #[test]
    fn decode_sup_bytes() {
        let bytes_utf16 = vec![0xD800, 0xDC00, 0xDBFF, 0xDFFF];
        let bytes_wtf8 = wtf8::decode_ill_formed_utf16(bytes_utf16.into_iter());
        assert_eq!(bytes_wtf8, b"\xF0\x90\x80\x80\xF4\x8F\xBF\xBF");
    }

    #[test]
    fn decode_ascii_string() {
        let string = "foo bar baz";
        let string_wtf8 = wtf8::decode_ill_formed_utf16(string.encode_utf16());
        let string_utf8 = string.as_bytes();
        assert_eq!(string_wtf8, string_utf8);
    }

    #[test]
    fn decode_unicode_string() {
        let string = "zażółć gęślą jaźń";
        let string_wtf8 = wtf8::decode_ill_formed_utf16(string.encode_utf16());
        let string_utf8 = string.as_bytes();
        assert_eq!(string_wtf8, string_utf8);
    }
}
<|endoftext|>


<commit_before>use std::num::{ Float, from_f64 };
use std::mem::size_of;

use music;
use reader;
use wave;

pub fn quantize<T>(input: f64) -> T where T: FromPrimitive {
    let quantization_levels = 2.0.powf(size_of::<T>() as f64 * 8.0) - 1.0;
    // Convert from [-1, 1] to take up full quantization range
    from_f64::<T>(input * (quantization_levels / 2.0)).expect("failed to quantize to given type")
}

pub fn quantize_samples<T>(input: Vec<f64>) -> Vec<T> where T: FromPrimitive {
    input.iter().map(|s| { quantize::<T>(*s) }).collect()
}

pub fn generate<F>(x: f64, f: F) -> f64 where F: Fn<(f64, ), f64> {
    f(x)
}

pub fn make_samples<F>(length: f64, sample_rate: uint, waveform: F) -> Vec<f64> where F: Fn<(f64, ), f64>+Copy {
    let num_samples = (sample_rate as f64 * length).floor() as uint;
    let mut samples: Vec<f64> = Vec::with_capacity(num_samples);

    for i in range(0u, num_samples) {
        let t = i as f64 / sample_rate as f64;
        samples.push(generate(t, waveform));
    }

    samples
}

// This is really awful, is there a more elegant way to do this?
// TODO: Split out volume normalization into a function
pub fn make_samples_from_midi(sample_rate: uint, bpm: f64, filename: &str) -> Vec<f64> {
    let song = reader::read_midi(filename).unwrap();
    let length = (60.0 * song.max_time as f64) / (bpm * song.time_unit as f64);

    let mut notes_on_for_ticks: Vec<Vec<u8>> = Vec::new();
    for _ in range(0, song.max_time) {
        let notes_on_for_tick: Vec<u8> = Vec::new();
        notes_on_for_ticks.push(notes_on_for_tick);
    }

    for track in song.tracks.iter() {
        for i in range(0, track.messages.len()) {
            let event = track.messages[i];
            if event.message_type == 9 {
                let from_tick = event.time;
                let note = event.value1;

                let mut to_tick = song.max_time;
                for j in range(i, track.messages.len()) {
                    let event_to = track.messages[j];
                    if event_to.message_type == 8 && event_to.value1 == note {
                        to_tick = event_to.time;
                        break;
                    }
                }

                for tick in range(from_tick, to_tick) {
                    notes_on_for_ticks[tick].push(note);
                }
            }
        }
    }

    let midi_frequency_function = |t: f64| -> f64 {
        let tick = (t * bpm * song.time_unit as f64 / 60.0) as uint;
        let mut out = 0.0;

        if tick < notes_on_for_ticks.len() {
            for note in notes_on_for_ticks[tick].iter() {
                let frequency = music::note_midi(440.0, *note as uint);
                out += wave::SineWave(frequency)(t)
            }
        }

        out
    };

    let num_samples = (sample_rate as f64 * length).floor() as uint;
    let mut samples: Vec<f64> = Vec::with_capacity(num_samples);

    for i in range(0u, num_samples) {
        let t = i as f64 / sample_rate as f64;
        samples.push(midi_frequency_function(t));
    }

    // Kill clipping
    let peak = samples.iter().fold(0.0f64, |acc, sample| {
        if acc > *sample { acc } else { *sample }
    });

    samples = samples.iter().map(|el| {
        *el / peak
    }).collect();

    samples
}
<commit_msg>Move volume peak normalisation into function<commit_after>use std::num::{ Float, FloatMath, from_f64 };
use std::mem::size_of;

use music;
use reader;
use wave;

pub fn quantize<T>(input: f64) -> T where T: FromPrimitive {
    let quantization_levels = 2.0.powf(size_of::<T>() as f64 * 8.0) - 1.0;
    // Convert from [-1, 1] to take up full quantization range
    from_f64::<T>(input * (quantization_levels / 2.0)).expect("failed to quantize to given type")
}

pub fn quantize_samples<T>(input: Vec<f64>) -> Vec<T> where T: FromPrimitive {
    input.iter().map(|s| { quantize::<T>(*s) }).collect()
}

pub fn generate<F>(x: f64, f: F) -> f64 where F: Fn<(f64, ), f64> {
    f(x)
}

pub fn make_samples<F>(length: f64, sample_rate: uint, waveform: F) -> Vec<f64> where F: Fn<(f64, ), f64>+Copy {
    let num_samples = (sample_rate as f64 * length).floor() as uint;
    let mut samples: Vec<f64> = Vec::with_capacity(num_samples);

    for i in range(0u, num_samples) {
        let t = i as f64 / sample_rate as f64;
        samples.push(generate(t, waveform));
    }

    samples
}

pub fn normalize(samples: Vec<f64>) -> Vec<f64> {
    let peak = samples.iter().fold(0.0f64, |acc, sample| {
        acc.max(*sample)
    });

    samples.iter().map(|sample| {
        *sample / peak
    }).collect()
}

// This is really awful, is there a more elegant way to do this?
// TODO: Split out volume normalization into a function
pub fn make_samples_from_midi(sample_rate: uint, bpm: f64, filename: &str) -> Vec<f64> {
    let song = reader::read_midi(filename).unwrap();
    let length = (60.0 * song.max_time as f64) / (bpm * song.time_unit as f64);

    let mut notes_on_for_ticks: Vec<Vec<u8>> = Vec::new();
    for _ in range(0, song.max_time) {
        let notes_on_for_tick: Vec<u8> = Vec::new();
        notes_on_for_ticks.push(notes_on_for_tick);
    }

    for track in song.tracks.iter() {
        for i in range(0, track.messages.len()) {
            let event = track.messages[i];
            if event.message_type == 9 {
                let from_tick = event.time;
                let note = event.value1;

                let mut to_tick = song.max_time;
                for j in range(i, track.messages.len()) {
                    let event_to = track.messages[j];
                    if event_to.message_type == 8 && event_to.value1 == note {
                        to_tick = event_to.time;
                        break;
                    }
                }

                for tick in range(from_tick, to_tick) {
                    notes_on_for_ticks[tick].push(note);
                }
            }
        }
    }

    let midi_frequency_function = |t: f64| -> f64 {
        let tick = (t * bpm * song.time_unit as f64 / 60.0) as uint;
        let mut out = 0.0;

        if tick < notes_on_for_ticks.len() {
            for note in notes_on_for_ticks[tick].iter() {
                let frequency = music::note_midi(440.0, *note as uint);
                out += wave::SineWave(frequency)(t)
            }
        }

        out
    };

    let num_samples = (sample_rate as f64 * length).floor() as uint;
    let mut samples: Vec<f64> = Vec::with_capacity(num_samples);

    for i in range(0u, num_samples) {
        let t = i as f64 / sample_rate as f64;
        samples.push(midi_frequency_function(t));
    }

    normalize(samples)
}
<|endoftext|>


<commit_before>// TODO: data xrefs

use super::*;
use zydis;
use byteorder::{ByteOrder, LittleEndian};
use std::mem;

fn analyze_operand_xrefs(
    // used to determine valid addresses
    layout: &ModuleLayout,
    // used to detect RIP-relative addressing
    pc: zydis::enums::register::Register,
    // used to compute relative addresses
    rva: Rva,
    // used to compute relative addresses
    insn: &zydis::ffi::DecodedInstruction,
    // the operand to inspect
    op: &zydis::ffi::DecodedOperand,
) -> Result<Option<Rva>, Error> {
    match op.ty {
        zydis::enums::OperandType::Unused => Err(Error::NotImplemented("xref from unused register")),
        zydis::enums::OperandType::Register => {
            // like: CALL rbx
            // TODO: for now, don't index unresolved indirect branches
            Ok(None)
        }
        zydis::enums::OperandType::Memory => {
            // like: .text:0000000180001041 FF 15 D1 78 07 00      call    cs:__imp_RtlVirtualUnwind_0
            //           0x0000000000001041:                       call    [0x0000000000079980]
            if pc == op.mem.base && op.mem.disp.has_displacement && op.mem.scale == 0 {
                // RIP-relative
                // this is the default encoding on x64.
                // tools like IDA automatically compute and display the target.
                // CALL [RIP + 0x401000]
                if let zydis::enums::register::Register::NONE = op.mem.index {
                    let target =
                        (rva as i64
                        // TODO: cast from rva (u64) to i64 is lossy.
                        + op.mem.disp.displacement
                        + i64::from(insn.length)) as Rva;

                    if layout.is_rva_valid(target) {
                        Ok(Some(target))
                    } else {
                        // TODO: record this anomaly somewhere.
                        warn!("problem: invalid xref target: memory not in sections");
                        Ok(None)
                    }
                } else {
                    // unsupported
                    // like: CALL [RIP + 4*RCX + 0x401000] ??
                    println!("CALL [RIP + 4*RCX + 0x401000] ??");
                    Err(Error::NotImplemented("xref from RIP-relative, non-zero index memory"))
                }
            } else if op.mem.base == zydis::enums::register::Register::NONE {
                // like: CALL [0x401000] ??
                println!("TODO: other OperandType::Memory branch");
                Err(Error::NotImplemented("xref from non-RIP-relative memory"))
            } else {
                // like: CALL [rbx]
                // like: CALL [rbx + 0x10]
                // TODO: for now, don't index unresolved indirect branches
                Ok(None)
            }
        }
        zydis::enums::OperandType::Pointer => {
            println!("TODO: operand: pointer");
            Err(Error::NotImplemented("xref from pointer"))
        }
        zydis::enums::OperandType::Immediate => {
            if !op.imm.is_relative {
                println!("TODO: absolute immediate operand");
                Err(Error::NotImplemented("xref from absolute immediate"))
            } else {
                let imm = if op.imm.is_signed {
                    u64_i64(op.imm.value)
                } else {
                    op.imm.value as i64
                };

                // TODO: cast from rva (u64) to i64 is lossy.
                let dst = (rva as i64 + imm + i64::from(insn.length)) as Rva;
                if layout.is_rva_valid(dst) {
                    Ok(Some(dst))
                } else {
                    // TODO: record this anomaly somewhere.
                    debug!("problem: invalid xref target: relative immediate not in sections");
                    Ok(None)
                }
            }
        }
    }
}

pub fn analyze_insn_xrefs(
    layout: &ModuleLayout,
    pc: zydis::enums::register::Register,
    rva: Rva,
    insn: &zydis::ffi::DecodedInstruction,
) -> Result<Vec<Xref>, Error> {
    match insn.mnemonic {
        // see InstructionCategory
        // syscall, sysexit, sysret
        // vmcall, vmmcall
        zydis::enums::mnemonic::Mnemonic::CALL => {
            // a CALL always has an operand, so assume this is ok.
            let op = Workspace::get_first_operand(insn).unwrap();

            let fallthrough = Xref {
                src: rva,
                dst: rva + u64::from(insn.length),
                typ: XrefType::Fallthrough,
            };

            match analyze_operand_xrefs(layout, pc, rva, insn, op)? {
                // TODO: fallthrough is not guaranteed if the function is noret
                Some(dst) => Ok(vec![Xref { src: rva, dst: dst, typ: XrefType::Call }, fallthrough]),
                None => Ok(vec![fallthrough]),
            }
        }
        zydis::enums::mnemonic::Mnemonic::RET
        | zydis::enums::mnemonic::Mnemonic::IRET
        | zydis::enums::mnemonic::Mnemonic::IRETD
        | zydis::enums::mnemonic::Mnemonic::IRETQ => Ok(vec![]),
        zydis::enums::mnemonic::Mnemonic::JMP => {
            // a JMP always has an operand, so assume this is ok.
            let op = Workspace::get_first_operand(insn).unwrap();

            match analyze_operand_xrefs(layout, pc, rva, insn, op)? {
                Some(dst) => Ok(vec![Xref { src: rva, dst: dst, typ: XrefType::UnconditionalJump }]),
                None => Ok(vec![]),
            }
        }
        zydis::enums::mnemonic::Mnemonic::JB
        | zydis::enums::mnemonic::Mnemonic::JBE
        | zydis::enums::mnemonic::Mnemonic::JCXZ
        | zydis::enums::mnemonic::Mnemonic::JECXZ
        | zydis::enums::mnemonic::Mnemonic::JKNZD
        | zydis::enums::mnemonic::Mnemonic::JKZD
        | zydis::enums::mnemonic::Mnemonic::JL
        | zydis::enums::mnemonic::Mnemonic::JLE
        | zydis::enums::mnemonic::Mnemonic::JNB
        | zydis::enums::mnemonic::Mnemonic::JNBE
        | zydis::enums::mnemonic::Mnemonic::JNL
        | zydis::enums::mnemonic::Mnemonic::JNLE
        | zydis::enums::mnemonic::Mnemonic::JNO
        | zydis::enums::mnemonic::Mnemonic::JNP
        | zydis::enums::mnemonic::Mnemonic::JNS
        | zydis::enums::mnemonic::Mnemonic::JNZ
        | zydis::enums::mnemonic::Mnemonic::JO
        | zydis::enums::mnemonic::Mnemonic::JP
        | zydis::enums::mnemonic::Mnemonic::JRCXZ
        | zydis::enums::mnemonic::Mnemonic::JS
        | zydis::enums::mnemonic::Mnemonic::JZ => {
            // a J* always has an operand, so assume this is ok.
            let op = Workspace::get_first_operand(insn).unwrap();

            let fallthrough = Xref {
                src: rva,
                dst: rva + u64::from(insn.length),
                typ: XrefType::Fallthrough,
            };

            match analyze_operand_xrefs(layout, pc, rva, insn, op)? {
                Some(dst) => Ok(vec![
                    Xref { src: rva, dst: dst, typ: XrefType::ConditionalJump, },
                    fallthrough,
                ]),
                None => Ok(vec![fallthrough]),
            }
        }
        zydis::enums::mnemonic::Mnemonic::CMOVB
        | zydis::enums::mnemonic::Mnemonic::CMOVBE
        | zydis::enums::mnemonic::Mnemonic::CMOVL
        | zydis::enums::mnemonic::Mnemonic::CMOVLE
        | zydis::enums::mnemonic::Mnemonic::CMOVNB
        | zydis::enums::mnemonic::Mnemonic::CMOVNBE
        | zydis::enums::mnemonic::Mnemonic::CMOVNL
        | zydis::enums::mnemonic::Mnemonic::CMOVNLE
        | zydis::enums::mnemonic::Mnemonic::CMOVNO
        | zydis::enums::mnemonic::Mnemonic::CMOVNP
        | zydis::enums::mnemonic::Mnemonic::CMOVNS
        | zydis::enums::mnemonic::Mnemonic::CMOVNZ
        | zydis::enums::mnemonic::Mnemonic::CMOVO
        | zydis::enums::mnemonic::Mnemonic::CMOVP
        | zydis::enums::mnemonic::Mnemonic::CMOVS
        | zydis::enums::mnemonic::Mnemonic::CMOVZ => Ok(vec![
            Xref {
                src: rva,
                dst: rva + u64::from(insn.length),
                typ: XrefType::Fallthrough
            },
            Xref {
                src: rva,
                dst: rva + u64::from(insn.length),
                typ: XrefType::ConditionalMove
            },
        ]),
        _ => Ok(vec![Xref {
            src: rva,
            dst: rva + u64::from(insn.length),
            typ: XrefType::Fallthrough,
        }]),
    }
}

pub fn find_insns(ws: &Workspace, predicate: fn(&Instruction) -> bool) -> Result<Vec<Rva>, Error> {
    let mut ret: Vec<Rva> = vec![];

        // empirically, doing this in parallel (par_extend/par_bridge) is not worth it.
        // the predicates are likely too simple,
        // so the majority of time is spent synchronizing.
        ret.extend(
            ws
            .iter_insns()
            .filter(|insn| predicate(insn))
            .map(|insn| match insn {
                Instruction::Valid {loc, ..} => loc.addr,
                Instruction::Invalid {loc, ..} => loc.addr,
            })
        );

    Ok(ret)
}

pub fn find_roots(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    find_insns(ws, |insn| match insn {
        Instruction::Valid{loc, ..} => loc.xrefs.to.is_empty(),
        _ => false,
    })
}

pub fn find_call_targets(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    find_insns(ws, |insn| match insn {
        Instruction::Valid{loc, ..} => loc.xrefs.to.iter().any(|xref| match xref.typ {
            XrefType::Call => true,
            _ => false,
        }),
        _ => false,
    })
}

pub fn find_branch_targets(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    find_insns(ws, |insn| match insn {
        Instruction::Valid{loc, ..} => loc.xrefs.to.iter().any(|xref| match xref.typ {
            XrefType::UnconditionalJump => true,
            XrefType::ConditionalJump => true,
            _ => false,
        }),
        _ => false,
    })
}

fn get_section_by_name<'a>(ws: &'a Workspace, name: &[u8]) -> Result<Option<&'a Section>, Error> {
    match ws.get_obj()? {
        Object::PE(pe) => {
            match pe.sections.iter()
                .enumerate()
                .map(|(index, section)| (index, section.name))
                .find(|(_, sec_name)| sec_name == name) {
                    Some((index, _)) => Ok(Some(&ws.sections[index])),
                    None => Ok(None),
            }
        },
        _ => Err(Error::NotImplemented("analyzer for non-PE module"))
    }
}

// pub fn find_runtime_functions
//  on pe64, entries in .pdata are struct RUNTIME_FUNCTION
//
//  00000000 FunctionStart   dd ?                    ; offset rva
//  00000004 FunctionEnd     dd ?                    ; offset rva pastend
//  00000008 UnwindInfo      dd ?                    ; offset rva
//
// ref: https://stackoverflow.com/a/9794688/87207
//
// for k32, the first three entries should be:
//  - 1010
//  - 1060
//  - 1140
pub fn find_runtime_functions(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    match ws.get_obj()? {
        Object::PE(_) => {
            let mut ret: Vec<Rva> = vec![];

            if let Some(sec) = get_section_by_name(ws, b".pdata\x00\x00")? {
                let layout = ws.get_layout()?;
                for runtime_function in sec.buf.chunks_exact(mem::size_of::<u32>() * 3) {
                    let rva = LittleEndian::read_u32(&runtime_function) as Rva;
                    if rva == 0 {
                        // must be end of (valid) .pdata section
                        break;
                    } else if layout.is_rva_valid(rva) {
                        ret.push(rva);
                    } else {
                        warn!("malformed .pdata section: entry points outside of image");
                        return Ok(vec![]);
                    }
                }
            }

            Ok(ret)
        },
        _ => Err(Error::NotImplemented("analyzer for non-PE module"))
    }
}

// pub fn find_fixups

// pub fn find_ptrs
//  relies on the image being loaded at the base address
//  also relies on loc vs insn to be figured out.

/// ```
/// use lancelot::*;
/// use lancelot::rsrc::*;
/// use lancelot::analysis::*;
/// let ws = get_workspace(Rsrc::TINY);
/// assert_eq!(0, find_entrypoints(&ws).unwrap().len());
///
/// //let ws = get_workspace(Rsrc::K32);
/// //assert_eq!(1630, find_entrypoints(&ws).unwrap().len());
/// ```
pub fn find_entrypoints(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    match ws.get_obj()? {
        Object::PE(pe) => {
            let mut ret: Vec<Rva> = vec![];
            if let Some(opt) = pe.header.optional_header {
                ret.push(opt.standard_fields.address_of_entry_point);
                ret.extend(
                    pe.exports.iter().map(|export| export.rva as Rva)
                );
            } else {
                info!("no optional header")
            }

            Ok(ret)
        },
        _ => Err(Error::NotImplemented("analyzer for non-PE module"))
    }
}

// pub fn find_sigs

// pub fn find_impossible_paths
// like:
//   MOV
//   MOV
//   INVALID
<commit_msg>analysis: xrefs: handle VA displacements by inspecting layout<commit_after>// TODO: data xrefs

use super::*;
use zydis;
use byteorder::{ByteOrder, LittleEndian};
use std::mem;

fn analyze_operand_xrefs(
    // used to determine valid addresses
    layout: &ModuleLayout,
    // used to detect RIP-relative addressing
    pc: zydis::enums::register::Register,
    // used to compute relative addresses
    rva: Rva,
    // used to compute relative addresses
    insn: &zydis::ffi::DecodedInstruction,
    // the operand to inspect
    op: &zydis::ffi::DecodedOperand,
) -> Result<Option<Rva>, Error> {
    match op.ty {
        zydis::enums::OperandType::Unused => Err(Error::NotImplemented("xref from unused register")),
        zydis::enums::OperandType::Register => {
            // like: CALL rbx
            // TODO: for now, don't index unresolved indirect branches
            Ok(None)
        }
        zydis::enums::OperandType::Memory => {
            // like: .text:0000000180001041 FF 15 D1 78 07 00      call    cs:__imp_RtlVirtualUnwind_0
            //           0x0000000000001041:                       call    [0x0000000000079980]
            if pc == op.mem.base && op.mem.disp.has_displacement && op.mem.scale == 0 {
                // RIP-relative
                // this is the default encoding on x64.
                // tools like IDA automatically compute and display the target.
                // CALL [RIP + 0x401000]
                if let zydis::enums::register::Register::NONE = op.mem.index {
                    let target =
                        (rva as i64
                        // TODO: cast from rva (u64) to i64 is lossy.
                        + op.mem.disp.displacement
                        + i64::from(insn.length)) as Rva;

                    if layout.is_rva_valid(target) {
                        Ok(Some(target))
                    } else {
                        // TODO: record this anomaly somewhere.
                        warn!("problem: invalid xref target: memory not in sections");
                        Ok(None)
                    }
                } else {
                    // unsupported
                    // like: CALL [RIP + 4*RCX + 0x401000] ??
                    println!("CALL [RIP + 4*RCX + 0x401000] ??");
                    Err(Error::NotImplemented("xref from RIP-relative, non-zero index memory"))
                }
            } else if op.mem.base == zydis::enums::register::Register::NONE &&
                      op.mem.disp.has_displacement &&
                      op.mem.scale == 0 {
                // like: .text:00401078 FF 15 40 80 40 00    call   DWORD PTR ds:0x408040
                // "mem": {
                //    "ty": "Register",
                //    "segment": "DS",
                //    "base": "NONE",
                //    "index": "NONE",
                //    "scale": 0,
                //    "disp": {
                //        "has_displacement": true,
                //        "displacement": 4227136
                //    }
                // },

                let target = op.mem.disp.displacement as Va;
                match layout.va2rva(target) {
                    Ok(target) => {
                        debug!("found RVA 0x{:x} from VA 0x{:x} using base address 0x{:x}",
                            target, op.mem.disp.displacement, layout.base_address);
                        Ok(Some(target))
                    }
                    Err(_) => {
                        warn!("problem: VA 0x{:x} not mapped using base address 0x{:x}", target, layout.base_address);
                        Ok(None)
                    }
                }
            } else {
                // like: CALL [rbx]
                // like: CALL [rbx + 0x10]
                // TODO: for now, don't index unresolved indirect branches
                Ok(None)
            }
        }
        zydis::enums::OperandType::Pointer => {
            println!("TODO: operand: pointer");
            Err(Error::NotImplemented("xref from pointer"))
        }
        zydis::enums::OperandType::Immediate => {
            if !op.imm.is_relative {
                println!("TODO: absolute immediate operand");
                Err(Error::NotImplemented("xref from absolute immediate"))
            } else {
                let imm = if op.imm.is_signed {
                    u64_i64(op.imm.value)
                } else {
                    op.imm.value as i64
                };

                // TODO: cast from rva (u64) to i64 is lossy.
                let dst = (rva as i64 + imm + i64::from(insn.length)) as Rva;
                if layout.is_rva_valid(dst) {
                    Ok(Some(dst))
                } else {
                    // TODO: record this anomaly somewhere.
                    debug!("problem: invalid xref target: relative immediate not in sections");
                    Ok(None)
                }
            }
        }
    }
}

pub fn analyze_insn_xrefs(
    layout: &ModuleLayout,
    pc: zydis::enums::register::Register,
    rva: Rva,
    insn: &zydis::ffi::DecodedInstruction,
) -> Result<Vec<Xref>, Error> {
    match insn.mnemonic {
        // see InstructionCategory
        // syscall, sysexit, sysret
        // vmcall, vmmcall
        zydis::enums::mnemonic::Mnemonic::CALL => {
            // a CALL always has an operand, so assume this is ok.
            let op = Workspace::get_first_operand(insn).unwrap();

            let fallthrough = Xref {
                src: rva,
                dst: rva + u64::from(insn.length),
                typ: XrefType::Fallthrough,
            };

            match analyze_operand_xrefs(layout, pc, rva, insn, op)? {
                // TODO: fallthrough is not guaranteed if the function is noret
                Some(dst) => Ok(vec![Xref { src: rva, dst: dst, typ: XrefType::Call }, fallthrough]),
                None => Ok(vec![fallthrough]),
            }
        }
        zydis::enums::mnemonic::Mnemonic::RET
        | zydis::enums::mnemonic::Mnemonic::IRET
        | zydis::enums::mnemonic::Mnemonic::IRETD
        | zydis::enums::mnemonic::Mnemonic::IRETQ => Ok(vec![]),
        zydis::enums::mnemonic::Mnemonic::JMP => {
            // a JMP always has an operand, so assume this is ok.
            let op = Workspace::get_first_operand(insn).unwrap();

            match analyze_operand_xrefs(layout, pc, rva, insn, op)? {
                Some(dst) => Ok(vec![Xref { src: rva, dst: dst, typ: XrefType::UnconditionalJump }]),
                None => Ok(vec![]),
            }
        }
        zydis::enums::mnemonic::Mnemonic::JB
        | zydis::enums::mnemonic::Mnemonic::JBE
        | zydis::enums::mnemonic::Mnemonic::JCXZ
        | zydis::enums::mnemonic::Mnemonic::JECXZ
        | zydis::enums::mnemonic::Mnemonic::JKNZD
        | zydis::enums::mnemonic::Mnemonic::JKZD
        | zydis::enums::mnemonic::Mnemonic::JL
        | zydis::enums::mnemonic::Mnemonic::JLE
        | zydis::enums::mnemonic::Mnemonic::JNB
        | zydis::enums::mnemonic::Mnemonic::JNBE
        | zydis::enums::mnemonic::Mnemonic::JNL
        | zydis::enums::mnemonic::Mnemonic::JNLE
        | zydis::enums::mnemonic::Mnemonic::JNO
        | zydis::enums::mnemonic::Mnemonic::JNP
        | zydis::enums::mnemonic::Mnemonic::JNS
        | zydis::enums::mnemonic::Mnemonic::JNZ
        | zydis::enums::mnemonic::Mnemonic::JO
        | zydis::enums::mnemonic::Mnemonic::JP
        | zydis::enums::mnemonic::Mnemonic::JRCXZ
        | zydis::enums::mnemonic::Mnemonic::JS
        | zydis::enums::mnemonic::Mnemonic::JZ => {
            // a J* always has an operand, so assume this is ok.
            let op = Workspace::get_first_operand(insn).unwrap();

            let fallthrough = Xref {
                src: rva,
                dst: rva + u64::from(insn.length),
                typ: XrefType::Fallthrough,
            };

            match analyze_operand_xrefs(layout, pc, rva, insn, op)? {
                Some(dst) => Ok(vec![
                    Xref { src: rva, dst: dst, typ: XrefType::ConditionalJump, },
                    fallthrough,
                ]),
                None => Ok(vec![fallthrough]),
            }
        }
        zydis::enums::mnemonic::Mnemonic::CMOVB
        | zydis::enums::mnemonic::Mnemonic::CMOVBE
        | zydis::enums::mnemonic::Mnemonic::CMOVL
        | zydis::enums::mnemonic::Mnemonic::CMOVLE
        | zydis::enums::mnemonic::Mnemonic::CMOVNB
        | zydis::enums::mnemonic::Mnemonic::CMOVNBE
        | zydis::enums::mnemonic::Mnemonic::CMOVNL
        | zydis::enums::mnemonic::Mnemonic::CMOVNLE
        | zydis::enums::mnemonic::Mnemonic::CMOVNO
        | zydis::enums::mnemonic::Mnemonic::CMOVNP
        | zydis::enums::mnemonic::Mnemonic::CMOVNS
        | zydis::enums::mnemonic::Mnemonic::CMOVNZ
        | zydis::enums::mnemonic::Mnemonic::CMOVO
        | zydis::enums::mnemonic::Mnemonic::CMOVP
        | zydis::enums::mnemonic::Mnemonic::CMOVS
        | zydis::enums::mnemonic::Mnemonic::CMOVZ => Ok(vec![
            Xref {
                src: rva,
                dst: rva + u64::from(insn.length),
                typ: XrefType::Fallthrough
            },
            Xref {
                src: rva,
                dst: rva + u64::from(insn.length),
                typ: XrefType::ConditionalMove
            },
        ]),
        _ => Ok(vec![Xref {
            src: rva,
            dst: rva + u64::from(insn.length),
            typ: XrefType::Fallthrough,
        }]),
    }
}

pub fn find_insns(ws: &Workspace, predicate: fn(&Instruction) -> bool) -> Result<Vec<Rva>, Error> {
    let mut ret: Vec<Rva> = vec![];

        // empirically, doing this in parallel (par_extend/par_bridge) is not worth it.
        // the predicates are likely too simple,
        // so the majority of time is spent synchronizing.
        ret.extend(
            ws
            .iter_insns()
            .filter(|insn| predicate(insn))
            .map(|insn| match insn {
                Instruction::Valid {loc, ..} => loc.addr,
                Instruction::Invalid {loc, ..} => loc.addr,
            })
        );

    Ok(ret)
}

pub fn find_roots(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    find_insns(ws, |insn| match insn {
        Instruction::Valid{loc, ..} => loc.xrefs.to.is_empty(),
        _ => false,
    })
}

pub fn find_call_targets(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    find_insns(ws, |insn| match insn {
        Instruction::Valid{loc, ..} => loc.xrefs.to.iter().any(|xref| match xref.typ {
            XrefType::Call => true,
            _ => false,
        }),
        _ => false,
    })
}

pub fn find_branch_targets(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    find_insns(ws, |insn| match insn {
        Instruction::Valid{loc, ..} => loc.xrefs.to.iter().any(|xref| match xref.typ {
            XrefType::UnconditionalJump => true,
            XrefType::ConditionalJump => true,
            _ => false,
        }),
        _ => false,
    })
}

fn get_section_by_name<'a>(ws: &'a Workspace, name: &[u8]) -> Result<Option<&'a Section>, Error> {
    match ws.get_obj()? {
        Object::PE(pe) => {
            match pe.sections.iter()
                .enumerate()
                .map(|(index, section)| (index, section.name))
                .find(|(_, sec_name)| sec_name == name) {
                    Some((index, _)) => Ok(Some(&ws.sections[index])),
                    None => Ok(None),
            }
        },
        _ => Err(Error::NotImplemented("analyzer for non-PE module"))
    }
}

// pub fn find_runtime_functions
//  on pe64, entries in .pdata are struct RUNTIME_FUNCTION
//
//  00000000 FunctionStart   dd ?                    ; offset rva
//  00000004 FunctionEnd     dd ?                    ; offset rva pastend
//  00000008 UnwindInfo      dd ?                    ; offset rva
//
// ref: https://stackoverflow.com/a/9794688/87207
//
// for k32, the first three entries should be:
//  - 1010
//  - 1060
//  - 1140
pub fn find_runtime_functions(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    match ws.get_obj()? {
        Object::PE(_) => {
            let mut ret: Vec<Rva> = vec![];

            if let Some(sec) = get_section_by_name(ws, b".pdata\x00\x00")? {
                let layout = ws.get_layout()?;
                for runtime_function in sec.buf.chunks_exact(mem::size_of::<u32>() * 3) {
                    let rva = LittleEndian::read_u32(&runtime_function) as Rva;
                    if rva == 0 {
                        // must be end of (valid) .pdata section
                        break;
                    } else if layout.is_rva_valid(rva) {
                        ret.push(rva);
                    } else {
                        warn!("malformed .pdata section: entry points outside of image");
                        return Ok(vec![]);
                    }
                }
            }

            Ok(ret)
        },
        _ => Err(Error::NotImplemented("analyzer for non-PE module"))
    }
}

// pub fn find_fixups

// pub fn find_ptrs
//  relies on the image being loaded at the base address
//  also relies on loc vs insn to be figured out.

/// ```
/// use lancelot::*;
/// use lancelot::rsrc::*;
/// use lancelot::analysis::*;
/// let ws = get_workspace(Rsrc::TINY);
/// assert_eq!(0, find_entrypoints(&ws).unwrap().len());
///
/// //let ws = get_workspace(Rsrc::K32);
/// //assert_eq!(1630, find_entrypoints(&ws).unwrap().len());
/// ```
pub fn find_entrypoints(ws: &Workspace) -> Result<Vec<Rva>, Error> {
    match ws.get_obj()? {
        Object::PE(pe) => {
            let mut ret: Vec<Rva> = vec![];
            if let Some(opt) = pe.header.optional_header {
                ret.push(opt.standard_fields.address_of_entry_point);
                ret.extend(
                    pe.exports.iter().map(|export| export.rva as Rva)
                );
            } else {
                info!("no optional header")
            }

            Ok(ret)
        },
        _ => Err(Error::NotImplemented("analyzer for non-PE module"))
    }
}

// pub fn find_sigs

// pub fn find_impossible_paths
// like:
//   MOV
//   MOV
//   INVALID
<|endoftext|>


<commit_before>use Data;
use test::examplelib::EXAMPLELIB;
use Func;
use Lib;
use Symbol;

#[test]
fn load_examplelib() {
    unsafe {
        Lib::new(EXAMPLELIB).unwrap();
    }
}

#[test]
fn add_2_numbers() {
    unsafe {
        let lib = Lib::new(EXAMPLELIB).unwrap();
        let add_u32s: Func<extern "C" fn(u32, u32) -> u32> = lib.find_func("add_u32s").unwrap();
        let result = add_u32s.get()(1, 1);
        assert_eq!(result, 2);
    }
}

#[test]
fn check_test_value() {
    unsafe {
        let lib = Lib::new(EXAMPLELIB).unwrap();
        let test_value: Data<u32> = lib.find_data("TEST_VALUE").unwrap();
        assert_eq!(*test_value.get(), 100);
    }
}
<commit_msg>Added some tests which show how to use structs.<commit_after>use Data;
use Func;
use FuncArc;
use FuncUnsafe;
use Lib;
use LibArc;
use LibUnsafe;
use std::mem;
use std::cell::RefCell;
use std::ops::DerefMut;
use Symbol;
use test::examplelib::EXAMPLELIB;

#[test]
fn load_examplelib() {
    unsafe {
        Lib::new(EXAMPLELIB).unwrap();
    }
}

#[test]
fn add_2_numbers() {
    unsafe {
        let lib = Lib::new(EXAMPLELIB).unwrap();
        let add_u32s: Func<extern "C" fn(u32, u32) -> u32> = lib.find_func("add_u32s").unwrap();
        let result = add_u32s.get()(1, 1);
        assert_eq!(result, 2);
    }
}

#[test]
fn check_test_value() {
    unsafe {
        let lib = Lib::new(EXAMPLELIB).unwrap();
        let test_value: Data<u32> = lib.find_data("TEST_VALUE").unwrap();
        assert_eq!(*test_value.get(), 100);
    }
}

struct SafeDll<'a> {
    _lib: Lib,
    _func: RefCell<Option<Func<'a, extern "C" fn(u32, u32) ->  u32>>>,
}

#[test]
fn create_struct_safe() {
    unsafe {
        let lib = Lib::new(EXAMPLELIB).unwrap();
        let dll = SafeDll { _lib: lib, _func: RefCell::new(None), };
        let add_u32s: Func<extern "C" fn(u32, u32) -> u32> = dll._lib.find_func("add_u32s").unwrap();
        mem::replace(dll._func.borrow_mut().deref_mut(), Some(add_u32s));
    }
}

struct UnsafeDll {
    _lib: LibUnsafe,
    _func: FuncUnsafe<extern "C" fn(u32, u32) -> u32>,
}

#[test]
fn create_struct_unsafe() {
    unsafe {
        let lib = LibUnsafe::new(EXAMPLELIB).unwrap();
        let add_u32s: FuncUnsafe<extern "C" fn(u32, u32) -> u32> = lib.find_func("add_u32s").unwrap();
        let _dll = UnsafeDll { _lib: lib, _func: add_u32s, };
    }
}

struct ArcDll {
    _lib: LibArc,
    _func: FuncArc<extern "C" fn(u32, u32) -> u32>,
}

#[test]
fn create_struct_arc() {
    unsafe {
        let lib = LibArc::new(EXAMPLELIB).unwrap();
        let add_u32s: FuncArc<extern "C" fn(u32, u32) -> u32> = lib.find_func("add_u32s").unwrap();
        let _dll = ArcDll { _lib: lib, _func: add_u32s, };
    }
}
<|endoftext|>


<commit_before>//! Folds constants in SIR programs.
//!
//! This transform walks each SIR function and attempts to evaluate expressions at runtime,
//! deleting definitions and variables which evaluate to simple constants.
//!
//! The transform in conservative in two ways:
//!
//! 1. It will not remove symbols required by terminators, in case a code generator relies on the
//! existence of these symbols.
//!
//! 2. It will not remove symbols which are used as parameters to a function (though this can
//!    probably be changed to support a simple inteprocedural constant simplification).

use std::collections::BTreeMap;

use error::*;
use sir::*;

extern crate fnv;

/// Evaluates a binary operation over two literals, returning a new literal.
// TODO(shoumik): Maybe we should just implement this as a method or as traits (e.g., Add, Sub,
// etc.) on LiteralKind? Not every literal implements every binop, however...
fn evaluate_binop(kind: BinOpKind, left: LiteralKind, right: LiteralKind) -> WeldResult<LiteralKind> {
    use ast::LiteralKind::*;
    use ast::BinOpKind::*;
    let result = match kind {
        // Just support the basics for now.
        Add => {
            match (left, right) {
                (I8Literal(l), I8Literal(r)) => I8Literal(l + r),
                (I16Literal(l), I16Literal(r)) => I16Literal(l + r),
                (I32Literal(l), I32Literal(r)) => I32Literal(l + r),
                (I64Literal(l), I64Literal(r)) => I64Literal(l + r),
                (F32Literal(l), F32Literal(r)) => F32Literal((f32::from_bits(l) + f32::from_bits(r)).to_bits()),
                (F64Literal(l), F64Literal(r)) => F64Literal((f64::from_bits(l) + f64::from_bits(r)).to_bits()),
                _ => {
                    return weld_err!("Mismatched types in evaluate_binop");
                }
            }
        }
        Subtract => {
            match (left, right) {
                (I8Literal(l), I8Literal(r)) => I8Literal(l - r),
                (I16Literal(l), I16Literal(r)) => I16Literal(l - r),
                (I32Literal(l), I32Literal(r)) => I32Literal(l - r),
                (I64Literal(l), I64Literal(r)) => I64Literal(l - r),
                (F32Literal(l), F32Literal(r)) => F32Literal((f32::from_bits(l) - f32::from_bits(r)).to_bits()),
                (F64Literal(l), F64Literal(r)) => F64Literal((f64::from_bits(l) - f64::from_bits(r)).to_bits()),
                _ => {
                    return weld_err!("Mismatched types in evaluate_binop");
                }
            }
        }
        Multiply => {
            match (left, right) {
                (I8Literal(l), I8Literal(r)) => I8Literal(l * r),
                (I16Literal(l), I16Literal(r)) => I16Literal(l * r),
                (I32Literal(l), I32Literal(r)) => I32Literal(l * r),
                (I64Literal(l), I64Literal(r)) => I64Literal(l * r),
                (F32Literal(l), F32Literal(r)) => F32Literal((f32::from_bits(l) * f32::from_bits(r)).to_bits()),
                (F64Literal(l), F64Literal(r)) => F64Literal((f64::from_bits(l) * f64::from_bits(r)).to_bits()),
                _ => {
                    return weld_err!("Mismatched types in evaluate_binop");
                }
            }
        }
        Divide => {
            match (left, right) {
                (I8Literal(l), I8Literal(r)) => I8Literal(l / r),
                (I16Literal(l), I16Literal(r)) => I16Literal(l / r),
                (I32Literal(l), I32Literal(r)) => I32Literal(l / r),
                (I64Literal(l), I64Literal(r)) => I64Literal(l / r),
                (F32Literal(l), F32Literal(r)) => F32Literal((f32::from_bits(l) / f32::from_bits(r)).to_bits()),
                (F64Literal(l), F64Literal(r)) => F64Literal((f64::from_bits(l) / f64::from_bits(r)).to_bits()),
                _ => {
                    return weld_err!("Mismatched types in evaluate_binop");
                }
            }
        }
        _ => {
            return weld_err!("Unsupported binary operation in evaluate_binop");
        }
    };
    Ok(result)
}

pub fn fold_constants(prog: &mut SirProgram) -> WeldResult<()> {

    let ref mut parameters = fnv::FnvHashSet::default();

    // Collect all the Symbols passed between functions. We will keep the
    // definitions of these symbols intact (even if we assign simple literal
    // values to them).
    //
    // TODO(shoumik): We could also run an interprocedural version of this transform
    // that eliminates passing these, instead propagating the computed value through
    // the program.
    for func in prog.funcs.iter() {
        parameters.extend(func.params.iter().map(&|(k, _)| k).cloned());
        // Add these because code gen expects them to be present even if they are unused.
        for block in func.blocks.iter() {
            parameters.extend(block.terminator.children().cloned());
        }
    }

    for func in prog.funcs.iter_mut() {
        fold_constants_in_function(func, parameters)?;
    }
    Ok(())
}

fn fold_constants_in_function(func: &mut SirFunction, global_params: &fnv::FnvHashSet<Symbol>) -> WeldResult<()> {
    use sir::StatementKind::*;

    // Maps a Symbol to a known literal value.
    let mut values: fnv::FnvHashMap<Symbol, LiteralKind> = fnv::FnvHashMap::default();
    // Set of used symbols in a function.
    let mut used_symbols = fnv::FnvHashSet::default();

    for block in func.blocks.iter_mut() {
        for statement in block.statements.iter_mut() {
            let replaced = match statement.kind {
                // Literal value - add it to the map
                AssignLiteral(ref lit) => {
                    values.insert(statement.output.clone().unwrap(), *lit);
                    true
                },
                // Aliases
                Assign(ref sym) => {
                    // contains_key instead of get to avoid double mutable borrow
                    if values.contains_key(sym) {
                        let value = *values.get(sym).unwrap();
                        values.insert(statement.output.clone().unwrap(), value);
                        true
                    } else {
                        false
                    }
                }
                BinOp { ref op, ref left, ref right } if values.contains_key(left) && values.contains_key(right) => {
                    let left_val = *values.get(left).unwrap();
                    let right_val = *values.get(right).unwrap();
                    // If this throws an error, it just means that we don't support evaluating the
                    // expression right now.
                    if let Ok(result) = evaluate_binop(*op, left_val, right_val) {
                        values.insert(statement.output.clone().unwrap(), result);
                        true
                    } else {
                        false
                    }
                }
                // Unsupported for now
                _ => false
            };

            if replaced {
                // If this is true, we computed a constant value for this statement. Replace it
                // with that constant value.
                let output_sym = statement.output.clone().unwrap();
                let kind = AssignLiteral(*values.get(&output_sym).unwrap());
                let new_statement = Statement::new(Some(output_sym), kind);
                *statement = new_statement;
            } else {
                // We didn't replace this statement with an AssignLiteral, so record the symbols
                // which it uses.
                used_symbols.extend(statement.kind.children().cloned());
            }
        }
    }

    // Delete unused symbols. We'll only delete the locals, since deleting the parameters will
    // require modifying global state.
    for block in func.blocks.iter_mut() {
        block.statements.retain(|ref s| {
            // Keep the statement that defines the symbol if it is used on the RHS of some
            // expression, and remove it otherwise.
            if let Some(ref sym) = s.output {
                used_symbols.contains(sym) || global_params.contains(sym)
            } else {
                // If there's no output symbol (i.e., it's a `Merge`), retain it.
                true
            }
        });
    }

    // TODO(shoumik): is there a better way of doing this...
    let mut locals = BTreeMap::new();
    for (k, v) in func.locals.iter() {
        if used_symbols.contains(k) || global_params.contains(k) {
            locals.insert(k.clone(), v.clone());
        }
    }
    func.locals = locals;
    Ok(())
}
<commit_msg>fold_constants bugfix for symbols assigned multiple times (#298)<commit_after>//! Folds constants in SIR programs.
//!
//! This transform walks each SIR function and attempts to evaluate expressions at runtime,
//! deleting definitions and variables which evaluate to simple constants.
//!
//! The transform in conservative in two ways:
//!
//! 1. It will not remove symbols required by terminators, in case a code generator relies on the
//! existence of these symbols.
//!
//! 2. It will not remove symbols which are used as parameters to a function (though this can
//!    probably be changed to support a simple inteprocedural constant simplification).

use std::collections::BTreeMap;

use error::*;
use sir::*;

extern crate fnv;

/// Evaluates a binary operation over two literals, returning a new literal.
// TODO(shoumik): Maybe we should just implement this as a method or as traits (e.g., Add, Sub,
// etc.) on LiteralKind? Not every literal implements every binop, however...
fn evaluate_binop(kind: BinOpKind, left: LiteralKind, right: LiteralKind) -> WeldResult<LiteralKind> {
    use ast::LiteralKind::*;
    use ast::BinOpKind::*;
    let result = match kind {
        // Just support the basics for now.
        Add => {
            match (left, right) {
                (I8Literal(l), I8Literal(r)) => I8Literal(l + r),
                (I16Literal(l), I16Literal(r)) => I16Literal(l + r),
                (I32Literal(l), I32Literal(r)) => I32Literal(l + r),
                (I64Literal(l), I64Literal(r)) => I64Literal(l + r),
                (F32Literal(l), F32Literal(r)) => F32Literal((f32::from_bits(l) + f32::from_bits(r)).to_bits()),
                (F64Literal(l), F64Literal(r)) => F64Literal((f64::from_bits(l) + f64::from_bits(r)).to_bits()),
                _ => {
                    return weld_err!("Mismatched types in evaluate_binop");
                }
            }
        }
        Subtract => {
            match (left, right) {
                (I8Literal(l), I8Literal(r)) => I8Literal(l - r),
                (I16Literal(l), I16Literal(r)) => I16Literal(l - r),
                (I32Literal(l), I32Literal(r)) => I32Literal(l - r),
                (I64Literal(l), I64Literal(r)) => I64Literal(l - r),
                (F32Literal(l), F32Literal(r)) => F32Literal((f32::from_bits(l) - f32::from_bits(r)).to_bits()),
                (F64Literal(l), F64Literal(r)) => F64Literal((f64::from_bits(l) - f64::from_bits(r)).to_bits()),
                _ => {
                    return weld_err!("Mismatched types in evaluate_binop");
                }
            }
        }
        Multiply => {
            match (left, right) {
                (I8Literal(l), I8Literal(r)) => I8Literal(l * r),
                (I16Literal(l), I16Literal(r)) => I16Literal(l * r),
                (I32Literal(l), I32Literal(r)) => I32Literal(l * r),
                (I64Literal(l), I64Literal(r)) => I64Literal(l * r),
                (F32Literal(l), F32Literal(r)) => F32Literal((f32::from_bits(l) * f32::from_bits(r)).to_bits()),
                (F64Literal(l), F64Literal(r)) => F64Literal((f64::from_bits(l) * f64::from_bits(r)).to_bits()),
                _ => {
                    return weld_err!("Mismatched types in evaluate_binop");
                }
            }
        }
        Divide => {
            match (left, right) {
                (I8Literal(l), I8Literal(r)) => I8Literal(l / r),
                (I16Literal(l), I16Literal(r)) => I16Literal(l / r),
                (I32Literal(l), I32Literal(r)) => I32Literal(l / r),
                (I64Literal(l), I64Literal(r)) => I64Literal(l / r),
                (F32Literal(l), F32Literal(r)) => F32Literal((f32::from_bits(l) / f32::from_bits(r)).to_bits()),
                (F64Literal(l), F64Literal(r)) => F64Literal((f64::from_bits(l) / f64::from_bits(r)).to_bits()),
                _ => {
                    return weld_err!("Mismatched types in evaluate_binop");
                }
            }
        }
        _ => {
            return weld_err!("Unsupported binary operation in evaluate_binop");
        }
    };
    Ok(result)
}

pub fn fold_constants(prog: &mut SirProgram) -> WeldResult<()> {

    let ref mut parameters = fnv::FnvHashSet::default();

    // Collect all the Symbols passed between functions. We will keep the
    // definitions of these symbols intact (even if we assign simple literal
    // values to them).
    //
    // TODO(shoumik): We could also run an interprocedural version of this transform
    // that eliminates passing these, instead propagating the computed value through
    // the program.
    for func in prog.funcs.iter() {
        parameters.extend(func.params.iter().map(&|(k, _)| k).cloned());
        // Add these because code gen expects them to be present even if they are unused.
        for block in func.blocks.iter() {
            parameters.extend(block.terminator.children().cloned());
        }
    }

    for func in prog.funcs.iter_mut() {
        fold_constants_in_function(func, parameters)?;
    }
    Ok(())
}

fn fold_constants_in_function(func: &mut SirFunction, global_params: &fnv::FnvHashSet<Symbol>) -> WeldResult<()> {
    use sir::StatementKind::*;

    let mut assignment_counts: fnv::FnvHashMap<Symbol, i32> = fnv::FnvHashMap::default();
    for block in func.blocks.iter_mut() {
        for statement in block.statements.iter_mut() {
            if statement.output.is_some() {
                let assignment_count = assignment_counts.entry(statement.output.clone().unwrap()).or_insert(0);
                *assignment_count += 1;
            }
        }
    }
    // Maps a Symbol to a known literal value.
    let mut values: fnv::FnvHashMap<Symbol, LiteralKind> = fnv::FnvHashMap::default();
    // Set of used symbols in a function.
    let mut used_symbols = fnv::FnvHashSet::default();

    for block in func.blocks.iter_mut() {
        for statement in block.statements.iter_mut() {
            let replacement_lit = match statement.kind {
                // Literal value - add it to the map
                AssignLiteral(ref lit) => {
                    let output_sym = statement.output.clone().unwrap();
                    if *assignment_counts.get(&output_sym).unwrap() == 1 {
                        values.insert(output_sym, *lit);
                        Some(*lit)
                    } else {
                        None
                    }
                },
                // Aliases
                Assign(ref sym) => {
                    // contains_key instead of get to avoid double mutable borrow
                    if values.contains_key(sym) {
                        let output_sym = statement.output.clone().unwrap();
                        let value = *values.get(sym).unwrap();
                        if *assignment_counts.get(&output_sym).unwrap() == 1 {
                            values.insert(output_sym, value);
                        }
                        Some(value)
                    } else {
                        None
                    }
                }
                BinOp { ref op, ref left, ref right } if values.contains_key(left) && values.contains_key(right) => {
                    let left_val = *values.get(left).unwrap();
                    let right_val = *values.get(right).unwrap();
                    // If this throws an error, it just means that we don't support evaluating the
                    // expression right now.
                    if let Ok(result) = evaluate_binop(*op, left_val, right_val) {
                        let output_sym = statement.output.clone().unwrap();
                        if *assignment_counts.get(&output_sym).unwrap() == 1 {
                            values.insert(output_sym, result);
                        }
                        Some(result)
                    } else {
                        None
                    }
                }
                // Unsupported for now
                _ => None
            };

            if replacement_lit.is_some() {
                // If this is true, we computed a constant value for this statement. Replace it
                // with that constant value.
                let kind = AssignLiteral(replacement_lit.unwrap());
                let new_statement = Statement::new(statement.output.clone(), kind);
                *statement = new_statement;
            } else {
                // We didn't replace this statement with an AssignLiteral, so record the symbols
                // which it uses.
                used_symbols.extend(statement.kind.children().cloned());
            }
        }
    }

    // Delete unused symbols. We'll only delete the locals, since deleting the parameters will
    // require modifying global state.
    for block in func.blocks.iter_mut() {
        block.statements.retain(|ref s| {
            // Keep the statement that defines the symbol if it is used on the RHS of some
            // expression, and remove it otherwise.
            if let Some(ref sym) = s.output {
                used_symbols.contains(sym) || global_params.contains(sym)
            } else {
                // If there's no output symbol (i.e., it's a `Merge`), retain it.
                true
            }
        });
    }

    // TODO(shoumik): is there a better way of doing this...
    let mut locals = BTreeMap::new();
    for (k, v) in func.locals.iter() {
        if used_symbols.contains(k) || global_params.contains(k) {
            locals.insert(k.clone(), v.clone());
        }
    }
    func.locals = locals;
    Ok(())
}
<|endoftext|>