<commit_before>//! # The Rust core allocation and collections library
//!
//! This library provides smart pointers and collections for managing
//! heap-allocated values.
//!
//! This library, like libcore, normally doesn’t need to be used directly
//! since its contents are re-exported in the [`std` crate](../std/index.html).
//! Crates that use the `#![no_std]` attribute however will typically
//! not depend on `std`, so they’d use this crate instead.
//!
//! ## Boxed values
//!
//! The [`Box`] type is a smart pointer type. There can only be one owner of a
//! [`Box`], and the owner can decide to mutate the contents, which live on the
//! heap.
//!
//! This type can be sent among threads efficiently as the size of a `Box` value
//! is the same as that of a pointer. Tree-like data structures are often built
//! with boxes because each node often has only one owner, the parent.
//!
//! ## Reference counted pointers
//!
//! The [`Rc`] type is a non-threadsafe reference-counted pointer type intended
//! for sharing memory within a thread. An [`Rc`] pointer wraps a type, `T`, and
//! only allows access to `&T`, a shared reference.
//!
//! This type is useful when inherited mutability (such as using [`Box`]) is too
//! constraining for an application, and is often paired with the [`Cell`] or
//! [`RefCell`] types in order to allow mutation.
//!
//! ## Atomically reference counted pointers
//!
//! The [`Arc`] type is the threadsafe equivalent of the [`Rc`] type. It
//! provides all the same functionality of [`Rc`], except it requires that the
//! contained type `T` is shareable. Additionally, [`Arc<T>`][`Arc`] is itself
//! sendable while [`Rc<T>`][`Rc`] is not.
//!
//! This type allows for shared access to the contained data, and is often
//! paired with synchronization primitives such as mutexes to allow mutation of
//! shared resources.
//!
//! ## Collections
//!
//! Implementations of the most common general purpose data structures are
//! defined in this library. They are re-exported through the
//! [standard collections library](../std/collections/index.html).
//!
//! ## Heap interfaces
//!
//! The [`alloc`](alloc/index.html) module defines the low-level interface to the
//! default global allocator. It is not compatible with the libc allocator API.
//!
//! [`Arc`]: sync
//! [`Box`]: boxed
//! [`Cell`]: core::cell
//! [`Rc`]: rc
//! [`RefCell`]: core::cell

#![allow(unused_attributes)]
#![stable(feature = "alloc", since = "1.36.0")]
#![doc(
    html_root_url = "https://doc.rust-lang.org/nightly/",
    html_playground_url = "https://play.rust-lang.org/",
    issue_tracker_base_url = "https://github.com/rust-lang/rust/issues/",
    test(no_crate_inject, attr(allow(unused_variables), deny(warnings)))
)]
#![no_std]
#![needs_allocator]
#![warn(deprecated_in_future)]
#![warn(missing_docs)]
#![warn(missing_debug_implementations)]
#![allow(explicit_outlives_requirements)]
#![deny(unsafe_op_in_unsafe_fn)]
#![feature(rustc_allow_const_fn_unstable)]
#![cfg_attr(not(test), feature(generator_trait))]
#![cfg_attr(test, feature(test))]
#![cfg_attr(test, feature(new_uninit))]
#![feature(allocator_api)]
#![feature(array_chunks)]
#![feature(array_methods)]
#![feature(array_windows)]
#![feature(allow_internal_unstable)]
#![feature(arbitrary_self_types)]
#![feature(async_stream)]
#![feature(box_patterns)]
#![feature(box_syntax)]
#![feature(cfg_sanitize)]
#![feature(cfg_target_has_atomic)]
#![feature(coerce_unsized)]
#![cfg_attr(not(no_global_oom_handling), feature(const_btree_new))]
#![feature(const_fn_trait_bound)]
#![feature(cow_is_borrowed)]
#![feature(const_cow_is_borrowed)]
#![feature(destructuring_assignment)]
#![feature(dispatch_from_dyn)]
#![feature(core_intrinsics)]
#![feature(dropck_eyepatch)]
#![feature(exact_size_is_empty)]
#![feature(exclusive_range_pattern)]
#![feature(extend_one)]
#![feature(fmt_internals)]
#![feature(fn_traits)]
#![feature(fundamental)]
#![feature(inplace_iteration)]
// Technically, this is a bug in rustdoc: rustdoc sees the documentation on `#[lang = slice_alloc]`
// blocks is for `&[T]`, which also has documentation using this feature in `core`, and gets mad
// that the feature-gate isn't enabled. Ideally, it wouldn't check for the feature gate for docs
// from other crates, but since this can only appear for lang items, it doesn't seem worth fixing.
#![feature(intra_doc_pointers)]
#![feature(iter_zip)]
#![feature(lang_items)]
#![feature(layout_for_ptr)]
#![feature(maybe_uninit_ref)]
#![feature(negative_impls)]
#![feature(never_type)]
#![feature(nll)]
#![feature(nonnull_slice_from_raw_parts)]
#![feature(auto_traits)]
#![feature(option_result_unwrap_unchecked)]
#![feature(pattern)]
#![feature(ptr_internals)]
#![feature(rustc_attrs)]
#![feature(receiver_trait)]
#![feature(min_specialization)]
#![feature(set_ptr_value)]
#![feature(slice_ptr_get)]
#![feature(slice_ptr_len)]
#![feature(slice_range)]
#![feature(staged_api)]
#![feature(str_internals)]
#![feature(trusted_len)]
#![feature(unboxed_closures)]
#![feature(unicode_internals)]
#![feature(unsize)]
#![feature(unsized_fn_params)]
#![feature(allocator_internals)]
#![feature(slice_partition_dedup)]
#![feature(maybe_uninit_extra, maybe_uninit_slice, maybe_uninit_uninit_array)]
#![feature(alloc_layout_extra)]
#![feature(trusted_random_access)]
#![cfg_attr(bootstrap, feature(try_trait))]
#![cfg_attr(not(bootstrap), feature(try_trait_v2))]
#![feature(min_type_alias_impl_trait)]
#![feature(associated_type_bounds)]
#![feature(slice_group_by)]
#![feature(decl_macro)]
#![feature(bindings_after_at)]
// Allow testing this library

#[cfg(test)]
#[macro_use]
extern crate std;
#[cfg(test)]
extern crate test;

// Module with internal macros used by other modules (needs to be included before other modules).
#[macro_use]
mod macros;

// Heaps provided for low-level allocation strategies

pub mod alloc;

// Primitive types using the heaps above

// Need to conditionally define the mod from `boxed.rs` to avoid
// duplicating the lang-items when building in test cfg; but also need
// to allow code to have `use boxed::Box;` declarations.
#[cfg(not(test))]
pub mod boxed;
#[cfg(test)]
mod boxed {
    pub use std::boxed::Box;
}
pub mod borrow;
pub mod collections;
pub mod fmt;
pub mod prelude;
pub mod raw_vec;
pub mod rc;
pub mod slice;
pub mod str;
pub mod string;
#[cfg(target_has_atomic = "ptr")]
pub mod sync;
#[cfg(all(not(no_global_oom_handling), target_has_atomic = "ptr"))]
pub mod task;
#[cfg(test)]
mod tests;
pub mod vec;

#[doc(hidden)]
#[unstable(feature = "liballoc_internals", issue = "none", reason = "implementation detail")]
pub mod __export {
    pub use core::format_args;
}
<commit_msg>removed ref to bindings_after_at<commit_after>//! # The Rust core allocation and collections library
//!
//! This library provides smart pointers and collections for managing
//! heap-allocated values.
//!
//! This library, like libcore, normally doesn’t need to be used directly
//! since its contents are re-exported in the [`std` crate](../std/index.html).
//! Crates that use the `#![no_std]` attribute however will typically
//! not depend on `std`, so they’d use this crate instead.
//!
//! ## Boxed values
//!
//! The [`Box`] type is a smart pointer type. There can only be one owner of a
//! [`Box`], and the owner can decide to mutate the contents, which live on the
//! heap.
//!
//! This type can be sent among threads efficiently as the size of a `Box` value
//! is the same as that of a pointer. Tree-like data structures are often built
//! with boxes because each node often has only one owner, the parent.
//!
//! ## Reference counted pointers
//!
//! The [`Rc`] type is a non-threadsafe reference-counted pointer type intended
//! for sharing memory within a thread. An [`Rc`] pointer wraps a type, `T`, and
//! only allows access to `&T`, a shared reference.
//!
//! This type is useful when inherited mutability (such as using [`Box`]) is too
//! constraining for an application, and is often paired with the [`Cell`] or
//! [`RefCell`] types in order to allow mutation.
//!
//! ## Atomically reference counted pointers
//!
//! The [`Arc`] type is the threadsafe equivalent of the [`Rc`] type. It
//! provides all the same functionality of [`Rc`], except it requires that the
//! contained type `T` is shareable. Additionally, [`Arc<T>`][`Arc`] is itself
//! sendable while [`Rc<T>`][`Rc`] is not.
//!
//! This type allows for shared access to the contained data, and is often
//! paired with synchronization primitives such as mutexes to allow mutation of
//! shared resources.
//!
//! ## Collections
//!
//! Implementations of the most common general purpose data structures are
//! defined in this library. They are re-exported through the
//! [standard collections library](../std/collections/index.html).
//!
//! ## Heap interfaces
//!
//! The [`alloc`](alloc/index.html) module defines the low-level interface to the
//! default global allocator. It is not compatible with the libc allocator API.
//!
//! [`Arc`]: sync
//! [`Box`]: boxed
//! [`Cell`]: core::cell
//! [`Rc`]: rc
//! [`RefCell`]: core::cell

#![allow(unused_attributes)]
#![stable(feature = "alloc", since = "1.36.0")]
#![doc(
    html_root_url = "https://doc.rust-lang.org/nightly/",
    html_playground_url = "https://play.rust-lang.org/",
    issue_tracker_base_url = "https://github.com/rust-lang/rust/issues/",
    test(no_crate_inject, attr(allow(unused_variables), deny(warnings)))
)]
#![no_std]
#![needs_allocator]
#![warn(deprecated_in_future)]
#![warn(missing_docs)]
#![warn(missing_debug_implementations)]
#![allow(explicit_outlives_requirements)]
#![deny(unsafe_op_in_unsafe_fn)]
#![feature(rustc_allow_const_fn_unstable)]
#![cfg_attr(not(test), feature(generator_trait))]
#![cfg_attr(test, feature(test))]
#![cfg_attr(test, feature(new_uninit))]
#![feature(allocator_api)]
#![feature(array_chunks)]
#![feature(array_methods)]
#![feature(array_windows)]
#![feature(allow_internal_unstable)]
#![feature(arbitrary_self_types)]
#![feature(async_stream)]
#![feature(box_patterns)]
#![feature(box_syntax)]
#![feature(cfg_sanitize)]
#![feature(cfg_target_has_atomic)]
#![feature(coerce_unsized)]
#![cfg_attr(not(no_global_oom_handling), feature(const_btree_new))]
#![feature(const_fn_trait_bound)]
#![feature(cow_is_borrowed)]
#![feature(const_cow_is_borrowed)]
#![feature(destructuring_assignment)]
#![feature(dispatch_from_dyn)]
#![feature(core_intrinsics)]
#![feature(dropck_eyepatch)]
#![feature(exact_size_is_empty)]
#![feature(exclusive_range_pattern)]
#![feature(extend_one)]
#![feature(fmt_internals)]
#![feature(fn_traits)]
#![feature(fundamental)]
#![feature(inplace_iteration)]
// Technically, this is a bug in rustdoc: rustdoc sees the documentation on `#[lang = slice_alloc]`
// blocks is for `&[T]`, which also has documentation using this feature in `core`, and gets mad
// that the feature-gate isn't enabled. Ideally, it wouldn't check for the feature gate for docs
// from other crates, but since this can only appear for lang items, it doesn't seem worth fixing.
#![feature(intra_doc_pointers)]
#![feature(iter_zip)]
#![feature(lang_items)]
#![feature(layout_for_ptr)]
#![feature(maybe_uninit_ref)]
#![feature(negative_impls)]
#![feature(never_type)]
#![feature(nll)]
#![feature(nonnull_slice_from_raw_parts)]
#![feature(auto_traits)]
#![feature(option_result_unwrap_unchecked)]
#![feature(pattern)]
#![feature(ptr_internals)]
#![feature(rustc_attrs)]
#![feature(receiver_trait)]
#![feature(min_specialization)]
#![feature(set_ptr_value)]
#![feature(slice_ptr_get)]
#![feature(slice_ptr_len)]
#![feature(slice_range)]
#![feature(staged_api)]
#![feature(str_internals)]
#![feature(trusted_len)]
#![feature(unboxed_closures)]
#![feature(unicode_internals)]
#![feature(unsize)]
#![feature(unsized_fn_params)]
#![feature(allocator_internals)]
#![feature(slice_partition_dedup)]
#![feature(maybe_uninit_extra, maybe_uninit_slice, maybe_uninit_uninit_array)]
#![feature(alloc_layout_extra)]
#![feature(trusted_random_access)]
#![cfg_attr(bootstrap, feature(try_trait))]
#![cfg_attr(not(bootstrap), feature(try_trait_v2))]
#![feature(min_type_alias_impl_trait)]
#![feature(associated_type_bounds)]
#![feature(slice_group_by)]
#![feature(decl_macro)]
// Allow testing this library

#[cfg(test)]
#[macro_use]
extern crate std;
#[cfg(test)]
extern crate test;

// Module with internal macros used by other modules (needs to be included before other modules).
#[macro_use]
mod macros;

// Heaps provided for low-level allocation strategies

pub mod alloc;

// Primitive types using the heaps above

// Need to conditionally define the mod from `boxed.rs` to avoid
// duplicating the lang-items when building in test cfg; but also need
// to allow code to have `use boxed::Box;` declarations.
#[cfg(not(test))]
pub mod boxed;
#[cfg(test)]
mod boxed {
    pub use std::boxed::Box;
}
pub mod borrow;
pub mod collections;
pub mod fmt;
pub mod prelude;
pub mod raw_vec;
pub mod rc;
pub mod slice;
pub mod str;
pub mod string;
#[cfg(target_has_atomic = "ptr")]
pub mod sync;
#[cfg(all(not(no_global_oom_handling), target_has_atomic = "ptr"))]
pub mod task;
#[cfg(test)]
mod tests;
pub mod vec;

#[doc(hidden)]
#[unstable(feature = "liballoc_internals", issue = "none", reason = "implementation detail")]
pub mod __export {
    pub use core::format_args;
}
<|endoftext|>


<commit_before>//! Structural const qualification.
//!
//! See the `Qualif` trait for more info.

use rustc_errors::ErrorReported;
use rustc_hir as hir;
use rustc_infer::infer::TyCtxtInferExt;
use rustc_middle::mir::*;
use rustc_middle::ty::{self, subst::SubstsRef, AdtDef, Ty};
use rustc_span::DUMMY_SP;
use rustc_trait_selection::traits::{
    self, ImplSource, Obligation, ObligationCause, SelectionContext,
};

use super::ConstCx;

pub fn in_any_value_of_ty(
    cx: &ConstCx<'_, 'tcx>,
    ty: Ty<'tcx>,
    error_occured: Option<ErrorReported>,
) -> ConstQualifs {
    ConstQualifs {
        has_mut_interior: HasMutInterior::in_any_value_of_ty(cx, ty),
        needs_drop: NeedsNonConstDrop::in_any_value_of_ty(cx, ty),
        custom_eq: CustomEq::in_any_value_of_ty(cx, ty),
        error_occured,
    }
}

/// A "qualif"(-ication) is a way to look for something "bad" in the MIR that would disqualify some
/// code for promotion or prevent it from evaluating at compile time.
///
/// Normally, we would determine what qualifications apply to each type and error when an illegal
/// operation is performed on such a type. However, this was found to be too imprecise, especially
/// in the presence of `enum`s. If only a single variant of an enum has a certain qualification, we
/// needn't reject code unless it actually constructs and operates on the qualified variant.
///
/// To accomplish this, const-checking and promotion use a value-based analysis (as opposed to a
/// type-based one). Qualifications propagate structurally across variables: If a local (or a
/// projection of a local) is assigned a qualified value, that local itself becomes qualified.
pub trait Qualif {
    /// The name of the file used to debug the dataflow analysis that computes this qualif.
    const ANALYSIS_NAME: &'static str;

    /// Whether this `Qualif` is cleared when a local is moved from.
    const IS_CLEARED_ON_MOVE: bool = false;

    /// Extracts the field of `ConstQualifs` that corresponds to this `Qualif`.
    fn in_qualifs(qualifs: &ConstQualifs) -> bool;

    /// Returns `true` if *any* value of the given type could possibly have this `Qualif`.
    ///
    /// This function determines `Qualif`s when we cannot do a value-based analysis. Since qualif
    /// propagation is context-insenstive, this includes function arguments and values returned
    /// from a call to another function.
    ///
    /// It also determines the `Qualif`s for primitive types.
    fn in_any_value_of_ty(cx: &ConstCx<'_, 'tcx>, ty: Ty<'tcx>) -> bool;

    /// Returns `true` if this `Qualif` is inherent to the given struct or enum.
    ///
    /// By default, `Qualif`s propagate into ADTs in a structural way: An ADT only becomes
    /// qualified if part of it is assigned a value with that `Qualif`. However, some ADTs *always*
    /// have a certain `Qualif`, regardless of whether their fields have it. For example, a type
    /// with a custom `Drop` impl is inherently `NeedsDrop`.
    ///
    /// Returning `true` for `in_adt_inherently` but `false` for `in_any_value_of_ty` is unsound.
    fn in_adt_inherently(
        cx: &ConstCx<'_, 'tcx>,
        adt: &'tcx AdtDef,
        substs: SubstsRef<'tcx>,
    ) -> bool;
}

/// Constant containing interior mutability (`UnsafeCell<T>`).
/// This must be ruled out to make sure that evaluating the constant at compile-time
/// and at *any point* during the run-time would produce the same result. In particular,
/// promotion of temporaries must not change program behavior; if the promoted could be
/// written to, that would be a problem.
pub struct HasMutInterior;

impl Qualif for HasMutInterior {
    const ANALYSIS_NAME: &'static str = "flow_has_mut_interior";

    fn in_qualifs(qualifs: &ConstQualifs) -> bool {
        qualifs.has_mut_interior
    }

    fn in_any_value_of_ty(cx: &ConstCx<'_, 'tcx>, ty: Ty<'tcx>) -> bool {
        !ty.is_freeze(cx.tcx.at(DUMMY_SP), cx.param_env)
    }

    fn in_adt_inherently(cx: &ConstCx<'_, 'tcx>, adt: &'tcx AdtDef, _: SubstsRef<'tcx>) -> bool {
        // Exactly one type, `UnsafeCell`, has the `HasMutInterior` qualif inherently.
        // It arises structurally for all other types.
        Some(adt.did) == cx.tcx.lang_items().unsafe_cell_type()
    }
}

/// Constant containing an ADT that implements `Drop`.
/// This must be ruled out (a) because we cannot run `Drop` during compile-time
/// as that might not be a `const fn`, and (b) because implicit promotion would
/// remove side-effects that occur as part of dropping that value.
pub struct NeedsNonConstDrop;

impl Qualif for NeedsNonConstDrop {
    const ANALYSIS_NAME: &'static str = "flow_needs_nonconst_drop";
    const IS_CLEARED_ON_MOVE: bool = true;

    fn in_qualifs(qualifs: &ConstQualifs) -> bool {
        qualifs.needs_drop
    }

    fn in_any_value_of_ty(cx: &ConstCx<'_, 'tcx>, ty: Ty<'tcx>) -> bool {
        let trait_ref = ty::TraitRef {
            def_id: cx.tcx.require_lang_item(hir::LangItem::Drop, None),
            substs: cx.tcx.mk_substs_trait(ty, &[]),
        };
        let obligation = Obligation::new(
            ObligationCause::dummy(),
            cx.param_env,
            ty::Binder::dummy(ty::TraitPredicate {
                trait_ref,
                constness: ty::BoundConstness::ConstIfConst,
            }),
        );

        let implsrc = cx.tcx.infer_ctxt().enter(|infcx| {
            let mut selcx = SelectionContext::with_constness(&infcx, hir::Constness::Const);
            selcx.select(&obligation)
        });
        match implsrc {
            Ok(Some(ImplSource::ConstDrop(_)))
            | Ok(Some(ImplSource::Param(_, ty::BoundConstness::ConstIfConst))) => false,
            _ => true,
        }
    }

    fn in_adt_inherently(cx: &ConstCx<'_, 'tcx>, adt: &'tcx AdtDef, _: SubstsRef<'tcx>) -> bool {
        adt.has_non_const_dtor(cx.tcx)
    }
}

/// A constant that cannot be used as part of a pattern in a `match` expression.
pub struct CustomEq;

impl Qualif for CustomEq {
    const ANALYSIS_NAME: &'static str = "flow_custom_eq";

    fn in_qualifs(qualifs: &ConstQualifs) -> bool {
        qualifs.custom_eq
    }

    fn in_any_value_of_ty(cx: &ConstCx<'_, 'tcx>, ty: Ty<'tcx>) -> bool {
        // If *any* component of a composite data type does not implement `Structural{Partial,}Eq`,
        // we know that at least some values of that type are not structural-match. I say "some"
        // because that component may be part of an enum variant (e.g.,
        // `Option::<NonStructuralMatchTy>::Some`), in which case some values of this type may be
        // structural-match (`Option::None`).
        let id = cx.tcx.hir().local_def_id_to_hir_id(cx.def_id());
        traits::search_for_structural_match_violation(id, cx.body.span, cx.tcx, ty).is_some()
    }

    fn in_adt_inherently(
        cx: &ConstCx<'_, 'tcx>,
        adt: &'tcx AdtDef,
        substs: SubstsRef<'tcx>,
    ) -> bool {
        let ty = cx.tcx.mk_ty(ty::Adt(adt, substs));
        !ty.is_structural_eq_shallow(cx.tcx)
    }
}

// FIXME: Use `mir::visit::Visitor` for the `in_*` functions if/when it supports early return.

/// Returns `true` if this `Rvalue` contains qualif `Q`.
pub fn in_rvalue<Q, F>(cx: &ConstCx<'_, 'tcx>, in_local: &mut F, rvalue: &Rvalue<'tcx>) -> bool
where
    Q: Qualif,
    F: FnMut(Local) -> bool,
{
    match rvalue {
        Rvalue::ThreadLocalRef(_) | Rvalue::NullaryOp(..) => {
            Q::in_any_value_of_ty(cx, rvalue.ty(cx.body, cx.tcx))
        }

        Rvalue::Discriminant(place) | Rvalue::Len(place) => {
            in_place::<Q, _>(cx, in_local, place.as_ref())
        }

        Rvalue::Use(operand)
        | Rvalue::Repeat(operand, _)
        | Rvalue::UnaryOp(_, operand)
        | Rvalue::Cast(_, operand, _) => in_operand::<Q, _>(cx, in_local, operand),

        Rvalue::BinaryOp(_, box (lhs, rhs)) | Rvalue::CheckedBinaryOp(_, box (lhs, rhs)) => {
            in_operand::<Q, _>(cx, in_local, lhs) || in_operand::<Q, _>(cx, in_local, rhs)
        }

        Rvalue::Ref(_, _, place) | Rvalue::AddressOf(_, place) => {
            // Special-case reborrows to be more like a copy of the reference.
            if let Some((place_base, ProjectionElem::Deref)) = place.as_ref().last_projection() {
                let base_ty = place_base.ty(cx.body, cx.tcx).ty;
                if let ty::Ref(..) = base_ty.kind() {
                    return in_place::<Q, _>(cx, in_local, place_base);
                }
            }

            in_place::<Q, _>(cx, in_local, place.as_ref())
        }

        Rvalue::Aggregate(kind, operands) => {
            // Return early if we know that the struct or enum being constructed is always
            // qualified.
            if let AggregateKind::Adt(def, _, substs, ..) = **kind {
                if Q::in_adt_inherently(cx, def, substs) {
                    return true;
                }
            }

            // Otherwise, proceed structurally...
            operands.iter().any(|o| in_operand::<Q, _>(cx, in_local, o))
        }
    }
}

/// Returns `true` if this `Place` contains qualif `Q`.
pub fn in_place<Q, F>(cx: &ConstCx<'_, 'tcx>, in_local: &mut F, place: PlaceRef<'tcx>) -> bool
where
    Q: Qualif,
    F: FnMut(Local) -> bool,
{
    let mut place = place;
    while let Some((place_base, elem)) = place.last_projection() {
        match elem {
            ProjectionElem::Index(index) if in_local(index) => return true,

            ProjectionElem::Deref
            | ProjectionElem::Field(_, _)
            | ProjectionElem::ConstantIndex { .. }
            | ProjectionElem::Subslice { .. }
            | ProjectionElem::Downcast(_, _)
            | ProjectionElem::Index(_) => {}
        }

        let base_ty = place_base.ty(cx.body, cx.tcx);
        let proj_ty = base_ty.projection_ty(cx.tcx, elem).ty;
        if !Q::in_any_value_of_ty(cx, proj_ty) {
            return false;
        }

        place = place_base;
    }

    assert!(place.projection.is_empty());
    in_local(place.local)
}

/// Returns `true` if this `Operand` contains qualif `Q`.
pub fn in_operand<Q, F>(cx: &ConstCx<'_, 'tcx>, in_local: &mut F, operand: &Operand<'tcx>) -> bool
where
    Q: Qualif,
    F: FnMut(Local) -> bool,
{
    let constant = match operand {
        Operand::Copy(place) | Operand::Move(place) => {
            return in_place::<Q, _>(cx, in_local, place.as_ref());
        }

        Operand::Constant(c) => c,
    };

    // Check the qualifs of the value of `const` items.
    if let Some(ct) = constant.literal.const_for_ty() {
        if let ty::ConstKind::Unevaluated(ty::Unevaluated { def, substs_: _, promoted }) = ct.val {
            assert!(promoted.is_none());
            // Don't peek inside trait associated constants.
            if cx.tcx.trait_of_item(def.did).is_none() {
                let qualifs = if let Some((did, param_did)) = def.as_const_arg() {
                    cx.tcx.at(constant.span).mir_const_qualif_const_arg((did, param_did))
                } else {
                    cx.tcx.at(constant.span).mir_const_qualif(def.did)
                };

                if !Q::in_qualifs(&qualifs) {
                    return false;
                }

                // Just in case the type is more specific than
                // the definition, e.g., impl associated const
                // with type parameters, take it into account.
            }
        }
    }
    // Otherwise use the qualifs of the type.
    Q::in_any_value_of_ty(cx, constant.literal.ty())
}
<commit_msg>do not require lang item<commit_after>//! Structural const qualification.
//!
//! See the `Qualif` trait for more info.

use rustc_errors::ErrorReported;
use rustc_hir as hir;
use rustc_infer::infer::TyCtxtInferExt;
use rustc_middle::mir::*;
use rustc_middle::ty::{self, subst::SubstsRef, AdtDef, Ty};
use rustc_span::DUMMY_SP;
use rustc_trait_selection::traits::{
    self, ImplSource, Obligation, ObligationCause, SelectionContext,
};

use super::ConstCx;

pub fn in_any_value_of_ty(
    cx: &ConstCx<'_, 'tcx>,
    ty: Ty<'tcx>,
    error_occured: Option<ErrorReported>,
) -> ConstQualifs {
    ConstQualifs {
        has_mut_interior: HasMutInterior::in_any_value_of_ty(cx, ty),
        needs_drop: NeedsNonConstDrop::in_any_value_of_ty(cx, ty),
        custom_eq: CustomEq::in_any_value_of_ty(cx, ty),
        error_occured,
    }
}

/// A "qualif"(-ication) is a way to look for something "bad" in the MIR that would disqualify some
/// code for promotion or prevent it from evaluating at compile time.
///
/// Normally, we would determine what qualifications apply to each type and error when an illegal
/// operation is performed on such a type. However, this was found to be too imprecise, especially
/// in the presence of `enum`s. If only a single variant of an enum has a certain qualification, we
/// needn't reject code unless it actually constructs and operates on the qualified variant.
///
/// To accomplish this, const-checking and promotion use a value-based analysis (as opposed to a
/// type-based one). Qualifications propagate structurally across variables: If a local (or a
/// projection of a local) is assigned a qualified value, that local itself becomes qualified.
pub trait Qualif {
    /// The name of the file used to debug the dataflow analysis that computes this qualif.
    const ANALYSIS_NAME: &'static str;

    /// Whether this `Qualif` is cleared when a local is moved from.
    const IS_CLEARED_ON_MOVE: bool = false;

    /// Extracts the field of `ConstQualifs` that corresponds to this `Qualif`.
    fn in_qualifs(qualifs: &ConstQualifs) -> bool;

    /// Returns `true` if *any* value of the given type could possibly have this `Qualif`.
    ///
    /// This function determines `Qualif`s when we cannot do a value-based analysis. Since qualif
    /// propagation is context-insenstive, this includes function arguments and values returned
    /// from a call to another function.
    ///
    /// It also determines the `Qualif`s for primitive types.
    fn in_any_value_of_ty(cx: &ConstCx<'_, 'tcx>, ty: Ty<'tcx>) -> bool;

    /// Returns `true` if this `Qualif` is inherent to the given struct or enum.
    ///
    /// By default, `Qualif`s propagate into ADTs in a structural way: An ADT only becomes
    /// qualified if part of it is assigned a value with that `Qualif`. However, some ADTs *always*
    /// have a certain `Qualif`, regardless of whether their fields have it. For example, a type
    /// with a custom `Drop` impl is inherently `NeedsDrop`.
    ///
    /// Returning `true` for `in_adt_inherently` but `false` for `in_any_value_of_ty` is unsound.
    fn in_adt_inherently(
        cx: &ConstCx<'_, 'tcx>,
        adt: &'tcx AdtDef,
        substs: SubstsRef<'tcx>,
    ) -> bool;
}

/// Constant containing interior mutability (`UnsafeCell<T>`).
/// This must be ruled out to make sure that evaluating the constant at compile-time
/// and at *any point* during the run-time would produce the same result. In particular,
/// promotion of temporaries must not change program behavior; if the promoted could be
/// written to, that would be a problem.
pub struct HasMutInterior;

impl Qualif for HasMutInterior {
    const ANALYSIS_NAME: &'static str = "flow_has_mut_interior";

    fn in_qualifs(qualifs: &ConstQualifs) -> bool {
        qualifs.has_mut_interior
    }

    fn in_any_value_of_ty(cx: &ConstCx<'_, 'tcx>, ty: Ty<'tcx>) -> bool {
        !ty.is_freeze(cx.tcx.at(DUMMY_SP), cx.param_env)
    }

    fn in_adt_inherently(cx: &ConstCx<'_, 'tcx>, adt: &'tcx AdtDef, _: SubstsRef<'tcx>) -> bool {
        // Exactly one type, `UnsafeCell`, has the `HasMutInterior` qualif inherently.
        // It arises structurally for all other types.
        Some(adt.did) == cx.tcx.lang_items().unsafe_cell_type()
    }
}

/// Constant containing an ADT that implements `Drop`.
/// This must be ruled out (a) because we cannot run `Drop` during compile-time
/// as that might not be a `const fn`, and (b) because implicit promotion would
/// remove side-effects that occur as part of dropping that value.
pub struct NeedsNonConstDrop;

impl Qualif for NeedsNonConstDrop {
    const ANALYSIS_NAME: &'static str = "flow_needs_nonconst_drop";
    const IS_CLEARED_ON_MOVE: bool = true;

    fn in_qualifs(qualifs: &ConstQualifs) -> bool {
        qualifs.needs_drop
    }

    fn in_any_value_of_ty(cx: &ConstCx<'_, 'tcx>, ty: Ty<'tcx>) -> bool {
        let drop_trait = if let Some(did) = cx.tcx.lang_items().drop_trait() {
            did
        } else {
            // there is no way to define a type that needs non-const drop
            // without having the lang item present.
            return false;
        };
        let trait_ref = ty::TraitRef {
            def_id: drop_trait,
            substs: cx.tcx.mk_substs_trait(ty, &[]),
        };
        let obligation = Obligation::new(
            ObligationCause::dummy(),
            cx.param_env,
            ty::Binder::dummy(ty::TraitPredicate {
                trait_ref,
                constness: ty::BoundConstness::ConstIfConst,
            }),
        );

        let implsrc = cx.tcx.infer_ctxt().enter(|infcx| {
            let mut selcx = SelectionContext::with_constness(&infcx, hir::Constness::Const);
            selcx.select(&obligation)
        });
        match implsrc {
            Ok(Some(ImplSource::ConstDrop(_)))
            | Ok(Some(ImplSource::Param(_, ty::BoundConstness::ConstIfConst))) => false,
            _ => true,
        }
    }

    fn in_adt_inherently(cx: &ConstCx<'_, 'tcx>, adt: &'tcx AdtDef, _: SubstsRef<'tcx>) -> bool {
        adt.has_non_const_dtor(cx.tcx)
    }
}

/// A constant that cannot be used as part of a pattern in a `match` expression.
pub struct CustomEq;

impl Qualif for CustomEq {
    const ANALYSIS_NAME: &'static str = "flow_custom_eq";

    fn in_qualifs(qualifs: &ConstQualifs) -> bool {
        qualifs.custom_eq
    }

    fn in_any_value_of_ty(cx: &ConstCx<'_, 'tcx>, ty: Ty<'tcx>) -> bool {
        // If *any* component of a composite data type does not implement `Structural{Partial,}Eq`,
        // we know that at least some values of that type are not structural-match. I say "some"
        // because that component may be part of an enum variant (e.g.,
        // `Option::<NonStructuralMatchTy>::Some`), in which case some values of this type may be
        // structural-match (`Option::None`).
        let id = cx.tcx.hir().local_def_id_to_hir_id(cx.def_id());
        traits::search_for_structural_match_violation(id, cx.body.span, cx.tcx, ty).is_some()
    }

    fn in_adt_inherently(
        cx: &ConstCx<'_, 'tcx>,
        adt: &'tcx AdtDef,
        substs: SubstsRef<'tcx>,
    ) -> bool {
        let ty = cx.tcx.mk_ty(ty::Adt(adt, substs));
        !ty.is_structural_eq_shallow(cx.tcx)
    }
}

// FIXME: Use `mir::visit::Visitor` for the `in_*` functions if/when it supports early return.

/// Returns `true` if this `Rvalue` contains qualif `Q`.
pub fn in_rvalue<Q, F>(cx: &ConstCx<'_, 'tcx>, in_local: &mut F, rvalue: &Rvalue<'tcx>) -> bool
where
    Q: Qualif,
    F: FnMut(Local) -> bool,
{
    match rvalue {
        Rvalue::ThreadLocalRef(_) | Rvalue::NullaryOp(..) => {
            Q::in_any_value_of_ty(cx, rvalue.ty(cx.body, cx.tcx))
        }

        Rvalue::Discriminant(place) | Rvalue::Len(place) => {
            in_place::<Q, _>(cx, in_local, place.as_ref())
        }

        Rvalue::Use(operand)
        | Rvalue::Repeat(operand, _)
        | Rvalue::UnaryOp(_, operand)
        | Rvalue::Cast(_, operand, _) => in_operand::<Q, _>(cx, in_local, operand),

        Rvalue::BinaryOp(_, box (lhs, rhs)) | Rvalue::CheckedBinaryOp(_, box (lhs, rhs)) => {
            in_operand::<Q, _>(cx, in_local, lhs) || in_operand::<Q, _>(cx, in_local, rhs)
        }

        Rvalue::Ref(_, _, place) | Rvalue::AddressOf(_, place) => {
            // Special-case reborrows to be more like a copy of the reference.
            if let Some((place_base, ProjectionElem::Deref)) = place.as_ref().last_projection() {
                let base_ty = place_base.ty(cx.body, cx.tcx).ty;
                if let ty::Ref(..) = base_ty.kind() {
                    return in_place::<Q, _>(cx, in_local, place_base);
                }
            }

            in_place::<Q, _>(cx, in_local, place.as_ref())
        }

        Rvalue::Aggregate(kind, operands) => {
            // Return early if we know that the struct or enum being constructed is always
            // qualified.
            if let AggregateKind::Adt(def, _, substs, ..) = **kind {
                if Q::in_adt_inherently(cx, def, substs) {
                    return true;
                }
            }

            // Otherwise, proceed structurally...
            operands.iter().any(|o| in_operand::<Q, _>(cx, in_local, o))
        }
    }
}

/// Returns `true` if this `Place` contains qualif `Q`.
pub fn in_place<Q, F>(cx: &ConstCx<'_, 'tcx>, in_local: &mut F, place: PlaceRef<'tcx>) -> bool
where
    Q: Qualif,
    F: FnMut(Local) -> bool,
{
    let mut place = place;
    while let Some((place_base, elem)) = place.last_projection() {
        match elem {
            ProjectionElem::Index(index) if in_local(index) => return true,

            ProjectionElem::Deref
            | ProjectionElem::Field(_, _)
            | ProjectionElem::ConstantIndex { .. }
            | ProjectionElem::Subslice { .. }
            | ProjectionElem::Downcast(_, _)
            | ProjectionElem::Index(_) => {}
        }

        let base_ty = place_base.ty(cx.body, cx.tcx);
        let proj_ty = base_ty.projection_ty(cx.tcx, elem).ty;
        if !Q::in_any_value_of_ty(cx, proj_ty) {
            return false;
        }

        place = place_base;
    }

    assert!(place.projection.is_empty());
    in_local(place.local)
}

/// Returns `true` if this `Operand` contains qualif `Q`.
pub fn in_operand<Q, F>(cx: &ConstCx<'_, 'tcx>, in_local: &mut F, operand: &Operand<'tcx>) -> bool
where
    Q: Qualif,
    F: FnMut(Local) -> bool,
{
    let constant = match operand {
        Operand::Copy(place) | Operand::Move(place) => {
            return in_place::<Q, _>(cx, in_local, place.as_ref());
        }

        Operand::Constant(c) => c,
    };

    // Check the qualifs of the value of `const` items.
    if let Some(ct) = constant.literal.const_for_ty() {
        if let ty::ConstKind::Unevaluated(ty::Unevaluated { def, substs_: _, promoted }) = ct.val {
            assert!(promoted.is_none());
            // Don't peek inside trait associated constants.
            if cx.tcx.trait_of_item(def.did).is_none() {
                let qualifs = if let Some((did, param_did)) = def.as_const_arg() {
                    cx.tcx.at(constant.span).mir_const_qualif_const_arg((did, param_did))
                } else {
                    cx.tcx.at(constant.span).mir_const_qualif(def.did)
                };

                if !Q::in_qualifs(&qualifs) {
                    return false;
                }

                // Just in case the type is more specific than
                // the definition, e.g., impl associated const
                // with type parameters, take it into account.
            }
        }
    }
    // Otherwise use the qualifs of the type.
    Q::in_any_value_of_ty(cx, constant.literal.ty())
}
<|endoftext|>


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

//! Helper library for accessing Oak storage services.

use crate::grpc;
use crate::proto::oak::storage::{
    StorageChannelDeleteRequest, StorageChannelDeleteResponse, StorageChannelReadRequest,
    StorageChannelReadResponse, StorageChannelWriteRequest, StorageChannelWriteResponse,
    StorageItem,
};
use log::info;

/// Default name for predefined Node config that corresponds to a storage
/// pseudo-Node.
pub const DEFAULT_CONFIG_NAME: &str = "storage";

/// Local representation of the connection to an external storage service.
pub struct Storage {
    client: crate::grpc::client::Client,
}

impl Storage {
    /// Create a default `Storage` instance assuming the default pre-defined
    /// name (`"storage"`) identifying storage Node config.
    pub fn default() -> Option<Storage> {
        Storage::new(DEFAULT_CONFIG_NAME)
    }

    /// Create a `Storage` instance using the given name identifying storage
    /// Node configuration.
    pub fn new(config: &str) -> Option<Storage> {
        crate::grpc::client::Client::new(config, "oak_main").map(|client| Storage { client })
    }

    fn execute_operation<Req, Res>(
        &mut self,
        grpc_method_name: &str,
        operation_request: Req,
    ) -> grpc::Result<Res>
    where
        Req: prost::Message,
        Res: prost::Message + Default,
    {
        info!("StorageChannelRequest: {:?}", operation_request);
        crate::grpc::invoke_grpc_method(
            grpc_method_name,
            &operation_request,
            None,
            &self.client.invocation_sender,
        )
    }

    /// Read the value associated with the given `name` from the storage
    /// instance identified by `name`.
    pub fn read(&mut self, storage_name: &[u8], name: &[u8]) -> grpc::Result<Vec<u8>> {
        let read_request = StorageChannelReadRequest {
            storage_name: storage_name.to_owned(),
            transaction_id: vec![],
            item: Some(StorageItem {
                name: name.to_owned(),
                value: vec![],
                label: None,
            }),
        };

        // TODO(#757): Automatically generate boilerplate from the proto
        // definition.
        self.execute_operation::<StorageChannelReadRequest, StorageChannelReadResponse>(
            "/oak.storage.StorageService/Read",
            read_request,
        )
        .map(|r| r.item.unwrap_or_default().value.to_vec())
    }

    /// Set the value associated with the given `name` from the storage instance
    /// identified by `name`.
    pub fn write(&mut self, storage_name: &[u8], name: &[u8], value: &[u8]) -> grpc::Result<()> {
        // TODO(#449): Set policy for item.
        let write_request = StorageChannelWriteRequest {
            storage_name: storage_name.to_owned(),
            transaction_id: vec![],
            item: Some(StorageItem {
                name: name.to_owned(),
                value: value.to_owned(),
                label: None,
            }),
        };

        // TODO(#757): Automatically generate boilerplate from the proto definition.
        self.execute_operation::<StorageChannelWriteRequest, StorageChannelWriteResponse>(
            "/oak.storage.StorageService/Write",
            write_request,
        )
        .map(|_| ())
    }

    /// Delete the value associated with the given `name` from the storage
    /// instance identified by `name`.
    pub fn delete(&mut self, storage_name: &[u8], name: &[u8]) -> grpc::Result<()> {
        let delete_request = StorageChannelDeleteRequest {
            storage_name: storage_name.to_owned(),
            transaction_id: vec![],
            item: Some(StorageItem {
                name: name.to_owned(),
                value: vec![],
                label: None,
            }),
        };

        // TODO(#757): Automatically generate boilerplate from the proto definition.
        self.execute_operation::<StorageChannelDeleteRequest, StorageChannelDeleteResponse>(
            "/oak.storage.StorageService/Delete",
            delete_request,
        )
        .map(|_| ())
    }
}
<commit_msg>sdk/storage: use generic gRPC client helpers (#835)<commit_after>//
// Copyright 2019 The Project Oak Authors
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
//

//! Helper library for accessing Oak storage services.

use crate::grpc;
use crate::proto::oak::storage::{
    StorageChannelDeleteRequest, StorageChannelReadRequest, StorageChannelWriteRequest,
    StorageItem, StorageServiceClient,
};

/// Default name for predefined Node config that corresponds to a storage
/// pseudo-Node.
pub const DEFAULT_CONFIG_NAME: &str = "storage";

/// Local representation of the connection to an external storage service.
pub struct Storage {
    client: StorageServiceClient,
}

impl Storage {
    /// Create a default `Storage` instance assuming the default pre-defined
    /// name (`"storage"`) identifying storage Node config.
    pub fn default() -> Option<Storage> {
        Storage::new(DEFAULT_CONFIG_NAME)
    }

    /// Create a `Storage` instance using the given name identifying storage
    /// Node configuration.
    pub fn new(config: &str) -> Option<Storage> {
        crate::grpc::client::Client::new(config, "oak_main").map(|client| Storage {
            client: StorageServiceClient(client),
        })
    }

    /// Read the value associated with the given `name` from the storage
    /// instance identified by `name`.
    pub fn read(&mut self, storage_name: &[u8], name: &[u8]) -> grpc::Result<Vec<u8>> {
        let read_request = StorageChannelReadRequest {
            storage_name: storage_name.to_owned(),
            transaction_id: vec![],
            item: Some(StorageItem {
                name: name.to_owned(),
                value: vec![],
                label: None,
            }),
        };

        self.client
            .read(read_request)
            .map(|r| r.item.unwrap_or_default().value.to_vec())
    }

    /// Set the value associated with the given `name` from the storage instance
    /// identified by `name`.
    pub fn write(&mut self, storage_name: &[u8], name: &[u8], value: &[u8]) -> grpc::Result<()> {
        // TODO(#449): Set policy for item.
        let write_request = StorageChannelWriteRequest {
            storage_name: storage_name.to_owned(),
            transaction_id: vec![],
            item: Some(StorageItem {
                name: name.to_owned(),
                value: value.to_owned(),
                label: None,
            }),
        };

        self.client.write(write_request).map(|_| ())
    }

    /// Delete the value associated with the given `name` from the storage
    /// instance identified by `name`.
    pub fn delete(&mut self, storage_name: &[u8], name: &[u8]) -> grpc::Result<()> {
        let delete_request = StorageChannelDeleteRequest {
            storage_name: storage_name.to_owned(),
            transaction_id: vec![],
            item: Some(StorageItem {
                name: name.to_owned(),
                value: vec![],
                label: None,
            }),
        };

        self.client.delete(delete_request).map(|_| ())
    }
}
<|endoftext|>


<commit_before>use std::f64;
use std::option::Option;
use rand::Rng;
use distribution::{Gamma, Distribution, Univariate, Continuous};
use result;

pub struct ChiSquared {
    k: f64,
    g: Gamma
}

impl ChiSquared {
    pub fn new(k: f64) -> result::Result<ChiSquared> {
        match Gamma::new(k / 2.0, 0.5) {
            Ok(g) => Ok(ChiSquared{k: k, g: g}),
            Err(e) => Err(e)
        }
    }

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

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

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

impl Distribution for ChiSquared {
    fn sample<R: Rng>(&self, r: &mut R) -> f64 {
        self.g.sample(r)
    }
}

impl Univariate for ChiSquared {
    fn mean(&self) -> f64 {
        self.g.mean()
    }

    fn variance(&self) -> f64 {
        self.g.variance()
    }

    fn std_dev(&self) -> f64 {
        self.g.std_dev()
    }

    fn entropy(&self) -> f64 {
        self.g.entropy()
    }

    fn skewness(&self) -> f64 {
        self.g.skewness()
    }

    fn median(&self) -> Option<f64> {
        None
    }

    fn cdf(&self, x: f64) -> result::Result<f64> {
        self.g.cdf(x)
    }
}

impl Continuous for ChiSquared {
    fn mode(&self) -> f64 {
        self.g.mode()
    }

    fn min(&self) -> f64 {
        0.0
    }

    fn max(&self) -> f64 {
        f64::INFINITY
    }

    fn pdf(&self, x: f64) -> f64 {
        self.g.pdf(x)
    }

    fn ln_pdf(&self, x: f64) -> f64 {
        self.g.ln_pdf(x)
    }
}<commit_msg>derived traits for chi squared<commit_after>use std::f64;
use std::option::Option;
use rand::Rng;
use distribution::{Gamma, Distribution, Univariate, Continuous};
use result;

pub struct ChiSquared {
    k: f64,
    g: Gamma
}

#[derive(Debug, Clone, PartialEq)]
impl ChiSquared {
    pub fn new(k: f64) -> result::Result<ChiSquared> {
        match Gamma::new(k / 2.0, 0.5) {
            Ok(g) => Ok(ChiSquared{k: k, g: g}),
            Err(e) => Err(e)
        }
    }

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

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

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

impl Distribution for ChiSquared {
    fn sample<R: Rng>(&self, r: &mut R) -> f64 {
        self.g.sample(r)
    }
}

impl Univariate for ChiSquared {
    fn mean(&self) -> f64 {
        self.g.mean()
    }

    fn variance(&self) -> f64 {
        self.g.variance()
    }

    fn std_dev(&self) -> f64 {
        self.g.std_dev()
    }

    fn entropy(&self) -> f64 {
        self.g.entropy()
    }

    fn skewness(&self) -> f64 {
        self.g.skewness()
    }

    fn median(&self) -> Option<f64> {
        None
    }

    fn cdf(&self, x: f64) -> result::Result<f64> {
        self.g.cdf(x)
    }
}

impl Continuous for ChiSquared {
    fn mode(&self) -> f64 {
        self.g.mode()
    }

    fn min(&self) -> f64 {
        0.0
    }

    fn max(&self) -> f64 {
        f64::INFINITY
    }

    fn pdf(&self, x: f64) -> f64 {
        self.g.pdf(x)
    }

    fn ln_pdf(&self, x: f64) -> f64 {
        self.g.ln_pdf(x)
    }
}<|endoftext|>


<commit_before>use super::frontend_prelude::*;

use crate::models::ApiToken;
use crate::schema::api_tokens;
use crate::util::read_fill;
use crate::views::EncodableApiTokenWithToken;

use conduit::{Body, Response};
use serde_json as json;

/// Handles the `GET /me/tokens` route.
pub fn list(req: &mut dyn RequestExt) -> EndpointResult {
    let authenticated_user = req.authenticate()?.forbid_api_token_auth()?;
    let conn = req.db_write()?;
    let user = authenticated_user.user();

    let tokens: Vec<ApiToken> = ApiToken::belonging_to(&user)
        .filter(api_tokens::revoked.eq(false))
        .order(api_tokens::created_at.desc())
        .load(&*conn)?;

    Ok(req.json(&json!({ "api_tokens": tokens })))
}

/// Handles the `PUT /me/tokens` route.
pub fn new(req: &mut dyn RequestExt) -> EndpointResult {
    /// The incoming serialization format for the `ApiToken` model.
    #[derive(Deserialize, Serialize)]
    struct NewApiToken {
        name: String,
    }

    /// The incoming serialization format for the `ApiToken` model.
    #[derive(Deserialize, Serialize)]
    struct NewApiTokenRequest {
        api_token: NewApiToken,
    }

    let max_size = 2000;
    let length = req
        .content_length()
        .ok_or_else(|| bad_request("missing header: Content-Length"))?;

    if length > max_size {
        return Err(bad_request(&format!("max content length is: {max_size}")));
    }

    let mut json = vec![0; length as usize];
    read_fill(req.body(), &mut json)?;

    let json =
        String::from_utf8(json).map_err(|_| bad_request(&"json body was not valid utf-8"))?;

    let new: NewApiTokenRequest = json::from_str(&json)
        .map_err(|e| bad_request(&format!("invalid new token request: {e:?}")))?;

    let name = &new.api_token.name;
    if name.is_empty() {
        return Err(bad_request("name must have a value"));
    }

    let authenticated_user = req.authenticate()?;
    if authenticated_user.api_token_id().is_some() {
        return Err(bad_request(
            "cannot use an API token to create a new API token",
        ));
    }

    let conn = req.db_write()?;
    let user = authenticated_user.user();

    let max_token_per_user = 500;
    let count: i64 = ApiToken::belonging_to(&user).count().get_result(&*conn)?;
    if count >= max_token_per_user {
        return Err(bad_request(&format!(
            "maximum tokens per user is: {}",
            max_token_per_user
        )));
    }

    let api_token = ApiToken::insert(&*conn, user.id, name)?;
    let api_token = EncodableApiTokenWithToken::from(api_token);

    Ok(req.json(&json!({ "api_token": api_token })))
}

/// Handles the `DELETE /me/tokens/:id` route.
pub fn revoke(req: &mut dyn RequestExt) -> EndpointResult {
    let id = req.params()["id"]
        .parse::<i32>()
        .map_err(|e| bad_request(&format!("invalid token id: {e:?}")))?;

    let authenticated_user = req.authenticate()?;
    let conn = req.db_write()?;
    let user = authenticated_user.user();
    diesel::update(ApiToken::belonging_to(&user).find(id))
        .set(api_tokens::revoked.eq(true))
        .execute(&*conn)?;

    Ok(req.json(&json!({})))
}

/// Handles the `DELETE /tokens/current` route.
pub fn revoke_current(req: &mut dyn RequestExt) -> EndpointResult {
    let authenticated_user = req.authenticate()?;
    let api_token_id = authenticated_user
        .api_token_id()
        .ok_or_else(|| bad_request("token not provided"))?;

    let conn = req.db_write()?;
    diesel::update(api_tokens::table.filter(api_tokens::id.eq(api_token_id)))
        .set(api_tokens::revoked.eq(true))
        .execute(&*conn)?;

    Ok(Response::builder().status(204).body(Body::empty()).unwrap())
}
<commit_msg>Use `db_read_prefer_primary()` for `GET /me/tokens` endpoint<commit_after>use super::frontend_prelude::*;

use crate::models::ApiToken;
use crate::schema::api_tokens;
use crate::util::read_fill;
use crate::views::EncodableApiTokenWithToken;

use conduit::{Body, Response};
use serde_json as json;

/// Handles the `GET /me/tokens` route.
pub fn list(req: &mut dyn RequestExt) -> EndpointResult {
    let authenticated_user = req.authenticate()?.forbid_api_token_auth()?;
    let conn = req.db_read_prefer_primary()?;
    let user = authenticated_user.user();

    let tokens: Vec<ApiToken> = ApiToken::belonging_to(&user)
        .filter(api_tokens::revoked.eq(false))
        .order(api_tokens::created_at.desc())
        .load(&*conn)?;

    Ok(req.json(&json!({ "api_tokens": tokens })))
}

/// Handles the `PUT /me/tokens` route.
pub fn new(req: &mut dyn RequestExt) -> EndpointResult {
    /// The incoming serialization format for the `ApiToken` model.
    #[derive(Deserialize, Serialize)]
    struct NewApiToken {
        name: String,
    }

    /// The incoming serialization format for the `ApiToken` model.
    #[derive(Deserialize, Serialize)]
    struct NewApiTokenRequest {
        api_token: NewApiToken,
    }

    let max_size = 2000;
    let length = req
        .content_length()
        .ok_or_else(|| bad_request("missing header: Content-Length"))?;

    if length > max_size {
        return Err(bad_request(&format!("max content length is: {max_size}")));
    }

    let mut json = vec![0; length as usize];
    read_fill(req.body(), &mut json)?;

    let json =
        String::from_utf8(json).map_err(|_| bad_request(&"json body was not valid utf-8"))?;

    let new: NewApiTokenRequest = json::from_str(&json)
        .map_err(|e| bad_request(&format!("invalid new token request: {e:?}")))?;

    let name = &new.api_token.name;
    if name.is_empty() {
        return Err(bad_request("name must have a value"));
    }

    let authenticated_user = req.authenticate()?;
    if authenticated_user.api_token_id().is_some() {
        return Err(bad_request(
            "cannot use an API token to create a new API token",
        ));
    }

    let conn = req.db_write()?;
    let user = authenticated_user.user();

    let max_token_per_user = 500;
    let count: i64 = ApiToken::belonging_to(&user).count().get_result(&*conn)?;
    if count >= max_token_per_user {
        return Err(bad_request(&format!(
            "maximum tokens per user is: {}",
            max_token_per_user
        )));
    }

    let api_token = ApiToken::insert(&*conn, user.id, name)?;
    let api_token = EncodableApiTokenWithToken::from(api_token);

    Ok(req.json(&json!({ "api_token": api_token })))
}

/// Handles the `DELETE /me/tokens/:id` route.
pub fn revoke(req: &mut dyn RequestExt) -> EndpointResult {
    let id = req.params()["id"]
        .parse::<i32>()
        .map_err(|e| bad_request(&format!("invalid token id: {e:?}")))?;

    let authenticated_user = req.authenticate()?;
    let conn = req.db_write()?;
    let user = authenticated_user.user();
    diesel::update(ApiToken::belonging_to(&user).find(id))
        .set(api_tokens::revoked.eq(true))
        .execute(&*conn)?;

    Ok(req.json(&json!({})))
}

/// Handles the `DELETE /tokens/current` route.
pub fn revoke_current(req: &mut dyn RequestExt) -> EndpointResult {
    let authenticated_user = req.authenticate()?;
    let api_token_id = authenticated_user
        .api_token_id()
        .ok_or_else(|| bad_request("token not provided"))?;

    let conn = req.db_write()?;
    diesel::update(api_tokens::table.filter(api_tokens::id.eq(api_token_id)))
        .set(api_tokens::revoked.eq(true))
        .execute(&*conn)?;

    Ok(Response::builder().status(204).body(Body::empty()).unwrap())
}
<|endoftext|>


<commit_before>// pest. The Elegant Parser
// Copyright (C) 2017  Dragoș Tiselice
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

//! # pest. The Elegant Parser
//!
//! pest is a [PEG](https://en.wikipedia.org/wiki/Parsing_expression_grammar) parser built with
//! *simplicity* and *speed* in mind.
//!
//! ## Input & Parser
//!
//! pest works mainly through two `trait`s: `Input` & `Parser`. `Input` defines the capabilities of
//! an input source, while `Parser` provides an interface to the parsing functionality. Since
//! `Parser` is a `trait`, parsing needs to be defined either though the `#[derive(Parser)]`
//! attribute, or manually through the [`Position API`](inputs/struct.Position.html). The use of the
//! `derive` is highly encouraged since this is the only way you can make use of pest's PEG
//! grammar, while `manual parser definition` can be used where highly specific or efficient parsing
//! is required.
//!
//! ## `#[derive(Parser)]`
//!
//! pest comes with a procedural macro crate--`pest_derive`--which needs to be included in
//! `Cargo.toml` in order to enable the `derive`.
//!
//! ```toml
//! pest_derive = "*"
//! ```
//!
//! ## `.pest` files
//!
//! Grammar definitions reside in custom `.pest` files located in the `src` directory. Their path is
//! relative to `src` and is specified between the `derive` attribute and an empty `struct` that
//! `Parser` will be derived on.
//!
//! Because of a limitation in procedural macros, there is no way for Cargo to know that a module
//! needs to be recompiled based on the file that the procedural macro is opening. This leads to the
//! case where modifying a `.pest` file without touching the file where the `derive` is does not
//! recompile it if it already has a working binary in the cache. To avoid this issue, the grammar
//! file can be included in a dummy `const` definition while debugging.
//!
//! ```ignore
//! const _GRAMMAR: &'static = include_str!("path/to/my_grammar.pest"); // relative to this file
//!
//! #[derive(Parser)]
//! #[grammar = "path/to/my_grammar.pest"] // relative to src
//! struct MyParser;
//! ```
//!
//! The grammar of `.pest` files is documented in the
//! [`pest_derive` crate](https://docs.rs/pest_derive/#Grammar).

#![doc(html_root_url = "https://docs.rs/pest")]

use std::fmt::Debug;
use std::hash::Hash;

mod error;
pub mod inputs;
pub mod iterators;
mod macros;
mod parser;
mod parser_state;
pub mod prec_climber;
mod token;

/// A `trait` which parser rules must implement.
///
/// This trait is set up so that any struct that implements all of its required traits will
/// automatically implement this trait as well.
///
/// This is essentially a [trait alias](https://github.com/rust-lang/rfcs/pull/1733). When trait
/// aliases are implemented, this may be replaced by one.
pub trait RuleType: Copy + Debug + Eq + Hash + Ord {}
impl<T: Copy + Debug + Eq + Hash + Ord> RuleType for T {}

pub use error::Error;
pub use parser::Parser;
pub use parser_state::{state, ParserState};
pub use token::Token;
<commit_msg>Fixed wrong code block<commit_after>// pest. The Elegant Parser
// Copyright (C) 2017  Dragoș Tiselice
//
// This Source Code Form is subject to the terms of the Mozilla Public
// License, v. 2.0. If a copy of the MPL was not distributed with this
// file, You can obtain one at http://mozilla.org/MPL/2.0/.

//! # pest. The Elegant Parser
//!
//! pest is a [PEG](https://en.wikipedia.org/wiki/Parsing_expression_grammar) parser built with
//! *simplicity* and *speed* in mind.
//!
//! ## Input & Parser
//!
//! pest works mainly through two `trait`s: `Input` & `Parser`. `Input` defines the capabilities of
//! an input source, while `Parser` provides an interface to the parsing functionality. Since
//! `Parser` is a `trait`, parsing needs to be defined either though the `#[derive(Parser)]`
//! attribute, or manually through the [`Position API`](inputs/struct.Position.html). The use of the
//! `derive` is highly encouraged since this is the only way you can make use of pest's PEG
//! grammar, while manual parser definition can be used where highly specific or efficient parsing
//! is required.
//!
//! ## `#[derive(Parser)]`
//!
//! pest comes with a procedural macro crate--`pest_derive`--which needs to be included in
//! `Cargo.toml` in order to enable the `derive`.
//!
//! ```toml
//! pest_derive = "*"
//! ```
//!
//! ## `.pest` files
//!
//! Grammar definitions reside in custom `.pest` files located in the `src` directory. Their path is
//! relative to `src` and is specified between the `derive` attribute and an empty `struct` that
//! `Parser` will be derived on.
//!
//! Because of a limitation in procedural macros, there is no way for Cargo to know that a module
//! needs to be recompiled based on the file that the procedural macro is opening. This leads to the
//! case where modifying a `.pest` file without touching the file where the `derive` is does not
//! recompile it if it already has a working binary in the cache. To avoid this issue, the grammar
//! file can be included in a dummy `const` definition while debugging.
//!
//! ```ignore
//! const _GRAMMAR: &'static = include_str!("path/to/my_grammar.pest"); // relative to this file
//!
//! #[derive(Parser)]
//! #[grammar = "path/to/my_grammar.pest"] // relative to src
//! struct MyParser;
//! ```
//!
//! The grammar of `.pest` files is documented in the
//! [`pest_derive` crate](https://docs.rs/pest_derive/#Grammar).

#![doc(html_root_url = "https://docs.rs/pest")]

use std::fmt::Debug;
use std::hash::Hash;

mod error;
pub mod inputs;
pub mod iterators;
mod macros;
mod parser;
mod parser_state;
pub mod prec_climber;
mod token;

/// A `trait` which parser rules must implement.
///
/// This trait is set up so that any struct that implements all of its required traits will
/// automatically implement this trait as well.
///
/// This is essentially a [trait alias](https://github.com/rust-lang/rfcs/pull/1733). When trait
/// aliases are implemented, this may be replaced by one.
pub trait RuleType: Copy + Debug + Eq + Hash + Ord {}
impl<T: Copy + Debug + Eq + Hash + Ord> RuleType for T {}

pub use error::Error;
pub use parser::Parser;
pub use parser_state::{state, ParserState};
pub use token::Token;
<|endoftext|>


<commit_before>use anyhow::Result;
use clap::Parser;
use regex::Regex;
use std::collections::HashMap;
use std::fs::File;
use std::io::{self, BufRead};
use std::path::PathBuf;
use std::time::Instant;

use polars::frame::DataFrame;
use polars::prelude::*;

#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
    /// List of input files. Ideally with headers
    #[clap(short, long, multiple_values(true))]
    input: Vec<PathBuf>,

    /// An output file
    #[clap(short, long)]
    output: String,

    /// The column to select. If files have no headers, use "column_N"
    #[clap(short, long, multiple_values(true))]
    select: Vec<String>,

    /// The size of chunks to process at once
    #[clap(short, long, default_value_t = 100)]
    chunksize: usize,
}

fn load_data(path: &PathBuf, delim: u8) -> Result<DataFrame> {
    let mut df: DataFrame = CsvReader::from_path(path)?
        .has_header(true)
        .with_delimiter(delim)
        .with_null_values(Some(NullValues::AllColumns("NA".to_string())))
        .infer_schema(None)
        .finish()
        .unwrap();

    // hstack the experiment name (derived from the filename) into the DataFrame
    let iter_exp =
        vec![path.file_name().unwrap().to_str().unwrap().split('-').collect::<Vec<&str>>()[0..=2]
            .join("-")]
        .into_iter();
    let exp_col: Series = Series::new(
        "experiment",
        iter_exp.flat_map(|n| std::iter::repeat(n).take(df.height())).collect::<Vec<String>>(),
    );
    // exp_col.rename("experiment");
    df.hstack_mut(&[exp_col]).unwrap();

    Ok(df)
}

fn baseline_get_median_gt_zero(str_val: &Series) -> Series {
    let lists = str_val.utf8().unwrap().into_iter().map(|x| {
        x.unwrap().split(',').into_iter().map(|y| y.parse::<f64>().unwrap()).collect::<Vec<f64>>()
    });

    let medians: Vec<bool> =
        lists.into_iter().map(|x| Series::from_iter(x).median().unwrap() > 0.0).collect();
    Series::from_iter(medians)
}

fn filter_input(input: &mut DataFrame, baseline: bool) -> DataFrame {
    if baseline {
        filter_baseline(input)
    } else {
        filter_differential(input)
    }
}

fn filter_baseline(input: &mut DataFrame) -> DataFrame {
    // This is a baseline experiment.
    // Find columns starting with lower case g, then apply the function to convert to
    // medain and select greater than zero
    let mut meas = Vec::<String>::new();

    for col in input.get_column_names_owned() {
        if col.starts_with('g') {
            input.apply(&col, baseline_get_median_gt_zero).unwrap();
            meas.push(col);
        }
    }
    // Selection should now have all the gN column names in it
    input
        .clone()
        .lazy()
        .filter(any_exprs(meas.into_iter().map(|x| col(&x)).collect::<Vec<Expr>>()))
        .collect()
        .unwrap()
}

fn filter_differential(input: &DataFrame) -> DataFrame {
    // find the p value and log fold columns
    let pv_regex = Regex::new(r".*p-value.*").unwrap();
    let log_fold_regex = Regex::new(r".*log2.*").unwrap();
    let mut exprs: Vec<Expr> = Vec::new();
    let mut pv: Expr;
    let mut lf: Expr;
    for column in input.get_column_names_owned() {
        if pv_regex.is_match(&column) {
            pv = col(&column).is_not_null();
            exprs.push(pv);
        } else if log_fold_regex.is_match(&column) {
            lf = col(&column).neq(lit(0));
            exprs.push(lf);
        }
    }

    input.clone().lazy().filter(all_exprs(exprs)).collect().unwrap()
}

fn load_chunk(
    paths: &mut Vec<PathBuf>,
    select: &mut Vec<String>,
) -> Result<DataFrame, anyhow::Error> {
    /*
    Work with the chunks of input to reduce them into a single dataframe of the genes we want from
    that particular set of experiments. Should return a single dataframe
    */
    let mut output: Option<DataFrame> = None;
    while let Some(infile) = paths.pop() {
        // Load the input Csv
        let mut input: DataFrame = if infile.extension().unwrap() == "tsv" {
            load_data(&infile, b'\t')
        } else {
            load_data(&infile, b',')
        }
        .unwrap_or_else(|error| match error.downcast_ref::<PolarsError>() {
            Some(PolarsError::Io(_string)) => panic!("Input file does not exist! {:?}", infile),
            _ => panic!("An error occurred! {:?}", error),
        });

        // Rename columns to remove . in the names. Now also remove spaces
        let mut new_cols = Vec::new();
        let mut c: u8 = 0;
        for nm in input.get_column_names().iter() {
            let mut new_name = nm.replace('.', "").replace(' ', "");
            if new_cols.contains(&new_name) {
                c += 1;
                new_name += &c.to_string(); // Hopefully avoid duplicate names
            }
            new_cols.push(new_name);
        }

        if new_cols != input.get_column_names() {
            input.set_column_names(&new_cols).unwrap_or_else(|error| {
                panic!("Failed to set column names for some reason {:?}", error);
            });
        }

        // Need to detect which type of experiment it is - differential or baseline
        let type_re = Regex::new(r"tpms").unwrap();
        let filtered = filter_input(&mut input, type_re.is_match(infile.to_str().unwrap()));
        // Drop everything from the input except the genes
        let genes: DataFrame = filtered.select(select.iter()).unwrap_or_else(|error| match error {
            PolarsError::NotFound(_string) => {
                panic!(
                    "{:?} was not found in the header, does the file have a header?\n{:?}",
                    select,
                    input.get_column_names()
                );
            },
            _ => panic!("Error selecting column from input: {:?}", error),
        });

        output = match output {
            None => {
                let output_df = genes.clone();
                Some(output_df)
            },
            Some(mut output_df) => {
                // genes =/= output, so we are handling a new file
                output_df.extend(&genes).unwrap_or_else(|error| {
                    panic!("Failed to extend the output dataframe. {:?}", error)
                });
                output_df =
                    output_df.unique(None, UniqueKeepStrategy::First).unwrap_or_else(|error| {
                        panic!("Failed to parse selected column for uniqueness. {:?}", error)
                    });
                Some(output_df)
            },
        }
    }
    Ok(output.unwrap())
}

fn get_taxid(paths: &mut Vec<PathBuf>, lookup_table: &mut HashMap<String, String>) {
    let tax_regex = Regex::new(r".*Taxon_([0-9]{4,})").unwrap();

    while let Some(path) = paths.pop() {
        let exp_name =
            path.file_name().unwrap().to_str().unwrap().split('-').collect::<Vec<&str>>()[0..=2]
                .join("-")
                .replace(".condensed", "");

        let file = File::open(path).unwrap();
        let lines = io::BufReader::new(file).lines();
        for line in lines {
            match line {
                Ok(line) => {
                    let captures = tax_regex.captures(line.as_str());
                    match captures {
                        Some(caps) => {
                            let taxid = caps.get(1).map_or("", |m| m.as_str());
                            lookup_table.insert(exp_name, taxid.to_string());
                            break;
                        },
                        None => {},
                    };
                },
                Err(_idc) => {},
            }
        }
    }
}

fn main() {
    let timer = Instant::now();

    let mut args = Args::parse();

    // Always select the experiment
    args.select.push("experiment".to_string());

    // separate expression data from taxa data
    let mut taxa_files: Vec<PathBuf> = Vec::new();
    let mut expr_files: Vec<PathBuf> = Vec::new();

    let re = Regex::new(r"sdrf.tsv").unwrap();

    for infile in args.input {
        if re.is_match(infile.to_str().unwrap()) {
            // File is experiment setup data, parse for taxa
            taxa_files.push(infile);
        } else {
            // Should only be expression data, parse for genes
            expr_files.push(infile);
        }
    }

    println!("{} taxa files to parse", taxa_files.len());
    println!("{} expression files to parse", expr_files.len());

    let n_files = expr_files.len();

    let mut dataframes: Vec<DataFrame> = Vec::new();

    let chunked_input: Vec<Vec<PathBuf>> =
        expr_files.chunks(args.chunksize).map(|x| x.to_vec()).collect();
    // Read everything into a big vector
    let mut n_chunks = 0;
    for mut files_chunk in chunked_input {
        dataframes.push(
            load_chunk(&mut files_chunk, &mut args.select)
                .unwrap_or_else(|_error| panic!("Something wrong in one of the reads, aborting")),
        );
        n_chunks += 1;
        println!("Done {} files", n_chunks * args.chunksize);
    }

    let mut output: DataFrame = dataframes[0].clone();

    for item in dataframes.iter() {
        output.extend(item).unwrap_or_else(|error| {
            panic!("Unable to extend the DataFrame! Out of memory? {:?}", error)
        });
        output.unique(None, UniqueKeepStrategy::First).unwrap();
    }

    // Now parse the experiment files to get taxa
    let chunked_taxa: Vec<Vec<PathBuf>> =
        taxa_files.chunks(args.chunksize).map(|x| x.to_vec()).collect();

    let mut lookup_table: HashMap<String, String> = HashMap::new();

    for mut tax_chunk in chunked_taxa {
        get_taxid(&mut tax_chunk, &mut lookup_table);
    }

    // Now we have the experiment - taxid lookup_table, we just need to add the column to the output df

    let expt_col = output.select(["experiment"]).unwrap();
    let mut tax_ids: Vec<String> = Vec::with_capacity(expt_col.height());
    for ex in expt_col.iter() {
        let mut warn_flag: bool = true;
        let uhy = ex.utf8().unwrap();
        for (idx, x) in uhy.into_iter().enumerate() {
            if lookup_table.contains_key(x.unwrap()) {
                tax_ids.insert(idx, lookup_table.get(x.unwrap()).unwrap().to_string());
            } else {
                tax_ids.insert(idx, String::new());
                if warn_flag {
                    println!("Experiment {} does not name a taxon", x.unwrap());
                    warn_flag = false;
                }
            }
        }
    }

    output.hstack_mut(&[Series::new("taxid", tax_ids.into_iter())]).unwrap();

    let out_stream: File = File::create(args.output).unwrap();
    CsvWriter::new(out_stream)
        .has_header(true)
        .finish(&mut output)
        .unwrap_or_else(|error| panic!("Something wrong writing file: {:?}", error));

    println!("Processed {} in {} seconds", n_files, timer.elapsed().as_secs());
}
<commit_msg>Workaround stack overflow for very wide entries like E-MTAB-2770<commit_after>use anyhow::Result;
use clap::Parser;
use regex::Regex;
use std::collections::HashMap;
use std::fs::File;
use std::io::{self, BufRead};
use std::path::PathBuf;
use std::time::Instant;

use polars::frame::DataFrame;
use polars::prelude::*;

#[derive(Parser, Debug)]
#[clap(author, version, about, long_about = None)]
struct Args {
    /// List of input files. Ideally with headers
    #[clap(short, long, multiple_values(true))]
    input: Vec<PathBuf>,

    /// An output file
    #[clap(short, long)]
    output: String,

    /// The column to select. If files have no headers, use "column_N"
    #[clap(short, long, multiple_values(true))]
    select: Vec<String>,

    /// The size of chunks to process at once
    #[clap(short, long, default_value_t = 100)]
    chunksize: usize,
}

fn load_data(path: &PathBuf, delim: u8) -> Result<DataFrame> {
    let mut df: DataFrame = CsvReader::from_path(path)?
        .has_header(true)
        .with_delimiter(delim)
        .with_null_values(Some(NullValues::AllColumns("NA".to_string())))
        .infer_schema(None)
        .finish()
        .unwrap();

    // hstack the experiment name (derived from the filename) into the DataFrame
    let iter_exp =
        vec![path.file_name().unwrap().to_str().unwrap().split('-').collect::<Vec<&str>>()[0..=2]
            .join("-")]
        .into_iter();
    let exp_col: Series = Series::new(
        "experiment",
        iter_exp.flat_map(|n| std::iter::repeat(n).take(df.height())).collect::<Vec<String>>(),
    );
    // exp_col.rename("experiment");
    df.hstack_mut(&[exp_col]).unwrap();

    Ok(df)
}

fn baseline_get_median_gt_zero(str_val: &Series) -> Series {
    let lists = str_val.utf8().unwrap().into_iter().map(|x| {
        x.unwrap().split(',').into_iter().map(|y| y.parse::<f64>().unwrap()).collect::<Vec<f64>>()
    });

    let medians: Vec<bool> =
        lists.into_iter().map(|x| Series::from_iter(x).median().unwrap() > 0.0).collect();
    Series::from_iter(medians)
}

fn filter_input(input: &mut DataFrame, baseline: bool) -> DataFrame {
    if baseline {
        filter_baseline(input)
    } else {
        filter_differential(input)
    }
}

fn filter_baseline(input: &mut DataFrame) -> DataFrame {
    // This is a baseline experiment.
    // Find columns starting with lower case g, then apply the function to convert to
    // medain and select greater than zero
    let mut meas = Vec::<Expr>::new();

    for column in input.get_column_names_owned() {
        if column.starts_with('g') {
            input.apply(&column, baseline_get_median_gt_zero).unwrap();
            meas.push(col(&column));
        }
    }
    // println!("{:?}", meas);
    // Selection should now have all the gN column names in it
    input
        .clone()
        .lazy()
        .filter(any_exprs(&meas[0..meas.len()/2]))
        .filter(any_exprs(&meas[meas.len()/2 .. meas.len()])) // If we try to do the whole thing at once, we get a stack overflow
        .collect()
        .unwrap()
}

fn filter_differential(input: &DataFrame) -> DataFrame {
    // find the p value and log fold columns
    let pv_regex = Regex::new(r".*p-value.*").unwrap();
    let log_fold_regex = Regex::new(r".*log2.*").unwrap();
    let mut exprs: Vec<Expr> = Vec::new();
    let mut pv: Expr;
    let mut lf: Expr;
    for column in input.get_column_names_owned() {
        if pv_regex.is_match(&column) {
            pv = col(&column).is_not_null();
            exprs.push(pv);
        } else if log_fold_regex.is_match(&column) {
            lf = col(&column).neq(lit(0));
            exprs.push(lf);
        }
    }

    input.clone().lazy().filter(all_exprs(exprs)).collect().unwrap()
}

fn load_chunk(
    paths: &mut Vec<PathBuf>,
    select: &mut Vec<String>,
) -> Result<DataFrame, anyhow::Error> {
    /*
    Work with the chunks of input to reduce them into a single dataframe of the genes we want from
    that particular set of experiments. Should return a single dataframe
    */
    let mut output: Option<DataFrame> = None;
    while let Some(infile) = paths.pop() {
        // Load the input Csv
        let mut input: DataFrame = if infile.extension().unwrap() == "tsv" {
            load_data(&infile, b'\t')
        } else {
            load_data(&infile, b',')
        }
        .unwrap_or_else(|error| match error.downcast_ref::<PolarsError>() {
            Some(PolarsError::Io(_string)) => panic!("Input file does not exist! {:?}", infile),
            _ => panic!("An error occurred! {:?}", error),
        });

        // Rename columns to remove . in the names. Now also remove spaces
        let mut new_cols = Vec::new();
        let mut c: u8 = 0;
        for nm in input.get_column_names().iter() {
            let mut new_name = nm.replace('.', "").replace(' ', "");
            if new_cols.contains(&new_name) {
                c += 1;
                new_name += &c.to_string(); // Hopefully avoid duplicate names
            }
            new_cols.push(new_name);
        }

        if new_cols != input.get_column_names() {
            input.set_column_names(&new_cols).unwrap_or_else(|error| {
                panic!("Failed to set column names for some reason {:?}", error);
            });
        }

        // Need to detect which type of experiment it is - differential or baseline
        let type_re = Regex::new(r"tpms").unwrap();
        let filtered = filter_input(&mut input, type_re.is_match(infile.to_str().unwrap()));
        // Drop everything from the input except the genes
        let genes: DataFrame = filtered.select(select.iter()).unwrap_or_else(|error| match error {
            PolarsError::NotFound(_string) => {
                panic!(
                    "{:?} was not found in the header, does the file have a header?\n{:?}",
                    select,
                    input.get_column_names()
                );
            },
            _ => panic!("Error selecting column from input: {:?}", error),
        });

        output = match output {
            None => {
                let output_df = genes.clone();
                Some(output_df)
            },
            Some(mut output_df) => {
                // genes =/= output, so we are handling a new file
                output_df.extend(&genes).unwrap_or_else(|error| {
                    panic!("Failed to extend the output dataframe. {:?}", error)
                });
                output_df =
                    output_df.unique(None, UniqueKeepStrategy::First).unwrap_or_else(|error| {
                        panic!("Failed to parse selected column for uniqueness. {:?}", error)
                    });
                Some(output_df)
            },
        }
    }
    Ok(output.unwrap())
}

fn get_taxid(paths: &mut Vec<PathBuf>, lookup_table: &mut HashMap<String, String>) {
    let tax_regex = Regex::new(r".*Taxon_([0-9]{4,})").unwrap();

    while let Some(path) = paths.pop() {
        let exp_name =
            path.file_name().unwrap().to_str().unwrap().split('-').collect::<Vec<&str>>()[0..=2]
                .join("-")
                .replace(".condensed", "");

        let file = File::open(path).unwrap();
        let lines = io::BufReader::new(file).lines();
        for line in lines {
            match line {
                Ok(line) => {
                    let captures = tax_regex.captures(line.as_str());
                    match captures {
                        Some(caps) => {
                            let taxid = caps.get(1).map_or("", |m| m.as_str());
                            lookup_table.insert(exp_name, taxid.to_string());
                            break;
                        },
                        None => {},
                    };
                },
                Err(_idc) => {},
            }
        }
    }
}

fn main() {
    let timer = Instant::now();

    let mut args = Args::parse();

    // Always select the experiment
    args.select.push("experiment".to_string());

    // separate expression data from taxa data
    let mut taxa_files: Vec<PathBuf> = Vec::new();
    let mut expr_files: Vec<PathBuf> = Vec::new();

    let re = Regex::new(r"sdrf.tsv").unwrap();

    for infile in args.input {
        if re.is_match(infile.to_str().unwrap()) {
            // File is experiment setup data, parse for taxa
            taxa_files.push(infile);
        } else {
            // Should only be expression data, parse for genes
            expr_files.push(infile);
        }
    }

    println!("{} taxa files to parse", taxa_files.len());
    println!("{} expression files to parse", expr_files.len());

    let n_files = expr_files.len();

    let mut dataframes: Vec<DataFrame> = Vec::new();

    let chunked_input: Vec<Vec<PathBuf>> =
        expr_files.chunks(args.chunksize).map(|x| x.to_vec()).collect();
    // Read everything into a big vector
    let mut n_chunks = 0;
    for mut files_chunk in chunked_input {
        dataframes.push(
            load_chunk(&mut files_chunk, &mut args.select)
                .unwrap_or_else(|_error| panic!("Something wrong in one of the reads, aborting")),
        );
        n_chunks += 1;
        println!("Done {} files", n_chunks * args.chunksize);
    }

    let mut output: DataFrame = dataframes[0].clone();

    for item in dataframes.iter() {
        output.extend(item).unwrap_or_else(|error| {
            panic!("Unable to extend the DataFrame! Out of memory? {:?}", error)
        });
        output.unique(None, UniqueKeepStrategy::First).unwrap();
    }

    // Now parse the experiment files to get taxa
    let chunked_taxa: Vec<Vec<PathBuf>> =
        taxa_files.chunks(args.chunksize).map(|x| x.to_vec()).collect();

    let mut lookup_table: HashMap<String, String> = HashMap::new();

    for mut tax_chunk in chunked_taxa {
        get_taxid(&mut tax_chunk, &mut lookup_table);
    }

    // Now we have the experiment - taxid lookup_table, we just need to add the column to the output df

    let expt_col = output.select(["experiment"]).unwrap();
    let mut tax_ids: Vec<String> = Vec::with_capacity(expt_col.height());
    for ex in expt_col.iter() {
        let mut warn_flag: bool = true;
        let uhy = ex.utf8().unwrap();
        for (idx, x) in uhy.into_iter().enumerate() {
            if lookup_table.contains_key(x.unwrap()) {
                tax_ids.insert(idx, lookup_table.get(x.unwrap()).unwrap().to_string());
            } else {
                tax_ids.insert(idx, String::new());
                if warn_flag {
                    println!("Experiment {} does not name a taxon", x.unwrap());
                    warn_flag = false;
                }
            }
        }
    }

    output.hstack_mut(&[Series::new("taxid", tax_ids.into_iter())]).unwrap();

    let out_stream: File = File::create(args.output).unwrap();
    CsvWriter::new(out_stream)
        .has_header(true)
        .finish(&mut output)
        .unwrap_or_else(|error| panic!("Something wrong writing file: {:?}", error));

    println!("Processed {} in {} seconds", n_files, timer.elapsed().as_secs());
}
<|endoftext|>