<commit_before>extern crate libc;
extern crate std;

use super::liblua;
use super::Index;
use super::Lua;
use super::Pushable;
use super::Readable;

macro_rules! integer_impl(
    ($t:ident) => (
        impl Pushable for $t {
            fn push_to_lua(&self, lua: &mut Lua) {
                unsafe { liblua::lua_pushinteger(lua.lua, *self as liblua::lua_Integer) }
            }
        }
        impl Readable for $t {
            fn read_from_lua(lua: &mut Lua, index: i32) -> Option<$t> {
                let success: libc::c_int = unsafe { std::mem::uninitialized() };
                let val = unsafe { liblua::lua_tointegerx(lua.lua, index, &success) };
                match success {
                    0 => None,
                    _ => Some(val as $t)
                }
            }
        }
        impl Index for $t {
        }
    );
)

integer_impl!(int)
integer_impl!(i8)
integer_impl!(i16)
integer_impl!(i32)
//integer_impl!(i64)   // data loss

macro_rules! unsigned_impl(
    ($t:ident) => (
        impl Pushable for $t {
            fn push_to_lua(&self, lua: &mut Lua) {
                unsafe { liblua::lua_pushunsigned(lua.lua, *self as liblua::lua_Unsigned) }
            }
        }
        impl Readable for $t {
            fn read_from_lua(lua: &mut Lua, index: i32) -> Option<$t> {
                let success: libc::c_int = unsafe { std::mem::uninitialized() };
                let val = unsafe { liblua::lua_tounsignedx(lua.lua, index, &success) };
                match success {
                    0 => None,
                    _ => Some(val as $t)
                }
            }
        }
        impl Index for $t {
        }
    );
)

unsigned_impl!(uint)
unsigned_impl!(u8)
unsigned_impl!(u16)
unsigned_impl!(u32)
//unsigned_impl!(u64)   // data loss

macro_rules! numeric_impl(
    ($t:ident) => (
        impl Pushable for $t {
            fn push_to_lua(&self, lua: &mut Lua) {
                unsafe { liblua::lua_pushnumber(lua.lua, *self as f64) }
            }
        }
        impl Readable for $t {
            fn read_from_lua(lua: &mut Lua, index: i32) -> Option<$t> {
                let success: libc::c_int = unsafe { std::mem::uninitialized() };
                let val = unsafe { liblua::lua_tonumberx(lua.lua, index, &success) };
                match success {
                    0 => None,
                    _ => Some(val as $t)
                }
            }
        }
        impl Index for $t {
        }
    );
)

numeric_impl!(f32)
numeric_impl!(f64)

impl Pushable for std::string::String {
    fn push_to_lua(&self, lua: &mut Lua) {
        unsafe { liblua::lua_pushstring(lua.lua, self.to_c_str().unwrap()) }
    }
}

impl Readable for String {
    fn read_from_lua(lua: &mut Lua, index: i32) -> Option<std::string::String> {
        let mut size: libc::size_t = unsafe { std::mem::uninitialized() };
        let cStrRaw = unsafe { liblua::lua_tolstring(lua.lua, index, &mut size) };
        if cStrRaw.is_null() {
            return None;
        }

        unsafe { std::c_str::CString::new(cStrRaw, false) }.as_str().map(|s| s.to_string())
    }
}

impl Index for String {
}

impl<'a> Pushable for &'a str {
    fn push_to_lua(&self, lua: &mut Lua) {
        unsafe { liblua::lua_pushstring(lua.lua, self.to_c_str().unwrap()) }
    }
}

impl Pushable for bool {
    fn push_to_lua(&self, lua: &mut Lua) {
        unsafe { liblua::lua_pushboolean(lua.lua, self.clone() as libc::c_int) }
    }
}

impl Readable for bool {
    fn read_from_lua(lua: &mut Lua, index: i32) -> Option<bool> {
        if unsafe { liblua::lua_isboolean(lua.lua, index) } != true {
            return None;
        }

        Some(unsafe { liblua::lua_toboolean(lua.lua, index) != 0 })
    }
}

impl Index for bool {
}

#[cfg(test)]
mod tests {
    #[test]
    fn read_ints() {
        let mut lua = super::super::Lua::new();

        lua.set("a", 2).unwrap();

        let x: int = lua.get("a").unwrap();
        assert_eq!(x, 2);

        let y: i8 = lua.get("a").unwrap();
        assert_eq!(y, 2);

        let z: i16 = lua.get("a").unwrap();
        assert_eq!(z, 2);

        let w: i32 = lua.get("a").unwrap();
        assert_eq!(w, 2);

        let a: uint = lua.get("a").unwrap();
        assert_eq!(a, 2);

        let b: u8 = lua.get("a").unwrap();
        assert_eq!(b, 2);

        let c: u16 = lua.get("a").unwrap();
        assert_eq!(c, 2);

        let d: u32 = lua.get("a").unwrap();
        assert_eq!(d, 2);
    }

    #[test]
    fn write_ints() {
        // TODO:

        let mut lua = super::super::Lua::new();

        lua.set("a", 2).unwrap();
        let x: int = lua.get("a").unwrap();
        assert_eq!(x, 2);
    }

    #[test]
    fn readwrite_floats() {
        let mut lua = super::super::Lua::new();

        lua.set("a", 2.51234 as f32).unwrap();
        lua.set("b", 3.4123456789 as f64).unwrap();

        let x: f32 = lua.get("a").unwrap();
        assert!(x - 2.51234 < 0.000001);

        let y: f64 = lua.get("a").unwrap();
        assert!(y - 2.51234 < 0.000001);

        let z: f32 = lua.get("b").unwrap();
        assert!(z - 3.4123456789 < 0.000001);

        let w: f64 = lua.get("b").unwrap();
        assert!(w - 3.4123456789 < 0.000001);
    }

    #[test]
    fn readwrite_bools() {
        let mut lua = super::super::Lua::new();

        lua.set("a", true).unwrap();
        lua.set("b", false).unwrap();

        let x: bool = lua.get("a").unwrap();
        assert_eq!(x, true);

        let y: bool = lua.get("b").unwrap();
        assert_eq!(y, false);
    }

    #[test]
    fn readwrite_strings() {
        let mut lua = super::super::Lua::new();

        lua.set("a", "hello").unwrap();
        lua.set("b", "hello".to_string()).unwrap();

        let x: String = lua.get("a").unwrap();
        assert_eq!(x.as_slice(), "hello");

        let y: String = lua.get("b").unwrap();
        assert_eq!(y.as_slice(), "hello");
    }

    #[test]
    fn int_to_string() {
        let mut lua = super::super::Lua::new();

        lua.set("a", 2).unwrap();

        let x: String = lua.get("a").unwrap();
        assert_eq!(x.as_slice(), "2");
    }

    #[test]
    fn string_to_int() {
        let mut lua = super::super::Lua::new();

        lua.set("a", "2").unwrap();
        lua.set("b", "aaa").unwrap();

        let x: int = lua.get("a").unwrap();
        assert_eq!(x, 2);

        let y: Option<int> = lua.get("b");
        assert!(y.is_none());
    }
}
<commit_msg>Add pushable and readable implementations for nil<commit_after>extern crate libc;
extern crate std;

use super::liblua;
use super::Index;
use super::Lua;
use super::Pushable;
use super::Readable;

macro_rules! integer_impl(
    ($t:ident) => (
        impl Pushable for $t {
            fn push_to_lua(&self, lua: &mut Lua) {
                unsafe { liblua::lua_pushinteger(lua.lua, *self as liblua::lua_Integer) }
            }
        }
        impl Readable for $t {
            fn read_from_lua(lua: &mut Lua, index: i32) -> Option<$t> {
                let success: libc::c_int = unsafe { std::mem::uninitialized() };
                let val = unsafe { liblua::lua_tointegerx(lua.lua, index, &success) };
                match success {
                    0 => None,
                    _ => Some(val as $t)
                }
            }
        }
        impl Index for $t {
        }
    );
)

integer_impl!(int)
integer_impl!(i8)
integer_impl!(i16)
integer_impl!(i32)
//integer_impl!(i64)   // data loss

macro_rules! unsigned_impl(
    ($t:ident) => (
        impl Pushable for $t {
            fn push_to_lua(&self, lua: &mut Lua) {
                unsafe { liblua::lua_pushunsigned(lua.lua, *self as liblua::lua_Unsigned) }
            }
        }
        impl Readable for $t {
            fn read_from_lua(lua: &mut Lua, index: i32) -> Option<$t> {
                let success: libc::c_int = unsafe { std::mem::uninitialized() };
                let val = unsafe { liblua::lua_tounsignedx(lua.lua, index, &success) };
                match success {
                    0 => None,
                    _ => Some(val as $t)
                }
            }
        }
        impl Index for $t {
        }
    );
)

unsigned_impl!(uint)
unsigned_impl!(u8)
unsigned_impl!(u16)
unsigned_impl!(u32)
//unsigned_impl!(u64)   // data loss

macro_rules! numeric_impl(
    ($t:ident) => (
        impl Pushable for $t {
            fn push_to_lua(&self, lua: &mut Lua) {
                unsafe { liblua::lua_pushnumber(lua.lua, *self as f64) }
            }
        }
        impl Readable for $t {
            fn read_from_lua(lua: &mut Lua, index: i32) -> Option<$t> {
                let success: libc::c_int = unsafe { std::mem::uninitialized() };
                let val = unsafe { liblua::lua_tonumberx(lua.lua, index, &success) };
                match success {
                    0 => None,
                    _ => Some(val as $t)
                }
            }
        }
        impl Index for $t {
        }
    );
)

numeric_impl!(f32)
numeric_impl!(f64)

impl Pushable for std::string::String {
    fn push_to_lua(&self, lua: &mut Lua) {
        unsafe { liblua::lua_pushstring(lua.lua, self.to_c_str().unwrap()) }
    }
}

impl Readable for String {
    fn read_from_lua(lua: &mut Lua, index: i32) -> Option<std::string::String> {
        let mut size: libc::size_t = unsafe { std::mem::uninitialized() };
        let cStrRaw = unsafe { liblua::lua_tolstring(lua.lua, index, &mut size) };
        if cStrRaw.is_null() {
            return None;
        }

        unsafe { std::c_str::CString::new(cStrRaw, false) }.as_str().map(|s| s.to_string())
    }
}

impl Index for String {
}

impl<'a> Pushable for &'a str {
    fn push_to_lua(&self, lua: &mut Lua) {
        unsafe { liblua::lua_pushstring(lua.lua, self.to_c_str().unwrap()) }
    }
}

impl Pushable for bool {
    fn push_to_lua(&self, lua: &mut Lua) {
        unsafe { liblua::lua_pushboolean(lua.lua, self.clone() as libc::c_int) }
    }
}

impl Readable for bool {
    fn read_from_lua(lua: &mut Lua, index: i32) -> Option<bool> {
        if unsafe { liblua::lua_isboolean(lua.lua, index) } != true {
            return None;
        }

        Some(unsafe { liblua::lua_toboolean(lua.lua, index) != 0 })
    }
}

impl Index for bool {
}

impl Pushable for () {
    fn push_to_lua(&self, lua: &mut Lua) {
        unsafe { liblua::lua_pushnil(lua.lua) }
    }
}

impl Readable for () {
    fn read_from_lua(lua: &mut Lua, index: i32) -> Option<()> {
        match unsafe { liblua::lua_isnil(lua.lua, index) } {
            true => Some(()),
            _ => None
        }
    }
}

#[cfg(test)]
mod tests {
    #[test]
    fn read_ints() {
        let mut lua = super::super::Lua::new();

        lua.set("a", 2).unwrap();

        let x: int = lua.get("a").unwrap();
        assert_eq!(x, 2);

        let y: i8 = lua.get("a").unwrap();
        assert_eq!(y, 2);

        let z: i16 = lua.get("a").unwrap();
        assert_eq!(z, 2);

        let w: i32 = lua.get("a").unwrap();
        assert_eq!(w, 2);

        let a: uint = lua.get("a").unwrap();
        assert_eq!(a, 2);

        let b: u8 = lua.get("a").unwrap();
        assert_eq!(b, 2);

        let c: u16 = lua.get("a").unwrap();
        assert_eq!(c, 2);

        let d: u32 = lua.get("a").unwrap();
        assert_eq!(d, 2);
    }

    #[test]
    fn write_ints() {
        // TODO:

        let mut lua = super::super::Lua::new();

        lua.set("a", 2).unwrap();
        let x: int = lua.get("a").unwrap();
        assert_eq!(x, 2);
    }

    #[test]
    fn readwrite_floats() {
        let mut lua = super::super::Lua::new();

        lua.set("a", 2.51234 as f32).unwrap();
        lua.set("b", 3.4123456789 as f64).unwrap();

        let x: f32 = lua.get("a").unwrap();
        assert!(x - 2.51234 < 0.000001);

        let y: f64 = lua.get("a").unwrap();
        assert!(y - 2.51234 < 0.000001);

        let z: f32 = lua.get("b").unwrap();
        assert!(z - 3.4123456789 < 0.000001);

        let w: f64 = lua.get("b").unwrap();
        assert!(w - 3.4123456789 < 0.000001);
    }

    #[test]
    fn readwrite_bools() {
        let mut lua = super::super::Lua::new();

        lua.set("a", true).unwrap();
        lua.set("b", false).unwrap();

        let x: bool = lua.get("a").unwrap();
        assert_eq!(x, true);

        let y: bool = lua.get("b").unwrap();
        assert_eq!(y, false);
    }

    #[test]
    fn readwrite_strings() {
        let mut lua = super::super::Lua::new();

        lua.set("a", "hello").unwrap();
        lua.set("b", "hello".to_string()).unwrap();

        let x: String = lua.get("a").unwrap();
        assert_eq!(x.as_slice(), "hello");

        let y: String = lua.get("b").unwrap();
        assert_eq!(y.as_slice(), "hello");
    }

    #[test]
    fn int_to_string() {
        let mut lua = super::super::Lua::new();

        lua.set("a", 2).unwrap();

        let x: String = lua.get("a").unwrap();
        assert_eq!(x.as_slice(), "2");
    }

    #[test]
    fn string_to_int() {
        let mut lua = super::super::Lua::new();

        lua.set("a", "2").unwrap();
        lua.set("b", "aaa").unwrap();

        let x: int = lua.get("a").unwrap();
        assert_eq!(x, 2);

        let y: Option<int> = lua.get("b");
        assert!(y.is_none());
    }
}
<|endoftext|>


<commit_before>//! Codegen vtables and vtable accesses.
//!
//! See `rustc_codegen_ssa/src/meth.rs` for reference.
// FIXME dedup this logic between miri, cg_llvm and cg_clif

use crate::prelude::*;
use super::constant::pointer_for_allocation;

fn vtable_memflags() -> MemFlags {
    let mut flags = MemFlags::trusted(); // A vtable access is always aligned and will never trap.
    flags.set_readonly(); // A vtable is always read-only.
    flags
}

pub(crate) fn drop_fn_of_obj(fx: &mut FunctionCx<'_, '_, '_>, vtable: Value) -> Value {
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;
    fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (ty::COMMON_VTABLE_ENTRIES_DROPINPLACE * usize_size) as i32,
    )
}

pub(crate) fn size_of_obj(fx: &mut FunctionCx<'_, '_, '_>, vtable: Value) -> Value {
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;
    fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (ty::COMMON_VTABLE_ENTRIES_SIZE * usize_size) as i32,
    )
}

pub(crate) fn min_align_of_obj(fx: &mut FunctionCx<'_, '_, '_>, vtable: Value) -> Value {
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;
    fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (ty::COMMON_VTABLE_ENTRIES_SIZE * usize_size) as i32,
    )
}

pub(crate) fn get_ptr_and_method_ref<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    arg: CValue<'tcx>,
    idx: usize,
) -> (Value, Value) {
    let (ptr, vtable) = if let Abi::ScalarPair(_, _) = arg.layout().abi {
        arg.load_scalar_pair(fx)
    } else {
        let (ptr, vtable) = arg.try_to_ptr().unwrap();
        (ptr.get_addr(fx), vtable.unwrap())
    };

    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes();
    let func_ref = fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (idx * usize_size as usize) as i32,
    );
    (ptr, func_ref)
}

pub(crate) fn get_vtable<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    ty: Ty<'tcx>,
    trait_ref: Option<ty::PolyExistentialTraitRef<'tcx>>,
) -> Value {
    let vtable_ptr = if let Some(vtable_ptr) = fx.vtables.get(&(ty, trait_ref)) {
        *vtable_ptr
    } else {
        let vtable_alloc_id = fx.tcx.vtable_allocation(ty, trait_ref);
        let vtable_allocation = fx.tcx.global_alloc(vtable_alloc_id).unwrap_memory();
        let vtable_ptr = pointer_for_allocation(fx, vtable_allocation);

        fx.vtables.insert((ty, trait_ref), vtable_ptr);
        vtable_ptr
    };

    vtable_ptr.get_addr(fx)
}
<commit_msg>Support allocation failures when interperting MIR<commit_after>//! Codegen vtables and vtable accesses.
//!
//! See `rustc_codegen_ssa/src/meth.rs` for reference.
// FIXME dedup this logic between miri, cg_llvm and cg_clif

use crate::prelude::*;
use super::constant::pointer_for_allocation;

fn vtable_memflags() -> MemFlags {
    let mut flags = MemFlags::trusted(); // A vtable access is always aligned and will never trap.
    flags.set_readonly(); // A vtable is always read-only.
    flags
}

pub(crate) fn drop_fn_of_obj(fx: &mut FunctionCx<'_, '_, '_>, vtable: Value) -> Value {
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;
    fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (ty::COMMON_VTABLE_ENTRIES_DROPINPLACE * usize_size) as i32,
    )
}

pub(crate) fn size_of_obj(fx: &mut FunctionCx<'_, '_, '_>, vtable: Value) -> Value {
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;
    fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (ty::COMMON_VTABLE_ENTRIES_SIZE * usize_size) as i32,
    )
}

pub(crate) fn min_align_of_obj(fx: &mut FunctionCx<'_, '_, '_>, vtable: Value) -> Value {
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;
    fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (ty::COMMON_VTABLE_ENTRIES_SIZE * usize_size) as i32,
    )
}

pub(crate) fn get_ptr_and_method_ref<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    arg: CValue<'tcx>,
    idx: usize,
) -> (Value, Value) {
    let (ptr, vtable) = if let Abi::ScalarPair(_, _) = arg.layout().abi {
        arg.load_scalar_pair(fx)
    } else {
        let (ptr, vtable) = arg.try_to_ptr().unwrap();
        (ptr.get_addr(fx), vtable.unwrap())
    };

    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes();
    let func_ref = fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (idx * usize_size as usize) as i32,
    );
    (ptr, func_ref)
}

pub(crate) fn get_vtable<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    ty: Ty<'tcx>,
    trait_ref: Option<ty::PolyExistentialTraitRef<'tcx>>,
) -> Value {
    let vtable_ptr = if let Some(vtable_ptr) = fx.vtables.get(&(ty, trait_ref)) {
        *vtable_ptr
    } else {
        let vtable_alloc_id = match fx.tcx.vtable_allocation(ty, trait_ref) {
            Ok(alloc) => alloc,
            Err(_) => fx.tcx.sess().fatal("allocation of constant vtable failed"),
        };
        let vtable_allocation = fx.tcx.global_alloc(vtable_alloc_id).unwrap_memory();
        let vtable_ptr = pointer_for_allocation(fx, vtable_allocation);

        fx.vtables.insert((ty, trait_ref), vtable_ptr);
        vtable_ptr
    };

    vtable_ptr.get_addr(fx)
}
<|endoftext|>


<commit_before>//! Codegen vtables and vtable accesses.
//!
//! See librustc_codegen_llvm/meth.rs for reference
// FIXME dedup this logic between miri, cg_llvm and cg_clif

use crate::prelude::*;

const DROP_FN_INDEX: usize = 0;
const SIZE_INDEX: usize = 1;
const ALIGN_INDEX: usize = 2;

fn vtable_memflags() -> MemFlags {
    let mut flags = MemFlags::trusted(); // A vtable access is always aligned and will never trap.
    flags.set_readonly(); // A vtable is always read-only.
    flags
}

pub(crate) fn drop_fn_of_obj(fx: &mut FunctionCx<'_, '_, '_>, vtable: Value) -> Value {
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;
    fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (DROP_FN_INDEX * usize_size) as i32,
    )
}

pub(crate) fn size_of_obj(fx: &mut FunctionCx<'_, '_, '_>, vtable: Value) -> Value {
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;
    fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (SIZE_INDEX * usize_size) as i32,
    )
}

pub(crate) fn min_align_of_obj(fx: &mut FunctionCx<'_, '_, '_>, vtable: Value) -> Value {
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;
    fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (ALIGN_INDEX * usize_size) as i32,
    )
}

pub(crate) fn get_ptr_and_method_ref<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    arg: CValue<'tcx>,
    idx: usize,
) -> (Value, Value) {
    let (ptr, vtable) = if let Abi::ScalarPair(_, _) = arg.layout().abi {
        arg.load_scalar_pair(fx)
    } else {
        let (ptr, vtable) = arg.try_to_ptr().unwrap();
        (ptr.get_addr(fx), vtable.unwrap())
    };

    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes();
    let func_ref = fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        ((idx + 3) * usize_size as usize) as i32,
    );
    (ptr, func_ref)
}

pub(crate) fn get_vtable<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    layout: TyAndLayout<'tcx>,
    trait_ref: Option<ty::PolyExistentialTraitRef<'tcx>>,
) -> Value {
    let data_id = if let Some(data_id) = fx.vtables.get(&(layout.ty, trait_ref)) {
        *data_id
    } else {
        let data_id = build_vtable(fx, layout, trait_ref);
        fx.vtables.insert((layout.ty, trait_ref), data_id);
        data_id
    };

    let local_data_id = fx.module.declare_data_in_func(data_id, &mut fx.bcx.func);
    fx.bcx.ins().global_value(fx.pointer_type, local_data_id)
}

fn build_vtable<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    layout: TyAndLayout<'tcx>,
    trait_ref: Option<ty::PolyExistentialTraitRef<'tcx>>,
) -> DataId {
    let tcx = fx.tcx;
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;

    let drop_in_place_fn = import_function(
        tcx,
        fx.module,
        Instance::resolve_drop_in_place(tcx, layout.ty).polymorphize(fx.tcx),
    );

    let mut components: Vec<_> = vec![Some(drop_in_place_fn), None, None];

    let methods_root;
    let methods = if let Some(trait_ref) = trait_ref {
        methods_root = tcx.vtable_methods(trait_ref.with_self_ty(tcx, layout.ty));
        methods_root.iter()
    } else {
        (&[]).iter()
    };
    let methods = methods.cloned().map(|opt_mth| {
        opt_mth.map(|(def_id, substs)| {
            import_function(
                tcx,
                fx.module,
                Instance::resolve_for_vtable(tcx, ParamEnv::reveal_all(), def_id, substs)
                    .unwrap()
                    .polymorphize(fx.tcx),
            )
        })
    });
    components.extend(methods);

    let mut data_ctx = DataContext::new();
    let mut data = ::std::iter::repeat(0u8)
        .take(components.len() * usize_size)
        .collect::<Vec<u8>>()
        .into_boxed_slice();

    write_usize(fx.tcx, &mut data, SIZE_INDEX, layout.size.bytes());
    write_usize(fx.tcx, &mut data, ALIGN_INDEX, layout.align.abi.bytes());
    data_ctx.define(data);

    for (i, component) in components.into_iter().enumerate() {
        if let Some(func_id) = component {
            let func_ref = fx.module.declare_func_in_data(func_id, &mut data_ctx);
            data_ctx.write_function_addr((i * usize_size) as u32, func_ref);
        }
    }

    data_ctx.set_align(fx.tcx.data_layout.pointer_align.pref.bytes());

    let data_id = fx.module.declare_anonymous_data(false, false).unwrap();

    fx.module.define_data(data_id, &data_ctx).unwrap();

    data_id
}

fn write_usize(tcx: TyCtxt<'_>, buf: &mut [u8], idx: usize, num: u64) {
    let pointer_size =
        tcx.layout_of(ParamEnv::reveal_all().and(tcx.types.usize)).unwrap().size.bytes() as usize;
    let target = &mut buf[idx * pointer_size..(idx + 1) * pointer_size];

    match tcx.data_layout.endian {
        rustc_target::abi::Endian::Little => match pointer_size {
            4 => target.copy_from_slice(&(num as u32).to_le_bytes()),
            8 => target.copy_from_slice(&(num as u64).to_le_bytes()),
            _ => todo!("pointer size {} is not yet supported", pointer_size),
        },
        rustc_target::abi::Endian::Big => match pointer_size {
            4 => target.copy_from_slice(&(num as u32).to_be_bytes()),
            8 => target.copy_from_slice(&(num as u64).to_be_bytes()),
            _ => todo!("pointer size {} is not yet supported", pointer_size),
        },
    }
}
<commit_msg>Sync from rust b203b0d240b67916cfa77f640aedaf1c87d50f6d<commit_after>//! Codegen vtables and vtable accesses.
//!
//! See `rustc_codegen_ssa/src/meth.rs` for reference.
// FIXME dedup this logic between miri, cg_llvm and cg_clif

use crate::prelude::*;

const DROP_FN_INDEX: usize = 0;
const SIZE_INDEX: usize = 1;
const ALIGN_INDEX: usize = 2;

fn vtable_memflags() -> MemFlags {
    let mut flags = MemFlags::trusted(); // A vtable access is always aligned and will never trap.
    flags.set_readonly(); // A vtable is always read-only.
    flags
}

pub(crate) fn drop_fn_of_obj(fx: &mut FunctionCx<'_, '_, '_>, vtable: Value) -> Value {
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;
    fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (DROP_FN_INDEX * usize_size) as i32,
    )
}

pub(crate) fn size_of_obj(fx: &mut FunctionCx<'_, '_, '_>, vtable: Value) -> Value {
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;
    fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (SIZE_INDEX * usize_size) as i32,
    )
}

pub(crate) fn min_align_of_obj(fx: &mut FunctionCx<'_, '_, '_>, vtable: Value) -> Value {
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;
    fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        (ALIGN_INDEX * usize_size) as i32,
    )
}

pub(crate) fn get_ptr_and_method_ref<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    arg: CValue<'tcx>,
    idx: usize,
) -> (Value, Value) {
    let (ptr, vtable) = if let Abi::ScalarPair(_, _) = arg.layout().abi {
        arg.load_scalar_pair(fx)
    } else {
        let (ptr, vtable) = arg.try_to_ptr().unwrap();
        (ptr.get_addr(fx), vtable.unwrap())
    };

    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes();
    let func_ref = fx.bcx.ins().load(
        pointer_ty(fx.tcx),
        vtable_memflags(),
        vtable,
        ((idx + 3) * usize_size as usize) as i32,
    );
    (ptr, func_ref)
}

pub(crate) fn get_vtable<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    layout: TyAndLayout<'tcx>,
    trait_ref: Option<ty::PolyExistentialTraitRef<'tcx>>,
) -> Value {
    let data_id = if let Some(data_id) = fx.vtables.get(&(layout.ty, trait_ref)) {
        *data_id
    } else {
        let data_id = build_vtable(fx, layout, trait_ref);
        fx.vtables.insert((layout.ty, trait_ref), data_id);
        data_id
    };

    let local_data_id = fx.module.declare_data_in_func(data_id, &mut fx.bcx.func);
    fx.bcx.ins().global_value(fx.pointer_type, local_data_id)
}

fn build_vtable<'tcx>(
    fx: &mut FunctionCx<'_, '_, 'tcx>,
    layout: TyAndLayout<'tcx>,
    trait_ref: Option<ty::PolyExistentialTraitRef<'tcx>>,
) -> DataId {
    let tcx = fx.tcx;
    let usize_size = fx.layout_of(fx.tcx.types.usize).size.bytes() as usize;

    let drop_in_place_fn = import_function(
        tcx,
        fx.module,
        Instance::resolve_drop_in_place(tcx, layout.ty).polymorphize(fx.tcx),
    );

    let mut components: Vec<_> = vec![Some(drop_in_place_fn), None, None];

    let methods_root;
    let methods = if let Some(trait_ref) = trait_ref {
        methods_root = tcx.vtable_methods(trait_ref.with_self_ty(tcx, layout.ty));
        methods_root.iter()
    } else {
        (&[]).iter()
    };
    let methods = methods.cloned().map(|opt_mth| {
        opt_mth.map(|(def_id, substs)| {
            import_function(
                tcx,
                fx.module,
                Instance::resolve_for_vtable(tcx, ParamEnv::reveal_all(), def_id, substs)
                    .unwrap()
                    .polymorphize(fx.tcx),
            )
        })
    });
    components.extend(methods);

    let mut data_ctx = DataContext::new();
    let mut data = ::std::iter::repeat(0u8)
        .take(components.len() * usize_size)
        .collect::<Vec<u8>>()
        .into_boxed_slice();

    write_usize(fx.tcx, &mut data, SIZE_INDEX, layout.size.bytes());
    write_usize(fx.tcx, &mut data, ALIGN_INDEX, layout.align.abi.bytes());
    data_ctx.define(data);

    for (i, component) in components.into_iter().enumerate() {
        if let Some(func_id) = component {
            let func_ref = fx.module.declare_func_in_data(func_id, &mut data_ctx);
            data_ctx.write_function_addr((i * usize_size) as u32, func_ref);
        }
    }

    data_ctx.set_align(fx.tcx.data_layout.pointer_align.pref.bytes());

    let data_id = fx.module.declare_anonymous_data(false, false).unwrap();

    fx.module.define_data(data_id, &data_ctx).unwrap();

    data_id
}

fn write_usize(tcx: TyCtxt<'_>, buf: &mut [u8], idx: usize, num: u64) {
    let pointer_size =
        tcx.layout_of(ParamEnv::reveal_all().and(tcx.types.usize)).unwrap().size.bytes() as usize;
    let target = &mut buf[idx * pointer_size..(idx + 1) * pointer_size];

    match tcx.data_layout.endian {
        rustc_target::abi::Endian::Little => match pointer_size {
            4 => target.copy_from_slice(&(num as u32).to_le_bytes()),
            8 => target.copy_from_slice(&(num as u64).to_le_bytes()),
            _ => todo!("pointer size {} is not yet supported", pointer_size),
        },
        rustc_target::abi::Endian::Big => match pointer_size {
            4 => target.copy_from_slice(&(num as u32).to_be_bytes()),
            8 => target.copy_from_slice(&(num as u64).to_be_bytes()),
            _ => todo!("pointer size {} is not yet supported", pointer_size),
        },
    }
}
<|endoftext|>


<commit_before>//! Formats a DOM structure to a Writer
//!
//! ### Example
//! ```
//! use document::Package;
//! use document::writer::format_document;
//!
//! let package = Package::new();
//! let doc = package.as_document();
//!
//! let hello = doc.create_element("hello");
//! hello.set_attribute_value("planet", "Earth");
//! doc.root().append_child(hello);
//!
//! let mut output = std::io::stdio::stdout();
//! format_document(&doc, &mut output).ok().expect("unable to output XML");
//! ```
//!
//! ### Known issues
//!
//! Output is not escaped in any way,
//! it's very easy to create malformed XML!
//!
//! ### Potential options to support
//!
//! - Space before `/>`
//! - Single vs double quotes
//! - Fixed ordering of attributes

use std::collections::HashMap;
use std::collections::hash_map::Entry::{Occupied,Vacant};
use std::io::IoResult;

use self::Content::*;

use super::QName;

use super::dom4;
use super::dom4::ChildOfElement::*;
use super::dom4::ChildOfRoot::*;

struct PrefixMapping<'a> {
    defined_prefixes: HashMap<&'a str, String>,
    generated_prefix_count: uint,
}

impl<'a> PrefixMapping<'a> {
    fn new() -> PrefixMapping<'a> {
        PrefixMapping {
            defined_prefixes: HashMap::new(),
            generated_prefix_count: 0,
        }
    }

    fn prefix_for_uri(&mut self, namespace_uri: &'a str) -> &str {
        let prefix = match self.defined_prefixes.entry(namespace_uri) {
            Occupied(entry) => entry.into_mut(),
            Vacant(entry) => {
                let prefix = format!("autons{}", self.generated_prefix_count);
                // TODO: Make sure doesn't already exist
                self.generated_prefix_count += 1;
                entry.set(prefix)
            },
        };

        prefix.as_slice()
    }
}

enum Content<'d> {
    Element(dom4::Element<'d>),
    ElementEnd(dom4::Element<'d>),
    Text(dom4::Text<'d>),
    Comment(dom4::Comment<'d>),
    ProcessingInstruction(dom4::ProcessingInstruction<'d>),
}

fn format_qname<'d, W>(q: QName<'d>, mapping: &mut PrefixMapping<'d>, writer: &mut W) -> IoResult<()>
    where W: Writer
{
    if let Some(namespace_uri) = q.namespace_uri {
        let prefix = mapping.prefix_for_uri(namespace_uri);
        try!(writer.write_str(prefix));
        try!(writer.write_str(":"));
    }
    writer.write_str(q.local_part)
}

fn format_element<'d, W>(element: dom4::Element<'d>,
                         todo: &mut Vec<Content<'d>>,
                         mapping: &mut PrefixMapping<'d>,
                         writer: &mut W)
                         -> IoResult<()>
    where W: Writer
{
    try!(writer.write_str("<"));
    try!(format_qname(element.name(), mapping, writer));

    for attr in element.attributes().iter() {
        try!(writer.write_str(" "));
        try!(format_qname(attr.name(), mapping, writer));
        try!(write!(writer, "='{}'", attr.value()));
    }

    for (ns_uri, prefix) in mapping.defined_prefixes.iter() {
        try!(writer.write_str(" xmlns:"));
        try!(writer.write_str(prefix.as_slice()));
        try!(write!(writer, "='{}'", *ns_uri));
    }

    let mut children = element.children();
    if children.is_empty() {
        writer.write_str("/>")
    } else {
        try!(writer.write_str(">"));

        todo.push(ElementEnd(element));
        children.reverse();
        let x = children.into_iter().map(|c| match c {
            ElementCOE(element) => Element(element),
            TextCOE(t)          => Text(t),
            CommentCOE(c)       => Comment(c),
            ProcessingInstructionCOE(p)            => ProcessingInstruction(p),
        });
        todo.extend(x);

        Ok(())
    }
}

fn format_element_end<'d, W>(element: dom4::Element<'d>,
                             mapping: &mut PrefixMapping<'d>,
                             writer: &mut W)
                             -> IoResult<()>
    where W: Writer
{
    try!(writer.write_str("</"));
    try!(format_qname(element.name(), mapping, writer));
    writer.write_str(">")
}

fn format_comment<W>(comment: dom4::Comment, writer: &mut W) -> IoResult<()>
    where W: Writer
{
    write!(writer, "<!--{}-->", comment.text())
}

fn format_processing_instruction<W>(pi: dom4::ProcessingInstruction, writer: &mut W) -> IoResult<()>
    where W: Writer
{
    match pi.value() {
        None    => write!(writer, "<?{}?>", pi.target()),
        Some(v) => write!(writer, "<?{} {}?>", pi.target(), v),
    }
}

fn format_one<'d, W>(content: Content<'d>,
                     todo: &mut Vec<Content<'d>>,
                     mapping: &mut PrefixMapping<'d>,
                     writer: &mut W)
                     -> IoResult<()>
    where W: Writer
{
    match content {
        Element(e)               => format_element(e, todo, mapping, writer),
        ElementEnd(e)            => format_element_end(e, mapping, writer),
        Text(t)                  => writer.write_str(t.text().as_slice()),
        Comment(c)               => format_comment(c, writer),
        ProcessingInstruction(p) => format_processing_instruction(p, writer),
    }
}

fn format_body<W>(element: dom4::Element, writer: &mut W) -> IoResult<()>
    where W: Writer
{
    let mut todo = vec![Element(element)];
    let mut mapping = PrefixMapping::new();

    while ! todo.is_empty() {
        try!(format_one(todo.pop().unwrap(), &mut todo, &mut mapping, writer));
    }

    Ok(())
}

/// Formats a document into a Writer
pub fn format_document<'d, W>(doc: &'d dom4::Document<'d>, writer: &mut W) -> IoResult<()>
    where W: Writer
{
    try!(writer.write_str("<?xml version='1.0'?>"));

    for child in doc.root().children().into_iter() {
        try!(match child {
            ElementCOR(e) => format_body(e, writer),
            CommentCOR(c) => format_comment(c, writer),
            ProcessingInstructionCOR(p)      => format_processing_instruction(p, writer),
        })
    }

    Ok(())
}

#[cfg(test)]
mod test {
    use std::io::MemWriter;

    use super::super::{Package,QName};
    use super::super::dom4;
    use super::format_document;

    macro_rules! assert_str_eq(
        ($l:expr, $r:expr) => (assert_eq!($l.as_slice(), $r.as_slice()));
    )

    fn format_xml<'d>(doc: &'d dom4::Document<'d>) -> String {
        let mut w = MemWriter::new();
        format_document(doc, &mut w).ok().expect("Not formatted");
        String::from_utf8(w.into_inner()).ok().expect("Not a string")
    }

    #[test]
    fn top_element() {
        let p = Package::new();
        let d = p.as_document();
        let e = d.create_element("hello");
        d.root().append_child(e);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello/>");
    }

    #[test]
    fn element_with_namespace() {
        let p = Package::new();
        let d = p.as_document();
        let name = QName::with_namespace_uri(Some("namespace"), "local-part");
        let e = d.create_element(name);
        d.root().append_child(e);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><autons0:local-part xmlns:autons0='namespace'/>");
    }

    #[test]
    fn element_with_attributes() {
        let p = Package::new();
        let d = p.as_document();
        let e = d.create_element("hello");
        e.set_attribute_value("a", "b");
        d.root().append_child(e);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello a='b'/>");
    }

    #[test]
    fn attribute_with_namespace() {
        let p = Package::new();
        let d = p.as_document();
        let e = d.create_element("hello");
        let name = QName::with_namespace_uri(Some("namespace"), "a");
        e.set_attribute_value(name, "b");
        d.root().append_child(e);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello autons0:a='b' xmlns:autons0='namespace'/>");
    }

    #[test]
    fn nested_element() {
        let p = Package::new();
        let d = p.as_document();
        let hello = d.create_element("hello");
        let world = d.create_element("world");
        hello.append_child(world);
        d.root().append_child(hello);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello><world/></hello>");
    }

    #[test]
    fn nested_text() {
        let p = Package::new();
        let d = p.as_document();
        let hello = d.create_element("hello");
        let text = d.create_text("A fine day to you!");
        hello.append_child(text);
        d.root().append_child(hello);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello>A fine day to you!</hello>");
    }

    #[test]
    fn nested_comment() {
        let p = Package::new();
        let d = p.as_document();
        let hello = d.create_element("hello");
        let comment = d.create_comment(" Fill this in ");
        hello.append_child(comment);
        d.root().append_child(hello);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello><!-- Fill this in --></hello>");
    }

    #[test]
    fn nested_processing_instruction_without_value() {
        let p = Package::new();
        let d = p.as_document();
        let hello = d.create_element("hello");
        let pi = d.create_processing_instruction("display", None);
        hello.append_child(pi);
        d.root().append_child(hello);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello><?display?></hello>");
    }

    #[test]
    fn nested_processing_instruction_with_value() {
        let p = Package::new();
        let d = p.as_document();
        let hello = d.create_element("hello");
        let pi = d.create_processing_instruction("display", Some("screen"));
        hello.append_child(pi);
        d.root().append_child(hello);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello><?display screen?></hello>");
    }

    #[test]
    fn top_level_comment() {
        let p = Package::new();
        let d = p.as_document();
        let comment = d.create_comment(" Fill this in ");
        d.root().append_child(comment);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><!-- Fill this in -->");
    }

    #[test]
    fn top_level_processing_instruction() {
        let p = Package::new();
        let d = p.as_document();
        let pi = d.create_processing_instruction("display", None);
        d.root().append_child(pi);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><?display?>");
    }
}
<commit_msg>Track multiple levels of namespace prefixes<commit_after>//! Formats a DOM structure to a Writer
//!
//! ### Example
//! ```
//! use document::Package;
//! use document::writer::format_document;
//!
//! let package = Package::new();
//! let doc = package.as_document();
//!
//! let hello = doc.create_element("hello");
//! hello.set_attribute_value("planet", "Earth");
//! doc.root().append_child(hello);
//!
//! let mut output = std::io::stdio::stdout();
//! format_document(&doc, &mut output).ok().expect("unable to output XML");
//! ```
//!
//! ### Known issues
//!
//! Output is not escaped in any way,
//! it's very easy to create malformed XML!
//!
//! ### Potential options to support
//!
//! - Space before `/>`
//! - Single vs double quotes
//! - Fixed ordering of attributes

use std::collections::HashMap;
use std::collections::hash_map;
use std::collections::hash_map::Entry;
use std::io::IoResult;

use self::Content::*;

use super::QName;

use super::dom4;
use super::dom4::ChildOfElement::*;
use super::dom4::ChildOfRoot::*;

struct PrefixMapping<'a> {
    defined_prefixes: Vec<HashMap<&'a str, String>>,
    generated_prefix_count: uint,
}

impl<'a> PrefixMapping<'a> {
    fn new() -> PrefixMapping<'a> {
        PrefixMapping {
            defined_prefixes: vec![HashMap::new()],
            generated_prefix_count: 0,
        }
    }

    fn push_scope(&mut self) {
        self.defined_prefixes.push(HashMap::new());
    }

    fn pop_scope(&mut self) {
        self.defined_prefixes.pop();
    }

    fn prefixes_in_current_scope(&self) -> hash_map::Entries<&'a str, String> {
        self.defined_prefixes.last().unwrap().iter()
    }

    fn recursive_lookup(&self, namespace_uri: &'a str) -> Option<&str> {
        for scope in self.defined_prefixes.iter().rev() {
            if let Some(prefix) = scope.get(namespace_uri) {
                return Some(prefix.as_slice());
            }
        }

        None
    }

    fn prefix_for_uri(&mut self, namespace_uri: &'a str) -> &str {
        if let None = self.recursive_lookup(namespace_uri) {
            let current_scope = self.defined_prefixes.last_mut().unwrap();

            let prefix = format!("autons{}", self.generated_prefix_count);
            self.generated_prefix_count += 1;

            match current_scope.entry(namespace_uri) {
                Entry::Occupied(_) => panic!("Duplicate prefix"),
                Entry::Vacant(entry) => entry.set(prefix),
            };
        }

        // This is inefficient, but the current borrow checker doesn't
        // understand early returns well enough.
        self.recursive_lookup(namespace_uri).unwrap().as_slice()
    }
}

enum Content<'d> {
    Element(dom4::Element<'d>),
    ElementEnd(dom4::Element<'d>),
    Text(dom4::Text<'d>),
    Comment(dom4::Comment<'d>),
    ProcessingInstruction(dom4::ProcessingInstruction<'d>),
}

fn format_qname<'d, W>(q: QName<'d>, mapping: &mut PrefixMapping<'d>, writer: &mut W) -> IoResult<()>
    where W: Writer
{
    if let Some(namespace_uri) = q.namespace_uri {
        let prefix = mapping.prefix_for_uri(namespace_uri);
        try!(writer.write_str(prefix));
        try!(writer.write_str(":"));
    }
    writer.write_str(q.local_part)
}

fn format_element<'d, W>(element: dom4::Element<'d>,
                         todo: &mut Vec<Content<'d>>,
                         mapping: &mut PrefixMapping<'d>,
                         writer: &mut W)
                         -> IoResult<()>
    where W: Writer
{
    try!(writer.write_str("<"));
    try!(format_qname(element.name(), mapping, writer));

    for attr in element.attributes().iter() {
        try!(writer.write_str(" "));
        try!(format_qname(attr.name(), mapping, writer));
        try!(write!(writer, "='{}'", attr.value()));
    }

    for (ns_uri, prefix) in mapping.prefixes_in_current_scope() {
        try!(writer.write_str(" xmlns:"));
        try!(writer.write_str(prefix.as_slice()));
        try!(write!(writer, "='{}'", *ns_uri));
    }

    let mut children = element.children();
    if children.is_empty() {
        writer.write_str("/>")
    } else {
        try!(writer.write_str(">"));

        todo.push(ElementEnd(element));
        children.reverse();
        let x = children.into_iter().map(|c| match c {
            ElementCOE(element) => Element(element),
            TextCOE(t)          => Text(t),
            CommentCOE(c)       => Comment(c),
            ProcessingInstructionCOE(p)            => ProcessingInstruction(p),
        });
        todo.extend(x);

        Ok(())
    }
}

fn format_element_end<'d, W>(element: dom4::Element<'d>,
                             mapping: &mut PrefixMapping<'d>,
                             writer: &mut W)
                             -> IoResult<()>
    where W: Writer
{
    try!(writer.write_str("</"));
    try!(format_qname(element.name(), mapping, writer));
    writer.write_str(">")
}

fn format_comment<W>(comment: dom4::Comment, writer: &mut W) -> IoResult<()>
    where W: Writer
{
    write!(writer, "<!--{}-->", comment.text())
}

fn format_processing_instruction<W>(pi: dom4::ProcessingInstruction, writer: &mut W) -> IoResult<()>
    where W: Writer
{
    match pi.value() {
        None    => write!(writer, "<?{}?>", pi.target()),
        Some(v) => write!(writer, "<?{} {}?>", pi.target(), v),
    }
}

fn format_one<'d, W>(content: Content<'d>,
                     todo: &mut Vec<Content<'d>>,
                     mapping: &mut PrefixMapping<'d>,
                     writer: &mut W)
                     -> IoResult<()>
    where W: Writer
{
    match content {
        Element(e)               => {
            mapping.push_scope();
            format_element(e, todo, mapping, writer)
        },
        ElementEnd(e)            => {
            let r = format_element_end(e, mapping, writer);
            mapping.pop_scope();
            r
        },
        Text(t)                  => writer.write_str(t.text().as_slice()),
        Comment(c)               => format_comment(c, writer),
        ProcessingInstruction(p) => format_processing_instruction(p, writer),
    }
}

fn format_body<W>(element: dom4::Element, writer: &mut W) -> IoResult<()>
    where W: Writer
{
    let mut todo = vec![Element(element)];
    let mut mapping = PrefixMapping::new();

    while ! todo.is_empty() {
        try!(format_one(todo.pop().unwrap(), &mut todo, &mut mapping, writer));
    }

    Ok(())
}

/// Formats a document into a Writer
pub fn format_document<'d, W>(doc: &'d dom4::Document<'d>, writer: &mut W) -> IoResult<()>
    where W: Writer
{
    try!(writer.write_str("<?xml version='1.0'?>"));

    for child in doc.root().children().into_iter() {
        try!(match child {
            ElementCOR(e) => format_body(e, writer),
            CommentCOR(c) => format_comment(c, writer),
            ProcessingInstructionCOR(p)      => format_processing_instruction(p, writer),
        })
    }

    Ok(())
}

#[cfg(test)]
mod test {
    use std::io::MemWriter;

    use super::super::{Package,QName};
    use super::super::dom4;
    use super::format_document;

    macro_rules! assert_str_eq(
        ($l:expr, $r:expr) => (assert_eq!($l.as_slice(), $r.as_slice()));
    )

    fn format_xml<'d>(doc: &'d dom4::Document<'d>) -> String {
        let mut w = MemWriter::new();
        format_document(doc, &mut w).ok().expect("Not formatted");
        String::from_utf8(w.into_inner()).ok().expect("Not a string")
    }

    #[test]
    fn top_element() {
        let p = Package::new();
        let d = p.as_document();
        let e = d.create_element("hello");
        d.root().append_child(e);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello/>");
    }

    #[test]
    fn element_with_namespace() {
        let p = Package::new();
        let d = p.as_document();
        let name = QName::with_namespace_uri(Some("namespace"), "local-part");
        let e = d.create_element(name);
        d.root().append_child(e);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><autons0:local-part xmlns:autons0='namespace'/>");
    }

    #[test]
    fn element_with_attributes() {
        let p = Package::new();
        let d = p.as_document();
        let e = d.create_element("hello");
        e.set_attribute_value("a", "b");
        d.root().append_child(e);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello a='b'/>");
    }

    #[test]
    fn attribute_with_namespace() {
        let p = Package::new();
        let d = p.as_document();
        let e = d.create_element("hello");
        let name = QName::with_namespace_uri(Some("namespace"), "a");
        e.set_attribute_value(name, "b");
        d.root().append_child(e);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello autons0:a='b' xmlns:autons0='namespace'/>");
    }

    #[test]
    fn nested_element() {
        let p = Package::new();
        let d = p.as_document();
        let hello = d.create_element("hello");
        let world = d.create_element("world");
        hello.append_child(world);
        d.root().append_child(hello);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello><world/></hello>");
    }

    #[test]
    fn nested_element_with_namespaces() {
        let p = Package::new();
        let d = p.as_document();
        let outer_name = QName::with_namespace_uri(Some("outer"), "hello");
        let inner_name = QName::with_namespace_uri(Some("inner"), "world");
        let hello = d.create_element(outer_name);
        let world = d.create_element(inner_name);
        hello.append_child(world);
        d.root().append_child(hello);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><autons0:hello xmlns:autons0='outer'><autons1:world xmlns:autons1='inner'/></autons0:hello>");
    }

    #[test]
    fn nested_element_with_namespaces_with_reused_namespaces() {
        let p = Package::new();
        let d = p.as_document();
        let outer_name = QName::with_namespace_uri(Some("ns"), "hello");
        let inner_name = QName::with_namespace_uri(Some("ns"), "world");
        let hello = d.create_element(outer_name);
        let world = d.create_element(inner_name);
        hello.append_child(world);
        d.root().append_child(hello);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><autons0:hello xmlns:autons0='ns'><autons0:world/></autons0:hello>");
    }

    #[test]
    fn nested_text() {
        let p = Package::new();
        let d = p.as_document();
        let hello = d.create_element("hello");
        let text = d.create_text("A fine day to you!");
        hello.append_child(text);
        d.root().append_child(hello);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello>A fine day to you!</hello>");
    }

    #[test]
    fn nested_comment() {
        let p = Package::new();
        let d = p.as_document();
        let hello = d.create_element("hello");
        let comment = d.create_comment(" Fill this in ");
        hello.append_child(comment);
        d.root().append_child(hello);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello><!-- Fill this in --></hello>");
    }

    #[test]
    fn nested_processing_instruction_without_value() {
        let p = Package::new();
        let d = p.as_document();
        let hello = d.create_element("hello");
        let pi = d.create_processing_instruction("display", None);
        hello.append_child(pi);
        d.root().append_child(hello);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello><?display?></hello>");
    }

    #[test]
    fn nested_processing_instruction_with_value() {
        let p = Package::new();
        let d = p.as_document();
        let hello = d.create_element("hello");
        let pi = d.create_processing_instruction("display", Some("screen"));
        hello.append_child(pi);
        d.root().append_child(hello);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><hello><?display screen?></hello>");
    }

    #[test]
    fn top_level_comment() {
        let p = Package::new();
        let d = p.as_document();
        let comment = d.create_comment(" Fill this in ");
        d.root().append_child(comment);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><!-- Fill this in -->");
    }

    #[test]
    fn top_level_processing_instruction() {
        let p = Package::new();
        let d = p.as_document();
        let pi = d.create_processing_instruction("display", None);
        d.root().append_child(pi);

        let xml = format_xml(&d);
        assert_str_eq!(xml, "<?xml version='1.0'?><?display?>");
    }
}
<|endoftext|>


<commit_before>/*! Example IRC bot

    This sample bot uses the library to connect to Freenode, under the nickname "rustirclib###"
    (where ### is a random number). It then joins the channel ##rustirclib, says hello, and prints
    to standard output any messages sent to the channel.

    Any incoming message of the form "rustirclib###: some message" will elicit a generic response.

    Any incoming message of the form "!rustirclib### quit" will cause the bot to shut down.
 */

#![crate_id = "github.com/kballard/rust-irclib#ircbot:0.1"]
#![crate_type = "bin"]

#![allow(deprecated_owned_vector)]

extern crate irc;
extern crate rand;

use irc::conn::{Conn, Line, Event, IRCCmd, IRCCode, IRCAction};

use std::str;
use rand::Rng;

fn main() {
    let mut opts = irc::conn::Options::new("chat.freenode.net", irc::conn::DefaultPort);

    let nick = format!("rustirclib{}", rand::task_rng().gen_range(100u, 1000u));
    opts.nick = nick.as_slice();
    match irc::conn::connect(opts, (), |c,e,_| handler(c,e)) {
        Ok(()) => println!("Exiting..."),
        Err(err) => println!("Connection error: {}", err)
    }
}

fn handler(conn: &mut Conn, event: Event) {
    match event {
        irc::conn::Connected => println!("Connected"),
        irc::conn::Disconnected => println!("Disconnected"),
        irc::conn::LineReceived(line) => {
            match line {
                Line{command: IRCCode(1), ..} => {
                    println!("Logged in");
                    // we've logged in
                    conn.join(bytes!("##rustirclib"), [])
                }
                Line{command: IRCCmd(cmd), args, prefix: prefix } => match cmd.as_slice() {
                    "JOIN" if prefix.is_some() => {
                        let prefix = prefix.unwrap();
                        if prefix.nick() != conn.me().nick() {
                            return;
                        }
                        if args.is_empty() {
                            let line = Line{command: IRCCmd("JOIN".into_maybe_owned()),
                                            args: args, prefix: Some(prefix)};
                            println!("ERROR: Invalid JOIN message received: {}", line_desc(&line));
                            return;
                        }
                        let chan = args.move_iter().next().unwrap();
                        conn.privmsg(chan.as_slice(), bytes!("Hello"));
                        let chan = str::from_utf8_lossy(chan.as_slice());
                        println!("JOINED: {}", chan);
                    }
                    "PRIVMSG" | "NOTICE" => {
                        let (src, dst, msg) = match prefix {
                            Some(_) if args.len() == 2 => {
                                let mut args = args;
                                let (dst, msg) = (args.swap_remove(0).unwrap(),
                                                  args.move_iter().next().unwrap());
                                (prefix.as_ref().unwrap().nick(), dst, msg)
                            }
                            _ => {
                                print!("ERROR: Unexpected {} line: ", cmd);
                                let line = Line{command: IRCCmd(cmd), args: args, prefix: prefix};
                                println!("{}", line_desc(&line));
                                return;
                            }
                        };
                        let dsts = str::from_utf8_lossy(dst.as_slice());
                        let srcs = str::from_utf8_lossy(src.as_slice());
                        let msgs = str::from_utf8_lossy(msg.as_slice());
                        println!("<-- {}({}) {}: {}", cmd, dsts, srcs, msgs);
                        handle_privmsg(conn, msg.as_slice(), src, dst.as_slice())
                    }
                    _ => ()
                },
                Line{command: IRCAction(dst), args, prefix } => {
                    let (src, msg) = match prefix {
                        Some(_) if args.len() == 1 => {
                            let msg = args.move_iter().next().unwrap();
                            (prefix.as_ref().unwrap().nick(), msg)
                        }
                        _ => {
                            let line = Line{command: IRCAction(dst), args: args, prefix: prefix};
                            println!("ERROR: Unexpected ACTION line: {}", line_desc(&line));
                            return;
                        }
                    };
                    let dst = str::from_utf8_lossy(dst.as_slice());
                    let src = str::from_utf8_lossy(src.as_slice());
                    let msg = str::from_utf8_lossy(msg.as_slice());
                    println!("<-- PRIVMSG({}) {} {}", dst, src, msg);
                }
                _ => ()
            }
        }
    }
}

fn handle_privmsg(conn: &mut Conn, msg: &[u8], src: &[u8], dst: &[u8]) {
    enum MsgType<'a> {
        DirectedMessage,
        CommandMessage(&'a [u8]),
        NormalMessage
    }
    let msgtype = {
        let nick = conn.me().nick();
        if msg.starts_with(nick) && msg.slice_from(nick.len()).starts_with(bytes!(": ")) &&
        msg.len() > nick.len() + 2 {
            DirectedMessage
        } else if msg.starts_with(bytes!("!")) && msg.slice_from(1).starts_with(nick) &&
                msg.slice_from(nick.len()+1).starts_with(bytes!(" ")) {
            let args = msg.slice_from(conn.me().nick().len()+2);
            let mut argit = args.split(|&b| b == ' ' as u8).skip_while(|v| v.is_empty());
            match argit.next() {
                Some(arg) => CommandMessage(arg),
                None => CommandMessage(bytes!(""))
            }
        } else {
            NormalMessage
        }
    };
    match msgtype {
        DirectedMessage => {
            let reply = if dst == conn.me().nick() { src } else { dst };
            let msg = Vec::from_slice(src).append(bytes!(": Hello"));
            conn.privmsg(reply, msg.as_slice());
            let src = str::from_utf8(conn.me().nick()).unwrap_or("(invalid utf8)");
            let reply = str::from_utf8(reply).unwrap_or("(invalid utf8)");
            let msg = StrBuf::from_utf8(msg).unwrap_or_else(|_| "(invalid utf8)".to_strbuf());
            println!("--> PRIVMSG({}) {}: {}", reply, src, msg);
        }
        CommandMessage(cmd) if cmd == bytes!("quit") => {
            println!("Quitting...");
            conn.quit([]);
        }
        _ => ()
    }
}

fn line_desc(line: &Line) -> StrBuf {
    let raw = line.to_raw();
    str::from_utf8_lossy(raw.as_slice()).into_owned()
}
<commit_msg>s/StrBuf/String/<commit_after>/*! Example IRC bot

    This sample bot uses the library to connect to Freenode, under the nickname "rustirclib###"
    (where ### is a random number). It then joins the channel ##rustirclib, says hello, and prints
    to standard output any messages sent to the channel.

    Any incoming message of the form "rustirclib###: some message" will elicit a generic response.

    Any incoming message of the form "!rustirclib### quit" will cause the bot to shut down.
 */

#![crate_id = "github.com/kballard/rust-irclib#ircbot:0.1"]
#![crate_type = "bin"]

#![allow(deprecated_owned_vector)]

extern crate irc;
extern crate rand;

use irc::conn::{Conn, Line, Event, IRCCmd, IRCCode, IRCAction};

use std::str;
use rand::Rng;

fn main() {
    let mut opts = irc::conn::Options::new("chat.freenode.net", irc::conn::DefaultPort);

    let nick = format!("rustirclib{}", rand::task_rng().gen_range(100u, 1000u));
    opts.nick = nick.as_slice();
    match irc::conn::connect(opts, (), |c,e,_| handler(c,e)) {
        Ok(()) => println!("Exiting..."),
        Err(err) => println!("Connection error: {}", err)
    }
}

fn handler(conn: &mut Conn, event: Event) {
    match event {
        irc::conn::Connected => println!("Connected"),
        irc::conn::Disconnected => println!("Disconnected"),
        irc::conn::LineReceived(line) => {
            match line {
                Line{command: IRCCode(1), ..} => {
                    println!("Logged in");
                    // we've logged in
                    conn.join(bytes!("##rustirclib"), [])
                }
                Line{command: IRCCmd(cmd), args, prefix: prefix } => match cmd.as_slice() {
                    "JOIN" if prefix.is_some() => {
                        let prefix = prefix.unwrap();
                        if prefix.nick() != conn.me().nick() {
                            return;
                        }
                        if args.is_empty() {
                            let line = Line{command: IRCCmd("JOIN".into_maybe_owned()),
                                            args: args, prefix: Some(prefix)};
                            println!("ERROR: Invalid JOIN message received: {}", line_desc(&line));
                            return;
                        }
                        let chan = args.move_iter().next().unwrap();
                        conn.privmsg(chan.as_slice(), bytes!("Hello"));
                        let chan = str::from_utf8_lossy(chan.as_slice());
                        println!("JOINED: {}", chan);
                    }
                    "PRIVMSG" | "NOTICE" => {
                        let (src, dst, msg) = match prefix {
                            Some(_) if args.len() == 2 => {
                                let mut args = args;
                                let (dst, msg) = (args.swap_remove(0).unwrap(),
                                                  args.move_iter().next().unwrap());
                                (prefix.as_ref().unwrap().nick(), dst, msg)
                            }
                            _ => {
                                print!("ERROR: Unexpected {} line: ", cmd);
                                let line = Line{command: IRCCmd(cmd), args: args, prefix: prefix};
                                println!("{}", line_desc(&line));
                                return;
                            }
                        };
                        let dsts = str::from_utf8_lossy(dst.as_slice());
                        let srcs = str::from_utf8_lossy(src.as_slice());
                        let msgs = str::from_utf8_lossy(msg.as_slice());
                        println!("<-- {}({}) {}: {}", cmd, dsts, srcs, msgs);
                        handle_privmsg(conn, msg.as_slice(), src, dst.as_slice())
                    }
                    _ => ()
                },
                Line{command: IRCAction(dst), args, prefix } => {
                    let (src, msg) = match prefix {
                        Some(_) if args.len() == 1 => {
                            let msg = args.move_iter().next().unwrap();
                            (prefix.as_ref().unwrap().nick(), msg)
                        }
                        _ => {
                            let line = Line{command: IRCAction(dst), args: args, prefix: prefix};
                            println!("ERROR: Unexpected ACTION line: {}", line_desc(&line));
                            return;
                        }
                    };
                    let dst = str::from_utf8_lossy(dst.as_slice());
                    let src = str::from_utf8_lossy(src.as_slice());
                    let msg = str::from_utf8_lossy(msg.as_slice());
                    println!("<-- PRIVMSG({}) {} {}", dst, src, msg);
                }
                _ => ()
            }
        }
    }
}

fn handle_privmsg(conn: &mut Conn, msg: &[u8], src: &[u8], dst: &[u8]) {
    enum MsgType<'a> {
        DirectedMessage,
        CommandMessage(&'a [u8]),
        NormalMessage
    }
    let msgtype = {
        let nick = conn.me().nick();
        if msg.starts_with(nick) && msg.slice_from(nick.len()).starts_with(bytes!(": ")) &&
        msg.len() > nick.len() + 2 {
            DirectedMessage
        } else if msg.starts_with(bytes!("!")) && msg.slice_from(1).starts_with(nick) &&
                msg.slice_from(nick.len()+1).starts_with(bytes!(" ")) {
            let args = msg.slice_from(conn.me().nick().len()+2);
            let mut argit = args.split(|&b| b == ' ' as u8).skip_while(|v| v.is_empty());
            match argit.next() {
                Some(arg) => CommandMessage(arg),
                None => CommandMessage(bytes!(""))
            }
        } else {
            NormalMessage
        }
    };
    match msgtype {
        DirectedMessage => {
            let reply = if dst == conn.me().nick() { src } else { dst };
            let msg = Vec::from_slice(src).append(bytes!(": Hello"));
            conn.privmsg(reply, msg.as_slice());
            let src = str::from_utf8(conn.me().nick()).unwrap_or("(invalid utf8)");
            let reply = str::from_utf8(reply).unwrap_or("(invalid utf8)");
            let msg = String::from_utf8(msg).unwrap_or_else(|_| "(invalid utf8)".to_strbuf());
            println!("--> PRIVMSG({}) {}: {}", reply, src, msg);
        }
        CommandMessage(cmd) if cmd == bytes!("quit") => {
            println!("Quitting...");
            conn.quit([]);
        }
        _ => ()
    }
}

fn line_desc(line: &Line) -> String {
    let raw = line.to_raw();
    str::from_utf8_lossy(raw.as_slice()).into_owned()
}
<|endoftext|>


<commit_before>use std::collections::BTreeSet;

use anyhow::Result;

use crate::{
    analysis::{
        cfg,
        cfg::{flow::Flow, ChangeBatch, CFG},
        dis,
        dis::{zydis, Target},
    },
    aspace::AddressSpace,
    module::Module,
    VA,
};

// TODO: add cfg_check_noret(module, cfg, va) that optionally marks as noret, if
// valid.

// with the given function address,
// either as the target of a direct or indirect call,
// consider it to be non-returning (such as ExitProcess).
//
// recursively consider its callers,
// possibly pruning instructions after the calls,
// and possibly considering those functions as noret, too.
//
// returns the set of functions newly recognized as noret.
pub fn cfg_mark_noret(module: &Module, cfg: &mut CFG, va: VA) -> Result<BTreeSet<VA>> {
    let mut ret: BTreeSet<VA> = Default::default();
    let mut batch: ChangeBatch = Default::default();

    // the given address is the target to either direct or indirect calls (import).
    // for each of these, remove any fallthrough flows from that call instruction.
    // then, rebuild the CFG.
    let mut callers: Vec<VA> = Default::default();
    for flow in cfg
        .flows
        .flows_by_dst
        .get(&va)
        .unwrap_or(&Default::default())
        .clone()
        .into_iter()
    {
        let src = match flow {
            Flow::Call(Target::Direct(src)) => src,
            Flow::Call(Target::Indirect(src)) => src,
            _ => continue,
        };

        batch.prune_noret_call(src);
        callers.push(src);
    }
    cfg.commit(batch);

    let decoder = dis::get_disassembler(module).expect("invalid disassembler");

    // for each of the call instructions that flow to the given va,
    // search backwards for function starts.
    // these are basic blocks to which call instructions flow.
    //
    // then, search forwards to find "leaf" basic blocks.
    // these are basic blocks with no edge successors.
    // check to see if any of these blocks ends in a ret instruction.
    //
    // if none do, then this is a noret function, too.
    //
    // recurse.
    for call_insn in callers.into_iter() {
        // the basic block that ends with a call to noret function at given va.
        let leaf_block = cfg.basic_blocks.blocks_by_last_address[&call_insn];

        for head in cfg.get_reaches_to(leaf_block) {
            // if the head appears to be the start of a function,
            // by looking at if there are any call flows here.
            // TODO: augment with function database?
            //
            // this is a function that flows to the block ending with a noret call.
            if cfg.flows.flows_by_dst[&head.address]
                .iter()
                .find(|flow| matches!(flow, Flow::Call(_)))
                .is_some()
            {
                // are there any other exit points from this function?
                let is_noret = cfg
                    .get_reaches_from(head.address)
                    .filter(|block| cfg::empty(cfg::edges(&cfg.flows.flows_by_src[&block.address_of_last_insn])))
                    .find(|block| {
                        let mut insn_buf = [0u8; 16];
                        module
                            .address_space
                            .read_into(block.address_of_last_insn, &mut insn_buf)
                            .unwrap();
                        let insn = decoder
                            .decode(&insn_buf)
                            .expect("invalid instruction")
                            .expect("missing instruction");
                        matches!(insn.mnemonic, zydis::Mnemonic::RET)
                    })
                    .is_none();

                if is_noret {
                    log::debug!("noret function: {:#x}", head.address);
                    ret.insert(head.address);
                }
            }
        }
    }

    for &caller in ret.clone().iter() {
        ret.extend(cfg_mark_noret(module, cfg, caller)?);
    }

    return Ok(ret);
}
<commit_msg>noret: better handle tail calls<commit_after>use std::collections::BTreeSet;

use anyhow::Result;

use crate::{
    analysis::{
        cfg,
        cfg::{flow::Flow, ChangeBatch, CFG},
        dis,
        dis::{zydis, Target},
    },
    aspace::AddressSpace,
    module::Module,
    VA,
};

// TODO: add cfg_check_noret(module, cfg, va) that optionally marks as noret, if
// valid.

// with the given function address,
// either as the target of a direct or indirect call,
// consider it to be non-returning (such as ExitProcess).
//
// recursively consider its callers,
// possibly pruning instructions after the calls,
// and possibly considering those functions as noret, too.
//
// returns the set of functions newly recognized as noret.
pub fn cfg_mark_noret(module: &Module, cfg: &mut CFG, va: VA) -> Result<BTreeSet<VA>> {
    let mut ret: BTreeSet<VA> = Default::default();
    let mut batch: ChangeBatch = Default::default();

    // the given address is the target to either direct or indirect calls (import).
    // for each of these, remove any fallthrough flows from that call instruction.
    // then, rebuild the CFG.
    let mut callers: Vec<VA> = Default::default();
    for flow in cfg
        .flows
        .flows_by_dst
        .get(&va)
        .unwrap_or(&Default::default())
        .clone()
        .into_iter()
    {
        let src = match flow {
            Flow::Call(Target::Direct(src)) => src,
            Flow::Call(Target::Indirect(src)) => src,
            // tail call
            Flow::UnconditionalJump(Target::Direct(src)) => src,
            _ => continue,
        };

        batch.prune_noret_call(src);
        callers.push(src);
    }
    cfg.commit(batch);

    let decoder = dis::get_disassembler(module).expect("invalid disassembler");

    // for each of the call instructions that flow to the given va,
    // search backwards for function starts.
    // these are basic blocks to which call instructions flow.
    //
    // then, search forwards to find "leaf" basic blocks.
    // these are basic blocks with no edge successors.
    // check to see if any of these blocks ends in a ret instruction.
    //
    // if none do, then this is a noret function, too.
    //
    // recurse.
    for call_insn in callers.into_iter() {
        // the basic block that ends with a call to noret function at given va.
        let leaf_block = cfg.basic_blocks.blocks_by_last_address[&call_insn];

        for head in cfg.get_reaches_to(leaf_block) {
            // if the head appears to be the start of a function,
            // by looking at if there are any call flows here.
            // TODO: augment with function database?
            //
            // this is a function that flows to the block ending with a noret call.
            if cfg.flows.flows_by_dst[&head.address]
                .iter()
                .find(|flow| matches!(flow, Flow::Call(_)))
                .is_some()
            {
                // are there any other exit points from this function?
                let is_noret = cfg
                    .get_reaches_from(head.address)
                    .filter(|block| cfg::empty(cfg::edges(&cfg.flows.flows_by_src[&block.address_of_last_insn])))
                    .find(|block| {
                        let mut insn_buf = [0u8; 16];
                        module
                            .address_space
                            .read_into(block.address_of_last_insn, &mut insn_buf)
                            .unwrap();
                        let insn = decoder
                            .decode(&insn_buf)
                            .expect("invalid instruction")
                            .expect("missing instruction");
                        matches!(insn.mnemonic, zydis::Mnemonic::RET)
                    })
                    .is_none();

                if is_noret {
                    log::debug!("noret function: {:#x}", head.address);
                    ret.insert(head.address);
                }
            }
        }
    }

    for &caller in ret.clone().iter() {
        ret.extend(cfg_mark_noret(module, cfg, caller)?);
    }

    return Ok(ret);
}
<|endoftext|>


<commit_before>/// Provides implementations of `From<$a> for $b` and `From<$b> for $a` that transmutes the value.
macro_rules! from_transmute {
    { unsafe $a:ty => $b:ty } => {
        from_transmute!{ @impl $a => $b }
        from_transmute!{ @impl $b => $a }
    };
    { @impl $from:ty => $to:ty } => {
        impl core::convert::From<$from> for $to {
            #[inline]
            fn from(value: $from) -> $to {
                unsafe { core::mem::transmute(value) }
            }
        }
    };
}

/// Provides implementations of `From<$generic> for core::arch::{x86, x86_64}::$intel` and
/// vice-versa that transmutes the value.
macro_rules! from_transmute_x86 {
    { unsafe $generic:ty => $intel:ident } => {
        #[cfg(target_arch = "x86")]
        from_transmute! { unsafe $generic => core::arch::x86::$intel }

        #[cfg(target_arch = "x86_64")]
        from_transmute! { unsafe $generic => core::arch::x86_64::$intel }
    }
}

/// Implements common traits on the specified vector `$name`, holding multiple `$lanes` of `$type`.
macro_rules! impl_vector {
    { $name:ident, $type:ty } => {
        impl<const LANES: usize> $name<LANES> where Self: crate::LanesAtMost64 {
            /// Construct a SIMD vector by setting all lanes to the given value.
            pub const fn splat(value: $type) -> Self {
                Self([value; LANES])
            }

            /// Returns a slice containing the entire SIMD vector.
            pub const fn as_slice(&self) -> &[$type] {
                &self.0
            }

            /// Returns a mutable slice containing the entire SIMD vector.
            pub fn as_mut_slice(&mut self) -> &mut [$type] {
                &mut self.0
            }

            /// Converts an array to a SIMD vector.
            pub const fn from_array(array: [$type; LANES]) -> Self {
                Self(array)
            }

            /// Converts a SIMD vector to an array.
            pub const fn to_array(self) -> [$type; LANES] {
                // workaround for rust-lang/rust#80108
                // TODO fix this
                #[cfg(target_arch = "wasm32")]
                {
                    let mut arr = [self.0[0]; LANES];
                    let mut i = 0;
                    while i < LANES {
                        arr[i] = self.0[i];
                        i += 1;
                    }
                    arr
                }

                #[cfg(not(target_arch = "wasm32"))]
                {
                    self.0
                }
            }
        }

        impl<const LANES: usize> Copy for $name<LANES> where Self: crate::LanesAtMost64 {}

        impl<const LANES: usize> Clone for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn clone(&self) -> Self {
                *self
            }
        }

        impl<const LANES: usize> Default for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn default() -> Self {
                Self::splat(<$type>::default())
            }
        }

        impl<const LANES: usize> PartialEq for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn eq(&self, other: &Self) -> bool {
                // TODO use SIMD equality
                self.to_array() == other.to_array()
            }
        }

        impl<const LANES: usize> PartialOrd for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
                // TODO use SIMD equalitya
                self.to_array().partial_cmp(other.as_ref())
            }
        }

        // array references
        impl<const LANES: usize> AsRef<[$type; LANES]> for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn as_ref(&self) -> &[$type; LANES] {
                &self.0
            }
        }

        impl<const LANES: usize> AsMut<[$type; LANES]> for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn as_mut(&mut self) -> &mut [$type; LANES] {
                &mut self.0
            }
        }

        // slice references
        impl<const LANES: usize> AsRef<[$type]> for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn as_ref(&self) -> &[$type] {
                &self.0
            }
        }

        impl<const LANES: usize> AsMut<[$type]> for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn as_mut(&mut self) -> &mut [$type] {
                &mut self.0
            }
        }

        // vector/array conversion
        impl<const LANES: usize> From<[$type; LANES]> for $name<LANES> where Self: crate::LanesAtMost64 {
            fn from(array: [$type; LANES]) -> Self {
                Self(array)
            }
        }

        impl <const LANES: usize> From<$name<LANES>> for [$type; LANES] where $name<LANES>: crate::LanesAtMost64 {
            fn from(vector: $name<LANES>) -> Self {
                vector.to_array()
            }
        }

        // splat
        impl<const LANES: usize> From<$type> for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn from(value: $type) -> Self {
                Self::splat(value)
            }
        }

        impl_shuffle_2pow_lanes!{ $name }
    }
}

/// Implements additional integer traits (Eq, Ord, Hash) on the specified vector `$name`, holding multiple `$lanes` of `$type`.
macro_rules! impl_integer_vector {
    { $name:ident, $type:ty } => {
        impl_vector! { $name, $type }

        impl<const LANES: usize> Eq for $name<LANES> where Self: crate::LanesAtMost64 {}

        impl<const LANES: usize> Ord for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn cmp(&self, other: &Self) -> core::cmp::Ordering {
                // TODO use SIMD cmp
                self.to_array().cmp(other.as_ref())
            }
        }

        impl<const LANES: usize> core::hash::Hash for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn hash<H>(&self, state: &mut H)
            where
                H: core::hash::Hasher
            {
                self.as_slice().hash(state)
            }
        }
    }
}

/// Implements inherent methods for a float vector `$name` containing multiple
/// `$lanes` of float `$type`, which uses `$bits_ty` as its binary
/// representation. Called from `define_float_vector!`.
macro_rules! impl_float_vector {
    { $name:ident, $type:ty, $bits_ty:ident } => {
        impl_vector! { $name, $type }

        impl<const LANES: usize> $name<LANES>
        where
            Self: crate::LanesAtMost64,
            crate::$bits_ty<LANES>: crate::LanesAtMost64,
        {
            /// Raw transmutation to an unsigned integer vector type with the
            /// same size and number of lanes.
            #[inline]
            pub fn to_bits(self) -> crate::$bits_ty<LANES> {
                assert_eq!(core::mem::size_of::<Self>(), core::mem::size_of::<crate::$bits_ty<LANES>>());
                unsafe { core::mem::transmute_copy(&self) }
            }

            /// Raw transmutation from an unsigned integer vector type with the
            /// same size and number of lanes.
            #[inline]
            pub fn from_bits(bits: crate::$bits_ty<LANES>) -> Self {
                assert_eq!(core::mem::size_of::<Self>(), core::mem::size_of::<crate::$bits_ty<LANES>>());
                unsafe { core::mem::transmute_copy(&bits) }
            }

            /// Produces a vector where every lane has the absolute value of the
            /// equivalently-indexed lane in `self`.
            #[inline]
            pub fn abs(self) -> Self {
                let no_sign = crate::$bits_ty::splat(!0 >> 1);
                Self::from_bits(self.to_bits() & no_sign)
            }
        }
    };
}
<commit_msg>Remove From<Scalar> for SimdTy impl<commit_after>/// Provides implementations of `From<$a> for $b` and `From<$b> for $a` that transmutes the value.
macro_rules! from_transmute {
    { unsafe $a:ty => $b:ty } => {
        from_transmute!{ @impl $a => $b }
        from_transmute!{ @impl $b => $a }
    };
    { @impl $from:ty => $to:ty } => {
        impl core::convert::From<$from> for $to {
            #[inline]
            fn from(value: $from) -> $to {
                unsafe { core::mem::transmute(value) }
            }
        }
    };
}

/// Provides implementations of `From<$generic> for core::arch::{x86, x86_64}::$intel` and
/// vice-versa that transmutes the value.
macro_rules! from_transmute_x86 {
    { unsafe $generic:ty => $intel:ident } => {
        #[cfg(target_arch = "x86")]
        from_transmute! { unsafe $generic => core::arch::x86::$intel }

        #[cfg(target_arch = "x86_64")]
        from_transmute! { unsafe $generic => core::arch::x86_64::$intel }
    }
}

/// Implements common traits on the specified vector `$name`, holding multiple `$lanes` of `$type`.
macro_rules! impl_vector {
    { $name:ident, $type:ty } => {
        impl<const LANES: usize> $name<LANES> where Self: crate::LanesAtMost64 {
            /// Construct a SIMD vector by setting all lanes to the given value.
            pub const fn splat(value: $type) -> Self {
                Self([value; LANES])
            }

            /// Returns a slice containing the entire SIMD vector.
            pub const fn as_slice(&self) -> &[$type] {
                &self.0
            }

            /// Returns a mutable slice containing the entire SIMD vector.
            pub fn as_mut_slice(&mut self) -> &mut [$type] {
                &mut self.0
            }

            /// Converts an array to a SIMD vector.
            pub const fn from_array(array: [$type; LANES]) -> Self {
                Self(array)
            }

            /// Converts a SIMD vector to an array.
            pub const fn to_array(self) -> [$type; LANES] {
                // workaround for rust-lang/rust#80108
                // TODO fix this
                #[cfg(target_arch = "wasm32")]
                {
                    let mut arr = [self.0[0]; LANES];
                    let mut i = 0;
                    while i < LANES {
                        arr[i] = self.0[i];
                        i += 1;
                    }
                    arr
                }

                #[cfg(not(target_arch = "wasm32"))]
                {
                    self.0
                }
            }
        }

        impl<const LANES: usize> Copy for $name<LANES> where Self: crate::LanesAtMost64 {}

        impl<const LANES: usize> Clone for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn clone(&self) -> Self {
                *self
            }
        }

        impl<const LANES: usize> Default for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn default() -> Self {
                Self::splat(<$type>::default())
            }
        }

        impl<const LANES: usize> PartialEq for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn eq(&self, other: &Self) -> bool {
                // TODO use SIMD equality
                self.to_array() == other.to_array()
            }
        }

        impl<const LANES: usize> PartialOrd for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn partial_cmp(&self, other: &Self) -> Option<core::cmp::Ordering> {
                // TODO use SIMD equalitya
                self.to_array().partial_cmp(other.as_ref())
            }
        }

        // array references
        impl<const LANES: usize> AsRef<[$type; LANES]> for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn as_ref(&self) -> &[$type; LANES] {
                &self.0
            }
        }

        impl<const LANES: usize> AsMut<[$type; LANES]> for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn as_mut(&mut self) -> &mut [$type; LANES] {
                &mut self.0
            }
        }

        // slice references
        impl<const LANES: usize> AsRef<[$type]> for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn as_ref(&self) -> &[$type] {
                &self.0
            }
        }

        impl<const LANES: usize> AsMut<[$type]> for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn as_mut(&mut self) -> &mut [$type] {
                &mut self.0
            }
        }

        // vector/array conversion
        impl<const LANES: usize> From<[$type; LANES]> for $name<LANES> where Self: crate::LanesAtMost64 {
            fn from(array: [$type; LANES]) -> Self {
                Self(array)
            }
        }

        impl <const LANES: usize> From<$name<LANES>> for [$type; LANES] where $name<LANES>: crate::LanesAtMost64 {
            fn from(vector: $name<LANES>) -> Self {
                vector.to_array()
            }
        }

        impl_shuffle_2pow_lanes!{ $name }
    }
}

/// Implements additional integer traits (Eq, Ord, Hash) on the specified vector `$name`, holding multiple `$lanes` of `$type`.
macro_rules! impl_integer_vector {
    { $name:ident, $type:ty } => {
        impl_vector! { $name, $type }

        impl<const LANES: usize> Eq for $name<LANES> where Self: crate::LanesAtMost64 {}

        impl<const LANES: usize> Ord for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn cmp(&self, other: &Self) -> core::cmp::Ordering {
                // TODO use SIMD cmp
                self.to_array().cmp(other.as_ref())
            }
        }

        impl<const LANES: usize> core::hash::Hash for $name<LANES> where Self: crate::LanesAtMost64 {
            #[inline]
            fn hash<H>(&self, state: &mut H)
            where
                H: core::hash::Hasher
            {
                self.as_slice().hash(state)
            }
        }
    }
}

/// Implements inherent methods for a float vector `$name` containing multiple
/// `$lanes` of float `$type`, which uses `$bits_ty` as its binary
/// representation. Called from `define_float_vector!`.
macro_rules! impl_float_vector {
    { $name:ident, $type:ty, $bits_ty:ident } => {
        impl_vector! { $name, $type }

        impl<const LANES: usize> $name<LANES>
        where
            Self: crate::LanesAtMost64,
            crate::$bits_ty<LANES>: crate::LanesAtMost64,
        {
            /// Raw transmutation to an unsigned integer vector type with the
            /// same size and number of lanes.
            #[inline]
            pub fn to_bits(self) -> crate::$bits_ty<LANES> {
                assert_eq!(core::mem::size_of::<Self>(), core::mem::size_of::<crate::$bits_ty<LANES>>());
                unsafe { core::mem::transmute_copy(&self) }
            }

            /// Raw transmutation from an unsigned integer vector type with the
            /// same size and number of lanes.
            #[inline]
            pub fn from_bits(bits: crate::$bits_ty<LANES>) -> Self {
                assert_eq!(core::mem::size_of::<Self>(), core::mem::size_of::<crate::$bits_ty<LANES>>());
                unsafe { core::mem::transmute_copy(&bits) }
            }

            /// Produces a vector where every lane has the absolute value of the
            /// equivalently-indexed lane in `self`.
            #[inline]
            pub fn abs(self) -> Self {
                let no_sign = crate::$bits_ty::splat(!0 >> 1);
                Self::from_bits(self.to_bits() & no_sign)
            }
        }
    };
}
<|endoftext|>