<commit_before>
use size_hint;

use std::cmp::Ordering;
use std::mem::replace;

macro_rules! clone_fields {
    ($name:ident, $base:expr, $($field:ident),+) => (
        $name {
            $(
                $field : $base . $field .clone()
            ),*
        }
    );
}

/// Head element and Tail iterator pair
///
/// `PartialEq`, `Eq`, `PartialOrd` and `Ord` are implemented by comparing sequences based on
/// first items (which are guaranteed to exist).
///
/// The meanings of `PartialOrd` and `Ord` are reversed so as to turn the heap used in
/// `KMerge` into a min-heap.
struct HeadTail<I>
    where I: Iterator
{
    head: I::Item,
    tail: I,
}

impl<I> HeadTail<I>
    where I: Iterator
{
    /// Constructs a `HeadTail` from an `Iterator`. Returns `None` if the `Iterator` is empty.
    fn new(mut it: I) -> Option<HeadTail<I>> {
        let head = it.next();
        head.map(|h| {
            HeadTail {
                head: h,
                tail: it,
            }
        })
    }

    /// Get the next element and update `head`, returning the old head in `Some`.
    ///
    /// Returns `None` when the tail is exhausted (only `head` then remains).
    fn next(&mut self) -> Option<I::Item> {
        if let Some(next) = self.tail.next() {
            Some(replace(&mut self.head, next))
        } else {
            None
        }
    }

    /// Hints at the size of the sequence, same as the `Iterator` method.
    fn size_hint(&self) -> (usize, Option<usize>) {
        size_hint::add_scalar(self.tail.size_hint(), 1)
    }
}

impl<I> Clone for HeadTail<I>
    where I: Iterator + Clone,
          I::Item: Clone
{
    fn clone(&self) -> Self {
        clone_fields!(HeadTail, self, head, tail)
    }
}

impl<I> PartialEq for HeadTail<I>
    where I: Iterator,
          I::Item: PartialEq
{
    fn eq(&self, other: &HeadTail<I>) -> bool {
        self.head.eq(&other.head)
    }
}

impl<I> Eq for HeadTail<I>
    where I: Iterator,
          I::Item: Eq
{}

impl<I> PartialOrd for HeadTail<I>
    where I: Iterator,
          I::Item: PartialOrd
{
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        other.head.partial_cmp(&self.head)
    }

    fn lt(&self, other: &Self) -> bool {
        other.head.lt(&self.head)
    }

    fn le(&self, other: &Self) -> bool {
        other.head.le(&self.head)
    }

    fn gt(&self, other: &Self) -> bool {
        other.head.gt(&self.head)
    }

    fn ge(&self, other: &Self) -> bool {
        other.head.ge(&self.head)
    }
}

impl<I> Ord for HeadTail<I>
    where I: Iterator,
          I::Item: Ord
{
    fn cmp(&self, other: &Self) -> Ordering {
        other.head.cmp(&self.head)
    }
}

/// Make `data` a heap (max-heap w.r.t T's Ord).
fn heapify<T: Ord>(data: &mut [T]) {
    for i in 0..data.len() / 2 {
        sift_down(data, i);
    }
}

/// Sift down element at `index` (`heap` is a max-heap wrt T's Ord).
fn sift_down<T: Ord>(heap: &mut [T], index: usize) {
    debug_assert!(index <= heap.len());
    let mut pos = index;
    let mut child = 2 * pos + 1;
    // the `pos` conditional is to avoid a bounds check
    while pos < heap.len() && child < heap.len() {
        let right = child + 1;

        // pick the bigger of the two children
        if right < heap.len() && heap[child] < heap[right] {
            child = right;
        }

        // sift down is done if we are already in order
        if heap[pos] >= heap[child] {
            return;
        }
        heap.swap(pos, child);
        pos = child;
        child = 2 * pos + 1;
    }
}

/// An iterator adaptor that merges an abitrary number of base iterators in ascending order.
/// If all base iterators are sorted (ascending), the result is sorted.
///
/// Iterator element type is `I::Item`.
///
/// See [*.kmerge()*](trait.Itertools.html#method.kmerge) for more information.
pub struct KMerge<I>
    where I: Iterator
{
    heap: Vec<HeadTail<I>>,
}

/// Create a `KMerge` iterator.
pub fn kmerge_new<I>(iter: I) -> KMerge<<I::Item as IntoIterator>::IntoIter>
    where I: Iterator,
          I::Item: IntoIterator,
          <<I as Iterator>::Item as IntoIterator>::Item: Ord
{
    let (lower, _) = iter.size_hint();
    let mut heap = Vec::with_capacity(lower);
    heap.extend(iter.filter_map(|it| HeadTail::new(it.into_iter())));
    heapify(&mut heap);
    KMerge { heap: heap }
}

impl<I> Clone for KMerge<I>
    where I: Iterator + Clone,
          I::Item: Clone
{
    fn clone(&self) -> KMerge<I> {
        clone_fields!(KMerge, self, heap)
    }
}

impl<I> Iterator for KMerge<I>
    where I: Iterator,
          I::Item: Ord
{
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        if self.heap.is_empty() {
            return None;
        }
        let result = if let Some(next) = self.heap[0].next() {
            next
        } else {
            self.heap.swap_remove(0).head
        };
        sift_down(&mut self.heap, 0);
        Some(result)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        if !self.heap.is_empty() {
            let mut it = self.heap.iter();
            let sh0 = it.next().unwrap().size_hint();
            it.fold(sh0, |acc, it| size_hint::add(acc, it.size_hint()))
        } else {
            (0, Some(0))
        }
    }
}

<commit_msg>kmerge: simplify size_hint<commit_after>
use size_hint;
use Itertools;

use std::cmp::Ordering;
use std::mem::replace;

macro_rules! clone_fields {
    ($name:ident, $base:expr, $($field:ident),+) => (
        $name {
            $(
                $field : $base . $field .clone()
            ),*
        }
    );
}

/// Head element and Tail iterator pair
///
/// `PartialEq`, `Eq`, `PartialOrd` and `Ord` are implemented by comparing sequences based on
/// first items (which are guaranteed to exist).
///
/// The meanings of `PartialOrd` and `Ord` are reversed so as to turn the heap used in
/// `KMerge` into a min-heap.
struct HeadTail<I>
    where I: Iterator
{
    head: I::Item,
    tail: I,
}

impl<I> HeadTail<I>
    where I: Iterator
{
    /// Constructs a `HeadTail` from an `Iterator`. Returns `None` if the `Iterator` is empty.
    fn new(mut it: I) -> Option<HeadTail<I>> {
        let head = it.next();
        head.map(|h| {
            HeadTail {
                head: h,
                tail: it,
            }
        })
    }

    /// Get the next element and update `head`, returning the old head in `Some`.
    ///
    /// Returns `None` when the tail is exhausted (only `head` then remains).
    fn next(&mut self) -> Option<I::Item> {
        if let Some(next) = self.tail.next() {
            Some(replace(&mut self.head, next))
        } else {
            None
        }
    }

    /// Hints at the size of the sequence, same as the `Iterator` method.
    fn size_hint(&self) -> (usize, Option<usize>) {
        size_hint::add_scalar(self.tail.size_hint(), 1)
    }
}

impl<I> Clone for HeadTail<I>
    where I: Iterator + Clone,
          I::Item: Clone
{
    fn clone(&self) -> Self {
        clone_fields!(HeadTail, self, head, tail)
    }
}

impl<I> PartialEq for HeadTail<I>
    where I: Iterator,
          I::Item: PartialEq
{
    fn eq(&self, other: &HeadTail<I>) -> bool {
        self.head.eq(&other.head)
    }
}

impl<I> Eq for HeadTail<I>
    where I: Iterator,
          I::Item: Eq
{}

impl<I> PartialOrd for HeadTail<I>
    where I: Iterator,
          I::Item: PartialOrd
{
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        other.head.partial_cmp(&self.head)
    }

    fn lt(&self, other: &Self) -> bool {
        other.head.lt(&self.head)
    }

    fn le(&self, other: &Self) -> bool {
        other.head.le(&self.head)
    }

    fn gt(&self, other: &Self) -> bool {
        other.head.gt(&self.head)
    }

    fn ge(&self, other: &Self) -> bool {
        other.head.ge(&self.head)
    }
}

impl<I> Ord for HeadTail<I>
    where I: Iterator,
          I::Item: Ord
{
    fn cmp(&self, other: &Self) -> Ordering {
        other.head.cmp(&self.head)
    }
}

/// Make `data` a heap (max-heap w.r.t T's Ord).
fn heapify<T: Ord>(data: &mut [T]) {
    for i in 0..data.len() / 2 {
        sift_down(data, i);
    }
}

/// Sift down element at `index` (`heap` is a max-heap wrt T's Ord).
fn sift_down<T: Ord>(heap: &mut [T], index: usize) {
    debug_assert!(index <= heap.len());
    let mut pos = index;
    let mut child = 2 * pos + 1;
    // the `pos` conditional is to avoid a bounds check
    while pos < heap.len() && child < heap.len() {
        let right = child + 1;

        // pick the bigger of the two children
        if right < heap.len() && heap[child] < heap[right] {
            child = right;
        }

        // sift down is done if we are already in order
        if heap[pos] >= heap[child] {
            return;
        }
        heap.swap(pos, child);
        pos = child;
        child = 2 * pos + 1;
    }
}

/// An iterator adaptor that merges an abitrary number of base iterators in ascending order.
/// If all base iterators are sorted (ascending), the result is sorted.
///
/// Iterator element type is `I::Item`.
///
/// See [*.kmerge()*](trait.Itertools.html#method.kmerge) for more information.
pub struct KMerge<I>
    where I: Iterator
{
    heap: Vec<HeadTail<I>>,
}

/// Create a `KMerge` iterator.
pub fn kmerge_new<I>(iter: I) -> KMerge<<I::Item as IntoIterator>::IntoIter>
    where I: Iterator,
          I::Item: IntoIterator,
          <<I as Iterator>::Item as IntoIterator>::Item: Ord
{
    let (lower, _) = iter.size_hint();
    let mut heap = Vec::with_capacity(lower);
    heap.extend(iter.filter_map(|it| HeadTail::new(it.into_iter())));
    heapify(&mut heap);
    KMerge { heap: heap }
}

impl<I> Clone for KMerge<I>
    where I: Iterator + Clone,
          I::Item: Clone
{
    fn clone(&self) -> KMerge<I> {
        clone_fields!(KMerge, self, heap)
    }
}

impl<I> Iterator for KMerge<I>
    where I: Iterator,
          I::Item: Ord
{
    type Item = I::Item;

    fn next(&mut self) -> Option<Self::Item> {
        if self.heap.is_empty() {
            return None;
        }
        let result = if let Some(next) = self.heap[0].next() {
            next
        } else {
            self.heap.swap_remove(0).head
        };
        sift_down(&mut self.heap, 0);
        Some(result)
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        self.heap.iter()
                 .map(|i| i.size_hint())
                 .fold1(size_hint::add)
                 .unwrap_or((0, Some(0)))
    }
}

<|endoftext|>


<commit_before>use smallvec::smallvec;

use crate::traits::{Obligation, ObligationCause, PredicateObligation};
use rustc_data_structures::fx::FxHashSet;
use rustc_middle::ty::outlives::Component;
use rustc_middle::ty::{self, ToPolyTraitRef, ToPredicate, TyCtxt, WithConstness};
use rustc_span::Span;

pub fn anonymize_predicate<'tcx>(
    tcx: TyCtxt<'tcx>,
    pred: ty::Predicate<'tcx>,
) -> ty::Predicate<'tcx> {
    let kind = pred.kind();
    let new = match kind {
        &ty::PredicateKind::Trait(ref data, constness) => {
            ty::PredicateKind::Trait(tcx.anonymize_late_bound_regions(data), constness)
        }

        ty::PredicateKind::RegionOutlives(data) => {
            ty::PredicateKind::RegionOutlives(tcx.anonymize_late_bound_regions(data))
        }

        ty::PredicateKind::TypeOutlives(data) => {
            ty::PredicateKind::TypeOutlives(tcx.anonymize_late_bound_regions(data))
        }

        ty::PredicateKind::Projection(data) => {
            ty::PredicateKind::Projection(tcx.anonymize_late_bound_regions(data))
        }

        &ty::PredicateKind::WellFormed(data) => ty::PredicateKind::WellFormed(data),

        &ty::PredicateKind::ObjectSafe(data) => ty::PredicateKind::ObjectSafe(data),

        &ty::PredicateKind::ClosureKind(closure_def_id, closure_substs, kind) => {
            ty::PredicateKind::ClosureKind(closure_def_id, closure_substs, kind)
        }

        ty::PredicateKind::Subtype(data) => {
            ty::PredicateKind::Subtype(tcx.anonymize_late_bound_regions(data))
        }

        &ty::PredicateKind::ConstEvaluatable(def_id, substs) => {
            ty::PredicateKind::ConstEvaluatable(def_id, substs)
        }

        ty::PredicateKind::ConstEquate(c1, c2) => ty::PredicateKind::ConstEquate(c1, c2),
    };

    if new != *kind { new.to_predicate(tcx) } else { pred }
}

struct PredicateSet<'tcx> {
    tcx: TyCtxt<'tcx>,
    set: FxHashSet<ty::Predicate<'tcx>>,
}

impl PredicateSet<'tcx> {
    fn new(tcx: TyCtxt<'tcx>) -> Self {
        Self { tcx, set: Default::default() }
    }

    fn insert(&mut self, pred: ty::Predicate<'tcx>) -> bool {
        // We have to be careful here because we want
        //
        //    for<'a> Foo<&'a i32>
        //
        // and
        //
        //    for<'b> Foo<&'b i32>
        //
        // to be considered equivalent. So normalize all late-bound
        // regions before we throw things into the underlying set.
        self.set.insert(anonymize_predicate(self.tcx, pred))
    }
}

impl Extend<ty::Predicate<'tcx>> for PredicateSet<'tcx> {
    fn extend<I: IntoIterator<Item = ty::Predicate<'tcx>>>(&mut self, iter: I) {
        for pred in iter {
            self.insert(pred);
        }
    }

    fn extend_one(&mut self, pred: ty::Predicate<'tcx>) {
        self.insert(pred);
    }

    fn extend_reserve(&mut self, additional: usize) {
        Extend::<ty::Predicate<'tcx>>::extend_reserve(&mut self.set, additional);
    }
}

///////////////////////////////////////////////////////////////////////////
// `Elaboration` iterator
///////////////////////////////////////////////////////////////////////////

/// "Elaboration" is the process of identifying all the predicates that
/// are implied by a source predicate. Currently, this basically means
/// walking the "supertraits" and other similar assumptions. For example,
/// if we know that `T: Ord`, the elaborator would deduce that `T: PartialOrd`
/// holds as well. Similarly, if we have `trait Foo: 'static`, and we know that
/// `T: Foo`, then we know that `T: 'static`.
pub struct Elaborator<'tcx> {
    stack: Vec<PredicateObligation<'tcx>>,
    visited: PredicateSet<'tcx>,
}

pub fn elaborate_trait_ref<'tcx>(
    tcx: TyCtxt<'tcx>,
    trait_ref: ty::PolyTraitRef<'tcx>,
) -> Elaborator<'tcx> {
    elaborate_predicates(tcx, std::iter::once(trait_ref.without_const().to_predicate(tcx)))
}

pub fn elaborate_trait_refs<'tcx>(
    tcx: TyCtxt<'tcx>,
    trait_refs: impl Iterator<Item = ty::PolyTraitRef<'tcx>>,
) -> Elaborator<'tcx> {
    let predicates = trait_refs.map(|trait_ref| trait_ref.without_const().to_predicate(tcx));
    elaborate_predicates(tcx, predicates)
}

pub fn elaborate_predicates<'tcx>(
    tcx: TyCtxt<'tcx>,
    predicates: impl Iterator<Item = ty::Predicate<'tcx>>,
) -> Elaborator<'tcx> {
    let obligations = predicates.map(|predicate| predicate_obligation(predicate, None)).collect();
    elaborate_obligations(tcx, obligations)
}

pub fn elaborate_obligations<'tcx>(
    tcx: TyCtxt<'tcx>,
    mut obligations: Vec<PredicateObligation<'tcx>>,
) -> Elaborator<'tcx> {
    let mut visited = PredicateSet::new(tcx);
    obligations.retain(|obligation| visited.insert(obligation.predicate));
    Elaborator { stack: obligations, visited }
}

fn predicate_obligation<'tcx>(
    predicate: ty::Predicate<'tcx>,
    span: Option<Span>,
) -> PredicateObligation<'tcx> {
    let cause = if let Some(span) = span {
        ObligationCause::dummy_with_span(span)
    } else {
        ObligationCause::dummy()
    };

    Obligation { cause, param_env: ty::ParamEnv::empty(), recursion_depth: 0, predicate }
}

impl Elaborator<'tcx> {
    pub fn filter_to_traits(self) -> FilterToTraits<Self> {
        FilterToTraits::new(self)
    }

    fn elaborate(&mut self, obligation: &PredicateObligation<'tcx>) {
        let tcx = self.visited.tcx;
        match obligation.predicate.kind() {
            ty::PredicateKind::Trait(ref data, _) => {
                // Get predicates declared on the trait.
                let predicates = tcx.super_predicates_of(data.def_id());

                let obligations = predicates.predicates.iter().map(|(pred, span)| {
                    predicate_obligation(
                        pred.subst_supertrait(tcx, &data.to_poly_trait_ref()),
                        Some(*span),
                    )
                });
                debug!("super_predicates: data={:?}", data);

                // Only keep those bounds that we haven't already seen.
                // This is necessary to prevent infinite recursion in some
                // cases. One common case is when people define
                // `trait Sized: Sized { }` rather than `trait Sized { }`.
                let visited = &mut self.visited;
                let obligations = obligations.filter(|o| visited.insert(o.predicate));

                self.stack.extend(obligations);
            }
            ty::PredicateKind::WellFormed(..) => {
                // Currently, we do not elaborate WF predicates,
                // although we easily could.
            }
            ty::PredicateKind::ObjectSafe(..) => {
                // Currently, we do not elaborate object-safe
                // predicates.
            }
            ty::PredicateKind::Subtype(..) => {
                // Currently, we do not "elaborate" predicates like `X <: Y`,
                // though conceivably we might.
            }
            ty::PredicateKind::Projection(..) => {
                // Nothing to elaborate in a projection predicate.
            }
            ty::PredicateKind::ClosureKind(..) => {
                // Nothing to elaborate when waiting for a closure's kind to be inferred.
            }
            ty::PredicateKind::ConstEvaluatable(..) => {
                // Currently, we do not elaborate const-evaluatable
                // predicates.
            }
            ty::PredicateKind::ConstEquate(..) => {
                // Currently, we do not elaborate const-equate
                // predicates.
            }
            ty::PredicateKind::RegionOutlives(..) => {
                // Nothing to elaborate from `'a: 'b`.
            }
            ty::PredicateKind::TypeOutlives(ref data) => {
                // We know that `T: 'a` for some type `T`. We can
                // often elaborate this. For example, if we know that
                // `[U]: 'a`, that implies that `U: 'a`. Similarly, if
                // we know `&'a U: 'b`, then we know that `'a: 'b` and
                // `U: 'b`.
                //
                // We can basically ignore bound regions here. So for
                // example `for<'c> Foo<'a,'c>: 'b` can be elaborated to
                // `'a: 'b`.

                // Ignore `for<'a> T: 'a` -- we might in the future
                // consider this as evidence that `T: 'static`, but
                // I'm a bit wary of such constructions and so for now
                // I want to be conservative. --nmatsakis
                let ty_max = data.skip_binder().0;
                let r_min = data.skip_binder().1;
                if r_min.is_late_bound() {
                    return;
                }

                let visited = &mut self.visited;
                let mut components = smallvec![];
                tcx.push_outlives_components(ty_max, &mut components);
                self.stack.extend(
                    components
                        .into_iter()
                        .filter_map(|component| match component {
                            Component::Region(r) => {
                                if r.is_late_bound() {
                                    None
                                } else {
                                    Some(ty::PredicateKind::RegionOutlives(ty::Binder::dummy(
                                        ty::OutlivesPredicate(r, r_min),
                                    )))
                                }
                            }

                            Component::Param(p) => {
                                let ty = tcx.mk_ty_param(p.index, p.name);
                                Some(ty::PredicateKind::TypeOutlives(ty::Binder::dummy(
                                    ty::OutlivesPredicate(ty, r_min),
                                )))
                            }

                            Component::UnresolvedInferenceVariable(_) => None,

                            Component::Projection(_) | Component::EscapingProjection(_) => {
                                // We can probably do more here. This
                                // corresponds to a case like `<T as
                                // Foo<'a>>::U: 'b`.
                                None
                            }
                        })
                        .map(|predicate_kind| predicate_kind.to_predicate(tcx))
                        .filter(|&predicate| visited.insert(predicate))
                        .map(|predicate| predicate_obligation(predicate, None)),
                );
            }
        }
    }
}

impl Iterator for Elaborator<'tcx> {
    type Item = PredicateObligation<'tcx>;

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.stack.len(), None)
    }

    fn next(&mut self) -> Option<Self::Item> {
        // Extract next item from top-most stack frame, if any.
        if let Some(obligation) = self.stack.pop() {
            self.elaborate(&obligation);
            Some(obligation)
        } else {
            None
        }
    }
}

///////////////////////////////////////////////////////////////////////////
// Supertrait iterator
///////////////////////////////////////////////////////////////////////////

pub type Supertraits<'tcx> = FilterToTraits<Elaborator<'tcx>>;

pub fn supertraits<'tcx>(
    tcx: TyCtxt<'tcx>,
    trait_ref: ty::PolyTraitRef<'tcx>,
) -> Supertraits<'tcx> {
    elaborate_trait_ref(tcx, trait_ref).filter_to_traits()
}

pub fn transitive_bounds<'tcx>(
    tcx: TyCtxt<'tcx>,
    bounds: impl Iterator<Item = ty::PolyTraitRef<'tcx>>,
) -> Supertraits<'tcx> {
    elaborate_trait_refs(tcx, bounds).filter_to_traits()
}

///////////////////////////////////////////////////////////////////////////
// Other
///////////////////////////////////////////////////////////////////////////

/// A filter around an iterator of predicates that makes it yield up
/// just trait references.
pub struct FilterToTraits<I> {
    base_iterator: I,
}

impl<I> FilterToTraits<I> {
    fn new(base: I) -> FilterToTraits<I> {
        FilterToTraits { base_iterator: base }
    }
}

impl<'tcx, I: Iterator<Item = PredicateObligation<'tcx>>> Iterator for FilterToTraits<I> {
    type Item = ty::PolyTraitRef<'tcx>;

    fn next(&mut self) -> Option<ty::PolyTraitRef<'tcx>> {
        while let Some(obligation) = self.base_iterator.next() {
            if let ty::PredicateKind::Trait(data, _) = obligation.predicate.kind() {
                return Some(data.to_poly_trait_ref());
            }
        }
        None
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let (_, upper) = self.base_iterator.size_hint();
        (0, upper)
    }
}
<commit_msg>anonymize_predicate<commit_after>use smallvec::smallvec;

use crate::traits::{Obligation, ObligationCause, PredicateObligation};
use rustc_data_structures::fx::FxHashSet;
use rustc_middle::ty::outlives::Component;
use rustc_middle::ty::{self, ToPolyTraitRef, ToPredicate, TyCtxt, WithConstness};
use rustc_span::Span;

pub fn anonymize_predicate<'tcx>(
    tcx: TyCtxt<'tcx>,
    pred: ty::Predicate<'tcx>,
) -> ty::Predicate<'tcx> {
    let kind = pred.kint(tcx);
    let new = match kind {
        ty::PredicateKint::ForAll(binder) => {
            ty::PredicateKint::ForAll(tcx.anonymize_late_bound_regions(binder))
        }
        &ty::PredicateKint::Trait(data, constness) => ty::PredicateKint::Trait(data, constness),

        &ty::PredicateKint::RegionOutlives(data) => ty::PredicateKint::RegionOutlives(data),

        &ty::PredicateKint::TypeOutlives(data) => ty::PredicateKint::TypeOutlives(data),

        &ty::PredicateKint::Projection(data) => ty::PredicateKint::Projection(data),

        &ty::PredicateKint::WellFormed(data) => ty::PredicateKint::WellFormed(data),

        &ty::PredicateKint::ObjectSafe(data) => ty::PredicateKint::ObjectSafe(data),

        &ty::PredicateKint::ClosureKind(closure_def_id, closure_substs, kind) => {
            ty::PredicateKint::ClosureKind(closure_def_id, closure_substs, kind)
        }

        &ty::PredicateKint::Subtype(data) => ty::PredicateKint::Subtype(data),

        &ty::PredicateKint::ConstEvaluatable(def_id, substs) => {
            ty::PredicateKint::ConstEvaluatable(def_id, substs)
        }

        &ty::PredicateKint::ConstEquate(c1, c2) => ty::PredicateKint::ConstEquate(c1, c2),
    };

    if new != *kind { new.to_predicate(tcx) } else { pred }
}

struct PredicateSet<'tcx> {
    tcx: TyCtxt<'tcx>,
    set: FxHashSet<ty::Predicate<'tcx>>,
}

impl PredicateSet<'tcx> {
    fn new(tcx: TyCtxt<'tcx>) -> Self {
        Self { tcx, set: Default::default() }
    }

    fn insert(&mut self, pred: ty::Predicate<'tcx>) -> bool {
        // We have to be careful here because we want
        //
        //    for<'a> Foo<&'a i32>
        //
        // and
        //
        //    for<'b> Foo<&'b i32>
        //
        // to be considered equivalent. So normalize all late-bound
        // regions before we throw things into the underlying set.
        self.set.insert(anonymize_predicate(self.tcx, pred))
    }
}

impl Extend<ty::Predicate<'tcx>> for PredicateSet<'tcx> {
    fn extend<I: IntoIterator<Item = ty::Predicate<'tcx>>>(&mut self, iter: I) {
        for pred in iter {
            self.insert(pred);
        }
    }

    fn extend_one(&mut self, pred: ty::Predicate<'tcx>) {
        self.insert(pred);
    }

    fn extend_reserve(&mut self, additional: usize) {
        Extend::<ty::Predicate<'tcx>>::extend_reserve(&mut self.set, additional);
    }
}

///////////////////////////////////////////////////////////////////////////
// `Elaboration` iterator
///////////////////////////////////////////////////////////////////////////

/// "Elaboration" is the process of identifying all the predicates that
/// are implied by a source predicate. Currently, this basically means
/// walking the "supertraits" and other similar assumptions. For example,
/// if we know that `T: Ord`, the elaborator would deduce that `T: PartialOrd`
/// holds as well. Similarly, if we have `trait Foo: 'static`, and we know that
/// `T: Foo`, then we know that `T: 'static`.
pub struct Elaborator<'tcx> {
    stack: Vec<PredicateObligation<'tcx>>,
    visited: PredicateSet<'tcx>,
}

pub fn elaborate_trait_ref<'tcx>(
    tcx: TyCtxt<'tcx>,
    trait_ref: ty::PolyTraitRef<'tcx>,
) -> Elaborator<'tcx> {
    elaborate_predicates(tcx, std::iter::once(trait_ref.without_const().to_predicate(tcx)))
}

pub fn elaborate_trait_refs<'tcx>(
    tcx: TyCtxt<'tcx>,
    trait_refs: impl Iterator<Item = ty::PolyTraitRef<'tcx>>,
) -> Elaborator<'tcx> {
    let predicates = trait_refs.map(|trait_ref| trait_ref.without_const().to_predicate(tcx));
    elaborate_predicates(tcx, predicates)
}

pub fn elaborate_predicates<'tcx>(
    tcx: TyCtxt<'tcx>,
    predicates: impl Iterator<Item = ty::Predicate<'tcx>>,
) -> Elaborator<'tcx> {
    let obligations = predicates.map(|predicate| predicate_obligation(predicate, None)).collect();
    elaborate_obligations(tcx, obligations)
}

pub fn elaborate_obligations<'tcx>(
    tcx: TyCtxt<'tcx>,
    mut obligations: Vec<PredicateObligation<'tcx>>,
) -> Elaborator<'tcx> {
    let mut visited = PredicateSet::new(tcx);
    obligations.retain(|obligation| visited.insert(obligation.predicate));
    Elaborator { stack: obligations, visited }
}

fn predicate_obligation<'tcx>(
    predicate: ty::Predicate<'tcx>,
    span: Option<Span>,
) -> PredicateObligation<'tcx> {
    let cause = if let Some(span) = span {
        ObligationCause::dummy_with_span(span)
    } else {
        ObligationCause::dummy()
    };

    Obligation { cause, param_env: ty::ParamEnv::empty(), recursion_depth: 0, predicate }
}

impl Elaborator<'tcx> {
    pub fn filter_to_traits(self) -> FilterToTraits<Self> {
        FilterToTraits::new(self)
    }

    fn elaborate(&mut self, obligation: &PredicateObligation<'tcx>) {
        let tcx = self.visited.tcx;
        match obligation.predicate.kind() {
            ty::PredicateKind::Trait(ref data, _) => {
                // Get predicates declared on the trait.
                let predicates = tcx.super_predicates_of(data.def_id());

                let obligations = predicates.predicates.iter().map(|(pred, span)| {
                    predicate_obligation(
                        pred.subst_supertrait(tcx, &data.to_poly_trait_ref()),
                        Some(*span),
                    )
                });
                debug!("super_predicates: data={:?}", data);

                // Only keep those bounds that we haven't already seen.
                // This is necessary to prevent infinite recursion in some
                // cases. One common case is when people define
                // `trait Sized: Sized { }` rather than `trait Sized { }`.
                let visited = &mut self.visited;
                let obligations = obligations.filter(|o| visited.insert(o.predicate));

                self.stack.extend(obligations);
            }
            ty::PredicateKind::WellFormed(..) => {
                // Currently, we do not elaborate WF predicates,
                // although we easily could.
            }
            ty::PredicateKind::ObjectSafe(..) => {
                // Currently, we do not elaborate object-safe
                // predicates.
            }
            ty::PredicateKind::Subtype(..) => {
                // Currently, we do not "elaborate" predicates like `X <: Y`,
                // though conceivably we might.
            }
            ty::PredicateKind::Projection(..) => {
                // Nothing to elaborate in a projection predicate.
            }
            ty::PredicateKind::ClosureKind(..) => {
                // Nothing to elaborate when waiting for a closure's kind to be inferred.
            }
            ty::PredicateKind::ConstEvaluatable(..) => {
                // Currently, we do not elaborate const-evaluatable
                // predicates.
            }
            ty::PredicateKind::ConstEquate(..) => {
                // Currently, we do not elaborate const-equate
                // predicates.
            }
            ty::PredicateKind::RegionOutlives(..) => {
                // Nothing to elaborate from `'a: 'b`.
            }
            ty::PredicateKind::TypeOutlives(ref data) => {
                // We know that `T: 'a` for some type `T`. We can
                // often elaborate this. For example, if we know that
                // `[U]: 'a`, that implies that `U: 'a`. Similarly, if
                // we know `&'a U: 'b`, then we know that `'a: 'b` and
                // `U: 'b`.
                //
                // We can basically ignore bound regions here. So for
                // example `for<'c> Foo<'a,'c>: 'b` can be elaborated to
                // `'a: 'b`.

                // Ignore `for<'a> T: 'a` -- we might in the future
                // consider this as evidence that `T: 'static`, but
                // I'm a bit wary of such constructions and so for now
                // I want to be conservative. --nmatsakis
                let ty_max = data.skip_binder().0;
                let r_min = data.skip_binder().1;
                if r_min.is_late_bound() {
                    return;
                }

                let visited = &mut self.visited;
                let mut components = smallvec![];
                tcx.push_outlives_components(ty_max, &mut components);
                self.stack.extend(
                    components
                        .into_iter()
                        .filter_map(|component| match component {
                            Component::Region(r) => {
                                if r.is_late_bound() {
                                    None
                                } else {
                                    Some(ty::PredicateKind::RegionOutlives(ty::Binder::dummy(
                                        ty::OutlivesPredicate(r, r_min),
                                    )))
                                }
                            }

                            Component::Param(p) => {
                                let ty = tcx.mk_ty_param(p.index, p.name);
                                Some(ty::PredicateKind::TypeOutlives(ty::Binder::dummy(
                                    ty::OutlivesPredicate(ty, r_min),
                                )))
                            }

                            Component::UnresolvedInferenceVariable(_) => None,

                            Component::Projection(_) | Component::EscapingProjection(_) => {
                                // We can probably do more here. This
                                // corresponds to a case like `<T as
                                // Foo<'a>>::U: 'b`.
                                None
                            }
                        })
                        .map(|predicate_kind| predicate_kind.to_predicate(tcx))
                        .filter(|&predicate| visited.insert(predicate))
                        .map(|predicate| predicate_obligation(predicate, None)),
                );
            }
        }
    }
}

impl Iterator for Elaborator<'tcx> {
    type Item = PredicateObligation<'tcx>;

    fn size_hint(&self) -> (usize, Option<usize>) {
        (self.stack.len(), None)
    }

    fn next(&mut self) -> Option<Self::Item> {
        // Extract next item from top-most stack frame, if any.
        if let Some(obligation) = self.stack.pop() {
            self.elaborate(&obligation);
            Some(obligation)
        } else {
            None
        }
    }
}

///////////////////////////////////////////////////////////////////////////
// Supertrait iterator
///////////////////////////////////////////////////////////////////////////

pub type Supertraits<'tcx> = FilterToTraits<Elaborator<'tcx>>;

pub fn supertraits<'tcx>(
    tcx: TyCtxt<'tcx>,
    trait_ref: ty::PolyTraitRef<'tcx>,
) -> Supertraits<'tcx> {
    elaborate_trait_ref(tcx, trait_ref).filter_to_traits()
}

pub fn transitive_bounds<'tcx>(
    tcx: TyCtxt<'tcx>,
    bounds: impl Iterator<Item = ty::PolyTraitRef<'tcx>>,
) -> Supertraits<'tcx> {
    elaborate_trait_refs(tcx, bounds).filter_to_traits()
}

///////////////////////////////////////////////////////////////////////////
// Other
///////////////////////////////////////////////////////////////////////////

/// A filter around an iterator of predicates that makes it yield up
/// just trait references.
pub struct FilterToTraits<I> {
    base_iterator: I,
}

impl<I> FilterToTraits<I> {
    fn new(base: I) -> FilterToTraits<I> {
        FilterToTraits { base_iterator: base }
    }
}

impl<'tcx, I: Iterator<Item = PredicateObligation<'tcx>>> Iterator for FilterToTraits<I> {
    type Item = ty::PolyTraitRef<'tcx>;

    fn next(&mut self) -> Option<ty::PolyTraitRef<'tcx>> {
        while let Some(obligation) = self.base_iterator.next() {
            if let ty::PredicateKind::Trait(data, _) = obligation.predicate.kind() {
                return Some(data.to_poly_trait_ref());
            }
        }
        None
    }

    fn size_hint(&self) -> (usize, Option<usize>) {
        let (_, upper) = self.base_iterator.size_hint();
        (0, upper)
    }
}
<|endoftext|>


<commit_before>use std::fmt;
use std::iter;
use std::cmp;
use super::stacks_and_queues::bag::Bag;

/// a weighted edge
#[derive(Clone, Copy)]
pub struct Edge {
    v: usize,
    w: usize,
    weight: f64
}

impl Edge {
    pub fn new(v: usize, w: usize, weight: f64) -> Edge {
        Edge { v: v, w: w, weight: weight }
    }

    pub fn weight(&self) -> f64 {
        self.weight
    }

    pub fn either(&self) -> usize {
        self.v
    }

    pub fn other(&self, vertex: usize) -> usize {
        if vertex == self.v {
            self.w
        } else if vertex == self.w {
            self.v
        } else {
            panic!("illegal endppint")
        }
    }

    fn swap(mut self) -> Edge {
        let v = self.v;
        self.v = self.w;
        self.w =v;
        self
    }
}

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

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

impl fmt::Display for Edge {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}-{} {:.5}", self.v, self.w, self.weight)
    }
}

pub struct EdgeWeightedGraph {
    v: usize,
    e: usize,
    adj: Vec<Bag<Edge>>
}

/// an edge-weighted graph
impl EdgeWeightedGraph {
    pub fn new(v: usize) -> EdgeWeightedGraph {
        EdgeWeightedGraph {
            v: v,
            e: 0,
            adj: iter::repeat(Bag::new()).take(v).collect::<Vec<Bag<Edge>>>()
        }
    }

    pub fn v(&self) -> usize {
        self.v
    }

    pub fn e(&self) -> usize {
        self.e
    }

    #[inline]
    fn validate_vertex(&self, v: usize) {
        assert!(v < self.v, "vertex is not between 0 and max V")
    }

    pub fn add_edge(&mut self, e: Edge) {
        let v = e.either();
        let w = e.other(v);
        self.validate_vertex(v);
        self.validate_vertex(w);
        self.adj[v].add(e.clone());
        if v != w {
            self.adj[w].add(e.swap());
        }
        self.e += 1
    }

    pub fn adj(&self, v: usize) -> ::std::vec::IntoIter<Edge> {
        self.validate_vertex(v);
        self.adj[v].iter().map(|e| e.clone()).collect::<Vec<Edge>>().into_iter()
    }

    pub fn degree(&self, v: usize) -> usize {
        self.validate_vertex(v);
        self.adj[v].len()
    }

    pub fn edges(&self) -> ::std::vec::IntoIter<Edge> {
        self.adj
            .iter()
            .flat_map(|adj| {
                adj.iter().map(|e| e.clone()).collect::<Vec<Edge>>().into_iter()
            })
            .filter(|ref e| e.either() <= e.other(e.either()) )
            .collect::<Vec<Edge>>()
            .into_iter()
    }

    // use fdp or neato
    pub fn to_dot(&self) -> String {
        let mut dot = String::new();

        dot.push_str("graph G {\n");
        for i in 0 .. self.v {
            dot.push_str(&format!("  {};\n", i));
        }

        for e in self.edges() {
            let v = e.either();
            let w = e.other(v);
            dot.push_str(&format!("  {} -- {} [len={}];\n",
                                  v, w, e.weight))
        }
        dot.push_str("}\n");
        dot

    }
}

#[test]
fn test_edge_weighted_graph() {
    let mut g = EdgeWeightedGraph::new(6);
    g.add_edge(Edge::new(0, 1, 7.0));
    g.add_edge(Edge::new(1, 2, 10.0));
    g.add_edge(Edge::new(0, 2, 9.0));
    g.add_edge(Edge::new(0, 5, 14.0));
    g.add_edge(Edge::new(1, 3, 15.0));
    g.add_edge(Edge::new(2, 5, 2.0));
    g.add_edge(Edge::new(2, 3, 11.0));
    g.add_edge(Edge::new(4, 5, 9.0));
    g.add_edge(Edge::new(3, 4, 6.0));
    g.add_edge(Edge::new(2, 2, 1.0));

    assert_eq!(10, g.edges().count());
    assert!(!g.to_dot().is_empty());
}
<commit_msg>minimum_spanning_trees: add KruskalMST<commit_after>use std::fmt;
use std::iter;
use std::cmp;
use super::stacks_and_queues::bag::Bag;
use super::stacks_and_queues::Queue;
use super::stacks_and_queues::resizing_array_queue::ResizingArrayQueue;
use super::priority_queues::MinPQ;
use super::priority_queues::binary_heaps::BinaryHeapMinPQ;
use super::union_find::UF;
use super::union_find::weighted_quick_union::UnionFind;

/// a weighted edge
#[derive(Clone, Copy)]
pub struct Edge {
    v: usize,
    w: usize,
    weight: f64
}

impl Edge {
    pub fn new(v: usize, w: usize, weight: f64) -> Edge {
        Edge { v: v, w: w, weight: weight }
    }

    pub fn weight(&self) -> f64 {
        self.weight
    }

    pub fn either(&self) -> usize {
        self.v
    }

    pub fn other(&self, vertex: usize) -> usize {
        if vertex == self.v {
            self.w
        } else if vertex == self.w {
            self.v
        } else {
            panic!("illegal endppint")
        }
    }

    fn swap(mut self) -> Edge {
        let v = self.v;
        self.v = self.w;
        self.w =v;
        self
    }
}

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

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

impl fmt::Display for Edge {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{}-{} {:.5}", self.v, self.w, self.weight)
    }
}

pub struct EdgeWeightedGraph {
    v: usize,
    e: usize,
    adj: Vec<Bag<Edge>>
}

/// an edge-weighted graph
impl EdgeWeightedGraph {
    pub fn new(v: usize) -> EdgeWeightedGraph {
        EdgeWeightedGraph {
            v: v,
            e: 0,
            adj: iter::repeat(Bag::new()).take(v).collect::<Vec<Bag<Edge>>>()
        }
    }

    pub fn v(&self) -> usize {
        self.v
    }

    pub fn e(&self) -> usize {
        self.e
    }

    #[inline]
    fn validate_vertex(&self, v: usize) {
        assert!(v < self.v, "vertex is not between 0 and max V")
    }

    pub fn add_edge(&mut self, e: Edge) {
        let v = e.either();
        let w = e.other(v);
        self.validate_vertex(v);
        self.validate_vertex(w);
        self.adj[v].add(e.clone());
        if v != w {
            self.adj[w].add(e.swap());
        }
        self.e += 1
    }

    pub fn adj(&self, v: usize) -> ::std::vec::IntoIter<Edge> {
        self.validate_vertex(v);
        self.adj[v].iter().map(|e| e.clone()).collect::<Vec<Edge>>().into_iter()
    }

    pub fn degree(&self, v: usize) -> usize {
        self.validate_vertex(v);
        self.adj[v].len()
    }

    pub fn edges(&self) -> ::std::vec::IntoIter<Edge> {
        self.adj
            .iter()
            .flat_map(|adj| {
                adj.iter().map(|e| e.clone()).collect::<Vec<Edge>>().into_iter()
            })
            .filter(|ref e| e.either() <= e.other(e.either()) )
            .collect::<Vec<Edge>>()
            .into_iter()
    }

    // use fdp or neato
    pub fn to_dot(&self) -> String {
        let mut dot = String::new();

        dot.push_str("graph G {\n");
        for i in 0 .. self.v {
            dot.push_str(&format!("  {};\n", i));
        }

        for e in self.edges() {
            let v = e.either();
            let w = e.other(v);
            dot.push_str(&format!("  {} -- {} [len={}];\n",
                                  v, w, e.weight))
        }
        dot.push_str("}\n");
        dot

    }
}

#[test]
fn test_edge_weighted_graph() {
    let mut g = EdgeWeightedGraph::new(6);
    g.add_edge(Edge::new(0, 1, 7.0));
    g.add_edge(Edge::new(1, 2, 10.0));
    g.add_edge(Edge::new(0, 2, 9.0));
    g.add_edge(Edge::new(0, 5, 14.0));
    g.add_edge(Edge::new(1, 3, 15.0));
    g.add_edge(Edge::new(2, 5, 2.0));
    g.add_edge(Edge::new(2, 3, 11.0));
    g.add_edge(Edge::new(4, 5, 9.0));
    g.add_edge(Edge::new(3, 4, 6.0));
    g.add_edge(Edge::new(2, 2, 1.0));

    assert_eq!(10, g.edges().count());
    assert!(!g.to_dot().is_empty());
}

#[allow(dead_code)]
pub struct KruskalMST<'a> {
    graph: &'a EdgeWeightedGraph,
    weight: f64,
    mst: ResizingArrayQueue<Edge>
}

impl<'a> KruskalMST<'a> {
    fn new<'b>(graph: &'b EdgeWeightedGraph) -> KruskalMST<'b> {
        let n = graph.v();
        let mut weight = 0f64;
        let mut mst = ResizingArrayQueue::<Edge>::new();
        let mut pq = BinaryHeapMinPQ::<Edge>::new();
        for e in graph.edges() {
            pq.insert(e);
        }
        let mut uf = UnionFind::new(n);

        while !pq.is_empty() && mst.size() < n - 1 {
            let e = pq.del_min().unwrap();
            let v = e.either();
            let w = e.other(v);
            if !uf.connected(v, w) {
                uf.union(v, w);
                weight += e.weight();
                mst.enqueue(e);
            }
        }

        KruskalMST {
            graph: graph,
            weight: weight,
            mst: mst
        }
    }

    pub fn edges(&self) -> ::std::vec::IntoIter<Edge> {
        self.mst.clone().into_iter().collect::<Vec<Edge>>().into_iter()
    }
}

impl EdgeWeightedGraph {
    pub fn kruskal_mst<'a>(&'a self) -> KruskalMST<'a> {
        KruskalMST::new(self)
    }
}

#[test]
fn test_edge_weighted_graph_kruskal_mst() {
    let mut g = EdgeWeightedGraph::new(6);
    g.add_edge(Edge::new(0, 1, 7.0));
    g.add_edge(Edge::new(1, 2, 10.0));
    g.add_edge(Edge::new(0, 2, 9.0));
    g.add_edge(Edge::new(0, 5, 14.0));
    g.add_edge(Edge::new(1, 3, 15.0));
    g.add_edge(Edge::new(2, 5, 2.0));
    g.add_edge(Edge::new(2, 3, 11.0));
    g.add_edge(Edge::new(4, 5, 9.0));
    g.add_edge(Edge::new(3, 4, 6.0));
    g.add_edge(Edge::new(2, 2, 1.0));

    assert_eq!(33.0, g.kruskal_mst().weight);
    assert_eq!(33.0, g.kruskal_mst().edges().map(|e| e.weight).sum());

}

/// data type for computing a minimum spanning tree in an edge-weighted graph
pub struct LazyPrimMST<'a> {
    graph: &'a EdgeWeightedGraph,
    weight: f64,
    mst: ResizingArrayQueue<Edge>,
    marked: Vec<bool>,
    pq: BinaryHeapMinPQ<Edge>
}
<|endoftext|>


<commit_before>//! An echo bot that just repeats stuff sent to it (either in a channel or as PRIVMSG).

use libtiny_client::{Client, Event, ServerInfo};
use libtiny_wire::{Cmd, Msg, MsgTarget, Pfx};

use futures_util::stream::StreamExt;
use std::process::exit;

fn main() {
    // echo <nick> <server> <port> [<chan_1> ... <chan_N>]
    let mut args_vec: Vec<String> = std::env::args().collect();
    if args_vec.len() < 4 {
        show_usage();
        exit(1);
    }

    let nick = args_vec.remove(1);
    let server = args_vec.remove(1);
    let port_str = args_vec.remove(1);
    let port = match port_str.parse::<u16>() {
        Ok(port) => port,
        Err(err) => {
            println!("Can't parse port: {:?}", port_str);
            println!("{}", err);
            exit(1);
        }
    };

    let chans = &args_vec[1..];

    let server_info = ServerInfo {
        addr: server,
        port,
        tls: false,
        pass: None,
        realname: "tiny echo bot".to_owned(),
        nicks: vec![nick],
        auto_join: chans.to_owned(),
        nickserv_ident: None,
        sasl_auth: None,
    };

    println!("{:?}", server_info);

    let mut executor = tokio::runtime::current_thread::Runtime::new().unwrap();

    executor.spawn(echo_bot_task(server_info));

    executor.run().unwrap();
}

fn show_usage() {
    println!("echo <nick> <server> <port> [<chan_1> .. <chan_N>]");
}

static NICK_SEP: [&str; 4] = [": ", ", ", ":", ","];

async fn echo_bot_task(server_info: ServerInfo) {
    let (mut client, mut rcv_ev) = Client::new(server_info, None);

    while let Some(ev) = rcv_ev.next().await {
        println!("Client event: {:?}", ev);
        if let Event::Msg(Msg {
            pfx: Some(Pfx::User { nick, .. }),
            cmd: Cmd::PRIVMSG { target, msg, .. },
        }) = ev
        {
            let echo_msg = match target {
                MsgTarget::User(_) => {
                    // Message is a PRIVMSG to us, just echo the whole message to the sender
                    Some((nick, msg))
                }
                MsgTarget::Chan(chan) => {
                    // Message was sent to a channel. Only echo if it's directed at us
                    let our_nick = client.get_nick();
                    if msg.starts_with(&our_nick) {
                        let mut msg = &msg[our_nick.len()..];
                        for nick_sep in NICK_SEP.iter() {
                            if msg.starts_with(nick_sep) {
                                msg = &msg[nick_sep.len()..];
                                break;
                            }
                        }
                        Some((chan, msg.to_owned()))
                    } else {
                        None
                    }
                }
            };

            if let Some((target, msg)) = echo_msg {
                client.privmsg(&target, &msg, false);
            }
        }
    }
}
<commit_msg>Fix echo bot<commit_after>//! An echo bot that just repeats stuff sent to it (either in a channel or as PRIVMSG).

use libtiny_client::{Client, Event, ServerInfo};
use libtiny_wire::{Cmd, Msg, MsgTarget, Pfx};

use futures_util::stream::StreamExt;
use std::process::exit;

fn main() {
    // echo <nick> <server> <port> [<chan_1> ... <chan_N>]
    let mut args_vec: Vec<String> = std::env::args().collect();
    if args_vec.len() < 4 {
        show_usage();
        exit(1);
    }

    let nick = args_vec.remove(1);
    let server = args_vec.remove(1);
    let port_str = args_vec.remove(1);
    let port = match port_str.parse::<u16>() {
        Ok(port) => port,
        Err(err) => {
            println!("Can't parse port: {:?}", port_str);
            println!("{}", err);
            exit(1);
        }
    };

    let chans = &args_vec[1..];

    let server_info = ServerInfo {
        addr: server,
        port,
        tls: false,
        pass: None,
        realname: "tiny echo bot".to_owned(),
        nicks: vec![nick],
        auto_join: chans.to_owned(),
        nickserv_ident: None,
        sasl_auth: None,
    };

    println!("{:?}", server_info);

    let mut executor = tokio::runtime::current_thread::Runtime::new().unwrap();

    executor.spawn(echo_bot_task(server_info));

    executor.run().unwrap();
}

fn show_usage() {
    println!("echo <nick> <server> <port> [<chan_1> .. <chan_N>]");
}

static NICK_SEP: [&str; 4] = [": ", ", ", ":", ","];

async fn echo_bot_task(server_info: ServerInfo) {
    let (mut client, mut rcv_ev) = Client::new(server_info);

    while let Some(ev) = rcv_ev.next().await {
        println!("Client event: {:?}", ev);
        if let Event::Msg(Msg {
            pfx: Some(Pfx::User { nick, .. }),
            cmd: Cmd::PRIVMSG { target, msg, .. },
        }) = ev
        {
            let echo_msg = match target {
                MsgTarget::User(_) => {
                    // Message is a PRIVMSG to us, just echo the whole message to the sender
                    Some((nick, msg))
                }
                MsgTarget::Chan(chan) => {
                    // Message was sent to a channel. Only echo if it's directed at us
                    let our_nick = client.get_nick();
                    if msg.starts_with(&our_nick) {
                        let mut msg = &msg[our_nick.len()..];
                        for nick_sep in NICK_SEP.iter() {
                            if msg.starts_with(nick_sep) {
                                msg = &msg[nick_sep.len()..];
                                break;
                            }
                        }
                        Some((chan, msg.to_owned()))
                    } else {
                        None
                    }
                }
            };

            if let Some((target, msg)) = echo_msg {
                client.privmsg(&target, &msg, false);
            }
        }
    }
}
<|endoftext|>


<commit_before>//! UDP relay local server

use std::{io, net::SocketAddr};

use async_trait::async_trait;
use log::{error, info, trace};

use crate::{
    config::RedirType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        redir::UdpSocketRedirExt,
        socks5::Address,
    },
};

use super::{
    association::{ProxyAssociation, ProxyAssociationManager, ProxySend},
    redir::sys::UdpRedirSocket,
    MAXIMUM_UDP_PAYLOAD_SIZE,
};

struct ProxyHandler {
    ty: RedirType,
    src_addr: SocketAddr,
    cache_key: String,
    assoc_map: ProxyAssociationManager<String>,
}

impl ProxyHandler {
    pub fn new(
        ty: RedirType,
        src_addr: SocketAddr,
        cache_key: String,
        assoc_map: ProxyAssociationManager<String>,
    ) -> io::Result<ProxyHandler> {
        Ok(ProxyHandler {
            ty,
            src_addr,
            cache_key,
            assoc_map,
        })
    }
}

#[async_trait]
impl ProxySend for ProxyHandler {
    async fn send_packet(&mut self, addr: Address, data: Vec<u8>) -> io::Result<()> {
        // Redirect only if the target is a SocketAddress
        if let Address::SocketAddress(ref dst_addr) = addr {
            // Create a socket binds to destination addr
            // This only works for systems that supports binding to non-local addresses
            let mut local_udp = UdpRedirSocket::bind(self.ty, dst_addr)?;

            local_udp.send_to(&data, &self.src_addr).await?;

            // Update LRU
            self.assoc_map.keep_alive(&self.cache_key).await;

            return Ok(());
        }

        Err(io::Error::new(
            io::ErrorKind::Other,
            "Address from remote is not a socket addr",
        ))
    }
}

/// Starts a UDP local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = context.config().local_addr.as_ref().expect("missing local config");
    let bind_addr = local_addr.bind_addr(&context).await?;

    let ty = context.config().udp_redir;

    // let l = create_socket(&bind_addr).await?;
    let mut l = UdpRedirSocket::bind(ty, &bind_addr)?;
    let local_addr = l.local_addr().expect("determine port bound to");

    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Udp).await;

    info!("shadowsocks UDP redirect listening on {}", local_addr);

    let assoc_manager = ProxyAssociationManager::new(context.config());

    let mut pkt_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];

    loop {
        let (recv_len, src, dst) = l.recv_from_redir(&mut pkt_buf).await?;

        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
        // Copy bytes, because udp_associate runs in another tokio Task
        let pkt = &pkt_buf[..recv_len];

        trace!(
            "received UDP packet from {}, destination {}, length {} bytes",
            src,
            dst,
            recv_len
        );

        if recv_len == 0 {
            // For windows, it will generate a ICMP Port Unreachable Message
            // https://docs.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-recvfrom
            // Which will result in recv_from return 0.
            //
            // It cannot be solved here, because `WSAGetLastError` is already set.
            //
            // See `relay::udprelay::utils::create_socket` for more detail.
            continue;
        }

        // Check destination should be proxied or not
        let target = Address::SocketAddress(dst);
        let is_bypassed = context.check_target_bypassed(&target).await;

        // Check or (re)create an association
        let cache_key = format!("{}-{}", src, dst);
        let cache_key_cloned = cache_key.clone();
        let res = assoc_manager
            .send_packet(cache_key, target, pkt.to_vec(), async {
                // Pick a server
                let server = balancer.pick_server();

                let sender = match ProxyHandler::new(ty, src, cache_key_cloned, assoc_manager.clone()) {
                    Ok(s) => s,
                    Err(err) => {
                        debug!("create UDP association for {} <-> {}, error: {}", src, dst, err);
                        continue;
                    }
                };

                if is_bypassed {
                    ProxyAssociation::associate_bypassed(src, server, sender).await
                } else {
                    ProxyAssociation::associate_proxied(src, server, sender).await
                }
            })
            .await;

        if let Err(err) = res {
            debug!("failed to create UDP association, {}", err);
        }
    }
}
<commit_msg>Fixed build redir_local<commit_after>//! UDP relay local server

use std::{io, net::SocketAddr};

use async_trait::async_trait;
use log::{debug, info, trace};

use crate::{
    config::RedirType,
    context::SharedContext,
    relay::{
        loadbalancing::server::{PlainPingBalancer, ServerType},
        redir::UdpSocketRedirExt,
        socks5::Address,
    },
};

use super::{
    association::{ProxyAssociation, ProxyAssociationManager, ProxySend},
    redir::sys::UdpRedirSocket,
    MAXIMUM_UDP_PAYLOAD_SIZE,
};

struct ProxyHandler {
    ty: RedirType,
    src_addr: SocketAddr,
    cache_key: String,
    assoc_map: ProxyAssociationManager<String>,
}

impl ProxyHandler {
    pub fn new(
        ty: RedirType,
        src_addr: SocketAddr,
        cache_key: String,
        assoc_map: ProxyAssociationManager<String>,
    ) -> io::Result<ProxyHandler> {
        Ok(ProxyHandler {
            ty,
            src_addr,
            cache_key,
            assoc_map,
        })
    }
}

#[async_trait]
impl ProxySend for ProxyHandler {
    async fn send_packet(&mut self, addr: Address, data: Vec<u8>) -> io::Result<()> {
        // Redirect only if the target is a SocketAddress
        if let Address::SocketAddress(ref dst_addr) = addr {
            // Create a socket binds to destination addr
            // This only works for systems that supports binding to non-local addresses
            let mut local_udp = UdpRedirSocket::bind(self.ty, dst_addr)?;

            local_udp.send_to(&data, &self.src_addr).await?;

            // Update LRU
            self.assoc_map.keep_alive(&self.cache_key).await;

            return Ok(());
        }

        Err(io::Error::new(
            io::ErrorKind::Other,
            "Address from remote is not a socket addr",
        ))
    }
}

/// Starts a UDP local server
pub async fn run(context: SharedContext) -> io::Result<()> {
    let local_addr = context.config().local_addr.as_ref().expect("missing local config");
    let bind_addr = local_addr.bind_addr(&context).await?;

    let ty = context.config().udp_redir;

    // let l = create_socket(&bind_addr).await?;
    let mut l = UdpRedirSocket::bind(ty, &bind_addr)?;
    let local_addr = l.local_addr().expect("determine port bound to");

    let balancer = PlainPingBalancer::new(context.clone(), ServerType::Udp).await;

    info!("shadowsocks UDP redirect listening on {}", local_addr);

    let assoc_manager = ProxyAssociationManager::new(context.config());

    let mut pkt_buf = vec![0u8; MAXIMUM_UDP_PAYLOAD_SIZE];

    loop {
        let (recv_len, src, dst) = l.recv_from_redir(&mut pkt_buf).await?;

        // Packet length is limited by MAXIMUM_UDP_PAYLOAD_SIZE, excess bytes will be discarded.
        // Copy bytes, because udp_associate runs in another tokio Task
        let pkt = &pkt_buf[..recv_len];

        trace!(
            "received UDP packet from {}, destination {}, length {} bytes",
            src,
            dst,
            recv_len
        );

        if recv_len == 0 {
            // For windows, it will generate a ICMP Port Unreachable Message
            // https://docs.microsoft.com/en-us/windows/win32/api/winsock2/nf-winsock2-recvfrom
            // Which will result in recv_from return 0.
            //
            // It cannot be solved here, because `WSAGetLastError` is already set.
            //
            // See `relay::udprelay::utils::create_socket` for more detail.
            continue;
        }

        // Check destination should be proxied or not
        let target = Address::SocketAddress(dst);
        let is_bypassed = context.check_target_bypassed(&target).await;

        // Check or (re)create an association
        let cache_key = format!("{}-{}", src, dst);
        let cache_key_cloned = cache_key.clone();
        let res = assoc_manager
            .send_packet(cache_key, target, pkt.to_vec(), async {
                // Pick a server
                let server = balancer.pick_server();

                let sender = match ProxyHandler::new(ty, src, cache_key_cloned, assoc_manager.clone()) {
                    Ok(s) => s,
                    Err(err) => {
                        debug!("create UDP association for {} <-> {}, error: {}", src, dst, err);
                        return Err(err);
                    }
                };

                if is_bypassed {
                    ProxyAssociation::associate_bypassed(src, server, sender).await
                } else {
                    ProxyAssociation::associate_proxied(src, server, sender).await
                }
            })
            .await;

        if let Err(err) = res {
            debug!("failed to create UDP association, {}", err);
        }
    }
}
<|endoftext|>


<commit_before>extern crate gcc;

use std::env;

fn main() {
    println!("cargo:rustc-link-lib=dylib=msi");
    println!("cargo:rustc-link-lib=dylib=user32");
    println!("cargo:rustc-link-lib=dylib=mincore");

    // Part of WIX SDK
    println!("cargo:rustc-link-lib=static=wcautil");
    println!("cargo:rustc-link-lib=static=dutil");

    let wix_path = env::var("WIX").expect("WIX must be installed, and 'WIX' environment variable must be set");

    // For the correct WIX library path, we need to know which VS version we are using.
    // We use the `gcc` crate to get the path to the correct `cl.exe`, and then try to
    // guess the version from the path components (we don't depend on the actual location
    // of the VS installation, but only on the internal directory structure).
    let config = gcc::Config::new();
    let compiler = config.get_compiler();
    let path = compiler.path();
    let vs_version = if path.to_string_lossy().contains("VC\\bin") {
        println!("cargo:warning=Guessing VS version: VS2015");
        "VS2015"
    } else if path.to_string_lossy().contains("VC\\Tools\\MSVC") {
        println!("cargo:warning=Guessing VS version: VS2017");
        "VS2017"
    } else {
        panic!("Unknown VS version with compiler path {:?}", path);
    };

    let target_arch = env::var("CARGO_CFG_TARGET_ARCH").expect("cannot read CARGO_CFG_TARGET_ARCH in build script");
    let target_arch = match target_arch.as_str() {
        "x86" => "x86",
        "x86_64" => "x64",
        other => panic!("Target architecture {} not supported by WIX.", other)
    };

    // Tell cargo about the WIX SDK path for `wcautil.lib` and `dutil.lib`
    println!("cargo:rustc-link-search=native={}SDK\\{}\\lib\\{}", wix_path, vs_version, target_arch);
}<commit_msg>Use the gcc crate's find_vs_version()<commit_after>extern crate gcc;

use std::env;
use gcc::windows_registry::{self, VsVers};

fn main() {
    println!("cargo:rustc-link-lib=dylib=msi");
    println!("cargo:rustc-link-lib=dylib=user32");
    println!("cargo:rustc-link-lib=dylib=mincore");

    // Part of WIX SDK
    println!("cargo:rustc-link-lib=static=wcautil");
    println!("cargo:rustc-link-lib=static=dutil");

    let wix_path = env::var("WIX").expect("WIX must be installed, and 'WIX' environment variable must be set");

    // For the correct WIX library path, we need to know which VS version we are using.
    // We use the `gcc` crate's functionality to do this, which should always match what rustc is doing.
    let vs_version = windows_registry::find_vs_version().expect("Cannot find VS version");
    let vs_version_string = match vs_version {
        VsVers::Vs14 => "VS2015",
        VsVers::Vs15 => "VS2017",
        VsVers::Vs12 => panic!("Unsupported VS version: Vs12"),
        _ => panic!("Unsupported VS version") // FIXME: should use {:?}, but `VsVers` does not yet implement `Debug`
    };

    println!("cargo:warning=Using WIX libraries for VS version: {}", vs_version_string);

    let target_arch = env::var("CARGO_CFG_TARGET_ARCH").expect("cannot read CARGO_CFG_TARGET_ARCH in build script");
    let target_arch = match target_arch.as_str() {
        "x86" => "x86",
        "x86_64" => "x64",
        other => panic!("Target architecture {} not supported by WIX.", other)
    };

    // Tell cargo about the WIX SDK path for `wcautil.lib` and `dutil.lib`
    println!("cargo:rustc-link-search=native={}SDK\\{}\\lib\\{}", wix_path, vs_version_string, target_arch);
}<|endoftext|>


<commit_before>use std::str::FromStr;

#[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct SyntaxScope {
    name: String
}

pub type SyntaxScopePath = Vec<SyntaxScope>;

#[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct SyntaxScopeSelector {
    path: SyntaxScopePath,
    exclude: Option<SyntaxScope>
}

#[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct SyntaxScopeSelectors {
    selectors: Vec<SyntaxScopeSelector>
}

#[derive(Debug)]
pub struct ParseSyntaxScopeError(char);

impl FromStr for SyntaxScope {
    type Err = ParseSyntaxScopeError;

    fn from_str(s: &str) -> Result<SyntaxScope, ParseSyntaxScopeError> {
        let mut name = s.trim().to_lowercase();
        for c in name.chars() {
            match c {
                'a' ... 'z' | '0' ... '9' | '-' | '.' => (),
                c => return Err(ParseSyntaxScopeError(c))
            }
        };
        Ok(SyntaxScope { name: name })
    }
}

impl FromStr for SyntaxScopeSelector {
    type Err = ParseSyntaxScopeError;

    fn from_str(mut s: &str) -> Result<SyntaxScopeSelector, ParseSyntaxScopeError> {
        let exclude = match s.find(" - ") {
            Some(index) => {
                let (path_str, exclude_str) = s.split_at(index);
                s = path_str;
                Some(try!(SyntaxScope::from_str(exclude_str)))
            },
            None => None
        };
        let mut path = Vec::new();
        for name in s.split_whitespace() {
            path.push(try!(SyntaxScope::from_str(name)))
        };
        Ok(SyntaxScopeSelector {
            path: path,
            exclude: exclude
        })
    }
}

impl FromStr for SyntaxScopeSelectors {
    type Err = ParseSyntaxScopeError;

    fn from_str(mut s: &str) -> Result<SyntaxScopeSelectors, ParseSyntaxScopeError> {
        let mut selectors = Vec::new();
        for selector in s.split(',') {
            selectors.push(try!(SyntaxScopeSelector::from_str(selector)))
        };
        Ok(SyntaxScopeSelectors {
            selectors: selectors
        })
    }
}
<commit_msg>implement rank and matching for syntax selectors<commit_after>use std::str::FromStr;

pub type Rank = u64;

#[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct SyntaxScope {
    name: String
}

pub type SyntaxScopePath = Vec<SyntaxScope>;

#[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct SyntaxScopeSelector {
    path: SyntaxScopePath,
    exclude: Option<SyntaxScope>
}

#[derive(Debug, Default, Clone, PartialEq, Eq, PartialOrd, Ord)]
pub struct SyntaxScopeSelectors {
    selectors: Vec<SyntaxScopeSelector>
}

#[derive(Debug)]
pub struct ParseSyntaxScopeError(char);

impl FromStr for SyntaxScope {
    type Err = ParseSyntaxScopeError;

    fn from_str(s: &str) -> Result<SyntaxScope, ParseSyntaxScopeError> {
        let mut name = s.trim().to_lowercase();
        for c in name.chars() {
            match c {
                'a' ... 'z' | '0' ... '9' | '-' | '.' => (),
                c => return Err(ParseSyntaxScopeError(c))
            }
        };
        Ok(SyntaxScope { name: name })
    }
}

impl FromStr for SyntaxScopeSelector {
    type Err = ParseSyntaxScopeError;

    fn from_str(mut s: &str) -> Result<SyntaxScopeSelector, ParseSyntaxScopeError> {
        let exclude = match s.find(" - ") {
            Some(index) => {
                let (path_str, exclude_str) = s.split_at(index);
                s = path_str;
                Some(try!(SyntaxScope::from_str(exclude_str)))
            },
            None => None
        };
        let mut path = Vec::new();
        for name in s.split_whitespace() {
            path.push(try!(SyntaxScope::from_str(name)))
        };
        Ok(SyntaxScopeSelector {
            path: path,
            exclude: exclude
        })
    }
}

impl FromStr for SyntaxScopeSelectors {
    type Err = ParseSyntaxScopeError;

    fn from_str(mut s: &str) -> Result<SyntaxScopeSelectors, ParseSyntaxScopeError> {
        let mut selectors = Vec::new();
        for selector in s.split(',') {
            selectors.push(try!(SyntaxScopeSelector::from_str(selector)))
        };
        Ok(SyntaxScopeSelectors {
            selectors: selectors
        })
    }
}

impl SyntaxScope {
    pub fn rank(&self) -> Rank {
        self.name.split('.').count() as u64 + 1
    }

    pub fn matched(&self, scope: &str) -> bool {
        scope.starts_with(&self.name) &&
        (self.name.len() == scope.len() || scope[self.name.len()..].starts_with('.'))
    }
}

impl SyntaxScopeSelector {
    pub fn rank(&self) -> Rank {
        let mut rank = 0;
        for scope in self.path.iter().rev() { // todo: do we need rev?
            rank <<= 4;
            rank += scope.rank();
        }
        rank
    }

    pub fn matched(&self, path: &[SyntaxScope]) -> bool {
        let mut iter = path.iter();
        for scope in &self.path {
            while let Some(s) = iter.next() {
                if scope.matched(&s.name) {
                    continue
                }
            }
            return false
        }
        true
    }
}
<|endoftext|>


<commit_before>use num_traits::ToPrimitive;
use num_bigint::BigInt;
use std::str::FromStr;
use std::cmp::max;

mod precision;
use self::precision::*;

pub mod asset_symbol;
pub use self::asset_symbol::*;

pub mod cmp;
pub use self::cmp::*;

pub mod error;
pub use self::error::*;

pub const MAX_STR_LEN: usize = 32;
pub const MAX_PRECISION: u8 = 8;

#[derive(Debug, Clone)]
pub struct Asset {
    pub amount: i64,
    pub decimals: u8,
    pub symbol: AssetSymbol
}

impl Asset {

    pub fn add(&self, other: &Self) -> Option<Self> {
        if self.symbol != other.symbol { return None }
        let decimals = max(self.decimals, other.decimals);
        let a = set_decimals_i64(self.amount, self.decimals, decimals)?;
        let b = set_decimals_i64(other.amount, other.decimals, decimals)?;
        Some(Asset { amount: a.checked_add(b)?, decimals, symbol: self.symbol })
    }

    pub fn sub(&self, other: &Self) -> Option<Self> {
        if self.symbol != other.symbol { return None }
        let decimals = max(self.decimals, other.decimals);
        let a = set_decimals_i64(self.amount, self.decimals, decimals)?;
        let b = set_decimals_i64(other.amount, other.decimals, decimals)?;
        Some(Asset { amount: a.checked_sub(b)?, decimals, symbol: self.symbol })
    }

    pub fn mul(&self, other: &Self, precision: u8) -> Option<Self> {
        if self.symbol != other.symbol { return None }
        let decimals = self.decimals + other.decimals;

        let mul = i128::from(self.amount).checked_mul(i128::from(other.amount))?;
        let final_mul = set_decimals_i128(mul, decimals, precision)?;
        if final_mul > i128::from(::std::i64::MAX) { return None }
        Some(Asset {
            amount: final_mul as i64,
            decimals: precision,
            symbol: self.symbol
        })
    }

    pub fn div(&self, other: &Self, precision: u8) -> Option<Self> {
        if self.symbol != other.symbol || other.amount == 0 { return None }
        let decimals = max(max(self.decimals, other.decimals), precision);
        let a = set_decimals_i64(self.amount, self.decimals, decimals * 2)?;
        let b = set_decimals_i64(other.amount, other.decimals, decimals)?;
        Some(Asset {
            amount: set_decimals_i64(a.checked_div(b)?, decimals, precision)?,
            decimals: precision,
            symbol: self.symbol
        })
    }

    pub fn pow(&self, num: u16, precision: u8) -> Option<Self> {
        if num == 0 {
            return Some(Asset {
                amount: set_decimals_i64(1, 0, precision)?,
                decimals: precision,
                symbol: self.symbol
            })
        }
        let big_zero = BigInt::from(0);
        let big_one = BigInt::from(1);

        let decimals = u16::from(self.decimals).checked_mul(num)?;
        let mut res = BigInt::from(1);
        {
            let mut base = BigInt::from(self.amount);
            let mut exp = BigInt::from(num);
            loop {
                if &exp & &big_one == big_one { res = &res * &base; }
                exp >>= 1;
                if exp == big_zero { break };
                base = &base * &base;
            }
        }

        res = set_decimals_big(&res, decimals, u16::from(precision));
        Some(Asset {
            amount: res.to_i64()?,
            decimals: precision,
            symbol: self.symbol
        })
    }

    pub fn set_decimals(&self, decimals: u8) -> Option<Self> {
        let amount = set_decimals_i64(self.amount, self.decimals, decimals)?;
        Some(Asset {
            amount,
            decimals,
            symbol: self.symbol
        })
    }

    pub fn to_str(&self) -> Box<str> {
        let mut s = self.amount.to_string();
        if self.decimals > 0 {
            let len = s.len();
            if len < self.decimals as usize {
                let start = if self.amount < 0 { 1 } else { 0 };
                let diff = self.decimals as usize - len + start;
                s.insert_str(start, "0.");
                s.insert_str(start + 2, &"0".repeat(diff));
            } else if len == self.decimals as usize {
                s.insert_str(0, "0.");
            } else {
                s.insert(len - (self.decimals as usize), '.');
            }
        }
        s.push(' ');
        s.push_str(self.symbol.as_str());
        s.into_boxed_str()
    }
}

impl FromStr for Asset {
    type Err = AssetError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.len() > MAX_STR_LEN {
            return Err(AssetError { kind: AssetErrorKind::StrTooLarge })
        }
        let mut split = s.trim().splitn(2, ' ');

        let amount: i64;
        let mut decimals: u8 = 0;
        match split.next() {
            Some(x) => {
                if let Some(pos) = x.find('.') {
                    let len = x.len() - 1;
                    if pos > 0 { decimals = (len - pos) as u8; }
                    else { decimals = len as u8; }

                    if decimals > MAX_PRECISION {
                        return Err(AssetError { kind: AssetErrorKind::InvalidAmount })
                    }

                    amount = match x.replace('.', "").parse() {
                        Ok(x) => x,
                        Err(_) => return Err(AssetError { kind: AssetErrorKind::InvalidAmount })
                    }
                } else {
                    amount = match x.parse() {
                        Ok(x) => x,
                        Err(_) => return Err(AssetError { kind: AssetErrorKind::InvalidAmount })
                    }
                }
            },
            None => return Err(AssetError { kind: AssetErrorKind::InvalidFormat })
        };

        let symbol = match split.next() {
            Some(x) => {
                match AssetSymbol::parse_str(x) {
                    Some(x) => x,
                    None => return Err(AssetError { kind: AssetErrorKind::InvalidAssetType })
                }
            },
            None => return Err(AssetError { kind: AssetErrorKind::InvalidFormat })
        };

        Ok(Asset { amount, decimals, symbol })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_valid_input() {
        let c = |asset: Asset, amount: &str, decimals: u8, symbol: AssetSymbol| {
            assert_eq!(asset.amount.to_string(), amount);
            assert_eq!(asset.decimals, decimals);
            assert_eq!(asset.symbol, symbol);
        };

        c(get_asset("1 GOLD"), "1", 0, AssetSymbol::GOLD);
        c(get_asset("1. GOLD"), "1", 0, AssetSymbol::GOLD);
        c(get_asset(".1 GOLD"), "1", 1, AssetSymbol::GOLD);
        c(get_asset("-.1 GOLD"), "-1", 1, AssetSymbol::GOLD);
        c(get_asset("0.1 GOLD"), "1", 1, AssetSymbol::GOLD);
        c(get_asset("1.0 SILVER"), "10", 1, AssetSymbol::SILVER);
        c(get_asset("0 SILVER"), "0", 0, AssetSymbol::SILVER);
        c(get_asset("-0.0 SILVER"), "0", 1, AssetSymbol::SILVER);
        c(get_asset("-1.0 SILVER"), "-10", 1, AssetSymbol::SILVER);
    }

    #[test]
    fn test_asset_to_str() {
        let c = |asset: Asset, s: &str| {
            assert_eq!(&*asset.to_str(), s);
        };
        c(get_asset("1.00001 GOLD"), "1.00001 GOLD");
        c(get_asset("0.00001 GOLD"), "0.00001 GOLD");
        c(get_asset("-0.00001 GOLD"), "-0.00001 GOLD");
        c(get_asset(".00001 GOLD"), "0.00001 GOLD");
        c(get_asset(".1 GOLD"), "0.1 GOLD");
        c(get_asset("1.0 GOLD"), "1.0 GOLD");
    }

    #[test]
    fn test_fail_parsing_invalid_input() {
        let c = |asset: &str, err: AssetErrorKind| {
            let e = Asset::from_str(asset).err().unwrap();
            assert_eq!(e.kind, err);
        };

        c("1e10 GOLD", AssetErrorKind::InvalidAmount);
        c("a100 GOLD", AssetErrorKind::InvalidAmount);
        c("100a GOLD", AssetErrorKind::InvalidAmount);

        c("1234567890123456789012345678 GOLD", AssetErrorKind::StrTooLarge);
        c("1", AssetErrorKind::InvalidFormat);

        c("1.0 GOLD a", AssetErrorKind::InvalidAssetType);
        c("1.0 gold", AssetErrorKind::InvalidAssetType);
    }

    #[test]
    fn test_set_precision() {
        let a = get_asset("1.5678 GOLD");
        assert_eq!(a.decimals, 4);
        assert_eq!(a.amount.to_string(), "15678");

        let a = a.set_decimals(2).unwrap();

        assert_eq!(a.decimals, 2);
        assert_eq!(a.amount.to_string(), "156");

        let a = a.set_decimals(0).unwrap();
        assert_eq!(a.decimals, 0);
        assert_eq!(a.amount.to_string(), "1");

        let a = a.set_decimals(4).unwrap();
        assert_eq!(a.decimals, 4);
        assert_eq!(a.amount.to_string(), "10000");
    }

    #[test]
    fn test_perform_arithmetic() {
        let c = |asset: &Asset, amount: &str| {
            assert_eq!(&*asset.to_str(), amount);
        };

        let a = get_asset("123.456 GOLD");
        c(&a.add(&get_asset("2.0 GOLD")).unwrap(), "125.456 GOLD");
        c(&a.add(&get_asset("-2.0 GOLD")).unwrap(), "121.456 GOLD");
        c(&a.add(&get_asset(".00000001 GOLD")).unwrap(), "123.45600001 GOLD");
        c(&a.sub(&get_asset("2.0 GOLD")).unwrap(), "121.456 GOLD");
        c(&a.sub(&get_asset("-2.0 GOLD")).unwrap(), "125.456 GOLD");
        c(&a.mul(&get_asset("100000.11111111 GOLD"), 8).unwrap(), "12345613.71733319 GOLD");
        c(&a.mul(&get_asset("-100000.11111111 GOLD"), 8).unwrap(), "-12345613.71733319 GOLD");
        c(&a.div(&get_asset("23 GOLD"), 3).unwrap(), "5.367 GOLD");
        c(&a.div(&get_asset("-23 GOLD"), 8).unwrap(), "-5.36765217 GOLD");
        c(&a.pow(2, 8).unwrap(), "15241.38393600 GOLD");
        c(&a.pow(3, 8).unwrap(), "1881640.29520281 GOLD");
        c(&a, "123.456 GOLD");

        c(&get_asset("1.0002 GOLD").pow(1000, 8).unwrap(), "1.22137833 GOLD");
        c(&get_asset("10 GOLD").div(&get_asset("2 GOLD"), 0).unwrap(), "5 GOLD");
        c(&get_asset("5 GOLD").div(&get_asset("10 GOLD"), 1).unwrap(), "0.5 GOLD");

        assert!(&a.div(&get_asset("0 GOLD"), 1).is_none());
    }

    #[test]
    fn test_invalid_arithmetic() {
        let a = &get_asset("10 GOLD");
        let b = &get_asset("10 SILVER");

        assert!(a.add(b).is_none());
        assert!(a.sub(b).is_none());
        assert!(a.div(b, 8).is_none());
        assert!(a.mul(b, 8).is_none());
    }

    fn get_asset(s: &str) -> Asset {
        Asset::from_str(s).unwrap()
    }
}

<commit_msg>Cleanup<commit_after>use num_traits::ToPrimitive;
use num_bigint::BigInt;
use std::str::FromStr;
use std::cmp::max;

mod precision;
use self::precision::*;

pub mod asset_symbol;
pub use self::asset_symbol::*;

pub mod cmp;
pub use self::cmp::*;

pub mod error;
pub use self::error::*;

pub const MAX_STR_LEN: usize = 32;
pub const MAX_PRECISION: u8 = 8;

#[derive(Debug, Clone)]
pub struct Asset {
    pub amount: i64,
    pub decimals: u8,
    pub symbol: AssetSymbol
}

impl Asset {

    pub fn add(&self, other: &Self) -> Option<Self> {
        if self.symbol != other.symbol { return None }
        let decimals = max(self.decimals, other.decimals);
        let a = set_decimals_i64(self.amount, self.decimals, decimals)?;
        let b = set_decimals_i64(other.amount, other.decimals, decimals)?;
        Some(Asset { amount: a.checked_add(b)?, decimals, symbol: self.symbol })
    }

    pub fn sub(&self, other: &Self) -> Option<Self> {
        if self.symbol != other.symbol { return None }
        let decimals = max(self.decimals, other.decimals);
        let a = set_decimals_i64(self.amount, self.decimals, decimals)?;
        let b = set_decimals_i64(other.amount, other.decimals, decimals)?;
        Some(Asset { amount: a.checked_sub(b)?, decimals, symbol: self.symbol })
    }

    pub fn mul(&self, other: &Self, precision: u8) -> Option<Self> {
        if self.symbol != other.symbol { return None }
        let decimals = self.decimals + other.decimals;

        let mul = i128::from(self.amount).checked_mul(i128::from(other.amount))?;
        let final_mul = set_decimals_i128(mul, decimals, precision)?;
        if final_mul > i128::from(::std::i64::MAX) { return None }
        Some(Asset {
            amount: final_mul as i64,
            decimals: precision,
            symbol: self.symbol
        })
    }

    pub fn div(&self, other: &Self, precision: u8) -> Option<Self> {
        if self.symbol != other.symbol || other.amount == 0 { return None }
        let decimals = max(max(self.decimals, other.decimals), precision);
        let a = set_decimals_i64(self.amount, self.decimals, decimals * 2)?;
        let b = set_decimals_i64(other.amount, other.decimals, decimals)?;
        Some(Asset {
            amount: set_decimals_i64(a.checked_div(b)?, decimals, precision)?,
            decimals: precision,
            symbol: self.symbol
        })
    }

    pub fn pow(&self, num: u16, precision: u8) -> Option<Self> {
        if num == 0 {
            return Some(Asset {
                amount: set_decimals_i64(1, 0, precision)?,
                decimals: precision,
                symbol: self.symbol
            })
        }
        let big_zero = BigInt::from(0);
        let big_one = BigInt::from(1);

        let decimals = u16::from(self.decimals).checked_mul(num)?;
        let mut res = BigInt::from(1);
        {
            let mut base = BigInt::from(self.amount);
            let mut exp = BigInt::from(num);
            loop {
                if &exp & &big_one == big_one { res = &res * &base; }
                exp >>= 1;
                if exp == big_zero { break };
                base = &base * &base;
            }
        }

        res = set_decimals_big(&res, decimals, u16::from(precision));
        Some(Asset {
            amount: res.to_i64()?,
            decimals: precision,
            symbol: self.symbol
        })
    }

    pub fn set_decimals(&self, decimals: u8) -> Option<Self> {
        let amount = set_decimals_i64(self.amount, self.decimals, decimals)?;
        Some(Asset {
            amount,
            decimals,
            symbol: self.symbol
        })
    }

    pub fn to_str(&self) -> Box<str> {
        let mut s = self.amount.to_string();
        if self.decimals > 0 {
            let len = s.len();
            if len < self.decimals as usize {
                let start = if self.amount < 0 { 1 } else { 0 };
                let diff = self.decimals as usize - len + start;
                s.insert_str(start, "0.");
                s.insert_str(start + 2, &"0".repeat(diff));
            } else if len == self.decimals as usize {
                s.insert_str(0, "0.");
            } else {
                s.insert(len - (self.decimals as usize), '.');
            }
        }
        s.push(' ');
        s.push_str(self.symbol.as_str());
        s.into_boxed_str()
    }
}

impl FromStr for Asset {
    type Err = AssetError;

    fn from_str(s: &str) -> Result<Self, Self::Err> {
        if s.len() > MAX_STR_LEN {
            return Err(AssetError { kind: AssetErrorKind::StrTooLarge })
        }
        let mut split = s.trim().splitn(2, ' ');

        let amount: i64;
        let mut decimals: u8 = 0;
        match split.next() {
            Some(x) => {
                if let Some(pos) = x.find('.') {
                    let len = x.len() - 1;
                    if pos > 0 { decimals = (len - pos) as u8; }
                    else { decimals = len as u8; }

                    if decimals > MAX_PRECISION {
                        return Err(AssetError { kind: AssetErrorKind::InvalidAmount })
                    }

                    amount = match x.replace('.', "").parse() {
                        Ok(x) => x,
                        Err(_) => return Err(AssetError { kind: AssetErrorKind::InvalidAmount })
                    }
                } else {
                    amount = match x.parse() {
                        Ok(x) => x,
                        Err(_) => return Err(AssetError { kind: AssetErrorKind::InvalidAmount })
                    }
                }
            },
            None => return Err(AssetError { kind: AssetErrorKind::InvalidFormat })
        };

        let symbol = match split.next() {
            Some(x) => {
                match AssetSymbol::parse_str(x) {
                    Some(x) => x,
                    None => return Err(AssetError { kind: AssetErrorKind::InvalidAssetType })
                }
            },
            None => return Err(AssetError { kind: AssetErrorKind::InvalidFormat })
        };

        Ok(Asset { amount, decimals, symbol })
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_parse_valid_input() {
        let c = |asset: Asset, amount: &str, decimals: u8, symbol: AssetSymbol| {
            assert_eq!(asset.amount.to_string(), amount);
            assert_eq!(asset.decimals, decimals);
            assert_eq!(asset.symbol, symbol);
        };

        c(get_asset("1 GOLD"), "1", 0, AssetSymbol::GOLD);
        c(get_asset("1. GOLD"), "1", 0, AssetSymbol::GOLD);
        c(get_asset(".1 GOLD"), "1", 1, AssetSymbol::GOLD);
        c(get_asset("-.1 GOLD"), "-1", 1, AssetSymbol::GOLD);
        c(get_asset("0.1 GOLD"), "1", 1, AssetSymbol::GOLD);
        c(get_asset("1.0 SILVER"), "10", 1, AssetSymbol::SILVER);
        c(get_asset("0 SILVER"), "0", 0, AssetSymbol::SILVER);
        c(get_asset("-0.0 SILVER"), "0", 1, AssetSymbol::SILVER);
        c(get_asset("-1.0 SILVER"), "-10", 1, AssetSymbol::SILVER);
    }

    #[test]
    fn test_asset_to_str() {
        let c = |asset: Asset, s: &str| {
            assert_eq!(&*asset.to_str(), s);
        };
        c(get_asset("1.00001 GOLD"), "1.00001 GOLD");
        c(get_asset("0.00001 GOLD"), "0.00001 GOLD");
        c(get_asset("-0.00001 GOLD"), "-0.00001 GOLD");
        c(get_asset(".00001 GOLD"), "0.00001 GOLD");
        c(get_asset(".1 GOLD"), "0.1 GOLD");
        c(get_asset("1.0 GOLD"), "1.0 GOLD");
    }

    #[test]
    fn test_fail_parsing_invalid_input() {
        let c = |asset: &str, err: AssetErrorKind| {
            let e = Asset::from_str(asset).err().unwrap();
            assert_eq!(e.kind, err);
        };

        c("1e10 GOLD", AssetErrorKind::InvalidAmount);
        c("a100 GOLD", AssetErrorKind::InvalidAmount);
        c("100a GOLD", AssetErrorKind::InvalidAmount);

        c("1234567890123456789012345678 GOLD", AssetErrorKind::StrTooLarge);
        c("1", AssetErrorKind::InvalidFormat);

        c("1.0 GOLD a", AssetErrorKind::InvalidAssetType);
        c("1.0 gold", AssetErrorKind::InvalidAssetType);
    }

    #[test]
    fn test_set_precision() {
        let a = get_asset("1.5678 GOLD");
        assert_eq!(a.decimals, 4);
        assert_eq!(a.amount.to_string(), "15678");

        let a = a.set_decimals(2).unwrap();

        assert_eq!(a.decimals, 2);
        assert_eq!(a.amount.to_string(), "156");

        let a = a.set_decimals(0).unwrap();
        assert_eq!(a.decimals, 0);
        assert_eq!(a.amount.to_string(), "1");

        let a = a.set_decimals(4).unwrap();
        assert_eq!(a.decimals, 4);
        assert_eq!(a.amount.to_string(), "10000");
    }

    #[test]
    fn test_perform_arithmetic() {
        let c = |asset: &Asset, amount: &str| {
            assert_eq!(&*asset.to_str(), amount);
        };

        let a = get_asset("123.456 GOLD");
        c(&a.add(&get_asset("2.0 GOLD")).unwrap(), "125.456 GOLD");
        c(&a.add(&get_asset("-2.0 GOLD")).unwrap(), "121.456 GOLD");
        c(&a.add(&get_asset(".00000001 GOLD")).unwrap(), "123.45600001 GOLD");
        c(&a.sub(&get_asset("2.0 GOLD")).unwrap(), "121.456 GOLD");
        c(&a.sub(&get_asset("-2.0 GOLD")).unwrap(), "125.456 GOLD");
        c(&a.mul(&get_asset("100000.11111111 GOLD"), 8).unwrap(), "12345613.71733319 GOLD");
        c(&a.mul(&get_asset("-100000.11111111 GOLD"), 8).unwrap(), "-12345613.71733319 GOLD");
        c(&a.div(&get_asset("23 GOLD"), 3).unwrap(), "5.367 GOLD");
        c(&a.div(&get_asset("-23 GOLD"), 8).unwrap(), "-5.36765217 GOLD");
        c(&a.pow(2, 8).unwrap(), "15241.38393600 GOLD");
        c(&a.pow(3, 8).unwrap(), "1881640.29520281 GOLD");
        c(&a, "123.456 GOLD");

        c(&get_asset("1.0002 GOLD").pow(1000, 8).unwrap(), "1.22137833 GOLD");
        c(&get_asset("10 GOLD").div(&get_asset("2 GOLD"), 0).unwrap(), "5 GOLD");
        c(&get_asset("5 GOLD").div(&get_asset("10 GOLD"), 1).unwrap(), "0.5 GOLD");

        assert!(&a.div(&get_asset("0 GOLD"), 1).is_none());
    }

    #[test]
    fn test_invalid_arithmetic() {
        let a = &get_asset("10 GOLD");
        let b = &get_asset("10 SILVER");

        assert!(a.add(b).is_none());
        assert!(a.sub(b).is_none());
        assert!(a.div(b, 8).is_none());
        assert!(a.mul(b, 8).is_none());
    }

    fn get_asset(s: &str) -> Asset {
        Asset::from_str(s).unwrap()
    }
}
<|endoftext|>


<commit_before>//! Memory layout of matrices

use lapack::c;
use ndarray::*;

use super::error::*;

pub type LDA = i32;
pub type LEN = i32;
pub type Col = i32;
pub type Row = i32;

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum MatrixLayout {
    C((Row, LDA)),
    F((Col, LDA)),
}

impl MatrixLayout {
    pub fn size(&self) -> (Row, Col) {
        match *self {
            MatrixLayout::C((row, lda)) => (row, lda),
            MatrixLayout::F((col, lda)) => (lda, col),
        }
    }

    pub fn resized(&self, row: Row, col: Col) -> MatrixLayout {
        match *self {
            MatrixLayout::C(_) => MatrixLayout::C((row, col)),
            MatrixLayout::F(_) => MatrixLayout::F((col, row)),
        }
    }

    pub fn lda(&self) -> LDA {
        match *self {
            MatrixLayout::C((_, lda)) | MatrixLayout::F((_, lda)) => lda,
        }
    }

    pub fn len(&self) -> LEN {
        match *self {
            MatrixLayout::C((row, _)) => row,
            MatrixLayout::F((col, _)) => col,
        }
    }

    pub fn lapacke_layout(&self) -> c::Layout {
        match *self {
            MatrixLayout::C(_) => c::Layout::RowMajor,
            MatrixLayout::F(_) => c::Layout::ColumnMajor,
        }
    }

    pub fn same_order(&self, other: &MatrixLayout) -> bool {
        self.lapacke_layout() == other.lapacke_layout()
    }

    pub fn as_shape(&self) -> Shape<Ix2> {
        match *self {
            MatrixLayout::C((row, col)) => (row as usize, col as usize).into_shape(),
            MatrixLayout::F((col, row)) => (row as usize, col as usize).f().into_shape(),
        }
    }

    pub fn toggle_order(&self) -> Self {
        match *self {
            MatrixLayout::C((row, col)) => MatrixLayout::F((col, row)),
            MatrixLayout::F((col, row)) => MatrixLayout::C((row, col)),
        }
    }
}

pub trait AllocatedArray {
    type Elem;
    fn layout(&self) -> Result<MatrixLayout>;
    fn square_layout(&self) -> Result<MatrixLayout>;
    /// Returns Ok iff the matrix is square (without computing the layout).
    fn ensure_square(&self) -> Result<()>;
    fn as_allocated(&self) -> Result<&[Self::Elem]>;
}

pub trait AllocatedArrayMut: AllocatedArray {
    fn as_allocated_mut(&mut self) -> Result<&mut [Self::Elem]>;
}

impl<A, S> AllocatedArray for ArrayBase<S, Ix2>
where
    S: Data<Elem = A>,
{
    type Elem = A;

    fn layout(&self) -> Result<MatrixLayout> {
        let shape = self.shape();
        let strides = self.strides();
        if shape[0] == strides[1] as usize {
            return Ok(MatrixLayout::F((self.cols() as i32, self.rows() as i32)));
        }
        if shape[1] == strides[0] as usize {
            return Ok(MatrixLayout::C((self.rows() as i32, self.cols() as i32)));
        }
        Err(StrideError::new(strides[0], strides[1]).into())
    }

    fn square_layout(&self) -> Result<MatrixLayout> {
        let l = self.layout()?;
        let (n, m) = l.size();
        if n == m {
            Ok(l)
        } else {
            Err(NotSquareError::new(n, m).into())
        }
    }

    fn ensure_square(&self) -> Result<()> {
        if self.is_square() {
            Ok(())
        } else {
            Err(
                NotSquareError::new(self.rows() as i32, self.cols() as i32).into(),
            )
        }
    }

    fn as_allocated(&self) -> Result<&[A]> {
        Ok(self.as_slice_memory_order().ok_or_else(MemoryContError::new)?)
    }
}

impl<A, S> AllocatedArrayMut for ArrayBase<S, Ix2>
where
    S: DataMut<Elem = A>,
{
    fn as_allocated_mut(&mut self) -> Result<&mut [A]> {
        Ok(self.as_slice_memory_order_mut().ok_or_else(
            MemoryContError::new,
        )?)
    }
}
<commit_msg>Fix format<commit_after>//! Memory layout of matrices

use lapack::c;
use ndarray::*;

use super::error::*;

pub type LDA = i32;
pub type LEN = i32;
pub type Col = i32;
pub type Row = i32;

#[derive(Debug, Clone, Copy, PartialEq)]
pub enum MatrixLayout {
    C((Row, LDA)),
    F((Col, LDA)),
}

impl MatrixLayout {
    pub fn size(&self) -> (Row, Col) {
        match *self {
            MatrixLayout::C((row, lda)) => (row, lda),
            MatrixLayout::F((col, lda)) => (lda, col),
        }
    }

    pub fn resized(&self, row: Row, col: Col) -> MatrixLayout {
        match *self {
            MatrixLayout::C(_) => MatrixLayout::C((row, col)),
            MatrixLayout::F(_) => MatrixLayout::F((col, row)),
        }
    }

    pub fn lda(&self) -> LDA {
        match *self {
            MatrixLayout::C((_, lda)) |
            MatrixLayout::F((_, lda)) => lda,
        }
    }

    pub fn len(&self) -> LEN {
        match *self {
            MatrixLayout::C((row, _)) => row,
            MatrixLayout::F((col, _)) => col,
        }
    }

    pub fn lapacke_layout(&self) -> c::Layout {
        match *self {
            MatrixLayout::C(_) => c::Layout::RowMajor,
            MatrixLayout::F(_) => c::Layout::ColumnMajor,
        }
    }

    pub fn same_order(&self, other: &MatrixLayout) -> bool {
        self.lapacke_layout() == other.lapacke_layout()
    }

    pub fn as_shape(&self) -> Shape<Ix2> {
        match *self {
            MatrixLayout::C((row, col)) => (row as usize, col as usize).into_shape(),
            MatrixLayout::F((col, row)) => (row as usize, col as usize).f().into_shape(),
        }
    }

    pub fn toggle_order(&self) -> Self {
        match *self {
            MatrixLayout::C((row, col)) => MatrixLayout::F((col, row)),
            MatrixLayout::F((col, row)) => MatrixLayout::C((row, col)),
        }
    }
}

pub trait AllocatedArray {
    type Elem;
    fn layout(&self) -> Result<MatrixLayout>;
    fn square_layout(&self) -> Result<MatrixLayout>;
    /// Returns Ok iff the matrix is square (without computing the layout).
    fn ensure_square(&self) -> Result<()>;
    fn as_allocated(&self) -> Result<&[Self::Elem]>;
}

pub trait AllocatedArrayMut: AllocatedArray {
    fn as_allocated_mut(&mut self) -> Result<&mut [Self::Elem]>;
}

impl<A, S> AllocatedArray for ArrayBase<S, Ix2>
where
    S: Data<Elem = A>,
{
    type Elem = A;

    fn layout(&self) -> Result<MatrixLayout> {
        let shape = self.shape();
        let strides = self.strides();
        if shape[0] == strides[1] as usize {
            return Ok(MatrixLayout::F((self.cols() as i32, self.rows() as i32)));
        }
        if shape[1] == strides[0] as usize {
            return Ok(MatrixLayout::C((self.rows() as i32, self.cols() as i32)));
        }
        Err(StrideError::new(strides[0], strides[1]).into())
    }

    fn square_layout(&self) -> Result<MatrixLayout> {
        let l = self.layout()?;
        let (n, m) = l.size();
        if n == m {
            Ok(l)
        } else {
            Err(NotSquareError::new(n, m).into())
        }
    }

    fn ensure_square(&self) -> Result<()> {
        if self.is_square() {
            Ok(())
        } else {
            Err(
                NotSquareError::new(self.rows() as i32, self.cols() as i32).into(),
            )
        }
    }

    fn as_allocated(&self) -> Result<&[A]> {
        Ok(self.as_slice_memory_order().ok_or_else(
            MemoryContError::new,
        )?)
    }
}

impl<A, S> AllocatedArrayMut for ArrayBase<S, Ix2>
where
    S: DataMut<Elem = A>,
{
    fn as_allocated_mut(&mut self) -> Result<&mut [A]> {
        Ok(self.as_slice_memory_order_mut().ok_or_else(
            MemoryContError::new,
        )?)
    }
}
<|endoftext|>