<commit_before>// Copyright 2021 Vector 35 Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use binaryninjacore_sys::*;

use std::result;

use crate::rc::*;
use crate::string::*;

pub type Result<R> = result::Result<R, ()>;

#[derive(PartialEq, Eq, Hash)]
pub struct BackgroundTask {
    pub(crate) handle: *mut BNBackgroundTask,
}

impl BackgroundTask {
    pub(crate) unsafe fn from_raw(handle: *mut BNBackgroundTask) -> Self {
        debug_assert!(!handle.is_null());

        Self { handle }
    }

    pub fn new<S: BnStrCompatible>(initial_text: S, can_cancel: bool) -> Result<Ref<Self>> {
        let text = initial_text.as_bytes_with_nul();

        let handle = unsafe { BNBeginBackgroundTask(text.as_ref().as_ptr() as *mut _, can_cancel) };

        if handle.is_null() {
            return Err(());
        }

        unsafe { Ok(Ref::new(Self { handle })) }
    }

    pub fn can_cancel(&self) -> bool {
        unsafe { BNCanCancelBackgroundTask(self.handle) }
    }

    pub fn is_cancelled(&self) -> bool {
        unsafe { BNIsBackgroundTaskCancelled(self.handle) }
    }

    pub fn is_finished(&self) -> bool {
        unsafe { BNIsBackgroundTaskFinished(self.handle) }
    }

    pub fn get_progress_text(&self) -> BnString {
        unsafe { BnString::from_raw(BNGetBackgroundTaskProgressText(self.handle)) }
    }

    pub fn cancel(&mut self) {
        unsafe { BNCancelBackgroundTask(self.handle) }
    }

    pub fn finish(&self) {
        unsafe { BNFinishBackgroundTask(self.handle) }
    }

    pub fn set_progress_text<S: BnStrCompatible>(&self, text: S) {
        let progress_text = text.as_bytes_with_nul();

        unsafe {
            BNSetBackgroundTaskProgressText(self.handle, progress_text.as_ref().as_ptr() as *mut _)
        }
    }

    pub fn running_tasks() -> Array<BackgroundTask> {
        unsafe {
            let mut count = 0;
            let handles = BNGetRunningBackgroundTasks(&mut count);

            Array::new(handles, count, ())
        }
    }
}

unsafe impl RefCountable for BackgroundTask {
    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
        Ref::new(Self {
            handle: BNNewBackgroundTaskReference(handle.handle),
        })
    }

    unsafe fn dec_ref(handle: &Self) {
        BNFreeBackgroundTask(handle.handle);
    }
}

unsafe impl CoreOwnedArrayProvider for BackgroundTask {
    type Raw = *mut BNBackgroundTask;
    type Context = ();

    unsafe fn free(raw: *mut *mut BNBackgroundTask, count: usize, _context: &()) {
        BNFreeBackgroundTaskList(raw, count);
    }
}

unsafe impl<'a> CoreOwnedArrayWrapper<'a> for BackgroundTask {
    type Wrapped = Guard<'a, BackgroundTask>;

    unsafe fn wrap_raw(
        raw: &'a *mut BNBackgroundTask,
        context: &'a (),
    ) -> Guard<'a, BackgroundTask> {
        Guard::new(BackgroundTask::from_raw(*raw), context)
    }
}

impl AsRef<BackgroundTask> for BackgroundTask {
    fn as_ref(&self) -> &Self {
        self
    }
}

impl ToOwned for BackgroundTask {
    type Owned = Ref<Self>;

    fn to_owned(&self) -> Self::Owned {
        unsafe { RefCountable::inc_ref(self) }
    }
}

unsafe impl Send for BackgroundTask {}
unsafe impl Sync for BackgroundTask {}
<commit_msg>More mut<commit_after>// Copyright 2021 Vector 35 Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

use binaryninjacore_sys::*;

use std::result;

use crate::rc::*;
use crate::string::*;

pub type Result<R> = result::Result<R, ()>;

#[derive(PartialEq, Eq, Hash)]
pub struct BackgroundTask {
    pub(crate) handle: *mut BNBackgroundTask,
}

impl BackgroundTask {
    pub(crate) unsafe fn from_raw(handle: *mut BNBackgroundTask) -> Self {
        debug_assert!(!handle.is_null());

        Self { handle }
    }

    pub fn new<S: BnStrCompatible>(initial_text: S, can_cancel: bool) -> Result<Ref<Self>> {
        let text = initial_text.as_bytes_with_nul();

        let handle = unsafe { BNBeginBackgroundTask(text.as_ref().as_ptr() as *mut _, can_cancel) };

        if handle.is_null() {
            return Err(());
        }

        unsafe { Ok(Ref::new(Self { handle })) }
    }

    pub fn can_cancel(&self) -> bool {
        unsafe { BNCanCancelBackgroundTask(self.handle) }
    }

    pub fn is_cancelled(&self) -> bool {
        unsafe { BNIsBackgroundTaskCancelled(self.handle) }
    }

    pub fn is_finished(&self) -> bool {
        unsafe { BNIsBackgroundTaskFinished(self.handle) }
    }

    pub fn get_progress_text(&self) -> BnString {
        unsafe { BnString::from_raw(BNGetBackgroundTaskProgressText(self.handle)) }
    }

    pub fn cancel(&mut self) {
        unsafe { BNCancelBackgroundTask(self.handle) }
    }

    pub fn finish(&mut self) {
        unsafe { BNFinishBackgroundTask(self.handle) }
    }

    pub fn set_progress_text<S: BnStrCompatible>(&mut self, text: S) {
        let progress_text = text.as_bytes_with_nul();

        unsafe {
            BNSetBackgroundTaskProgressText(self.handle, progress_text.as_ref().as_ptr() as *mut _)
        }
    }

    pub fn running_tasks() -> Array<BackgroundTask> {
        unsafe {
            let mut count = 0;
            let handles = BNGetRunningBackgroundTasks(&mut count);

            Array::new(handles, count, ())
        }
    }
}

unsafe impl RefCountable for BackgroundTask {
    unsafe fn inc_ref(handle: &Self) -> Ref<Self> {
        Ref::new(Self {
            handle: BNNewBackgroundTaskReference(handle.handle),
        })
    }

    unsafe fn dec_ref(handle: &Self) {
        BNFreeBackgroundTask(handle.handle);
    }
}

unsafe impl CoreOwnedArrayProvider for BackgroundTask {
    type Raw = *mut BNBackgroundTask;
    type Context = ();

    unsafe fn free(raw: *mut *mut BNBackgroundTask, count: usize, _context: &()) {
        BNFreeBackgroundTaskList(raw, count);
    }
}

unsafe impl<'a> CoreOwnedArrayWrapper<'a> for BackgroundTask {
    type Wrapped = Guard<'a, BackgroundTask>;

    unsafe fn wrap_raw(
        raw: &'a *mut BNBackgroundTask,
        context: &'a (),
    ) -> Guard<'a, BackgroundTask> {
        Guard::new(BackgroundTask::from_raw(*raw), context)
    }
}

impl AsRef<BackgroundTask> for BackgroundTask {
    fn as_ref(&self) -> &Self {
        self
    }
}

impl ToOwned for BackgroundTask {
    type Owned = Ref<Self>;

    fn to_owned(&self) -> Self::Owned {
        unsafe { RefCountable::inc_ref(self) }
    }
}

unsafe impl Send for BackgroundTask {}
unsafe impl Sync for BackgroundTask {}
<|endoftext|>


<commit_before>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use common::{Location, WithLocation};
use graphql_ir::{
    FragmentSpread, InlineFragment, LinkedField, Program, ScalarField, Selection, Transformed,
    TransformedValue, Transformer,
};

use interner::{Intern, StringKey};
use schema::{FieldID, InterfaceID, ObjectID, Type};
use std::collections::HashMap;
use std::sync::Arc;
///
/// A transform that adds an `id` field on any type that has an id field but
/// where there is no unaliased `id` selection.
///
pub fn generate_id_field<'s>(program: &Program<'s>) -> Program<'s> {
    let mut transform = GenerateIDFieldTransform::new(program);
    transform
        .transform_program(program)
        .replace_or_else(|| program.clone())
}

struct GenerateIDFieldTransform<'s> {
    program: &'s Program<'s>,
    id_name: StringKey,
    node_interface_id: InterfaceID,
    cache: HashMap<Type, Option<FieldID>>,
    node_id_field_id: FieldID,
}

impl<'s> Transformer for GenerateIDFieldTransform<'s> {
    const NAME: &'static str = "GenerateIDFieldTransform";
    const VISIT_ARGUMENTS: bool = false;
    const VISIT_DIRECTIVES: bool = false;

    fn transform_linked_field(&mut self, field: &LinkedField) -> Transformed<Arc<LinkedField>> {
        let selections = self.transform_selections(&field.selections);

        let next_selections = if self.has_unaliased_id_field(&field.selections) {
            selections
        } else {
            let schema = self.program.schema();
            let type_ = self
                .program
                .schema()
                .field(field.definition.item)
                .type_
                .inner();
            match type_ {
                Type::Object(id) => {
                    let object = schema.object(id);
                    if let Some(id_field_id) = self.get_id_field_id(type_, &object.fields) {
                        let mut next_selections =
                            selections.replace_or_else(|| field.selections.clone());
                        next_selections
                            .push(self.create_id_selection(field.definition.location, id_field_id));
                        TransformedValue::Replace(next_selections)
                    } else {
                        selections
                    }
                }
                Type::Interface(id) => {
                    let interface = schema.interface(id);
                    if let Some(id_field_id) = self.get_id_field_id(type_, &interface.fields) {
                        let mut next_selections =
                            selections.replace_or_else(|| field.selections.clone());
                        next_selections
                            .push(self.create_id_selection(field.definition.location, id_field_id));
                        TransformedValue::Replace(next_selections)
                    } else {
                        let mut inline_fragments = self.create_id_inline_fragments(
                            field.definition.location,
                            &interface.implementors,
                        );
                        if inline_fragments.is_empty() {
                            selections
                        } else {
                            if let TransformedValue::Replace(selections) = selections {
                                inline_fragments.extend(selections.into_iter())
                            } else {
                                inline_fragments.extend(field.selections.iter().cloned());
                            }
                            TransformedValue::Replace(inline_fragments)
                        }
                    }
                }
                Type::Union(id) => {
                    let union = schema.union(id);
                    let mut inline_fragments =
                        self.create_id_inline_fragments(field.definition.location, &union.members);

                    if inline_fragments.is_empty() {
                        selections
                    } else {
                        if let TransformedValue::Replace(selections) = selections {
                            inline_fragments.extend(selections.into_iter())
                        } else {
                            inline_fragments.extend(field.selections.iter().cloned());
                        }
                        TransformedValue::Replace(inline_fragments)
                    }
                }
                _ => selections,
            }
        };
        match next_selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => Transformed::Replace(Arc::new(LinkedField {
                alias: field.alias,
                definition: field.definition,
                arguments: field.arguments.clone(),
                directives: field.directives.clone(),
                selections,
            })),
        }
    }

    fn transform_fragment_spread(
        &mut self,
        _spread: &FragmentSpread,
    ) -> Transformed<Arc<FragmentSpread>> {
        Transformed::Keep
    }

    fn transform_scalar_field(&mut self, _field: &ScalarField) -> Transformed<Arc<ScalarField>> {
        Transformed::Keep
    }
}

impl<'s> GenerateIDFieldTransform<'s> {
    fn new(program: &'s Program<'s>) -> Self {
        let schema = program.schema();
        let node_interface_id = match schema.get_type("Node".intern()) {
            Some(Type::Interface(id)) => id,
            _ => panic!("Expected the schema to contain an interface named `Node`."),
        };
        let node_interface = schema.interface(node_interface_id);
        let id_key = "id".intern();
        let node_id_field_id = *node_interface
            .fields
            .iter()
            .find(|&&id| schema.field(id).name == id_key)
            .expect("Expected `Node` to contain a field named `id`.");
        Self {
            program,
            id_name: schema.field(node_id_field_id).name,
            node_interface_id,
            cache: Default::default(),
            node_id_field_id,
        }
    }

    fn has_unaliased_id_field(&self, selections: &[Selection]) -> bool {
        selections.iter().any(|x| match x {
            Selection::ScalarField(child) => {
                child.alias.is_none()
                    && self.program.schema().field(child.definition.item).name == self.id_name
            }
            _ => false,
        })
    }

    fn get_id_field_id(&mut self, type_: Type, fields: &[FieldID]) -> Option<FieldID> {
        match self.cache.get(&type_) {
            Some(result) => *result,
            None => {
                for id in fields {
                    let field = self.program.schema().field(*id);
                    if field.name == self.id_name
                        && self.program.schema().is_id(field.type_.inner())
                    {
                        let result = Some(*id);
                        self.cache.insert(type_, result);
                        return result;
                    }
                }
                self.cache.insert(type_, None);
                None
            }
        }
    }

    /// For interfaces and unions: generate a `... on Node { id }`
    /// fragment if *any* concrete type implements Node. Then generate a
    /// `... on PossibleType { id }` for every concrete type that does *not*
    /// implement `Node`
    fn create_id_inline_fragments(
        &mut self,
        location: Location,
        concrete_ids: &[ObjectID],
    ) -> Vec<Selection> {
        let mut next_selections = vec![];
        let mut should_generate_node = false;

        for object_id in concrete_ids {
            let object = self.program.schema().object(*object_id);
            let implements_node = object
                .interfaces
                .iter()
                .any(|interface_id| interface_id == &self.node_interface_id);
            if implements_node {
                should_generate_node = true;
            } else if let Some(id_field_id) =
                self.get_id_field_id(Type::Object(*object_id), &object.fields)
            {
                next_selections.push(Selection::InlineFragment(self.create_inline_id_fragment(
                    location,
                    Type::Object(*object_id),
                    id_field_id,
                )));
            }
        }

        if should_generate_node {
            next_selections.push(Selection::InlineFragment(self.create_inline_id_fragment(
                location,
                Type::Interface(self.node_interface_id),
                self.node_id_field_id,
            )));
        }

        next_selections
    }

    fn create_id_selection(&self, location: Location, id_field_id: FieldID) -> Selection {
        Selection::ScalarField(Arc::new(ScalarField {
            alias: None,
            definition: WithLocation::new(location, id_field_id),
            arguments: Default::default(),
            directives: Default::default(),
        }))
    }

    fn create_inline_id_fragment(
        &self,
        location: Location,
        type_: Type,
        id_field_id: FieldID,
    ) -> Arc<InlineFragment> {
        Arc::new(InlineFragment {
            type_condition: Some(type_),
            directives: Default::default(),
            selections: vec![self.create_id_selection(location, id_field_id)],
        })
    }
}
<commit_msg>make generate id transform a noop if no Node type is defined<commit_after>/*
 * Copyright (c) Facebook, Inc. and its affiliates.
 *
 * This source code is licensed under the MIT license found in the
 * LICENSE file in the root directory of this source tree.
 */

use common::{Location, WithLocation};
use graphql_ir::{
    FragmentSpread, InlineFragment, LinkedField, Program, ScalarField, Selection, Transformed,
    TransformedValue, Transformer,
};

use interner::{Intern, StringKey};
use schema::{FieldID, InterfaceID, ObjectID, Type};
use std::collections::HashMap;
use std::sync::Arc;

/// A transform that adds an `id` field on any type that has an id field but
/// where there is no unaliased `id` selection.
pub fn generate_id_field<'s>(program: &Program<'s>) -> Program<'s> {
    let mut transform = GenerateIDFieldTransform::new(program);
    transform
        .transform_program(program)
        .replace_or_else(|| program.clone())
}

struct GenerateIDFieldTransform<'s> {
    program: &'s Program<'s>,
    id_name: StringKey,
    node_interface: Option<NodeInterface>,
    cache: HashMap<Type, Option<FieldID>>,
}

/// If the schema defines a `Node` interface, this contains information on that interface.
struct NodeInterface {
    id: InterfaceID,
    id_field: FieldID,
}

impl<'s> Transformer for GenerateIDFieldTransform<'s> {
    const NAME: &'static str = "GenerateIDFieldTransform";
    const VISIT_ARGUMENTS: bool = false;
    const VISIT_DIRECTIVES: bool = false;

    fn transform_linked_field(&mut self, field: &LinkedField) -> Transformed<Arc<LinkedField>> {
        let selections = self.transform_selections(&field.selections);

        let next_selections = if self.has_unaliased_id_field(&field.selections) {
            selections
        } else {
            let schema = self.program.schema();
            let type_ = self
                .program
                .schema()
                .field(field.definition.item)
                .type_
                .inner();
            match type_ {
                Type::Object(id) => {
                    let object = schema.object(id);
                    if let Some(id_field_id) = self.get_id_field_id(type_, &object.fields) {
                        let mut next_selections =
                            selections.replace_or_else(|| field.selections.clone());
                        next_selections
                            .push(self.create_id_selection(field.definition.location, id_field_id));
                        TransformedValue::Replace(next_selections)
                    } else {
                        selections
                    }
                }
                Type::Interface(id) => {
                    let interface = schema.interface(id);
                    if let Some(id_field_id) = self.get_id_field_id(type_, &interface.fields) {
                        let mut next_selections =
                            selections.replace_or_else(|| field.selections.clone());
                        next_selections
                            .push(self.create_id_selection(field.definition.location, id_field_id));
                        TransformedValue::Replace(next_selections)
                    } else {
                        let mut inline_fragments = self.create_id_inline_fragments(
                            field.definition.location,
                            &interface.implementors,
                        );
                        if inline_fragments.is_empty() {
                            selections
                        } else {
                            if let TransformedValue::Replace(selections) = selections {
                                inline_fragments.extend(selections.into_iter())
                            } else {
                                inline_fragments.extend(field.selections.iter().cloned());
                            }
                            TransformedValue::Replace(inline_fragments)
                        }
                    }
                }
                Type::Union(id) => {
                    let union = schema.union(id);
                    let mut inline_fragments =
                        self.create_id_inline_fragments(field.definition.location, &union.members);

                    if inline_fragments.is_empty() {
                        selections
                    } else {
                        if let TransformedValue::Replace(selections) = selections {
                            inline_fragments.extend(selections.into_iter())
                        } else {
                            inline_fragments.extend(field.selections.iter().cloned());
                        }
                        TransformedValue::Replace(inline_fragments)
                    }
                }
                _ => selections,
            }
        };
        match next_selections {
            TransformedValue::Keep => Transformed::Keep,
            TransformedValue::Replace(selections) => Transformed::Replace(Arc::new(LinkedField {
                alias: field.alias,
                definition: field.definition,
                arguments: field.arguments.clone(),
                directives: field.directives.clone(),
                selections,
            })),
        }
    }

    fn transform_fragment_spread(
        &mut self,
        _spread: &FragmentSpread,
    ) -> Transformed<Arc<FragmentSpread>> {
        Transformed::Keep
    }

    fn transform_scalar_field(&mut self, _field: &ScalarField) -> Transformed<Arc<ScalarField>> {
        Transformed::Keep
    }
}

impl<'s> GenerateIDFieldTransform<'s> {
    fn new(program: &'s Program<'s>) -> Self {
        let id_name = "id".intern();

        let schema = program.schema();
        let node_interface = match schema.get_type("Node".intern()) {
            Some(Type::Interface(node_interface_id)) => {
                let node_interface = schema.interface(node_interface_id);
                let id_field = *node_interface
                    .fields
                    .iter()
                    .find(|&&id| schema.field(id).name == id_name)
                    .expect("Expected `Node` to contain a field named `id`.");

                Some(NodeInterface {
                    id: node_interface_id,
                    id_field,
                })
            }
            _ => None,
        };

        Self {
            program,
            id_name,
            node_interface,
            cache: Default::default(),
        }
    }

    fn has_unaliased_id_field(&self, selections: &[Selection]) -> bool {
        selections.iter().any(|x| match x {
            Selection::ScalarField(child) => {
                child.alias.is_none()
                    && self.program.schema().field(child.definition.item).name == self.id_name
            }
            _ => false,
        })
    }

    fn get_id_field_id(&mut self, type_: Type, fields: &[FieldID]) -> Option<FieldID> {
        match self.cache.get(&type_) {
            Some(result) => *result,
            None => {
                for id in fields {
                    let field = self.program.schema().field(*id);
                    if field.name == self.id_name
                        && self.program.schema().is_id(field.type_.inner())
                    {
                        let result = Some(*id);
                        self.cache.insert(type_, result);
                        return result;
                    }
                }
                self.cache.insert(type_, None);
                None
            }
        }
    }

    /// For interfaces and unions: generate a `... on Node { id }`
    /// fragment if *any* concrete type implements Node. Then generate a
    /// `... on PossibleType { id }` for every concrete type that does *not*
    /// implement `Node`
    fn create_id_inline_fragments(
        &mut self,
        location: Location,
        concrete_ids: &[ObjectID],
    ) -> Vec<Selection> {
        let mut next_selections = vec![];
        let mut should_generate_node = false;

        for object_id in concrete_ids {
            let object = self.program.schema().object(*object_id);
            let implements_node = if let Some(ref node_interface) = self.node_interface {
                object
                    .interfaces
                    .iter()
                    .any(|&interface_id| interface_id == node_interface.id)
            } else {
                false
            };
            if implements_node {
                should_generate_node = true;
            } else if let Some(id_field_id) =
                self.get_id_field_id(Type::Object(*object_id), &object.fields)
            {
                next_selections.push(Selection::InlineFragment(self.create_inline_id_fragment(
                    location,
                    Type::Object(*object_id),
                    id_field_id,
                )));
            }
        }

        if should_generate_node {
            // This should not happen because we can only set
            // `should_generate_node` to true, if this is Some.
            let node_interface = self.node_interface.as_ref().unwrap();
            next_selections.push(Selection::InlineFragment(self.create_inline_id_fragment(
                location,
                Type::Interface(node_interface.id),
                node_interface.id_field,
            )));
        }

        next_selections
    }

    fn create_id_selection(&self, location: Location, id_field_id: FieldID) -> Selection {
        Selection::ScalarField(Arc::new(ScalarField {
            alias: None,
            definition: WithLocation::new(location, id_field_id),
            arguments: Default::default(),
            directives: Default::default(),
        }))
    }

    fn create_inline_id_fragment(
        &self,
        location: Location,
        type_: Type,
        id_field_id: FieldID,
    ) -> Arc<InlineFragment> {
        Arc::new(InlineFragment {
            type_condition: Some(type_),
            directives: Default::default(),
            selections: vec![self.create_id_selection(location, id_field_id)],
        })
    }
}
<|endoftext|>


<commit_before>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! Rust helpers for Gecko's `nsCSSShadowItem`.

use app_units::Au;
use gecko::values::{convert_rgba_to_nscolor, convert_nscolor_to_rgba};
use gecko_bindings::structs::nsCSSShadowItem;
use values::computed::Color;
use values::computed::effects::{BoxShadow, SimpleShadow};

impl nsCSSShadowItem {
    /// Sets this item from the given box shadow.
    #[inline]
    pub fn set_from_box_shadow(&mut self, shadow: BoxShadow) {
        self.set_from_simple_shadow(shadow.base);
        self.mSpread = shadow.spread.0;
        self.mInset = shadow.inset;
    }

    /// Returns this item as a box shadow.
    #[inline]
    pub fn to_box_shadow(&self) -> BoxShadow {
        BoxShadow {
            base: self.extract_simple_shadow(),
            spread: Au(self.mSpread),
            inset: self.mInset,
        }
    }

    /// Sets this item from the given simple shadow.
    #[inline]
    pub fn set_from_simple_shadow(&mut self, shadow: SimpleShadow) {
        self.mXOffset = shadow.horizontal.0;
        self.mYOffset = shadow.vertical.0;
        self.mRadius = shadow.blur.value();
        self.mSpread = 0;
        self.mInset = false;
        if shadow.color.is_currentcolor() {
            // TODO handle currentColor
            // https://bugzilla.mozilla.org/show_bug.cgi?id=760345
            self.mHasColor = false;
            self.mColor = 0;
        } else {
            self.mHasColor = true;
            self.mColor = convert_rgba_to_nscolor(&shadow.color.color);
        }
    }

    /// Gets a simple shadow from this item.
    #[inline]
    fn extract_simple_shadow(&self) -> SimpleShadow {
        SimpleShadow {
            color: Color::rgba(convert_nscolor_to_rgba(self.mColor)),
            horizontal: Au(self.mXOffset),
            vertical: Au(self.mYOffset),
            blur: Au(self.mRadius).into(),
        }
    }

    /// Returns this item as a simple shadow.
    #[inline]
    pub fn to_simple_shadow(&self) -> SimpleShadow {
        debug_assert_eq!(self.mSpread, 0);
        debug_assert_eq!(self.mInset, false);
        self.extract_simple_shadow()
    }
}
<commit_msg>Convert gecko's currentcolor to servo one.<commit_after>/* This Source Code Form is subject to the terms of the Mozilla Public
 * License, v. 2.0. If a copy of the MPL was not distributed with this
 * file, You can obtain one at http://mozilla.org/MPL/2.0/. */

//! Rust helpers for Gecko's `nsCSSShadowItem`.

use app_units::Au;
use gecko::values::{convert_rgba_to_nscolor, convert_nscolor_to_rgba};
use gecko_bindings::structs::nsCSSShadowItem;
use values::computed::Color;
use values::computed::effects::{BoxShadow, SimpleShadow};

impl nsCSSShadowItem {
    /// Sets this item from the given box shadow.
    #[inline]
    pub fn set_from_box_shadow(&mut self, shadow: BoxShadow) {
        self.set_from_simple_shadow(shadow.base);
        self.mSpread = shadow.spread.0;
        self.mInset = shadow.inset;
    }

    /// Returns this item as a box shadow.
    #[inline]
    pub fn to_box_shadow(&self) -> BoxShadow {
        BoxShadow {
            base: self.extract_simple_shadow(),
            spread: Au(self.mSpread),
            inset: self.mInset,
        }
    }

    /// Sets this item from the given simple shadow.
    #[inline]
    pub fn set_from_simple_shadow(&mut self, shadow: SimpleShadow) {
        self.mXOffset = shadow.horizontal.0;
        self.mYOffset = shadow.vertical.0;
        self.mRadius = shadow.blur.value();
        self.mSpread = 0;
        self.mInset = false;
        if shadow.color.is_currentcolor() {
            // TODO handle currentColor
            // https://bugzilla.mozilla.org/show_bug.cgi?id=760345
            self.mHasColor = false;
            self.mColor = 0;
        } else {
            self.mHasColor = true;
            self.mColor = convert_rgba_to_nscolor(&shadow.color.color);
        }
    }

    #[inline]
    fn extract_color(&self) -> Color {
        if self.mHasColor {
            Color::rgba(convert_nscolor_to_rgba(self.mColor))
        } else {
            Color::currentcolor()
        }
    }

    /// Gets a simple shadow from this item.
    #[inline]
    fn extract_simple_shadow(&self) -> SimpleShadow {
        SimpleShadow {
            color: self.extract_color(),
            horizontal: Au(self.mXOffset),
            vertical: Au(self.mYOffset),
            blur: Au(self.mRadius).into(),
        }
    }

    /// Returns this item as a simple shadow.
    #[inline]
    pub fn to_simple_shadow(&self) -> SimpleShadow {
        debug_assert_eq!(self.mSpread, 0);
        debug_assert_eq!(self.mInset, false);
        self.extract_simple_shadow()
    }
}
<|endoftext|>


<commit_before>// Claxon -- A FLAC decoding library in Rust
// Copyright (C) 2014-2015 Ruud van Asseldonk
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License, version 3,
// as published by the Free Software Foundation.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

#![feature(path_ext)]

extern crate claxon;
extern crate hound;

use std::fs;
use std::io;
use std::path;

fn run_metaflac(fname: &path::Path) -> String {
    use std::process::Command;

    // Run metaflac on the specified file and print all streaminfo data.
    let output = Command::new("metaflac")
                         .arg("--show-min-blocksize")
                         .arg("--show-max-blocksize")
                         .arg("--show-min-framesize")
                         .arg("--show-max-framesize")
                         .arg("--show-sample-rate")
                         .arg("--show-channels")
                         .arg("--show-bps")
                         .arg("--show-total-samples")
                         .arg("--show-md5sum")
                         .arg(fname.to_str().expect("unsupported filename"))
                         .output().ok().expect("failed to run metaflac");
    String::from_utf8(output.stdout).ok().expect("metaflac wrote invalid UTF-8")
}

fn decode_file(fname: &path::Path) {
    use std::process::Command;

    // Run the the reference flac decoder on the file.
    let success = Command::new("flac")
                          .arg("--decode")
                          .arg(fname.to_str().expect("unsupported filename"))
                          .status().ok().expect("failed to run flac")
                          .success();
    assert!(success);
}

fn print_hex(seq: &[u8]) -> String {
    let vec: Vec<String> = seq.iter().map(|x| format!("{:0>2x}", *x)).collect();
    vec.concat()
}

fn read_streaminfo(fname: &path::Path) -> String {
    use claxon::FlacStream;

    // Use a buffered reader, this speeds up the test by 120%.
    let file = fs::File::open(fname).unwrap();
    let mut reader = io::BufReader::new(file);
    let stream = FlacStream::new(&mut reader).unwrap();
    let streaminfo = stream.streaminfo();

    // Format the streaminfo in the same way that metaflac prints it.
    format!("{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n",
            streaminfo.min_block_size,
            streaminfo.max_block_size,
            streaminfo.min_frame_size.unwrap_or(0),
            streaminfo.max_frame_size.unwrap_or(0),
            streaminfo.sample_rate,
            streaminfo.channels,
            streaminfo.bits_per_sample,
            streaminfo.samples.unwrap_or(0),
            print_hex(&streaminfo.md5sum)) // TODO implement LowerHex for &[u8] and submit a PR.
}

fn compare_metaflac(fname: &path::Path) {
    let metaflac = run_metaflac(fname);
    let streaminfo = read_streaminfo(fname);
    let mut mf_lines = metaflac.lines_any();
    let mut si_lines = streaminfo.lines_any();
    while let (Some(mf), Some(si)) = (mf_lines.next(), si_lines.next()) {
        if mf != si {
            println!("metaflac\n--------\n{}", metaflac);
            println!("streaminfo\n----------\n{}", streaminfo);
            panic!("metaflac disagrees on parsed streaminfo");
        }
    };
}

fn compare_decoded_stream(fname: &path::Path) {
    use std::fs::PathExt;

    let ref_fname = fname.with_extension("wav");
    if !ref_fname.exists() {
        // TODO: actually, we might only want to run the test for files that
        // do exist already. It is fine to run streaminfo or decode-only tests
        // on thousands of files, but decoding thousands of files to wav before
        // running the tests will take a lot of space. For now, while the
        // number of samples is closer to 10 than to 100, this is fine.
        decode_file(fname);
    }

    // If the reference file does exist after decoding, we can compare it to
    // how Claxon decodes it, sample by sample.
    if ref_fname.exists() {
        let mut ref_stream = hound::WavReader::open(ref_fname).unwrap();

        let try_file = fs::File::open(fname).unwrap();
        let mut try_reader = io::BufReader::new(try_file);
        // TODO: Claxon should take R by value, not by mut reference.
        let mut try_stream = claxon::FlacStream::new(&mut try_reader).unwrap();

        // The streaminfo test will ensure that things like bit depth and
        // sample rate match, only the actual samples are compared here.

        let mut ref_samples = ref_stream.samples::<i32>();

        let samples = try_stream.streaminfo().samples.unwrap();
        let n_channels = try_stream.streaminfo().channels;
        let mut blocks = try_stream.blocks::<i32>();
        let mut sample = 0u64;
        let mut b = 0u64;

        while sample < samples {
            let block = blocks.read_next().unwrap();
            let mut channels: Vec<_> = (0 .. n_channels)
                                       .map(|i| block.channel(i).iter().cloned())
                                       .collect();
            for i in (0 .. block.len()) {
                for ch in (0 .. n_channels as usize) {
                    let ref_sample = ref_samples.next().map(|r| r.ok().unwrap());
                    let try_sample = channels[ch].next();
                    if ref_sample != try_sample {
                        println!("disagreement at sample {} of block {} in channel {}: reference is {} but decoded is {}",
                                 i, b, ch, ref_sample.unwrap(), try_sample.unwrap());
                        panic!("decoding differs from reference decoder");
                    }
                }
            }
            sample = sample + block.len() as u64;
            b = b + 1;
        }
    }
}

fn for_each_test_sample<F: Fn(&path::Path)>(f: F) {
    use std::ffi::OsStr;
    use std::fs::PathExt;

    // Enumerate all the flac files in the testsamples directory, and execute
    // the function `f` for it.
    let dir = fs::read_dir("testsamples")
                 .ok().expect("failed to enumerate flac files");
    for path in dir {
        let path = path.ok().expect("failed to obtain path info").path();
        if path.is_file() && path.extension() == Some(OsStr::new("flac")) {
            print!("    comparing {} ...", path.to_str()
                                               .expect("unsupported filename"));
            f(&path);
            println!(" ok");
        }
    }
}

#[test]
fn verify_streaminfo_all() {
    for_each_test_sample(compare_metaflac);
}

#[test]
fn verify_decoded_stream_all() {
    for_each_test_sample(compare_decoded_stream);
}

// Hard-coded tests for the test samples added py populate.sh.
// This allows us to run these tests in parallel and follow progress without
// enabling --nocapture.

#[test]
fn verify_streaminfo_p0() {
    compare_metaflac(path::Path::new("testsamples/p0.flac"));
}

#[test]
fn verify_streaminfo_p1() {
    compare_metaflac(path::Path::new("testsamples/p1.flac"));
}

#[test]
fn verify_streaminfo_p2() {
    compare_metaflac(path::Path::new("testsamples/p2.flac"));
}

#[test]
fn verify_streaminfo_p3() {
    compare_metaflac(path::Path::new("testsamples/p3.flac"));
}

#[test]
fn verify_streaminfo_p4() {
    compare_metaflac(path::Path::new("testsamples/p4.flac"));
}

#[test]
fn verify_decoded_stream_p0() {
    compare_decoded_stream(path::Path::new("testsamples/p0.flac"));
}

#[test]
fn verify_decoded_stream_p1() {
    compare_decoded_stream(path::Path::new("testsamples/p1.flac"));
}

#[test]
fn verify_decoded_stream_p2() {
    compare_decoded_stream(path::Path::new("testsamples/p2.flac"));
}

#[test]
fn verify_decoded_stream_p3() {
    compare_decoded_stream(path::Path::new("testsamples/p3.flac"));
}

#[test]
fn verify_decoded_stream_p4() {
    compare_decoded_stream(path::Path::new("testsamples/p4.flac"));
}
<commit_msg>remove tests that test all test samples -- for now<commit_after>// Claxon -- A FLAC decoding library in Rust
// Copyright (C) 2014-2015 Ruud van Asseldonk
//
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU General Public License, version 3,
// as published by the Free Software Foundation.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU General Public License for more details.
//
// You should have received a copy of the GNU General Public License
// along with this program.  If not, see <http://www.gnu.org/licenses/>.

#![feature(path_ext)]

extern crate claxon;
extern crate hound;

use std::fs;
use std::io;
use std::path;

fn run_metaflac(fname: &path::Path) -> String {
    use std::process::Command;

    // Run metaflac on the specified file and print all streaminfo data.
    let output = Command::new("metaflac")
                         .arg("--show-min-blocksize")
                         .arg("--show-max-blocksize")
                         .arg("--show-min-framesize")
                         .arg("--show-max-framesize")
                         .arg("--show-sample-rate")
                         .arg("--show-channels")
                         .arg("--show-bps")
                         .arg("--show-total-samples")
                         .arg("--show-md5sum")
                         .arg(fname.to_str().expect("unsupported filename"))
                         .output().ok().expect("failed to run metaflac");
    String::from_utf8(output.stdout).ok().expect("metaflac wrote invalid UTF-8")
}

fn decode_file(fname: &path::Path) {
    use std::process::Command;

    // Run the the reference flac decoder on the file.
    let success = Command::new("flac")
                          .arg("--decode")
                          .arg(fname.to_str().expect("unsupported filename"))
                          .status().ok().expect("failed to run flac")
                          .success();
    assert!(success);
}

fn print_hex(seq: &[u8]) -> String {
    let vec: Vec<String> = seq.iter().map(|x| format!("{:0>2x}", *x)).collect();
    vec.concat()
}

fn read_streaminfo(fname: &path::Path) -> String {
    use claxon::FlacStream;

    // Use a buffered reader, this speeds up the test by 120%.
    let file = fs::File::open(fname).unwrap();
    let mut reader = io::BufReader::new(file);
    let stream = FlacStream::new(&mut reader).unwrap();
    let streaminfo = stream.streaminfo();

    // Format the streaminfo in the same way that metaflac prints it.
    format!("{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n{}\n",
            streaminfo.min_block_size,
            streaminfo.max_block_size,
            streaminfo.min_frame_size.unwrap_or(0),
            streaminfo.max_frame_size.unwrap_or(0),
            streaminfo.sample_rate,
            streaminfo.channels,
            streaminfo.bits_per_sample,
            streaminfo.samples.unwrap_or(0),
            print_hex(&streaminfo.md5sum)) // TODO implement LowerHex for &[u8] and submit a PR.
}

fn compare_metaflac(fname: &path::Path) {
    let metaflac = run_metaflac(fname);
    let streaminfo = read_streaminfo(fname);
    let mut mf_lines = metaflac.lines_any();
    let mut si_lines = streaminfo.lines_any();
    while let (Some(mf), Some(si)) = (mf_lines.next(), si_lines.next()) {
        if mf != si {
            println!("metaflac\n--------\n{}", metaflac);
            println!("streaminfo\n----------\n{}", streaminfo);
            panic!("metaflac disagrees on parsed streaminfo");
        }
    };
}

fn compare_decoded_stream(fname: &path::Path) {
    use std::fs::PathExt;

    let ref_fname = fname.with_extension("wav");
    if !ref_fname.exists() {
        // TODO: actually, we might only want to run the test for files that
        // do exist already. It is fine to run streaminfo or decode-only tests
        // on thousands of files, but decoding thousands of files to wav before
        // running the tests will take a lot of space. For now, while the
        // number of samples is closer to 10 than to 100, this is fine.
        decode_file(fname);
    }

    // If the reference file does exist after decoding, we can compare it to
    // how Claxon decodes it, sample by sample.
    if ref_fname.exists() {
        let mut ref_stream = hound::WavReader::open(ref_fname).unwrap();

        let try_file = fs::File::open(fname).unwrap();
        let mut try_reader = io::BufReader::new(try_file);
        // TODO: Claxon should take R by value, not by mut reference.
        let mut try_stream = claxon::FlacStream::new(&mut try_reader).unwrap();

        // The streaminfo test will ensure that things like bit depth and
        // sample rate match, only the actual samples are compared here.

        let mut ref_samples = ref_stream.samples::<i32>();

        let samples = try_stream.streaminfo().samples.unwrap();
        let n_channels = try_stream.streaminfo().channels;
        let mut blocks = try_stream.blocks::<i32>();
        let mut sample = 0u64;
        let mut b = 0u64;

        while sample < samples {
            let block = blocks.read_next().unwrap();
            let mut channels: Vec<_> = (0 .. n_channels)
                                       .map(|i| block.channel(i).iter().cloned())
                                       .collect();
            for i in (0 .. block.len()) {
                for ch in (0 .. n_channels as usize) {
                    let ref_sample = ref_samples.next().map(|r| r.ok().unwrap());
                    let try_sample = channels[ch].next();
                    if ref_sample != try_sample {
                        println!("disagreement at sample {} of block {} in channel {}: reference is {} but decoded is {}",
                                 i, b, ch, ref_sample.unwrap(), try_sample.unwrap());
                        panic!("decoding differs from reference decoder");
                    }
                }
            }
            sample = sample + block.len() as u64;
            b = b + 1;
        }
    }
}

// Hard-coded tests for the test samples added py populate.sh.
// This allows us to run these tests in parallel and follow progress without
// enabling --nocapture.

#[test]
fn verify_streaminfo_p0() {
    compare_metaflac(path::Path::new("testsamples/p0.flac"));
}

#[test]
fn verify_streaminfo_p1() {
    compare_metaflac(path::Path::new("testsamples/p1.flac"));
}

#[test]
fn verify_streaminfo_p2() {
    compare_metaflac(path::Path::new("testsamples/p2.flac"));
}

#[test]
fn verify_streaminfo_p3() {
    compare_metaflac(path::Path::new("testsamples/p3.flac"));
}

#[test]
fn verify_streaminfo_p4() {
    compare_metaflac(path::Path::new("testsamples/p4.flac"));
}

#[test]
fn verify_decoded_stream_p0() {
    compare_decoded_stream(path::Path::new("testsamples/p0.flac"));
}

#[test]
fn verify_decoded_stream_p1() {
    compare_decoded_stream(path::Path::new("testsamples/p1.flac"));
}

#[test]
fn verify_decoded_stream_p2() {
    compare_decoded_stream(path::Path::new("testsamples/p2.flac"));
}

#[test]
fn verify_decoded_stream_p3() {
    compare_decoded_stream(path::Path::new("testsamples/p3.flac"));
}

#[test]
fn verify_decoded_stream_p4() {
    compare_decoded_stream(path::Path::new("testsamples/p4.flac"));
}
<|endoftext|>


<commit_before>extern crate natural;

use natural::tf_idf::TfIdf;

#[test]
fn test_basic_usage() {
	let mut tf_idf = TfIdf::new();

	tf_idf.add("this document is about rust.");
	tf_idf.add("this document is about erlang.");
	tf_idf.add("this document is about erlang and rust.");
	tf_idf.add("this document is about rust. it has rust examples");

	assert_eq!(tf_idf.get("rust erlang"), 0.2993708f32 );
	assert!(tf_idf.get("rust") > tf_idf.get("erlang"));
}<commit_msg>fix test<commit_after>extern crate natural;

use natural::tf_idf::TfIdf;

#[test]
fn test_basic_usage() {
	let mut tf_idf = TfIdf::new();

	tf_idf.add("this document is about rust.");
	tf_idf.add("this document is about erlang.");
	tf_idf.add("this document is about erlang and rust.");
	tf_idf.add("this document is about rust. it has rust examples");

	assert_eq!(tf_idf.get("rust"), 0.2993708f32 );
	assert!(tf_idf.get("rust") > tf_idf.get("erlang"));
}<|endoftext|>


<commit_before>use rand::Rng;
use multi_objective::MultiObjective;
use domination::Domination;
use sort::FastNonDominatedSorter;
use crowding_distance::{crowding_distance_assignment, CrowdingDistanceAssignment};
use std::cmp::{self, Ordering};

pub trait SelectSolutions<T, F>
    where T: CrowdingDistanceAssignment<F>, F: MultiObjective + Domination,
{
    /// Select `n` out of the `solutions`, assigning rank and distance using the first `num_objectives`
    /// objectives. `rng` might not be used.
    fn select_solutions<R>(solutions: &mut [T],
                           select_n: usize,
                           num_objectives: usize,
                           rng: &mut R)
        where R: Rng;
}

/// Uses the approach NSGA takes.
pub struct SelectNSGA;

impl<T, F> SelectSolutions<T, F> for SelectNSGA
    where T: CrowdingDistanceAssignment<F>,
          F: MultiObjective + Domination
{
    fn select_solutions<R>(solutions: &mut [T],
                           select_n: usize,
                           num_objectives: usize,
                           _rng: &mut R)
        where R: Rng
    {
        // We can select at most `select_n` solutions.
        let mut missing = cmp::min(solutions.len(), select_n);

        // make sure all solutions are unselected
        for s in solutions.iter_mut() {
            s.unselect();
        }

        let pareto_fronts = FastNonDominatedSorter::new(solutions, &|s| s.fitness());

        for (rank, mut front) in pareto_fronts.enumerate() {
            assert!(missing > 0);

            // assign rank and crowding distance of those solutions in `front`.
            crowding_distance_assignment(solutions, &mut front, rank as u32, num_objectives);

            if front.len() <= missing {
                // whole front fits into result. XXX: should we sort?

                for i in front.into_iter() {
                    solutions[i].select();
                    assert!(missing > 0);
                    missing -= 1;
                }

                if missing == 0 {
                    break;
                }
            } else {
                // choose only best from this front, according to the crowding distance.
                front.sort_by(|&a, &b| {
                    debug_assert!(solutions[a].rank() == solutions[b].rank());
                    if solutions[a].has_better_rank_and_crowding(&solutions[b]) {
                        Ordering::Less
                    } else if solutions[b].has_better_rank_and_crowding(&solutions[a]) {
                        Ordering::Greater
                    } else {
                        Ordering::Equal
                    }
                });

                for i in front.into_iter().take(missing) {
                    assert!(missing > 0);
                    solutions[i].select();
                    missing -= 1;
                }

                assert!(missing == 0);
                break;
            }
        }
    }
}

/// Selection using NSGP as described in [1].
///
/// [1]: Multi-objective genetic programming with redundancy-regulations for automatic construction of image feature extractors.
///      Watchareeruetai, Ukrit and Matsumoto, Tetsuya and Takeuchi, Yoshinori and Hiroaki, KUDO and Ohnishi, Noboru.
///      2010
pub struct SelectNSGP;

impl<T, F> SelectSolutions<T, F> for SelectNSGP
    where T: CrowdingDistanceAssignment<F>,
          F: MultiObjective + Domination + Eq
{
    fn select_solutions<R>(solutions: &mut [T],
                           select_n: usize,
                           num_objectives: usize,
                           rng: &mut R)
        where R: Rng
    {
        // We can select at most `select_n` solutions.
        let mut missing = cmp::min(solutions.len(), select_n);

        // make sure all solutions are unselected
        for s in solutions.iter_mut() {
            s.unselect();
        }

        let pareto_fronts = FastNonDominatedSorter::new(solutions, &|s| s.fitness());

        let mut fronts_grouped: Vec<Vec<Vec<_>>> = Vec::new();
        for (rank, mut front) in pareto_fronts.enumerate() {

            // first assign rank and crowding distance of those solutions in `front`.
            crowding_distance_assignment(solutions, &mut front, rank as u32, num_objectives);

            // then group this front into individuals with similar (or equal) fitness.
            let mut groups: Vec<Vec<usize>> = Vec::new();

            for idx in front {
                // insert `idx` into a group.
                // test fitness of `idx` (solutions[idx]) against a member of each group.
                // if none fits, create a new group.

                let mut insert_in = None;
                for (i, grp) in groups.iter().enumerate() {
                    if solutions[idx].fitness().similar_to(solutions[grp[0]].fitness()) {
                        insert_in = Some(i);
                        break;
                    }
                }

                if let Some(i) = insert_in {
                    groups[i].push(idx);
                } else {
                    groups.push(vec![idx]);
                }
            }

            // update the crowding_distance within each group. simply take the average.
            // also record the number of points in the group within each individual's fitness. this
            // is required for sexual selection.
            for grp in groups.iter() {
                // use same average crowding distance for each point in the group
                assert!(grp.len() > 0);
                let avg: f64 = grp.iter().map(|&idx| solutions[idx].dist()).sum::<f64>() /
                               grp.len() as f64;
                for &idx in grp.iter() {
                    solutions[idx].set_dist(avg);
                    solutions[idx].set_crowd(grp.len());
                }
            }

            fronts_grouped.push(groups);
        }

        // Each pareto_front is now grouped into individuals of same (or similar) fitness.
        // Now go through each front, and in each front, choose exactly one individual
        // from each group. Then go to the next front etc. once we are through all
        // fronts, continue with the first and so on, until enough individuals are selected.

        'outer: while missing > 0 {
            for current_front in fronts_grouped.iter_mut() {
                // select up to `missing` solutions from current_front, but from each group only one
                // (which is choosen randomly).
                for grp in current_front.iter_mut() {
                    if missing > 0 {
                        let i = rng.gen_range(0, grp.len());
                        let idx = grp.swap_remove(i); // remove `i`th element from grp returning the solution index
                        assert!(solutions[idx].is_selected() == false);
                        solutions[idx].select();
                        missing -= 1;
                    } else {
                        break 'outer;
                    }
                }

                // only retain groups which are non-empty
                current_front.retain(|grp| grp.len() > 0);
            }
        }
    }
}

/// Select the best individual out of `k` randomly choosen.
/// This gives individuals with better fitness a higher chance to reproduce.
/// `n` is the total number of individuals.
///
/// NOTE: We are not using `sample(rng, 0..n, k)` as it is *very* expensive.
/// Instead we call `rng.gen_range()` k-times. The drawn items could be the same,
/// but the probability is very low if `n` is high compared to `k`.
#[inline]
pub fn tournament_selection_fast<R: Rng, F>(rng: &mut R,
                                            better_than: F,
                                            n: usize,
                                            k: usize)
                                            -> usize
    where F: Fn(usize, usize) -> bool
{
    assert!(n > 0);
    assert!(k > 0);
    assert!(n >= k);

    let mut best: usize = rng.gen_range(0, n);

    for _ in 1..k {
        let i = rng.gen_range(0, n);
        if better_than(i, best) {
            best = i;
        }
    }

    return best;
}
<commit_msg>Remove Eq trait dependence<commit_after>use rand::Rng;
use multi_objective::MultiObjective;
use domination::Domination;
use sort::FastNonDominatedSorter;
use crowding_distance::{crowding_distance_assignment, CrowdingDistanceAssignment};
use std::cmp::{self, Ordering};

pub trait SelectSolutions<T, F>
    where T: CrowdingDistanceAssignment<F>, F: MultiObjective + Domination,
{
    /// Select `n` out of the `solutions`, assigning rank and distance using the first `num_objectives`
    /// objectives. `rng` might not be used.
    fn select_solutions<R>(solutions: &mut [T],
                           select_n: usize,
                           num_objectives: usize,
                           rng: &mut R)
        where R: Rng;
}

/// Uses the approach NSGA takes.
pub struct SelectNSGA;

impl<T, F> SelectSolutions<T, F> for SelectNSGA
    where T: CrowdingDistanceAssignment<F>,
          F: MultiObjective + Domination
{
    fn select_solutions<R>(solutions: &mut [T],
                           select_n: usize,
                           num_objectives: usize,
                           _rng: &mut R)
        where R: Rng
    {
        // We can select at most `select_n` solutions.
        let mut missing = cmp::min(solutions.len(), select_n);

        // make sure all solutions are unselected
        for s in solutions.iter_mut() {
            s.unselect();
        }

        let pareto_fronts = FastNonDominatedSorter::new(solutions, &|s| s.fitness());

        for (rank, mut front) in pareto_fronts.enumerate() {
            assert!(missing > 0);

            // assign rank and crowding distance of those solutions in `front`.
            crowding_distance_assignment(solutions, &mut front, rank as u32, num_objectives);

            if front.len() <= missing {
                // whole front fits into result. XXX: should we sort?

                for i in front.into_iter() {
                    solutions[i].select();
                    assert!(missing > 0);
                    missing -= 1;
                }

                if missing == 0 {
                    break;
                }
            } else {
                // choose only best from this front, according to the crowding distance.
                front.sort_by(|&a, &b| {
                    debug_assert!(solutions[a].rank() == solutions[b].rank());
                    if solutions[a].has_better_rank_and_crowding(&solutions[b]) {
                        Ordering::Less
                    } else if solutions[b].has_better_rank_and_crowding(&solutions[a]) {
                        Ordering::Greater
                    } else {
                        Ordering::Equal
                    }
                });

                for i in front.into_iter().take(missing) {
                    assert!(missing > 0);
                    solutions[i].select();
                    missing -= 1;
                }

                assert!(missing == 0);
                break;
            }
        }
    }
}

/// Selection using NSGP as described in [1].
///
/// [1]: Multi-objective genetic programming with redundancy-regulations for automatic construction of image feature extractors.
///      Watchareeruetai, Ukrit and Matsumoto, Tetsuya and Takeuchi, Yoshinori and Hiroaki, KUDO and Ohnishi, Noboru.
///      2010
pub struct SelectNSGP;

impl<T, F> SelectSolutions<T, F> for SelectNSGP
    where T: CrowdingDistanceAssignment<F>,
          F: MultiObjective + Domination
{
    fn select_solutions<R>(solutions: &mut [T],
                           select_n: usize,
                           num_objectives: usize,
                           rng: &mut R)
        where R: Rng
    {
        // We can select at most `select_n` solutions.
        let mut missing = cmp::min(solutions.len(), select_n);

        // make sure all solutions are unselected
        for s in solutions.iter_mut() {
            s.unselect();
        }

        let pareto_fronts = FastNonDominatedSorter::new(solutions, &|s| s.fitness());

        let mut fronts_grouped: Vec<Vec<Vec<_>>> = Vec::new();
        for (rank, mut front) in pareto_fronts.enumerate() {

            // first assign rank and crowding distance of those solutions in `front`.
            crowding_distance_assignment(solutions, &mut front, rank as u32, num_objectives);

            // then group this front into individuals with similar (or equal) fitness.
            let mut groups: Vec<Vec<usize>> = Vec::new();

            for idx in front {
                // insert `idx` into a group.
                // test fitness of `idx` (solutions[idx]) against a member of each group.
                // if none fits, create a new group.

                let mut insert_in = None;
                for (i, grp) in groups.iter().enumerate() {
                    if solutions[idx].fitness().similar_to(solutions[grp[0]].fitness()) {
                        insert_in = Some(i);
                        break;
                    }
                }

                if let Some(i) = insert_in {
                    groups[i].push(idx);
                } else {
                    groups.push(vec![idx]);
                }
            }

            // update the crowding_distance within each group. simply take the average.
            // also record the number of points in the group within each individual's fitness. this
            // is required for sexual selection.
            for grp in groups.iter() {
                // use same average crowding distance for each point in the group
                assert!(grp.len() > 0);
                let avg: f64 = grp.iter().map(|&idx| solutions[idx].dist()).sum::<f64>() /
                               grp.len() as f64;
                for &idx in grp.iter() {
                    solutions[idx].set_dist(avg);
                    solutions[idx].set_crowd(grp.len());
                }
            }

            fronts_grouped.push(groups);
        }

        // Each pareto_front is now grouped into individuals of same (or similar) fitness.
        // Now go through each front, and in each front, choose exactly one individual
        // from each group. Then go to the next front etc. once we are through all
        // fronts, continue with the first and so on, until enough individuals are selected.

        'outer: while missing > 0 {
            for current_front in fronts_grouped.iter_mut() {
                // select up to `missing` solutions from current_front, but from each group only one
                // (which is choosen randomly).
                for grp in current_front.iter_mut() {
                    if missing > 0 {
                        let i = rng.gen_range(0, grp.len());
                        let idx = grp.swap_remove(i); // remove `i`th element from grp returning the solution index
                        assert!(solutions[idx].is_selected() == false);
                        solutions[idx].select();
                        missing -= 1;
                    } else {
                        break 'outer;
                    }
                }

                // only retain groups which are non-empty
                current_front.retain(|grp| grp.len() > 0);
            }
        }
    }
}

/// Select the best individual out of `k` randomly choosen.
/// This gives individuals with better fitness a higher chance to reproduce.
/// `n` is the total number of individuals.
///
/// NOTE: We are not using `sample(rng, 0..n, k)` as it is *very* expensive.
/// Instead we call `rng.gen_range()` k-times. The drawn items could be the same,
/// but the probability is very low if `n` is high compared to `k`.
#[inline]
pub fn tournament_selection_fast<R: Rng, F>(rng: &mut R,
                                            better_than: F,
                                            n: usize,
                                            k: usize)
                                            -> usize
    where F: Fn(usize, usize) -> bool
{
    assert!(n > 0);
    assert!(k > 0);
    assert!(n >= k);

    let mut best: usize = rng.gen_range(0, n);

    for _ in 1..k {
        let i = rng.gen_range(0, n);
        if better_than(i, best) {
            best = i;
        }
    }

    return best;
}
<|endoftext|>


<commit_before>// Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
// Licensed under the Apache License, Version 2.0 (see LICENSE).

#![deny(warnings)]
// Enable all clippy lints except for many of the pedantic ones. It's a shame this needs to be copied and pasted across crates, but there doesn't appear to be a way to include inner attributes from a common source.
#![deny(
  clippy::all,
  clippy::default_trait_access,
  clippy::expl_impl_clone_on_copy,
  clippy::if_not_else,
  clippy::needless_continue,
  clippy::unseparated_literal_suffix,
  clippy::used_underscore_binding
)]
// It is often more clear to show that nothing is being moved.
#![allow(clippy::match_ref_pats)]
// Subjective style.
#![allow(
  clippy::len_without_is_empty,
  clippy::redundant_field_names,
  clippy::too_many_arguments
)]
// Default isn't as big a deal as people seem to think it is.
#![allow(clippy::new_without_default, clippy::new_ret_no_self)]
// Arc<Mutex> can be more clear than needing to grok Orderings:
#![allow(clippy::mutex_atomic)]

use std::sync::{Arc, Weak};

use tokio::sync::{oneshot, watch};

///
/// A cancellable value computed by one sender, and broadcast to multiple receivers.
///
/// Supports canceling the work associated with the value either:
///   1. explicitly if the value is dropped
///   2. implicitly if all receivers go away
///
/// NB: This is currently a `tokio::sync::watch` (which supports the second case), plus a
/// separate cancellation signal via `tokio::sync::oneshot` (to support the first case).
///
#[derive(Debug)]
pub struct AsyncValue<T: Clone + Send + Sync + 'static> {
  item_receiver: Weak<watch::Receiver<Option<T>>>,
  abort_sender: Option<oneshot::Sender<T>>,
}

impl<T: Clone + Send + Sync + 'static> AsyncValue<T> {
  pub fn new() -> (AsyncValue<T>, AsyncValueSender<T>, AsyncValueReceiver<T>) {
    let (abort_sender, abort_receiver) = oneshot::channel();
    let (item_sender, item_receiver) = watch::channel(None);
    let item_receiver = Arc::new(item_receiver);
    (
      AsyncValue {
        item_receiver: Arc::downgrade(&item_receiver),
        abort_sender: Some(abort_sender),
      },
      AsyncValueSender {
        item_sender,
        abort_receiver,
      },
      AsyncValueReceiver { item_receiver },
    )
  }

  ///
  /// Returns an AsyncValueReceiver for this value if the associated work has not already been
  /// canceled.
  ///
  pub fn receiver(&self) -> Option<AsyncValueReceiver<T>> {
    self
      .item_receiver
      .upgrade()
      .map(|item_receiver| AsyncValueReceiver { item_receiver })
  }

  pub fn try_abort(&mut self, t: T) -> Result<(), T> {
    if let Some(abort_sender) = self.abort_sender.take() {
      abort_sender.send(t)
    } else {
      Ok(())
    }
  }
}

pub struct AsyncValueReceiver<T: Clone + Send + Sync + 'static> {
  item_receiver: Arc<watch::Receiver<Option<T>>>,
}

impl<T: Clone + Send + Sync + 'static> AsyncValueReceiver<T> {
  ///
  /// Returns a Future that will wait for the result of this value, or None if the work was
  /// canceled.
  ///
  pub async fn recv(&self) -> Option<T> {
    let mut item_receiver = (*self.item_receiver).clone();
    loop {
      if let Some(ref value) = *item_receiver.borrow() {
        return Some(value.clone());
      }

      #[allow(clippy::question_mark)]
      if item_receiver.changed().await.is_err() {
        return None;
      }
    }
  }
}

pub struct AsyncValueSender<T: Clone + Send + Sync + 'static> {
  item_sender: watch::Sender<Option<T>>,
  abort_receiver: oneshot::Receiver<T>,
}

impl<T: Clone + Send + Sync + 'static> AsyncValueSender<T> {
  pub fn send(self, item: T) {
    let _ = self.item_sender.send(Some(item));
  }

  pub async fn aborted(&mut self) -> Option<T> {
    tokio::select! {
      res = &mut self.abort_receiver => {
        match res {
          Ok(res) => {
            // Aborted with a value.
            Some(res)
          },
          Err(_) => {
            // Was dropped.
            None
          },
        }
      }
      _ = self.item_sender.closed() => { None }
    }
  }
}

#[cfg(test)]
mod tests;
<commit_msg>[internal] add comment re clippy issue (#14188)<commit_after>// Copyright 2020 Pants project contributors (see CONTRIBUTORS.md).
// Licensed under the Apache License, Version 2.0 (see LICENSE).

#![deny(warnings)]
// Enable all clippy lints except for many of the pedantic ones. It's a shame this needs to be copied and pasted across crates, but there doesn't appear to be a way to include inner attributes from a common source.
#![deny(
  clippy::all,
  clippy::default_trait_access,
  clippy::expl_impl_clone_on_copy,
  clippy::if_not_else,
  clippy::needless_continue,
  clippy::unseparated_literal_suffix,
  clippy::used_underscore_binding
)]
// It is often more clear to show that nothing is being moved.
#![allow(clippy::match_ref_pats)]
// Subjective style.
#![allow(
  clippy::len_without_is_empty,
  clippy::redundant_field_names,
  clippy::too_many_arguments
)]
// Default isn't as big a deal as people seem to think it is.
#![allow(clippy::new_without_default, clippy::new_ret_no_self)]
// Arc<Mutex> can be more clear than needing to grok Orderings:
#![allow(clippy::mutex_atomic)]

use std::sync::{Arc, Weak};

use tokio::sync::{oneshot, watch};

///
/// A cancellable value computed by one sender, and broadcast to multiple receivers.
///
/// Supports canceling the work associated with the value either:
///   1. explicitly if the value is dropped
///   2. implicitly if all receivers go away
///
/// NB: This is currently a `tokio::sync::watch` (which supports the second case), plus a
/// separate cancellation signal via `tokio::sync::oneshot` (to support the first case).
///
#[derive(Debug)]
pub struct AsyncValue<T: Clone + Send + Sync + 'static> {
  item_receiver: Weak<watch::Receiver<Option<T>>>,
  abort_sender: Option<oneshot::Sender<T>>,
}

impl<T: Clone + Send + Sync + 'static> AsyncValue<T> {
  pub fn new() -> (AsyncValue<T>, AsyncValueSender<T>, AsyncValueReceiver<T>) {
    let (abort_sender, abort_receiver) = oneshot::channel();
    let (item_sender, item_receiver) = watch::channel(None);
    let item_receiver = Arc::new(item_receiver);
    (
      AsyncValue {
        item_receiver: Arc::downgrade(&item_receiver),
        abort_sender: Some(abort_sender),
      },
      AsyncValueSender {
        item_sender,
        abort_receiver,
      },
      AsyncValueReceiver { item_receiver },
    )
  }

  ///
  /// Returns an AsyncValueReceiver for this value if the associated work has not already been
  /// canceled.
  ///
  pub fn receiver(&self) -> Option<AsyncValueReceiver<T>> {
    self
      .item_receiver
      .upgrade()
      .map(|item_receiver| AsyncValueReceiver { item_receiver })
  }

  pub fn try_abort(&mut self, t: T) -> Result<(), T> {
    if let Some(abort_sender) = self.abort_sender.take() {
      abort_sender.send(t)
    } else {
      Ok(())
    }
  }
}

pub struct AsyncValueReceiver<T: Clone + Send + Sync + 'static> {
  item_receiver: Arc<watch::Receiver<Option<T>>>,
}

impl<T: Clone + Send + Sync + 'static> AsyncValueReceiver<T> {
  ///
  /// Returns a Future that will wait for the result of this value, or None if the work was
  /// canceled.
  ///
  pub async fn recv(&self) -> Option<T> {
    let mut item_receiver = (*self.item_receiver).clone();
    loop {
      if let Some(ref value) = *item_receiver.borrow() {
        return Some(value.clone());
      }

      // TODO: Remove the `allow` once https://github.com/rust-lang/rust-clippy/issues/8281
      // is fixed upstream.
      #[allow(clippy::question_mark)]
      if item_receiver.changed().await.is_err() {
        return None;
      }
    }
  }
}

pub struct AsyncValueSender<T: Clone + Send + Sync + 'static> {
  item_sender: watch::Sender<Option<T>>,
  abort_receiver: oneshot::Receiver<T>,
}

impl<T: Clone + Send + Sync + 'static> AsyncValueSender<T> {
  pub fn send(self, item: T) {
    let _ = self.item_sender.send(Some(item));
  }

  pub async fn aborted(&mut self) -> Option<T> {
    tokio::select! {
      res = &mut self.abort_receiver => {
        match res {
          Ok(res) => {
            // Aborted with a value.
            Some(res)
          },
          Err(_) => {
            // Was dropped.
            None
          },
        }
      }
      _ = self.item_sender.closed() => { None }
    }
  }
}

#[cfg(test)]
mod tests;
<|endoftext|>


<commit_before>/*
    TODO:
    Write code for downloading the package, into a
    specific path.
*/
crate hyper;
use self::hyper::client;

macro_rules! custom_try {
    ($x:expr) => (match $x {
        Ok(x) => x,
        Err(why) => panic!("An error occured during package downloading.", why),
    });
}

fn download_package(url: &str) {
    let client = Client::new();
    let response = custom_try!(client.get(url).send());
    let content = Vec::new();
    custom_try!(response.read_to_end(&mut content));
    return content;
}<commit_msg>New todos<commit_after>/*
    TODO:
    1. Copy downloaded files to other places.
    2. Verify the downloaded package.
    3. Make it support repos someday.
*/
crate hyper;
use self::hyper::client;

macro_rules! custom_try {
    ($x:expr) => (match $x {
        Ok(x) => x,
        Err(why) => panic!("An error occured during package downloading.", why),
    });
}

fn download_package(url: &str) {
    let client = Client::new();
    let response = custom_try!(client.get(url).send());
    let content = Vec::new();
    custom_try!(response.read_to_end(&mut content));
    return content;
}<|endoftext|>


<commit_before>#![allow(clippy::upper_case_acronyms)]

use anyhow::{anyhow, Result};
use clap::clap_app;
use log::{debug, error, info};

use lancelot::{
    analysis::{
        cfg::{InstructionIndex, CFG},
        dis,
        dis::zydis,
    },
    aspace::AddressSpace,
    loader::pe::PE,
    util, RVA, VA,
};

fn handle_functions(pe: &PE) -> Result<()> {
    let functions = lancelot::analysis::pe::find_function_starts(pe)?;

    info!("found {} functions", functions.len());
    for va in functions.iter() {
        println!("{:#x}", va);
    }

    Ok(())
}

fn render_insn_buf(buf: &[u8], width: usize) -> String {
    let mut out = String::new();
    for (i, c) in hex::encode(buf).chars().enumerate() {
        out.push(c);

        if i % 2 == 1 {
            out.push(' ');
        }
    }

    let out = out.trim_end_matches(' ').to_string();
    if out.len() > width {
        out[..width].to_string()
    } else {
        out
    }
}

fn render_insn(insn: &zydis::ffi::DecodedInstruction, va: VA) -> String {
    let formatter = zydis::Formatter::new(zydis::FormatterStyle::INTEL).expect("formatter");
    let mut buffer = [0u8; 200];
    let mut buffer = zydis::OutputBuffer::new(&mut buffer[..]);

    formatter
        .format_instruction(insn, &mut buffer, Some(va), None)
        .expect("format");
    format!("{}", buffer)
}

fn handle_disassemble(pe: &PE, va: VA) -> Result<()> {
    let decoder = dis::get_disassembler(&pe.module)?;
    let mut insns: InstructionIndex = Default::default();
    insns.build_index(&pe.module, va)?;
    let cfg = CFG::from_instructions(insns)?;

    info!("found {} basic blocks", cfg.basic_blocks.blocks_by_address.len());
    for bb in cfg.basic_blocks.blocks_by_address.values() {
        // need to over-read the bb buffer, to account for the final instructions.
        let buf = pe
            .module
            .address_space
            .read_bytes(bb.address, bb.length as usize + 0x10)?;
        for (offset, insn) in dis::linear_disassemble(&decoder, &buf) {
            // because we over-read the bb buffer,
            // discard the instructions found after it.
            if offset >= bb.length as usize {
                break;
            }

            let va = bb.address + offset as RVA;

            let name = &pe
                .module
                .sections
                .iter()
                .find(|sec| sec.virtual_range.contains(&va))
                .unwrap()
                .name;

            if let Ok(Some(insn)) = insn {
                let insn_buf = &buf[offset..offset + insn.length as usize];
                println!(
                    "{}:{:016x}  {:15}  {}",
                    name,
                    va,
                    render_insn_buf(insn_buf, 15),
                    render_insn(&insn, va)
                );
            } else {
                println!("  {}:{:#x}: INVALID", name, va);
                break;
            }
        }
        println!();
    }

    Ok(())
}

fn parse_va(s: &str) -> Result<VA> {
    if s.starts_with("0x") {
        let without_prefix = s.trim_start_matches("0x");
        Ok(u64::from_str_radix(without_prefix, 16)?)
    } else {
        Ok(s.parse()?)
    }
}

fn _main() -> Result<()> {
    better_panic::install();

    let matches = clap::clap_app!(lancelot =>
        (author: "Willi Ballenthin <william.ballenthin@mandiant.com>")
        (about: "Binary analysis framework")
        (@arg verbose: -v --verbose +multiple "log verbose messages")
        (@arg quiet: -q --quiet "disable informational messages")
        (@subcommand functions =>
            (about: "find functions")
            (@arg input: +required "path to file to analyze"))
        (@subcommand disassemble =>
            (about: "disassemble function")
            (@arg input: +required "path to file to analyze")
            (@arg va: +required "VA of function")))
    .get_matches();

    // --quiet overrides --verbose
    let log_level = if matches.is_present("quiet") {
        log::LevelFilter::Error
    } else {
        match matches.occurrences_of("verbose") {
            0 => log::LevelFilter::Info,
            1 => log::LevelFilter::Debug,
            2 => log::LevelFilter::Trace,
            _ => log::LevelFilter::Trace,
        }
    };

    fern::Dispatch::new()
        .format(move |out, message, record| {
            out.finish(format_args!(
                "{} [{:5}] {} {}",
                chrono::Local::now().format("%Y-%m-%d %H:%M:%S"),
                record.level(),
                if log_level == log::LevelFilter::Trace {
                    record.target()
                } else {
                    ""
                },
                message
            ))
        })
        .level(log_level)
        .chain(std::io::stderr())
        .filter(|metadata| !metadata.target().starts_with("goblin::pe"))
        .apply()
        .expect("failed to configure logging");

    if let Some(matches) = matches.subcommand_matches("functions") {
        debug!("mode: find functions");

        let filename = matches.value_of("input").unwrap();
        debug!("input: {}", filename);

        let buf = util::read_file(filename)?;
        let pe = PE::from_bytes(&buf)?;

        handle_functions(&pe)
    } else if let Some(matches) = matches.subcommand_matches("disassemble") {
        debug!("mode: disassemble");

        let filename = matches.value_of("input").unwrap();
        debug!("input: {}", filename);

        let va = parse_va(matches.value_of("va").unwrap())?;

        let buf = util::read_file(filename)?;
        let pe = PE::from_bytes(&buf)?;

        handle_disassemble(&pe, va)
    } else {
        Err(anyhow!("SUBCOMMAND required"))
    }
}

fn main() {
    if let Err(e) = _main() {
        #[cfg(debug_assertions)]
        error!("{:?}", e);
        #[cfg(not(debug_assertions))]
        error!("{:}", e);
    }
}
<commit_msg>bin: lancelot: use reachable routine to restrict to current cfg<commit_after>#![allow(clippy::upper_case_acronyms)]

use anyhow::{anyhow, Result};
use clap::clap_app;
use log::{debug, error, info};

use lancelot::{
    analysis::{
        cfg::{InstructionIndex, CFG},
        dis,
        dis::zydis,
    },
    aspace::AddressSpace,
    loader::pe::PE,
    util, RVA, VA,
};

fn handle_functions(pe: &PE) -> Result<()> {
    let functions = lancelot::analysis::pe::find_function_starts(pe)?;

    info!("found {} functions", functions.len());
    for va in functions.iter() {
        println!("{:#x}", va);
    }

    Ok(())
}

fn render_insn_buf(buf: &[u8], width: usize) -> String {
    let mut out = String::new();
    for (i, c) in hex::encode(buf).chars().enumerate() {
        out.push(c);

        if i % 2 == 1 {
            out.push(' ');
        }
    }

    let out = out.trim_end_matches(' ').to_string();
    if out.len() > width {
        out[..width].to_string()
    } else {
        out
    }
}

fn render_insn(insn: &zydis::ffi::DecodedInstruction, va: VA) -> String {
    let formatter = zydis::Formatter::new(zydis::FormatterStyle::INTEL).expect("formatter");
    let mut buffer = [0u8; 200];
    let mut buffer = zydis::OutputBuffer::new(&mut buffer[..]);

    formatter
        .format_instruction(insn, &mut buffer, Some(va), None)
        .expect("format");
    format!("{}", buffer)
}

fn handle_disassemble(pe: &PE, va: VA) -> Result<()> {
    let decoder = dis::get_disassembler(&pe.module)?;
    let mut insns: InstructionIndex = Default::default();
    insns.build_index(&pe.module, va)?;
    let cfg = CFG::from_instructions(insns)?;

    let mut blocks = cfg.get_reachable_blocks(va).collect::<Vec<_>>();
    blocks.sort_unstable_by_key(|&bb| bb.address);
    info!("found {} basic blocks", blocks.len());

    for bb in blocks.into_iter() {
        // need to over-read the bb buffer, to account for the final instructions.
        let buf = pe
            .module
            .address_space
            .read_bytes(bb.address, bb.length as usize + 0x10)?;
        for (offset, insn) in dis::linear_disassemble(&decoder, &buf) {
            // because we over-read the bb buffer,
            // discard the instructions found after it.
            if offset >= bb.length as usize {
                break;
            }

            let va = bb.address + offset as RVA;

            let name = &pe
                .module
                .sections
                .iter()
                .find(|sec| sec.virtual_range.contains(&va))
                .unwrap()
                .name;

            if let Ok(Some(insn)) = insn {
                let insn_buf = &buf[offset..offset + insn.length as usize];
                println!(
                    "{}:{:016x}  {:15}  {}",
                    name,
                    va,
                    render_insn_buf(insn_buf, 15),
                    render_insn(&insn, va)
                );
            } else {
                println!("  {}:{:#x}: INVALID", name, va);
                break;
            }
        }
        println!();
    }

    Ok(())
}

fn parse_va(s: &str) -> Result<VA> {
    if s.starts_with("0x") {
        let without_prefix = s.trim_start_matches("0x");
        Ok(u64::from_str_radix(without_prefix, 16)?)
    } else {
        Ok(s.parse()?)
    }
}

fn _main() -> Result<()> {
    better_panic::install();

    let matches = clap::clap_app!(lancelot =>
        (author: "Willi Ballenthin <william.ballenthin@mandiant.com>")
        (about: "Binary analysis framework")
        (@arg verbose: -v --verbose +multiple "log verbose messages")
        (@arg quiet: -q --quiet "disable informational messages")
        (@subcommand functions =>
            (about: "find functions")
            (@arg input: +required "path to file to analyze"))
        (@subcommand disassemble =>
            (about: "disassemble function")
            (@arg input: +required "path to file to analyze")
            (@arg va: +required "VA of function")))
    .get_matches();

    // --quiet overrides --verbose
    let log_level = if matches.is_present("quiet") {
        log::LevelFilter::Error
    } else {
        match matches.occurrences_of("verbose") {
            0 => log::LevelFilter::Info,
            1 => log::LevelFilter::Debug,
            2 => log::LevelFilter::Trace,
            _ => log::LevelFilter::Trace,
        }
    };

    fern::Dispatch::new()
        .format(move |out, message, record| {
            out.finish(format_args!(
                "{} [{:5}] {} {}",
                chrono::Local::now().format("%Y-%m-%d %H:%M:%S"),
                record.level(),
                if log_level == log::LevelFilter::Trace {
                    record.target()
                } else {
                    ""
                },
                message
            ))
        })
        .level(log_level)
        .chain(std::io::stderr())
        .filter(|metadata| !metadata.target().starts_with("goblin::pe"))
        .apply()
        .expect("failed to configure logging");

    if let Some(matches) = matches.subcommand_matches("functions") {
        debug!("mode: find functions");

        let filename = matches.value_of("input").unwrap();
        debug!("input: {}", filename);

        let buf = util::read_file(filename)?;
        let pe = PE::from_bytes(&buf)?;

        handle_functions(&pe)
    } else if let Some(matches) = matches.subcommand_matches("disassemble") {
        debug!("mode: disassemble");

        let filename = matches.value_of("input").unwrap();
        debug!("input: {}", filename);

        let va = parse_va(matches.value_of("va").unwrap())?;

        let buf = util::read_file(filename)?;
        let pe = PE::from_bytes(&buf)?;

        handle_disassemble(&pe, va)
    } else {
        Err(anyhow!("SUBCOMMAND required"))
    }
}

fn main() {
    if let Err(e) = _main() {
        #[cfg(debug_assertions)]
        error!("{:?}", e);
        #[cfg(not(debug_assertions))]
        error!("{:}", e);
    }
}
<|endoftext|>