<commit_before>// Copyright 2012-2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use borrow_check::place_ext::PlaceExt;
use dataflow::indexes::BorrowIndex;
use dataflow::move_paths::MoveData;
use rustc::mir::traversal;
use rustc::mir::visit::{
    PlaceContext, Visitor, NonUseContext, MutatingUseContext, NonMutatingUseContext
};
use rustc::mir::{self, Location, Mir, Place, Local};
use rustc::ty::{Region, TyCtxt};
use rustc::util::nodemap::{FxHashMap, FxHashSet};
use rustc_data_structures::indexed_vec::IndexVec;
use rustc_data_structures::bit_set::BitSet;
use std::fmt;
use std::hash::Hash;
use std::ops::Index;

crate struct BorrowSet<'tcx> {
    /// The fundamental map relating bitvector indexes to the borrows
    /// in the MIR.
    crate borrows: IndexVec<BorrowIndex, BorrowData<'tcx>>,

    /// Each borrow is also uniquely identified in the MIR by the
    /// `Location` of the assignment statement in which it appears on
    /// the right hand side; we map each such location to the
    /// corresponding `BorrowIndex`.
    crate location_map: FxHashMap<Location, BorrowIndex>,

    /// Locations which activate borrows.
    /// NOTE: A given location may activate more than one borrow in the future
    /// when more general two-phase borrow support is introduced, but for now we
    /// only need to store one borrow index
    crate activation_map: FxHashMap<Location, Vec<BorrowIndex>>,

    /// Every borrow has a region; this maps each such regions back to
    /// its borrow-indexes.
    crate region_map: FxHashMap<Region<'tcx>, FxHashSet<BorrowIndex>>,

    /// Map from local to all the borrows on that local
    crate local_map: FxHashMap<mir::Local, FxHashSet<BorrowIndex>>,

    crate locals_state_at_exit: LocalsStateAtExit,
}

impl<'tcx> Index<BorrowIndex> for BorrowSet<'tcx> {
    type Output = BorrowData<'tcx>;

    fn index(&self, index: BorrowIndex) -> &BorrowData<'tcx> {
        &self.borrows[index]
    }
}

/// Location where a two phase borrow is activated, if a borrow
/// is in fact a two phase borrow.
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
crate enum TwoPhaseActivation {
    NotTwoPhase,
    NotActivated,
    ActivatedAt(Location),
}

#[derive(Debug)]
crate struct BorrowData<'tcx> {
    /// Location where the borrow reservation starts.
    /// In many cases, this will be equal to the activation location but not always.
    crate reserve_location: Location,
    /// Location where the borrow is activated.
    crate activation_location: TwoPhaseActivation,
    /// What kind of borrow this is
    crate kind: mir::BorrowKind,
    /// The region for which this borrow is live
    crate region: Region<'tcx>,
    /// Place from which we are borrowing
    crate borrowed_place: mir::Place<'tcx>,
    /// Place to which the borrow was stored
    crate assigned_place: mir::Place<'tcx>,
}

impl<'tcx> fmt::Display for BorrowData<'tcx> {
    fn fmt(&self, w: &mut fmt::Formatter) -> fmt::Result {
        let kind = match self.kind {
            mir::BorrowKind::Shared => "",
            mir::BorrowKind::Shallow => "shallow ",
            mir::BorrowKind::Unique => "uniq ",
            mir::BorrowKind::Mut { .. } => "mut ",
        };
        let region = self.region.to_string();
        let separator = if !region.is_empty() {
            " "
        } else {
            ""
        };
        write!(w, "&{}{}{}{:?}", region, separator, kind, self.borrowed_place)
    }
}

crate enum LocalsStateAtExit {
    AllAreInvalidated,
    SomeAreInvalidated { has_storage_dead_or_moved: BitSet<Local> }
}

impl LocalsStateAtExit {
    fn build(
        locals_are_invalidated_at_exit: bool,
        mir: &Mir<'tcx>,
        move_data: &MoveData<'tcx>
    ) -> Self {
        struct HasStorageDead(BitSet<Local>);

        impl<'tcx> Visitor<'tcx> for HasStorageDead {
            fn visit_local(&mut self, local: &Local, ctx: PlaceContext<'tcx>, _: Location) {
                if ctx == PlaceContext::NonUse(NonUseContext::StorageDead) {
                    self.0.insert(*local);
                }
            }
        }

        if locals_are_invalidated_at_exit {
            LocalsStateAtExit::AllAreInvalidated
        } else {
            let mut has_storage_dead = HasStorageDead(BitSet::new_empty(mir.local_decls.len()));
            has_storage_dead.visit_mir(mir);
            let mut has_storage_dead_or_moved = has_storage_dead.0;
            for move_out in &move_data.moves {
                if let Some(index) = move_data.base_local(move_out.path) {
                    has_storage_dead_or_moved.insert(index);

                }
            }
            LocalsStateAtExit::SomeAreInvalidated{ has_storage_dead_or_moved }
        }
    }
}

impl<'tcx> BorrowSet<'tcx> {
    pub fn build(
        tcx: TyCtxt<'_, '_, 'tcx>,
        mir: &Mir<'tcx>,
        locals_are_invalidated_at_exit: bool,
        move_data: &MoveData<'tcx>
    ) -> Self {

        let mut visitor = GatherBorrows {
            tcx,
            mir,
            idx_vec: IndexVec::new(),
            location_map: Default::default(),
            activation_map: Default::default(),
            region_map: Default::default(),
            local_map: Default::default(),
            pending_activations: Default::default(),
            locals_state_at_exit:
                LocalsStateAtExit::build(locals_are_invalidated_at_exit, mir, move_data),
        };

        for (block, block_data) in traversal::preorder(mir) {
            visitor.visit_basic_block_data(block, block_data);
        }

        BorrowSet {
            borrows: visitor.idx_vec,
            location_map: visitor.location_map,
            activation_map: visitor.activation_map,
            region_map: visitor.region_map,
            local_map: visitor.local_map,
            locals_state_at_exit: visitor.locals_state_at_exit,
        }
    }

    crate fn activations_at_location(&self, location: Location) -> &[BorrowIndex] {
        self.activation_map
            .get(&location)
            .map(|activations| &activations[..])
            .unwrap_or(&[])
    }
}

struct GatherBorrows<'a, 'gcx: 'tcx, 'tcx: 'a> {
    tcx: TyCtxt<'a, 'gcx, 'tcx>,
    mir: &'a Mir<'tcx>,
    idx_vec: IndexVec<BorrowIndex, BorrowData<'tcx>>,
    location_map: FxHashMap<Location, BorrowIndex>,
    activation_map: FxHashMap<Location, Vec<BorrowIndex>>,
    region_map: FxHashMap<Region<'tcx>, FxHashSet<BorrowIndex>>,
    local_map: FxHashMap<mir::Local, FxHashSet<BorrowIndex>>,

    /// When we encounter a 2-phase borrow statement, it will always
    /// be assigning into a temporary TEMP:
    ///
    ///    TEMP = &foo
    ///
    /// We add TEMP into this map with `b`, where `b` is the index of
    /// the borrow. When we find a later use of this activation, we
    /// remove from the map (and add to the "tombstone" set below).
    pending_activations: FxHashMap<mir::Local, BorrowIndex>,

    locals_state_at_exit: LocalsStateAtExit,
}

impl<'a, 'gcx, 'tcx> Visitor<'tcx> for GatherBorrows<'a, 'gcx, 'tcx> {
    fn visit_assign(
        &mut self,
        block: mir::BasicBlock,
        assigned_place: &mir::Place<'tcx>,
        rvalue: &mir::Rvalue<'tcx>,
        location: mir::Location,
    ) {
        if let mir::Rvalue::Ref(region, kind, ref borrowed_place) = *rvalue {
            if borrowed_place.ignore_borrow(
                self.tcx, self.mir, &self.locals_state_at_exit) {
                return;
            }

            let borrow = BorrowData {
                kind,
                region,
                reserve_location: location,
                activation_location: TwoPhaseActivation::NotTwoPhase,
                borrowed_place: borrowed_place.clone(),
                assigned_place: assigned_place.clone(),
            };
            let idx = self.idx_vec.push(borrow);
            self.location_map.insert(location, idx);

            self.insert_as_pending_if_two_phase(location, &assigned_place, region, kind, idx);

            insert(&mut self.region_map, &region, idx);
            if let Some(local) = borrowed_place.root_local() {
                insert(&mut self.local_map, &local, idx);
            }
        }

        return self.super_assign(block, assigned_place, rvalue, location);

        fn insert<'a, K, V>(map: &'a mut FxHashMap<K, FxHashSet<V>>, k: &K, v: V)
        where
            K: Clone + Eq + Hash,
            V: Eq + Hash,
        {
            map.entry(k.clone()).or_default().insert(v);
        }
    }

    fn visit_place(
        &mut self,
        place: &mir::Place<'tcx>,
        context: PlaceContext<'tcx>,
        location: Location,
    ) {
        self.super_place(place, context, location);

        // We found a use of some temporary TEMP...
        if let Place::Local(temp) = place {
            // ... check whether we (earlier) saw a 2-phase borrow like
            //
            //     TMP = &mut place
            if let Some(&borrow_index) = self.pending_activations.get(temp) {
                let borrow_data = &mut self.idx_vec[borrow_index];

                // Watch out: the use of TMP in the borrow itself
                // doesn't count as an activation. =)
                if borrow_data.reserve_location == location &&
                    context == PlaceContext::MutatingUse(MutatingUseContext::Store)
                {
                    return;
                }

                if let TwoPhaseActivation::ActivatedAt(other_location) =
                        borrow_data.activation_location {
                    span_bug!(
                        self.mir.source_info(location).span,
                        "found two uses for 2-phase borrow temporary {:?}: \
                         {:?} and {:?}",
                        temp,
                        location,
                        other_location,
                    );
                }

                // Otherwise, this is the unique later use
                // that we expect.
                borrow_data.activation_location = match context {
                    // The use of TMP in a shared borrow does not
                    // count as an actual activation.
                    PlaceContext::NonMutatingUse(NonMutatingUseContext::SharedBorrow(..)) |
                    PlaceContext::NonMutatingUse(NonMutatingUseContext::ShallowBorrow(..)) =>
                        TwoPhaseActivation::NotActivated,
                    _ => {
                        // Double check: This borrow is indeed a two-phase borrow (that is,
                        // we are 'transitioning' from `NotActivated` to `ActivatedAt`) and
                        // we've not found any other activations (checked above).
                        assert_eq!(
                            borrow_data.activation_location,
                            TwoPhaseActivation::NotActivated,
                            "never found an activation for this borrow!",
                        );

                        self.activation_map
                            .entry(location)
                            .or_default()
                            .push(borrow_index);
                        TwoPhaseActivation::ActivatedAt(location)
                    }
                };
            }
        }
    }

    fn visit_rvalue(&mut self, rvalue: &mir::Rvalue<'tcx>, location: mir::Location) {
        if let mir::Rvalue::Ref(region, kind, ref place) = *rvalue {
            // double-check that we already registered a BorrowData for this

            let borrow_index = self.location_map[&location];
            let borrow_data = &self.idx_vec[borrow_index];
            assert_eq!(borrow_data.reserve_location, location);
            assert_eq!(borrow_data.kind, kind);
            assert_eq!(borrow_data.region, region);
            assert_eq!(borrow_data.borrowed_place, *place);
        }

        return self.super_rvalue(rvalue, location);
    }

    fn visit_statement(
        &mut self,
        block: mir::BasicBlock,
        statement: &mir::Statement<'tcx>,
        location: Location,
    ) {
        return self.super_statement(block, statement, location);
    }
}

impl<'a, 'gcx, 'tcx> GatherBorrows<'a, 'gcx, 'tcx> {
    /// Returns true if the borrow represented by `kind` is
    /// allowed to be split into separate Reservation and
    /// Activation phases.
    fn allow_two_phase_borrow(&self, kind: mir::BorrowKind) -> bool {
        self.tcx.two_phase_borrows()
            && (kind.allows_two_phase_borrow()
                || self.tcx.sess.opts.debugging_opts.two_phase_beyond_autoref)
    }

    /// If this is a two-phase borrow, then we will record it
    /// as "pending" until we find the activating use.
    fn insert_as_pending_if_two_phase(
        &mut self,
        start_location: Location,
        assigned_place: &mir::Place<'tcx>,
        region: Region<'tcx>,
        kind: mir::BorrowKind,
        borrow_index: BorrowIndex,
    ) {
        debug!(
            "Borrows::insert_as_pending_if_two_phase({:?}, {:?}, {:?}, {:?})",
            start_location, assigned_place, region, borrow_index,
        );

        if !self.allow_two_phase_borrow(kind) {
            debug!("  -> {:?}", start_location);
            return;
        }

        // When we encounter a 2-phase borrow statement, it will always
        // be assigning into a temporary TEMP:
        //
        //    TEMP = &foo
        //
        // so extract `temp`.
        let temp = if let &mir::Place::Local(temp) = assigned_place {
            temp
        } else {
            span_bug!(
                self.mir.source_info(start_location).span,
                "expected 2-phase borrow to assign to a local, not `{:?}`",
                assigned_place,
            );
        };

        // Consider the borrow not activated to start. When we find an activation, we'll update
        // this field.
        {
            let borrow_data = &mut self.idx_vec[borrow_index];
            borrow_data.activation_location = TwoPhaseActivation::NotActivated;
        }

        // Insert `temp` into the list of pending activations. From
        // now on, we'll be on the lookout for a use of it. Note that
        // we are guaranteed that this use will come after the
        // assignment.
        let old_value = self.pending_activations.insert(temp, borrow_index);
        if let Some(old_index) = old_value {
            span_bug!(self.mir.source_info(start_location).span,
                      "found already pending activation for temp: {:?} \
                       at borrow_index: {:?} with associated data {:?}",
                      temp, old_index, self.idx_vec[old_index]);
        }
    }
}
<commit_msg>borrow_set: remove a helper function and a clone it uses<commit_after>// Copyright 2012-2017 The Rust Project Developers. See the COPYRIGHT
// file at the top-level directory of this distribution and at
// http://rust-lang.org/COPYRIGHT.
//
// Licensed under the Apache License, Version 2.0 <LICENSE-APACHE or
// http://www.apache.org/licenses/LICENSE-2.0> or the MIT license
// <LICENSE-MIT or http://opensource.org/licenses/MIT>, at your
// option. This file may not be copied, modified, or distributed
// except according to those terms.

use borrow_check::place_ext::PlaceExt;
use dataflow::indexes::BorrowIndex;
use dataflow::move_paths::MoveData;
use rustc::mir::traversal;
use rustc::mir::visit::{
    PlaceContext, Visitor, NonUseContext, MutatingUseContext, NonMutatingUseContext
};
use rustc::mir::{self, Location, Mir, Place, Local};
use rustc::ty::{Region, TyCtxt};
use rustc::util::nodemap::{FxHashMap, FxHashSet};
use rustc_data_structures::indexed_vec::IndexVec;
use rustc_data_structures::bit_set::BitSet;
use std::fmt;
use std::ops::Index;

crate struct BorrowSet<'tcx> {
    /// The fundamental map relating bitvector indexes to the borrows
    /// in the MIR.
    crate borrows: IndexVec<BorrowIndex, BorrowData<'tcx>>,

    /// Each borrow is also uniquely identified in the MIR by the
    /// `Location` of the assignment statement in which it appears on
    /// the right hand side; we map each such location to the
    /// corresponding `BorrowIndex`.
    crate location_map: FxHashMap<Location, BorrowIndex>,

    /// Locations which activate borrows.
    /// NOTE: A given location may activate more than one borrow in the future
    /// when more general two-phase borrow support is introduced, but for now we
    /// only need to store one borrow index
    crate activation_map: FxHashMap<Location, Vec<BorrowIndex>>,

    /// Every borrow has a region; this maps each such regions back to
    /// its borrow-indexes.
    crate region_map: FxHashMap<Region<'tcx>, FxHashSet<BorrowIndex>>,

    /// Map from local to all the borrows on that local
    crate local_map: FxHashMap<mir::Local, FxHashSet<BorrowIndex>>,

    crate locals_state_at_exit: LocalsStateAtExit,
}

impl<'tcx> Index<BorrowIndex> for BorrowSet<'tcx> {
    type Output = BorrowData<'tcx>;

    fn index(&self, index: BorrowIndex) -> &BorrowData<'tcx> {
        &self.borrows[index]
    }
}

/// Location where a two phase borrow is activated, if a borrow
/// is in fact a two phase borrow.
#[derive(Copy, Clone, PartialEq, Eq, Debug)]
crate enum TwoPhaseActivation {
    NotTwoPhase,
    NotActivated,
    ActivatedAt(Location),
}

#[derive(Debug)]
crate struct BorrowData<'tcx> {
    /// Location where the borrow reservation starts.
    /// In many cases, this will be equal to the activation location but not always.
    crate reserve_location: Location,
    /// Location where the borrow is activated.
    crate activation_location: TwoPhaseActivation,
    /// What kind of borrow this is
    crate kind: mir::BorrowKind,
    /// The region for which this borrow is live
    crate region: Region<'tcx>,
    /// Place from which we are borrowing
    crate borrowed_place: mir::Place<'tcx>,
    /// Place to which the borrow was stored
    crate assigned_place: mir::Place<'tcx>,
}

impl<'tcx> fmt::Display for BorrowData<'tcx> {
    fn fmt(&self, w: &mut fmt::Formatter) -> fmt::Result {
        let kind = match self.kind {
            mir::BorrowKind::Shared => "",
            mir::BorrowKind::Shallow => "shallow ",
            mir::BorrowKind::Unique => "uniq ",
            mir::BorrowKind::Mut { .. } => "mut ",
        };
        let region = self.region.to_string();
        let separator = if !region.is_empty() {
            " "
        } else {
            ""
        };
        write!(w, "&{}{}{}{:?}", region, separator, kind, self.borrowed_place)
    }
}

crate enum LocalsStateAtExit {
    AllAreInvalidated,
    SomeAreInvalidated { has_storage_dead_or_moved: BitSet<Local> }
}

impl LocalsStateAtExit {
    fn build(
        locals_are_invalidated_at_exit: bool,
        mir: &Mir<'tcx>,
        move_data: &MoveData<'tcx>
    ) -> Self {
        struct HasStorageDead(BitSet<Local>);

        impl<'tcx> Visitor<'tcx> for HasStorageDead {
            fn visit_local(&mut self, local: &Local, ctx: PlaceContext<'tcx>, _: Location) {
                if ctx == PlaceContext::NonUse(NonUseContext::StorageDead) {
                    self.0.insert(*local);
                }
            }
        }

        if locals_are_invalidated_at_exit {
            LocalsStateAtExit::AllAreInvalidated
        } else {
            let mut has_storage_dead = HasStorageDead(BitSet::new_empty(mir.local_decls.len()));
            has_storage_dead.visit_mir(mir);
            let mut has_storage_dead_or_moved = has_storage_dead.0;
            for move_out in &move_data.moves {
                if let Some(index) = move_data.base_local(move_out.path) {
                    has_storage_dead_or_moved.insert(index);

                }
            }
            LocalsStateAtExit::SomeAreInvalidated{ has_storage_dead_or_moved }
        }
    }
}

impl<'tcx> BorrowSet<'tcx> {
    pub fn build(
        tcx: TyCtxt<'_, '_, 'tcx>,
        mir: &Mir<'tcx>,
        locals_are_invalidated_at_exit: bool,
        move_data: &MoveData<'tcx>
    ) -> Self {

        let mut visitor = GatherBorrows {
            tcx,
            mir,
            idx_vec: IndexVec::new(),
            location_map: Default::default(),
            activation_map: Default::default(),
            region_map: Default::default(),
            local_map: Default::default(),
            pending_activations: Default::default(),
            locals_state_at_exit:
                LocalsStateAtExit::build(locals_are_invalidated_at_exit, mir, move_data),
        };

        for (block, block_data) in traversal::preorder(mir) {
            visitor.visit_basic_block_data(block, block_data);
        }

        BorrowSet {
            borrows: visitor.idx_vec,
            location_map: visitor.location_map,
            activation_map: visitor.activation_map,
            region_map: visitor.region_map,
            local_map: visitor.local_map,
            locals_state_at_exit: visitor.locals_state_at_exit,
        }
    }

    crate fn activations_at_location(&self, location: Location) -> &[BorrowIndex] {
        self.activation_map
            .get(&location)
            .map(|activations| &activations[..])
            .unwrap_or(&[])
    }
}

struct GatherBorrows<'a, 'gcx: 'tcx, 'tcx: 'a> {
    tcx: TyCtxt<'a, 'gcx, 'tcx>,
    mir: &'a Mir<'tcx>,
    idx_vec: IndexVec<BorrowIndex, BorrowData<'tcx>>,
    location_map: FxHashMap<Location, BorrowIndex>,
    activation_map: FxHashMap<Location, Vec<BorrowIndex>>,
    region_map: FxHashMap<Region<'tcx>, FxHashSet<BorrowIndex>>,
    local_map: FxHashMap<mir::Local, FxHashSet<BorrowIndex>>,

    /// When we encounter a 2-phase borrow statement, it will always
    /// be assigning into a temporary TEMP:
    ///
    ///    TEMP = &foo
    ///
    /// We add TEMP into this map with `b`, where `b` is the index of
    /// the borrow. When we find a later use of this activation, we
    /// remove from the map (and add to the "tombstone" set below).
    pending_activations: FxHashMap<mir::Local, BorrowIndex>,

    locals_state_at_exit: LocalsStateAtExit,
}

impl<'a, 'gcx, 'tcx> Visitor<'tcx> for GatherBorrows<'a, 'gcx, 'tcx> {
    fn visit_assign(
        &mut self,
        block: mir::BasicBlock,
        assigned_place: &mir::Place<'tcx>,
        rvalue: &mir::Rvalue<'tcx>,
        location: mir::Location,
    ) {
        if let mir::Rvalue::Ref(region, kind, ref borrowed_place) = *rvalue {
            if borrowed_place.ignore_borrow(
                self.tcx, self.mir, &self.locals_state_at_exit) {
                return;
            }

            let borrow = BorrowData {
                kind,
                region,
                reserve_location: location,
                activation_location: TwoPhaseActivation::NotTwoPhase,
                borrowed_place: borrowed_place.clone(),
                assigned_place: assigned_place.clone(),
            };
            let idx = self.idx_vec.push(borrow);
            self.location_map.insert(location, idx);

            self.insert_as_pending_if_two_phase(location, &assigned_place, region, kind, idx);

            self.region_map.entry(region).or_default().insert(idx);
            if let Some(local) = borrowed_place.root_local() {
                self.local_map.entry(local).or_default().insert(idx);
            }
        }

        self.super_assign(block, assigned_place, rvalue, location)
    }

    fn visit_place(
        &mut self,
        place: &mir::Place<'tcx>,
        context: PlaceContext<'tcx>,
        location: Location,
    ) {
        self.super_place(place, context, location);

        // We found a use of some temporary TEMP...
        if let Place::Local(temp) = place {
            // ... check whether we (earlier) saw a 2-phase borrow like
            //
            //     TMP = &mut place
            if let Some(&borrow_index) = self.pending_activations.get(temp) {
                let borrow_data = &mut self.idx_vec[borrow_index];

                // Watch out: the use of TMP in the borrow itself
                // doesn't count as an activation. =)
                if borrow_data.reserve_location == location &&
                    context == PlaceContext::MutatingUse(MutatingUseContext::Store)
                {
                    return;
                }

                if let TwoPhaseActivation::ActivatedAt(other_location) =
                        borrow_data.activation_location {
                    span_bug!(
                        self.mir.source_info(location).span,
                        "found two uses for 2-phase borrow temporary {:?}: \
                         {:?} and {:?}",
                        temp,
                        location,
                        other_location,
                    );
                }

                // Otherwise, this is the unique later use
                // that we expect.
                borrow_data.activation_location = match context {
                    // The use of TMP in a shared borrow does not
                    // count as an actual activation.
                    PlaceContext::NonMutatingUse(NonMutatingUseContext::SharedBorrow(..)) |
                    PlaceContext::NonMutatingUse(NonMutatingUseContext::ShallowBorrow(..)) =>
                        TwoPhaseActivation::NotActivated,
                    _ => {
                        // Double check: This borrow is indeed a two-phase borrow (that is,
                        // we are 'transitioning' from `NotActivated` to `ActivatedAt`) and
                        // we've not found any other activations (checked above).
                        assert_eq!(
                            borrow_data.activation_location,
                            TwoPhaseActivation::NotActivated,
                            "never found an activation for this borrow!",
                        );

                        self.activation_map
                            .entry(location)
                            .or_default()
                            .push(borrow_index);
                        TwoPhaseActivation::ActivatedAt(location)
                    }
                };
            }
        }
    }

    fn visit_rvalue(&mut self, rvalue: &mir::Rvalue<'tcx>, location: mir::Location) {
        if let mir::Rvalue::Ref(region, kind, ref place) = *rvalue {
            // double-check that we already registered a BorrowData for this

            let borrow_index = self.location_map[&location];
            let borrow_data = &self.idx_vec[borrow_index];
            assert_eq!(borrow_data.reserve_location, location);
            assert_eq!(borrow_data.kind, kind);
            assert_eq!(borrow_data.region, region);
            assert_eq!(borrow_data.borrowed_place, *place);
        }

        return self.super_rvalue(rvalue, location);
    }

    fn visit_statement(
        &mut self,
        block: mir::BasicBlock,
        statement: &mir::Statement<'tcx>,
        location: Location,
    ) {
        return self.super_statement(block, statement, location);
    }
}

impl<'a, 'gcx, 'tcx> GatherBorrows<'a, 'gcx, 'tcx> {
    /// Returns true if the borrow represented by `kind` is
    /// allowed to be split into separate Reservation and
    /// Activation phases.
    fn allow_two_phase_borrow(&self, kind: mir::BorrowKind) -> bool {
        self.tcx.two_phase_borrows()
            && (kind.allows_two_phase_borrow()
                || self.tcx.sess.opts.debugging_opts.two_phase_beyond_autoref)
    }

    /// If this is a two-phase borrow, then we will record it
    /// as "pending" until we find the activating use.
    fn insert_as_pending_if_two_phase(
        &mut self,
        start_location: Location,
        assigned_place: &mir::Place<'tcx>,
        region: Region<'tcx>,
        kind: mir::BorrowKind,
        borrow_index: BorrowIndex,
    ) {
        debug!(
            "Borrows::insert_as_pending_if_two_phase({:?}, {:?}, {:?}, {:?})",
            start_location, assigned_place, region, borrow_index,
        );

        if !self.allow_two_phase_borrow(kind) {
            debug!("  -> {:?}", start_location);
            return;
        }

        // When we encounter a 2-phase borrow statement, it will always
        // be assigning into a temporary TEMP:
        //
        //    TEMP = &foo
        //
        // so extract `temp`.
        let temp = if let &mir::Place::Local(temp) = assigned_place {
            temp
        } else {
            span_bug!(
                self.mir.source_info(start_location).span,
                "expected 2-phase borrow to assign to a local, not `{:?}`",
                assigned_place,
            );
        };

        // Consider the borrow not activated to start. When we find an activation, we'll update
        // this field.
        {
            let borrow_data = &mut self.idx_vec[borrow_index];
            borrow_data.activation_location = TwoPhaseActivation::NotActivated;
        }

        // Insert `temp` into the list of pending activations. From
        // now on, we'll be on the lookout for a use of it. Note that
        // we are guaranteed that this use will come after the
        // assignment.
        let old_value = self.pending_activations.insert(temp, borrow_index);
        if let Some(old_index) = old_value {
            span_bug!(self.mir.source_info(start_location).span,
                      "found already pending activation for temp: {:?} \
                       at borrow_index: {:?} with associated data {:?}",
                      temp, old_index, self.idx_vec[old_index]);
        }
    }
}
<|endoftext|>


<commit_before>#[macro_use]
extern crate redis_async;
#[macro_use]
extern crate serde_derive;

use std::io;

use actix::prelude::Addr;
use actix_files::{Files, NamedFile};
use actix_redis::{Command, RedisActor, RespValue};
use actix_web::{middleware, web, App, HttpResponse, HttpServer, Result};
use dotenv::dotenv;
use futures::future;
use futures::future::Future;
use serde::Serialize;
use uuid::Uuid;

#[derive(Serialize, Deserialize)]
pub struct SecretBody {
    content: String,
}

#[derive(Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct CreateSecretPostBody {
    content: String,
    expiration_seconds: usize,
}

#[derive(Serialize)]
#[serde(rename_all = "snake_case")]
pub struct CreateSecretResponseBody {
    id: Uuid,
}

#[derive(Serialize)]
#[serde(rename_all = "snake_case")]
pub struct ApiError<T: Serialize> {
    error: T,
}

impl<T: Serialize> ApiError<T> {
    pub fn from(err: T) -> ApiError<T> {
        ApiError { error: err }
    }
}

#[derive(Debug, Clone, Copy)]
struct AppConfig {
    max_secret_length: usize,
    max_secret_expiration: usize,
}

struct AppState {
    redis: Addr<RedisActor>,
    config: AppConfig,
}

impl AppState {
    fn create_secret(
        self: &AppState,
        id: Uuid,
        body: String,
        expiration_seconds: usize,
    ) -> impl Future<Item = bool, Error = actix_web::Error> {
        self.redis
            .send(Command(resp_array![
                "SETEX",
                id.to_string(),
                expiration_seconds.to_string(),
                body
            ]))
            .map_err(actix_web::Error::from)
            .and_then(|res: Result<RespValue, actix_redis::Error>| match res {
                Ok(RespValue::SimpleString(_)) => Ok(true),

                err => {
                    println!("unexpected response to creating secret: {:?}", err);
                    Err(actix_web::error::ErrorInternalServerError(
                        "failed to create secret",
                    ))
                }
            })
    }

    fn fetch_secret(
        self: &AppState,
        id: String,
    ) -> impl Future<Item = Option<String>, Error = actix_web::Error> {
        self.redis
            .send(Command(resp_array!["GET", id]))
            .map_err(actix_web::Error::from)
            .map(|res| match res {
                Ok(RespValue::BulkString(bytes)) => {
                    let str = String::from_utf8(bytes).expect("non-utf8 string in redis");
                    Some(str)
                }

                _ => None,
            })
    }

    fn consume_secret(
        self: &AppState,
        id: String,
    ) -> impl Future<Item = Option<String>, Error = actix_web::Error> {
        let remove_secret = self
            .redis
            .send(Command(resp_array!["DEL", id.clone()]))
            .map_err(actix_web::Error::from);

        self.fetch_secret(id)
            .and_then(|result| remove_secret.map(|_| result))
    }
}

fn create_secret(
    body: web::Json<CreateSecretPostBody>,
    state: web::Data<AppState>,
) -> impl Future<Item = HttpResponse, Error = actix_web::Error> {
    let body = body.into_inner();

    let id = Uuid::new_v4();
    let expiration = body.expiration_seconds.to_string();

    let validate = if body.content.len() > state.config.max_secret_length {
        Err("content too large")
    } else if body.expiration_seconds > state.config.max_secret_expiration {
        Err("expiration too long")
    } else {
        Ok(())
    };

    future::result(validate)
        .map_err(|err| HttpResponse::BadRequest().json(ApiError::from(err)))
        .map_err(actix_web::Error::from)
        .and_then(move |_| {
            println!("Writing: {:?}: {:?}", id.to_string(), expiration);

            // TODO: Possibly re-encrypt secret here? libsodium or something.

            state
                .create_secret(id, body.content, body.expiration_seconds)
                .map(move |_| HttpResponse::Ok().json(CreateSecretResponseBody { id }))
        })
}

fn get_secret(
    path: web::Path<String>,
    state: web::Data<AppState>,
) -> impl Future<Item = HttpResponse, Error = actix_web::Error> {
    let secret_id = path.to_string();
    println!("Fetching: {:?}", secret_id);

    state
        .consume_secret(secret_id)
        .map(|res: Option<String>| match res {
            Some(secret) => HttpResponse::Ok().json(SecretBody { content: secret }),
            None => HttpResponse::NotFound().json(ApiError::from("not found")),
        })
}

fn view_secret_page(
    path: web::Path<String>,
    state: web::Data<AppState>,
) -> impl Future<Item = NamedFile, Error = actix_web::Error> {
    let secret_id = path.to_string();

    state.fetch_secret(secret_id).and_then(|res| match res {
        Some(_) => Ok(NamedFile::open("./static/view.html")?),

        None => {
            let file = NamedFile::open("./static/404.html")?
                // TODO: (requires HTTP crate).set_status_code()
                ;

            Ok(file)
        }
    })
}

fn main() -> io::Result<()> {
    env_logger::init();
    dotenv().ok();

    let config = AppConfig {
        max_secret_length: std::env::var("MAX_SECRET_LENGTH")
            .unwrap_or("4096".to_string())
            .parse()
            .expect("unvalid integer value"),
        max_secret_expiration: std::env::var("MAX_SECRET_EXPIRATION")
            .unwrap_or("4096".to_string())
            .parse()
            .expect("unvalid integer value"),
    };

    let bind_addr = std::env::var("BIND_ADDR").unwrap_or("127.0.0.1:8080".to_string());

    println!("Booting with config: {:?}", config);

    HttpServer::new(move || {
        let state = AppState {
            redis: RedisActor::start("localhost:6379"),
            config: config,
        };

        App::new()
            .data(state)
            .wrap(middleware::Logger::default())
            .service(
                web::resource("/api/secret").route(web::post().to_async(create_secret)),
            )
            .service(
                web::resource("/api/secret/{secret_id}")
                    .route(web::get().to_async(get_secret)),
            )
            .service(
                web::resource("/view/{secret_id}")
                    .route(web::get().to_async(view_secret_page)),
            )
            .service(Files::new("/", "./static/").index_file("index.html"))
    })
    .bind(bind_addr)?
    .run()
}
<commit_msg>Fix incorrect command ordering<commit_after>#[macro_use]
extern crate redis_async;
#[macro_use]
extern crate serde_derive;

use std::io;

use actix::prelude::Addr;
use actix_files::{Files, NamedFile};
use actix_redis::{Command, RedisActor, RespValue};
use actix_web::{middleware, web, App, HttpResponse, HttpServer, Result};
use dotenv::dotenv;
use futures::future;
use futures::future::Future;
use serde::Serialize;
use uuid::Uuid;

#[derive(Serialize, Deserialize)]
pub struct SecretBody {
    content: String,
}

#[derive(Deserialize)]
#[serde(rename_all = "snake_case")]
pub struct CreateSecretPostBody {
    content: String,
    expiration_seconds: usize,
}

#[derive(Serialize)]
#[serde(rename_all = "snake_case")]
pub struct CreateSecretResponseBody {
    id: Uuid,
}

#[derive(Serialize)]
#[serde(rename_all = "snake_case")]
pub struct ApiError<T: Serialize> {
    error: T,
}

impl<T: Serialize> ApiError<T> {
    pub fn from(err: T) -> ApiError<T> {
        ApiError { error: err }
    }
}

#[derive(Debug, Clone, Copy)]
struct AppConfig {
    max_secret_length: usize,
    max_secret_expiration: usize,
}

struct AppState {
    redis: Addr<RedisActor>,
    config: AppConfig,
}

impl AppState {
    fn create_secret(
        self: &AppState,
        id: Uuid,
        body: String,
        expiration_seconds: usize,
    ) -> impl Future<Item = bool, Error = actix_web::Error> {
        self.redis
            .send(Command(resp_array![
                "SETEX",
                id.to_string(),
                expiration_seconds.to_string(),
                body
            ]))
            .map_err(actix_web::Error::from)
            .and_then(|res: Result<RespValue, actix_redis::Error>| match res {
                Ok(RespValue::SimpleString(_)) => Ok(true),

                err => {
                    println!("unexpected response to creating secret: {:?}", err);
                    Err(actix_web::error::ErrorInternalServerError(
                        "failed to create secret",
                    ))
                }
            })
    }

    fn fetch_secret(
        self: &AppState,
        id: String,
    ) -> impl Future<Item = Option<String>, Error = actix_web::Error> {
        self.redis
            .send(Command(resp_array!["GET", id]))
            .map_err(actix_web::Error::from)
            .map(|res| match res {
                Ok(RespValue::BulkString(bytes)) => {
                    let str = String::from_utf8(bytes).expect("non-utf8 string in redis");
                    Some(str)
                }

                _ => None,
            })
    }

    fn consume_secret(
        self: &AppState,
        id: String,
    ) -> impl Future<Item = Option<String>, Error = actix_web::Error> {
        let fetch_secret = self.fetch_secret(id.clone());
        let delete_secret = self
            .redis
            .send(Command(resp_array!["DEL", id]))
            .map_err(actix_web::Error::from);

        fetch_secret.and_then(|res| delete_secret.map(|_| res))
    }
}

fn create_secret(
    body: web::Json<CreateSecretPostBody>,
    state: web::Data<AppState>,
) -> impl Future<Item = HttpResponse, Error = actix_web::Error> {
    let body = body.into_inner();

    let id = Uuid::new_v4();
    let expiration = body.expiration_seconds.to_string();

    let validate = if body.content.len() > state.config.max_secret_length {
        Err("content too large")
    } else if body.expiration_seconds > state.config.max_secret_expiration {
        Err("expiration too long")
    } else {
        Ok(())
    };

    future::result(validate)
        .map_err(|err| HttpResponse::BadRequest().json(ApiError::from(err)))
        .map_err(actix_web::Error::from)
        .and_then(move |_| {
            println!("Writing: {:?}: {:?}", id.to_string(), expiration);

            // TODO: Possibly re-encrypt secret here? libsodium or something.

            state
                .create_secret(id, body.content, body.expiration_seconds)
                .map(move |_| HttpResponse::Ok().json(CreateSecretResponseBody { id }))
        })
}

fn get_secret(
    path: web::Path<String>,
    state: web::Data<AppState>,
) -> impl Future<Item = HttpResponse, Error = actix_web::Error> {
    let secret_id = path.to_string();
    println!("Fetching: {:?}", secret_id);

    state
        .consume_secret(secret_id)
        .map(|res: Option<String>| match res {
            Some(secret) => HttpResponse::Ok().json(SecretBody { content: secret }),
            None => HttpResponse::NotFound().json(ApiError::from("not found")),
        })
}

fn view_secret_page(
    path: web::Path<String>,
    state: web::Data<AppState>,
) -> impl Future<Item = NamedFile, Error = actix_web::Error> {
    let secret_id = path.to_string();

    state.fetch_secret(secret_id).and_then(|res| match res {
        Some(_) => Ok(NamedFile::open("./static/view.html")?),

        None => {
            let file = NamedFile::open("./static/404.html")?
                // TODO: (requires HTTP crate).set_status_code()
                ;

            Ok(file)
        }
    })
}

fn main() -> io::Result<()> {
    env_logger::init();
    dotenv().ok();

    let config = AppConfig {
        max_secret_length: std::env::var("MAX_SECRET_LENGTH")
            .unwrap_or("4096".to_string())
            .parse()
            .expect("unvalid integer value"),
        max_secret_expiration: std::env::var("MAX_SECRET_EXPIRATION")
            .unwrap_or("4096".to_string())
            .parse()
            .expect("unvalid integer value"),
    };

    let bind_addr = std::env::var("BIND_ADDR").unwrap_or("127.0.0.1:8080".to_string());

    println!("Booting with config: {:?}", config);

    HttpServer::new(move || {
        let state = AppState {
            redis: RedisActor::start("localhost:6379"),
            config: config,
        };

        App::new()
            .data(state)
            .wrap(middleware::Logger::default())
            .service(
                web::resource("/api/secret").route(web::post().to_async(create_secret)),
            )
            .service(
                web::resource("/api/secret/{secret_id}")
                    .route(web::get().to_async(get_secret)),
            )
            .service(
                web::resource("/view/{secret_id}")
                    .route(web::get().to_async(view_secret_page)),
            )
            .service(Files::new("/", "./static/").index_file("index.html"))
    })
    .bind(bind_addr)?
    .run()
}
<|endoftext|>


<commit_before>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! Implementation of [`StreamMultiplexer`] using the [`yamux`] protocol over TCP
//!
//! [`StreamMultiplexer`]: crate::multiplexing::StreamMultiplexer
//! [`yamux`]: https://github.com/hashicorp/yamux/blob/master/spec.md

use crate::{
    multiplexing::StreamMultiplexer,
    negotiate::{negotiate_inbound, negotiate_outbound_interactive},
    transport::ConnectionOrigin,
};
use futures::{
    compat::{Compat, Compat01As03},
    future::{self, Future},
    io::{AsyncRead, AsyncWrite},
    stream::Stream,
};
use futures_01::future::poll_fn as poll_fn_01;
use std::{
    fmt::Debug,
    io,
    pin::Pin,
    task::{Context, Poll},
};
use yamux;

/// Re-export `Mode` from the yamux crate
pub use yamux::Mode;

/// The substream type produced by the `Yamux` multiplexer
pub type StreamHandle<T> = Compat01As03<yamux::StreamHandle<Compat<T>>>;

const YAMUX_PROTOCOL_NAME: &[u8] = b"/yamux/1.0.0";

#[derive(Debug)]
pub struct Yamux<TSocket> {
    inner: yamux::Connection<Compat<TSocket>>,
}

const MAX_BUFFER_SIZE: u32 = 8 * 1024 * 1024; // 8MB
const RECEIVE_WINDOW: u32 = 4 * 1024 * 1024; // 4MB

impl<TSocket> Yamux<TSocket>
where
    TSocket: AsyncRead + AsyncWrite + Send + Debug + Unpin,
{
    pub fn new(socket: TSocket, mode: Mode) -> Self {
        let mut config = yamux::Config::default();
        // Use OnRead mode instead of OnReceive mode to provide back pressure to the sending side.
        // Caveat: the OnRead mode has the risk of deadlock, where both sides send data larger than
        // receive window and don't read before finishing writes. But it doesn't apply to our use
        // cases. Some of our streams are unidirectional, where only one end writes data, e.g.,
        // Direct Send. Some of our streams are bidirectional, but only one end write data at a
        // time, e.g., RPC. Some of our streams may write at the same time, but the frames are
        // shorter than the receive window, e.g., protocol negotiation.
        config.set_window_update_mode(yamux::WindowUpdateMode::OnRead);
        // Because OnRead mode increases the RTT of window update, bigger buffer size and receive
        // window size perform better.
        config.set_max_buffer_size(MAX_BUFFER_SIZE as usize);
        config
            .set_receive_window(RECEIVE_WINDOW)
            .expect("Invalid receive window size");
        let socket = Compat::new(socket);
        Self {
            inner: yamux::Connection::new(socket, config, mode),
        }
    }

    pub async fn upgrade_connection(socket: TSocket, origin: ConnectionOrigin) -> io::Result<Self> {
        // Perform protocol negotiation
        let (socket, proto) = match origin {
            ConnectionOrigin::Inbound => negotiate_inbound(socket, [YAMUX_PROTOCOL_NAME]).await?,
            ConnectionOrigin::Outbound => {
                negotiate_outbound_interactive(socket, [YAMUX_PROTOCOL_NAME]).await?
            }
        };

        assert_eq!(proto, YAMUX_PROTOCOL_NAME);

        let mode = match origin {
            ConnectionOrigin::Inbound => Mode::Server,
            ConnectionOrigin::Outbound => Mode::Client,
        };

        Ok(Yamux::new(socket, mode))
    }
}

impl<TSocket> StreamMultiplexer for Yamux<TSocket>
where
    TSocket: AsyncRead + AsyncWrite + Send + Debug + Unpin,
{
    type Substream = StreamHandle<TSocket>;
    type Listener = Listener<TSocket>;
    type Outbound = future::Ready<io::Result<Self::Substream>>;
    type Close = Close<TSocket>;

    fn listen_for_inbound(&self) -> Self::Listener {
        Listener::new(self.inner.clone())
    }

    fn open_outbound(&self) -> Self::Outbound {
        let output = match self.inner.open_stream() {
            Ok(Some(substream)) => Ok(Compat01As03::new(substream)),
            Ok(None) => Err(io::Error::new(
                io::ErrorKind::Other,
                "Unable to open substream; underlying connection is dead",
            )),
            Err(e) => Err(io::Error::new(io::ErrorKind::Other, e)),
        };

        future::ready(output)
    }

    fn close(&self) -> Self::Close {
        Close::new(self.inner.clone())
    }
}

#[derive(Debug)]
#[must_use = "futures do nothing unless polled"]
pub struct Close<TSocket> {
    inner: yamux::Connection<Compat<TSocket>>,
}

impl<TSocket> Close<TSocket>
where
    TSocket: AsyncRead + AsyncWrite + Unpin,
{
    pub fn new(connection: yamux::Connection<Compat<TSocket>>) -> Self {
        Self { inner: connection }
    }
}

impl<TSocket> Future for Close<TSocket>
where
    TSocket: AsyncRead + AsyncWrite + Unpin,
{
    type Output = io::Result<()>;

    fn poll(self: Pin<&mut Self>, context: &mut Context) -> Poll<Self::Output> {
        let mut close_fut = Compat01As03::new(poll_fn_01(|| {
            self.inner
                .close()
                .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
        }));
        Pin::new(&mut close_fut).poll(context)
    }
}

#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct Listener<TSocket> {
    inner: Compat01As03<yamux::Connection<Compat<TSocket>>>,
}

impl<TSocket> Listener<TSocket>
where
    TSocket: AsyncRead + AsyncWrite + Unpin,
{
    pub fn new(connection: yamux::Connection<Compat<TSocket>>) -> Self {
        Self {
            inner: Compat01As03::new(connection),
        }
    }
}

impl<TSocket> Stream for Listener<TSocket>
where
    TSocket: AsyncRead + AsyncWrite + Unpin,
{
    type Item = io::Result<Compat01As03<yamux::StreamHandle<Compat<TSocket>>>>;

    fn poll_next(mut self: Pin<&mut Self>, context: &mut Context) -> Poll<Option<Self::Item>> {
        match Pin::new(&mut self.inner).poll_next(context) {
            Poll::Pending => Poll::Pending,
            Poll::Ready(None) => Poll::Ready(None),
            Poll::Ready(Some(Err(e))) => {
                Poll::Ready(Some(Err(io::Error::new(io::ErrorKind::Other, e))))
            }
            Poll::Ready(Some(Ok(substream))) => Poll::Ready(Some(Ok(Compat01As03::new(substream)))),
        }
    }
}

#[cfg(test)]
mod test {
    use crate::multiplexing::{
        yamux::{Mode, Yamux},
        StreamMultiplexer,
    };
    use futures::{
        executor::block_on,
        future::join,
        io::{AsyncReadExt, AsyncWriteExt},
        stream::StreamExt,
    };
    use memsocket::MemorySocket;
    use std::io;

    #[test]
    fn open_substream() -> io::Result<()> {
        let (dialer, listener) = MemorySocket::new_pair();
        let msg = b"The Way of Kings";

        let dialer = async move {
            let muxer = Yamux::new(dialer, Mode::Client);

            let mut substream = muxer.open_outbound().await?;

            substream.write_all(msg).await?;
            substream.flush().await?;

            // Force return type of the async block
            let result: io::Result<()> = Ok(());
            result
        };

        let listener = async move {
            let muxer = Yamux::new(listener, Mode::Server);

            let (maybe_substream, _listener) = muxer.listen_for_inbound().into_future().await;
            let mut substream = maybe_substream
                .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "no substream"))??;

            let mut buf = Vec::new();
            substream.read_to_end(&mut buf).await?;
            assert_eq!(buf, msg);

            // Force return type of the async block
            let result: io::Result<()> = Ok(());
            result
        };

        let (dialer_result, listener_result) = block_on(join(dialer, listener));
        dialer_result?;
        listener_result?;
        Ok(())
    }

    #[test]
    fn close() -> io::Result<()> {
        let (dialer, listener) = MemorySocket::new_pair();
        let msg = b"Words of Radiance";

        let dialer = async move {
            let muxer = Yamux::new(dialer, Mode::Client);

            let mut substream = muxer.open_outbound().await?;

            substream.write_all(msg).await?;
            substream.flush().await?;

            let mut buf = Vec::new();
            substream.read_to_end(&mut buf).await?;
            assert_eq!(buf, b"");

            // Force return type of the async block
            let result: io::Result<()> = Ok(());
            result
        };

        let listener = async move {
            let muxer = Yamux::new(listener, Mode::Server);

            let (maybe_substream, _listener) = muxer.listen_for_inbound().into_future().await;
            let mut substream = maybe_substream
                .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "no substream"))??;

            let mut buf = vec![0; msg.len()];
            substream.read_exact(&mut buf).await?;
            assert_eq!(buf, msg);

            // Close the muxer and then try to write to it
            muxer.close().await?;

            let result = substream.write_all(b"ignored message").await;
            match result {
                Ok(()) => panic!("Write should have failed"),
                Err(e) => assert_eq!(e.kind(), io::ErrorKind::WriteZero),
            }

            // Force return type of the async block
            let result: io::Result<()> = Ok(());
            result
        };

        let (dialer_result, listener_result) = block_on(join(dialer, listener));
        dialer_result?;
        listener_result?;
        Ok(())
    }

    #[test]
    fn send_big_message() -> io::Result<()> {
        #[allow(non_snake_case)]
        let MiB: usize = 1 << 20;
        let msg_len = 16 * MiB;

        let (dialer, listener) = MemorySocket::new_pair();

        let dialer = async move {
            let muxer = Yamux::new(dialer, Mode::Client);
            let mut substream = muxer.open_outbound().await?;

            let msg = vec![0x55u8; msg_len];
            substream.write_all(msg.as_slice()).await?;

            let mut buf = vec![0u8; msg_len];
            substream.read_exact(&mut buf).await?;
            substream.close().await?;

            assert_eq!(buf.len(), msg_len);
            assert_eq!(buf, vec![0xAAu8; msg_len]);

            // Force return type of the async block
            let result: io::Result<Yamux<_>> = Ok(muxer);
            result
        };

        let listener = async move {
            let muxer = Yamux::new(listener, Mode::Server);
            let (maybe_substream, _listener) = muxer.listen_for_inbound().into_future().await;
            let mut substream = maybe_substream
                .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "no substream"))??;

            let mut buf = vec![0u8; msg_len];
            substream.read_exact(&mut buf).await?;
            assert_eq!(buf, vec![0x55u8; msg_len]);

            let msg = vec![0xAAu8; msg_len];
            substream.write_all(msg.as_slice()).await?;
            substream.close().await?;

            // Force return type of the async block
            let result: io::Result<Yamux<_>> = Ok(muxer);
            result
        };

        let (dialer_result, listener_result) = block_on(join(dialer, listener));
        let _ = dialer_result?;
        let _ = listener_result?;
        Ok(())
    }
}
<commit_msg>[network] Test multiplexing of yamux substreams within yamux substreams<commit_after>// Copyright (c) The Libra Core Contributors
// SPDX-License-Identifier: Apache-2.0

//! Implementation of [`StreamMultiplexer`] using the [`yamux`] protocol over TCP
//!
//! [`StreamMultiplexer`]: crate::multiplexing::StreamMultiplexer
//! [`yamux`]: https://github.com/hashicorp/yamux/blob/master/spec.md

use crate::{
    multiplexing::StreamMultiplexer,
    negotiate::{negotiate_inbound, negotiate_outbound_interactive},
    transport::ConnectionOrigin,
};
use futures::{
    compat::{Compat, Compat01As03},
    future::{self, Future},
    io::{AsyncRead, AsyncWrite},
    stream::Stream,
};
use futures_01::future::poll_fn as poll_fn_01;
use std::{
    fmt::Debug,
    io,
    pin::Pin,
    task::{Context, Poll},
};
use yamux;

/// Re-export `Mode` from the yamux crate
pub use yamux::Mode;

/// The substream type produced by the `Yamux` multiplexer
pub type StreamHandle<T> = Compat01As03<yamux::StreamHandle<Compat<T>>>;

const YAMUX_PROTOCOL_NAME: &[u8] = b"/yamux/1.0.0";

#[derive(Debug)]
pub struct Yamux<TSocket> {
    inner: yamux::Connection<Compat<TSocket>>,
}

const MAX_BUFFER_SIZE: u32 = 8 * 1024 * 1024; // 8MB
const RECEIVE_WINDOW: u32 = 4 * 1024 * 1024; // 4MB

impl<TSocket> Yamux<TSocket>
where
    TSocket: AsyncRead + AsyncWrite + Send + Debug + Unpin,
{
    pub fn new(socket: TSocket, mode: Mode) -> Self {
        let mut config = yamux::Config::default();
        // Use OnRead mode instead of OnReceive mode to provide back pressure to the sending side.
        // Caveat: the OnRead mode has the risk of deadlock, where both sides send data larger than
        // receive window and don't read before finishing writes. But it doesn't apply to our use
        // cases. Some of our streams are unidirectional, where only one end writes data, e.g.,
        // Direct Send. Some of our streams are bidirectional, but only one end write data at a
        // time, e.g., RPC. Some of our streams may write at the same time, but the frames are
        // shorter than the receive window, e.g., protocol negotiation.
        config.set_window_update_mode(yamux::WindowUpdateMode::OnRead);
        // Because OnRead mode increases the RTT of window update, bigger buffer size and receive
        // window size perform better.
        config.set_max_buffer_size(MAX_BUFFER_SIZE as usize);
        config
            .set_receive_window(RECEIVE_WINDOW)
            .expect("Invalid receive window size");
        let socket = Compat::new(socket);
        Self {
            inner: yamux::Connection::new(socket, config, mode),
        }
    }

    pub async fn upgrade_connection(socket: TSocket, origin: ConnectionOrigin) -> io::Result<Self> {
        // Perform protocol negotiation
        let (socket, proto) = match origin {
            ConnectionOrigin::Inbound => negotiate_inbound(socket, [YAMUX_PROTOCOL_NAME]).await?,
            ConnectionOrigin::Outbound => {
                negotiate_outbound_interactive(socket, [YAMUX_PROTOCOL_NAME]).await?
            }
        };

        assert_eq!(proto, YAMUX_PROTOCOL_NAME);

        let mode = match origin {
            ConnectionOrigin::Inbound => Mode::Server,
            ConnectionOrigin::Outbound => Mode::Client,
        };

        Ok(Yamux::new(socket, mode))
    }
}

impl<TSocket> StreamMultiplexer for Yamux<TSocket>
where
    TSocket: AsyncRead + AsyncWrite + Send + Debug + Unpin,
{
    type Substream = StreamHandle<TSocket>;
    type Listener = Listener<TSocket>;
    type Outbound = future::Ready<io::Result<Self::Substream>>;
    type Close = Close<TSocket>;

    fn listen_for_inbound(&self) -> Self::Listener {
        Listener::new(self.inner.clone())
    }

    fn open_outbound(&self) -> Self::Outbound {
        let output = match self.inner.open_stream() {
            Ok(Some(substream)) => Ok(Compat01As03::new(substream)),
            Ok(None) => Err(io::Error::new(
                io::ErrorKind::Other,
                "Unable to open substream; underlying connection is dead",
            )),
            Err(e) => Err(io::Error::new(io::ErrorKind::Other, e)),
        };

        future::ready(output)
    }

    fn close(&self) -> Self::Close {
        Close::new(self.inner.clone())
    }
}

#[derive(Debug)]
#[must_use = "futures do nothing unless polled"]
pub struct Close<TSocket> {
    inner: yamux::Connection<Compat<TSocket>>,
}

impl<TSocket> Close<TSocket>
where
    TSocket: AsyncRead + AsyncWrite + Unpin,
{
    pub fn new(connection: yamux::Connection<Compat<TSocket>>) -> Self {
        Self { inner: connection }
    }
}

impl<TSocket> Future for Close<TSocket>
where
    TSocket: AsyncRead + AsyncWrite + Unpin,
{
    type Output = io::Result<()>;

    fn poll(self: Pin<&mut Self>, context: &mut Context) -> Poll<Self::Output> {
        let mut close_fut = Compat01As03::new(poll_fn_01(|| {
            self.inner
                .close()
                .map_err(|e| io::Error::new(io::ErrorKind::Other, e))
        }));
        Pin::new(&mut close_fut).poll(context)
    }
}

#[derive(Debug)]
#[must_use = "streams do nothing unless polled"]
pub struct Listener<TSocket> {
    inner: Compat01As03<yamux::Connection<Compat<TSocket>>>,
}

impl<TSocket> Listener<TSocket>
where
    TSocket: AsyncRead + AsyncWrite + Unpin,
{
    pub fn new(connection: yamux::Connection<Compat<TSocket>>) -> Self {
        Self {
            inner: Compat01As03::new(connection),
        }
    }
}

impl<TSocket> Stream for Listener<TSocket>
where
    TSocket: AsyncRead + AsyncWrite + Unpin,
{
    type Item = io::Result<Compat01As03<yamux::StreamHandle<Compat<TSocket>>>>;

    fn poll_next(mut self: Pin<&mut Self>, context: &mut Context) -> Poll<Option<Self::Item>> {
        match Pin::new(&mut self.inner).poll_next(context) {
            Poll::Pending => Poll::Pending,
            Poll::Ready(None) => Poll::Ready(None),
            Poll::Ready(Some(Err(e))) => {
                Poll::Ready(Some(Err(io::Error::new(io::ErrorKind::Other, e))))
            }
            Poll::Ready(Some(Ok(substream))) => Poll::Ready(Some(Ok(Compat01As03::new(substream)))),
        }
    }
}

#[cfg(test)]
mod test {
    use crate::multiplexing::{
        yamux::{Mode, Yamux},
        StreamMultiplexer,
    };
    use futures::{
        executor::block_on,
        future::join,
        io::{AsyncReadExt, AsyncWriteExt},
        stream::StreamExt,
    };
    use memsocket::MemorySocket;
    use std::io;

    #[test]
    fn substream_within_substream() -> io::Result<()> {
        let (dialer, listener) = MemorySocket::new_pair();
        let msg = b"Too fast too furious";
        let msg2 = b"Mission Impossible";

        let dialer = async move {
            let muxer = Yamux::new(dialer, Mode::Client);
            // Open outer substream.
            let substream = muxer.open_outbound().await?;
            // Create multiplexer over substream.
            let muxer = Yamux::new(substream, Mode::Client);
            // Open new inner outbound substream over outer substream.
            let mut substream = muxer.open_outbound().await?;
            // Send data over inner substream.
            substream.write_all(msg).await?;
            substream.flush().await?;
            // We have to close the substream to unblock the read end if it uses `read_to_end`.
            substream.close().await?;
            // Listen for new inbound inner substreams over outer substream.
            let (maybe_substream, _) = muxer.listen_for_inbound().into_future().await;
            let mut substream = maybe_substream
                .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "no substream"))??;
            // Receive data over inner substream.
            let mut buf = Vec::new();
            substream.read_to_end(&mut buf).await?;
            assert_eq!(buf, msg2);
            // Force return type of the async block
            let result: io::Result<()> = Ok(());
            result
        };

        let listener = async move {
            let muxer = Yamux::new(listener, Mode::Server);
            // Listen for inbound outer substream.
            let (maybe_substream, _) = muxer.listen_for_inbound().into_future().await;
            let substream = maybe_substream
                .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "no substream"))??;
            // Create multiplexer over substream.
            let muxer = Yamux::new(substream, Mode::Server);
            // Listen for inbound inner substream over outer substream.
            let (maybe_substream, _listener) = muxer.listen_for_inbound().into_future().await;
            let mut substream = maybe_substream
                .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "no substream"))??;
            // Receive data over inner substream.
            let mut buf = Vec::new();
            substream.read_to_end(&mut buf).await?;
            assert_eq!(buf, msg);
            // Open new inner outbound substream over outer substream.
            let mut substream = muxer.open_outbound().await?;
            // Send data over inner substream.
            substream.write_all(msg2).await?;
            substream.flush().await?;
            // Force return type of the async block
            let result: io::Result<()> = Ok(());
            result
        };

        let (dialer_result, listener_result) = block_on(join(dialer, listener));
        dialer_result?;
        listener_result?;
        Ok(())
    }

    #[test]
    fn open_substream() -> io::Result<()> {
        let (dialer, listener) = MemorySocket::new_pair();
        let msg = b"The Way of Kings";

        let dialer = async move {
            let muxer = Yamux::new(dialer, Mode::Client);

            let mut substream = muxer.open_outbound().await?;

            substream.write_all(msg).await?;
            substream.flush().await?;

            // Force return type of the async block
            let result: io::Result<()> = Ok(());
            result
        };

        let listener = async move {
            let muxer = Yamux::new(listener, Mode::Server);

            let (maybe_substream, _listener) = muxer.listen_for_inbound().into_future().await;
            let mut substream = maybe_substream
                .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "no substream"))??;

            let mut buf = Vec::new();
            substream.read_to_end(&mut buf).await?;
            assert_eq!(buf, msg);

            // Force return type of the async block
            let result: io::Result<()> = Ok(());
            result
        };

        let (dialer_result, listener_result) = block_on(join(dialer, listener));
        dialer_result?;
        listener_result?;
        Ok(())
    }

    #[test]
    fn close() -> io::Result<()> {
        let (dialer, listener) = MemorySocket::new_pair();
        let msg = b"Words of Radiance";

        let dialer = async move {
            let muxer = Yamux::new(dialer, Mode::Client);

            let mut substream = muxer.open_outbound().await?;

            substream.write_all(msg).await?;
            substream.flush().await?;

            let mut buf = Vec::new();
            substream.read_to_end(&mut buf).await?;
            assert_eq!(buf, b"");

            // Force return type of the async block
            let result: io::Result<()> = Ok(());
            result
        };

        let listener = async move {
            let muxer = Yamux::new(listener, Mode::Server);

            let (maybe_substream, _listener) = muxer.listen_for_inbound().into_future().await;
            let mut substream = maybe_substream
                .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "no substream"))??;

            let mut buf = vec![0; msg.len()];
            substream.read_exact(&mut buf).await?;
            assert_eq!(buf, msg);

            // Close the muxer and then try to write to it
            muxer.close().await?;

            let result = substream.write_all(b"ignored message").await;
            match result {
                Ok(()) => panic!("Write should have failed"),
                Err(e) => assert_eq!(e.kind(), io::ErrorKind::WriteZero),
            }

            // Force return type of the async block
            let result: io::Result<()> = Ok(());
            result
        };

        let (dialer_result, listener_result) = block_on(join(dialer, listener));
        dialer_result?;
        listener_result?;
        Ok(())
    }

    #[test]
    fn send_big_message() -> io::Result<()> {
        #[allow(non_snake_case)]
        let MiB: usize = 1 << 20;
        let msg_len = 16 * MiB;

        let (dialer, listener) = MemorySocket::new_pair();

        let dialer = async move {
            let muxer = Yamux::new(dialer, Mode::Client);
            let mut substream = muxer.open_outbound().await?;

            let msg = vec![0x55u8; msg_len];
            substream.write_all(msg.as_slice()).await?;

            let mut buf = vec![0u8; msg_len];
            substream.read_exact(&mut buf).await?;
            substream.close().await?;

            assert_eq!(buf.len(), msg_len);
            assert_eq!(buf, vec![0xAAu8; msg_len]);

            // Force return type of the async block
            let result: io::Result<Yamux<_>> = Ok(muxer);
            result
        };

        let listener = async move {
            let muxer = Yamux::new(listener, Mode::Server);
            let (maybe_substream, _listener) = muxer.listen_for_inbound().into_future().await;
            let mut substream = maybe_substream
                .ok_or_else(|| io::Error::new(io::ErrorKind::Other, "no substream"))??;

            let mut buf = vec![0u8; msg_len];
            substream.read_exact(&mut buf).await?;
            assert_eq!(buf, vec![0x55u8; msg_len]);

            let msg = vec![0xAAu8; msg_len];
            substream.write_all(msg.as_slice()).await?;
            substream.close().await?;

            // Force return type of the async block
            let result: io::Result<Yamux<_>> = Ok(muxer);
            result
        };

        let (dialer_result, listener_result) = block_on(join(dialer, listener));
        let _ = dialer_result?;
        let _ = listener_result?;
        Ok(())
    }
}
<|endoftext|>


<commit_before>use clap::ArgMatches;
use diesel::expression::sql;
#[cfg(feature="postgres")]
use diesel::pg::PgConnection;
#[cfg(feature="sqlite")]
use diesel::sqlite::SqliteConnection;
#[cfg(feature="mysql")]
use diesel::mysql::MysqlConnection;
use diesel::types::Bool;
use diesel::*;

use database_error::{DatabaseError, DatabaseResult};

use std::error::Error;
use std::env;

enum Backend {
    #[cfg(feature="postgres")]
    Pg,
    #[cfg(feature="sqlite")]
    Sqlite,
    #[cfg(feature="mysql")]
    Mysql,
}

impl Backend {
    fn for_url(database_url: &str) -> Self {
        match database_url {
            #[cfg(feature="postgres")]
            _ if database_url.starts_with("postgres://") || database_url.starts_with("postgresql://") =>
                Backend::Pg,
            #[cfg(feature="mysql")]
            _ if database_url.starts_with("mysql://") =>
                Backend::Mysql,
            #[cfg(feature="sqlite")]
            _ => Backend::Sqlite,
            #[cfg(not(feature="sqlite"))]
            _ => {
                panic!("{:?} is not a valid database URL. It should start with \
                        `postgres://` or `mysql://`", database_url);
            }
        }
    }
}

pub enum InferConnection {
    #[cfg(feature="postgres")]
    Pg(PgConnection),
    #[cfg(feature="sqlite")]
    Sqlite(SqliteConnection),
    #[cfg(feature="mysql")]
    Mysql(MysqlConnection),
}

impl InferConnection {
    pub fn establish(database_url: &str) -> DatabaseResult<Self> {
        match Backend::for_url(database_url) {
            #[cfg(feature="postgres")]
            Backend::Pg => PgConnection::establish(database_url)
                .map(InferConnection::Pg),
            #[cfg(feature="sqlite")]
            Backend::Sqlite => SqliteConnection::establish(database_url)
                    .map(InferConnection::Sqlite),
            #[cfg(feature="mysql")]
            Backend::Mysql => MysqlConnection::establish(database_url)
                    .map(InferConnection::Mysql),
        }.map_err(Into::into)
    }
}

macro_rules! call_with_conn {
    (
        $database_url:expr,
        $($func:ident)::+
    ) => {
        call_with_conn!($database_url, $($func)::+ ())
    };

    (
        $database_url:expr,
        $($func:ident)::+ ($($args:expr),*)
    ) => {
        match ::database::InferConnection::establish(&$database_url).unwrap() {
            #[cfg(feature="postgres")]
            ::database::InferConnection::Pg(ref conn) => $($func)::+ (conn, $($args),*),
            #[cfg(feature="sqlite")]
            ::database::InferConnection::Sqlite(ref conn) => $($func)::+ (conn, $($args),*),
            #[cfg(feature="mysql")]
            ::database::InferConnection::Mysql(ref conn) => $($func)::+ (conn, $($args),*),
        }
    };
}

pub fn reset_database(args: &ArgMatches) -> DatabaseResult<()> {
    try!(drop_database(&database_url(args)));
    setup_database(args)
}

pub fn setup_database(args: &ArgMatches) -> DatabaseResult<()> {
    let database_url = database_url(args);

    try!(create_database_if_needed(&database_url));
    create_schema_table_and_run_migrations_if_needed(&database_url)
}

pub fn drop_database_command(args: &ArgMatches) -> DatabaseResult<()> {
    drop_database(&database_url(args))
}

/// Creates the database specified in the connection url. It returns an error
/// it it was unable to create the database.
fn create_database_if_needed(database_url: &str) -> DatabaseResult<()> {
    match Backend::for_url(database_url) {
        #[cfg(feature="postgres")]
        Backend::Pg => {
            if PgConnection::establish(database_url).is_err() {
                let (database, postgres_url) = change_database_of_url(database_url, "postgres");
                println!("Creating database: {}", database);
                let conn = try!(PgConnection::establish(&postgres_url));
                try!(conn.execute(&format!("CREATE DATABASE {}", database)));
            }
        },
        #[cfg(feature="sqlite")]
        Backend::Sqlite => {
            if !::std::path::Path::new(database_url).exists() {
                println!("Creating database: {}", database_url);
                try!(SqliteConnection::establish(database_url));
            }
        },
        #[cfg(feature="mysql")]
        Backend::Mysql => {
            if MysqlConnection::establish(database_url).is_err() {
                let (database, mysql_url) = change_database_of_url(database_url, "information_schema");
                println!("Creating database: {}", database);
                let conn = try!(MysqlConnection::establish(&mysql_url));
                try!(conn.execute(&format!("CREATE DATABASE {}", database)));
            }
        },
    }

    Ok(())
}

/// Creates the `__diesel_schema_migrations` table if it doesn't exist. If the
/// table didn't exist, it also runs any pending migrations. Returns a
/// `DatabaseError::ConnectionError` if it can't create the table, and exits
/// with a migration error if it can't run migrations.
fn create_schema_table_and_run_migrations_if_needed(database_url: &str)
    -> DatabaseResult<()>
{
    if !schema_table_exists(database_url).unwrap_or_else(handle_error) {
        try!(call_with_conn!(database_url, migrations::setup_database));
        call_with_conn!(database_url, migrations::run_pending_migrations).unwrap_or_else(handle_error);
    };
    Ok(())
}

/// Drops the database specified in the connection url. It returns an error
/// if it was unable to drop the database.
fn drop_database(database_url: &str) -> DatabaseResult<()> {
    match Backend::for_url(database_url) {
        #[cfg(feature="postgres")]
        Backend::Pg => {
            let (database, postgres_url) = change_database_of_url(database_url, "postgres");
            let conn = try!(PgConnection::establish(&postgres_url));
            if try!(pg_database_exists(&conn, &database)) {
                println!("Dropping database: {}", database);
                try!(conn.execute(&format!("DROP DATABASE IF EXISTS {}", database)));
            }
        },
        #[cfg(feature="sqlite")]
        Backend::Sqlite => {
            use std::path::Path;
            use std::fs;

            if Path::new(database_url).exists() {
                println!("Dropping database: {}", database_url);
                try!(fs::remove_file(&database_url));
            }
        },
        #[cfg(feature="mysql")]
        Backend::Mysql => {
            let (database, mysql_url) = change_database_of_url(database_url, "information_schema");
            let conn = try!(MysqlConnection::establish(&mysql_url));
            if try!(mysql_database_exists(&conn, &database)) {
                println!("Dropping database: {}", database);
                try!(conn.execute(&format!("DROP DATABASE IF EXISTS {}", database)));
            }
        },
    }
    Ok(())
}

#[cfg(feature="postgres")]
table! {
    pg_database (datname) {
        datname -> Text,
        datistemplate -> Bool,
    }
}

#[cfg(feature="postgres")]
fn pg_database_exists(conn: &PgConnection, database_name: &str) -> QueryResult<bool> {
    use self::pg_database::dsl::*;

    pg_database.select(datname) // here come dat name!!!! o shit waddup!!!!
        .filter(datname.eq(database_name))
        .filter(datistemplate.eq(false))
        .get_result::<String>(conn)
        .optional()
        .map(|x| x.is_some())
}

#[cfg(feature="mysql")]
table! {
    information_schema.schemata (schema_name) {
        schema_name -> Text,
    }
}

#[cfg(feature="mysql")]
fn mysql_database_exists(conn: &MysqlConnection, database_name: &str) -> QueryResult<bool> {
    use self::schemata::dsl::*;

    schemata.select(schema_name)
        .filter(schema_name.eq(database_name))
        .get_result::<String>(conn)
        .optional()
        .map(|x| x.is_some())
}

/// Returns true if the `__diesel_schema_migrations` table exists in the
/// database we connect to, returns false if it does not.
pub fn schema_table_exists(database_url: &str) -> DatabaseResult<bool> {
    match InferConnection::establish(database_url).unwrap() {
        #[cfg(feature="postgres")]
        InferConnection::Pg(conn) => {
            select(sql::<Bool>("EXISTS \
                    (SELECT 1 \
                     FROM information_schema.tables \
                     WHERE table_name = '__diesel_schema_migrations')"))
                .get_result(&conn)
        },
        #[cfg(feature="sqlite")]
        InferConnection::Sqlite(conn) => {
            select(sql::<Bool>("EXISTS \
                    (SELECT 1 \
                     FROM sqlite_master \
                     WHERE type = 'table' \
                     AND name = '__diesel_schema_migrations')"))
                .get_result(&conn)
        },
        #[cfg(feature="mysql")]
        InferConnection::Mysql(conn) => {
            select(sql::<Bool>("EXISTS \
                    (SELECT 1 \
                     FROM information_schema.tables \
                     WHERE table_name = '__diesel_schema_migrations'
                     AND table_schema = DATABASE())"))
                .get_result(&conn)
        },
    }.map_err(Into::into)
}

pub fn database_url(matches: &ArgMatches) -> String {
    matches.value_of("DATABASE_URL")
        .map(|s| s.into())
        .or_else(|| env::var("DATABASE_URL").ok())
        .unwrap_or_else(|| {
            handle_error(DatabaseError::DatabaseUrlMissing)
        })
}

#[cfg(any(feature="postgres", feature="mysql"))]
fn change_database_of_url(database_url: &str, default_database: &str) -> (String, String) {
    let mut split: Vec<&str> = database_url.split('/').collect();
    let database = split.pop().unwrap();
    let new_url = format!("{}/{}", split.join("/"), default_database);
    (database.to_owned(), new_url)
}

fn handle_error<E: Error, T>(error: E) -> T {
    println!("{}", error);
    ::std::process::exit(1);
}

#[cfg(all(test, any(feature="postgres", feature="mysql")))]
mod tests {
    use super::change_database_of_url;

    #[test]
    fn split_pg_connection_string_returns_postgres_url_and_database() {
        let database = "database".to_owned();
        let base_url = "postgresql://localhost:5432".to_owned();
        let database_url = format!("{}/{}", base_url, database);
        let postgres_url = format!("{}/{}", base_url, "postgres");
        assert_eq!((database, postgres_url), change_database_of_url(&database_url, "postgres"));
    }

    #[test]
    fn split_pg_connection_string_handles_user_and_password() {
        let database = "database".to_owned();
        let base_url = "postgresql://user:password@localhost:5432".to_owned();
        let database_url = format!("{}/{}", base_url, database);
        let postgres_url = format!("{}/{}", base_url, "postgres");
        assert_eq!((database, postgres_url), change_database_of_url(&database_url, "postgres"));
    }
}
<commit_msg>Nicer error when CLI can't parse database URL<commit_after>use clap::ArgMatches;
use diesel::expression::sql;
#[cfg(feature="postgres")]
use diesel::pg::PgConnection;
#[cfg(feature="sqlite")]
use diesel::sqlite::SqliteConnection;
#[cfg(feature="mysql")]
use diesel::mysql::MysqlConnection;
use diesel::types::Bool;
use diesel::*;

use database_error::{DatabaseError, DatabaseResult};

use std::error::Error;
use std::env;

enum Backend {
    #[cfg(feature="postgres")]
    Pg,
    #[cfg(feature="sqlite")]
    Sqlite,
    #[cfg(feature="mysql")]
    Mysql,
}

impl Backend {
    fn for_url(database_url: &str) -> Self {
        match database_url {
            #[cfg(feature="postgres")]
            _ if database_url.starts_with("postgres://") || database_url.starts_with("postgresql://") =>
                Backend::Pg,
            #[cfg(feature="mysql")]
            _ if database_url.starts_with("mysql://") =>
                Backend::Mysql,
            #[cfg(feature="sqlite")]
            _ => Backend::Sqlite,
            #[cfg(not(feature="sqlite"))]
            _ => {
                let mut available_schemas: Vec<&str> = Vec::new();

                // One of these will always be true, or you are compiling
                // diesel_cli without a backend. And why would you ever want to
                // do that?
                if cfg!(feature="postgres") {
                    available_schemes.push("`postgres://`");
                }
                if cfg!(feature="mysql") {
                    available_schemes.push("`mysql://`");
                }

                panic!("`{}` is not a valid database URL. It should start with {}",
                    database_url, available_schemes.join(" or "));
            },
        }
    }
}

pub enum InferConnection {
    #[cfg(feature="postgres")]
    Pg(PgConnection),
    #[cfg(feature="sqlite")]
    Sqlite(SqliteConnection),
    #[cfg(feature="mysql")]
    Mysql(MysqlConnection),
}

impl InferConnection {
    pub fn establish(database_url: &str) -> DatabaseResult<Self> {
        match Backend::for_url(database_url) {
            #[cfg(feature="postgres")]
            Backend::Pg => PgConnection::establish(database_url)
                .map(InferConnection::Pg),
            #[cfg(feature="sqlite")]
            Backend::Sqlite => SqliteConnection::establish(database_url)
                    .map(InferConnection::Sqlite),
            #[cfg(feature="mysql")]
            Backend::Mysql => MysqlConnection::establish(database_url)
                    .map(InferConnection::Mysql),
        }.map_err(Into::into)
    }
}

macro_rules! call_with_conn {
    (
        $database_url:expr,
        $($func:ident)::+
    ) => {
        call_with_conn!($database_url, $($func)::+ ())
    };

    (
        $database_url:expr,
        $($func:ident)::+ ($($args:expr),*)
    ) => {
        match ::database::InferConnection::establish(&$database_url).unwrap() {
            #[cfg(feature="postgres")]
            ::database::InferConnection::Pg(ref conn) => $($func)::+ (conn, $($args),*),
            #[cfg(feature="sqlite")]
            ::database::InferConnection::Sqlite(ref conn) => $($func)::+ (conn, $($args),*),
            #[cfg(feature="mysql")]
            ::database::InferConnection::Mysql(ref conn) => $($func)::+ (conn, $($args),*),
        }
    };
}

pub fn reset_database(args: &ArgMatches) -> DatabaseResult<()> {
    try!(drop_database(&database_url(args)));
    setup_database(args)
}

pub fn setup_database(args: &ArgMatches) -> DatabaseResult<()> {
    let database_url = database_url(args);

    try!(create_database_if_needed(&database_url));
    create_schema_table_and_run_migrations_if_needed(&database_url)
}

pub fn drop_database_command(args: &ArgMatches) -> DatabaseResult<()> {
    drop_database(&database_url(args))
}

/// Creates the database specified in the connection url. It returns an error
/// it it was unable to create the database.
fn create_database_if_needed(database_url: &str) -> DatabaseResult<()> {
    match Backend::for_url(database_url) {
        #[cfg(feature="postgres")]
        Backend::Pg => {
            if PgConnection::establish(database_url).is_err() {
                let (database, postgres_url) = change_database_of_url(database_url, "postgres");
                println!("Creating database: {}", database);
                let conn = try!(PgConnection::establish(&postgres_url));
                try!(conn.execute(&format!("CREATE DATABASE {}", database)));
            }
        },
        #[cfg(feature="sqlite")]
        Backend::Sqlite => {
            if !::std::path::Path::new(database_url).exists() {
                println!("Creating database: {}", database_url);
                try!(SqliteConnection::establish(database_url));
            }
        },
        #[cfg(feature="mysql")]
        Backend::Mysql => {
            if MysqlConnection::establish(database_url).is_err() {
                let (database, mysql_url) = change_database_of_url(database_url, "information_schema");
                println!("Creating database: {}", database);
                let conn = try!(MysqlConnection::establish(&mysql_url));
                try!(conn.execute(&format!("CREATE DATABASE {}", database)));
            }
        },
    }

    Ok(())
}

/// Creates the `__diesel_schema_migrations` table if it doesn't exist. If the
/// table didn't exist, it also runs any pending migrations. Returns a
/// `DatabaseError::ConnectionError` if it can't create the table, and exits
/// with a migration error if it can't run migrations.
fn create_schema_table_and_run_migrations_if_needed(database_url: &str)
    -> DatabaseResult<()>
{
    if !schema_table_exists(database_url).unwrap_or_else(handle_error) {
        try!(call_with_conn!(database_url, migrations::setup_database));
        call_with_conn!(database_url, migrations::run_pending_migrations).unwrap_or_else(handle_error);
    };
    Ok(())
}

/// Drops the database specified in the connection url. It returns an error
/// if it was unable to drop the database.
fn drop_database(database_url: &str) -> DatabaseResult<()> {
    match Backend::for_url(database_url) {
        #[cfg(feature="postgres")]
        Backend::Pg => {
            let (database, postgres_url) = change_database_of_url(database_url, "postgres");
            let conn = try!(PgConnection::establish(&postgres_url));
            if try!(pg_database_exists(&conn, &database)) {
                println!("Dropping database: {}", database);
                try!(conn.execute(&format!("DROP DATABASE IF EXISTS {}", database)));
            }
        },
        #[cfg(feature="sqlite")]
        Backend::Sqlite => {
            use std::path::Path;
            use std::fs;

            if Path::new(database_url).exists() {
                println!("Dropping database: {}", database_url);
                try!(fs::remove_file(&database_url));
            }
        },
        #[cfg(feature="mysql")]
        Backend::Mysql => {
            let (database, mysql_url) = change_database_of_url(database_url, "information_schema");
            let conn = try!(MysqlConnection::establish(&mysql_url));
            if try!(mysql_database_exists(&conn, &database)) {
                println!("Dropping database: {}", database);
                try!(conn.execute(&format!("DROP DATABASE IF EXISTS {}", database)));
            }
        },
    }
    Ok(())
}

#[cfg(feature="postgres")]
table! {
    pg_database (datname) {
        datname -> Text,
        datistemplate -> Bool,
    }
}

#[cfg(feature="postgres")]
fn pg_database_exists(conn: &PgConnection, database_name: &str) -> QueryResult<bool> {
    use self::pg_database::dsl::*;

    pg_database.select(datname) // here come dat name!!!! o shit waddup!!!!
        .filter(datname.eq(database_name))
        .filter(datistemplate.eq(false))
        .get_result::<String>(conn)
        .optional()
        .map(|x| x.is_some())
}

#[cfg(feature="mysql")]
table! {
    information_schema.schemata (schema_name) {
        schema_name -> Text,
    }
}

#[cfg(feature="mysql")]
fn mysql_database_exists(conn: &MysqlConnection, database_name: &str) -> QueryResult<bool> {
    use self::schemata::dsl::*;

    schemata.select(schema_name)
        .filter(schema_name.eq(database_name))
        .get_result::<String>(conn)
        .optional()
        .map(|x| x.is_some())
}

/// Returns true if the `__diesel_schema_migrations` table exists in the
/// database we connect to, returns false if it does not.
pub fn schema_table_exists(database_url: &str) -> DatabaseResult<bool> {
    match InferConnection::establish(database_url).unwrap() {
        #[cfg(feature="postgres")]
        InferConnection::Pg(conn) => {
            select(sql::<Bool>("EXISTS \
                    (SELECT 1 \
                     FROM information_schema.tables \
                     WHERE table_name = '__diesel_schema_migrations')"))
                .get_result(&conn)
        },
        #[cfg(feature="sqlite")]
        InferConnection::Sqlite(conn) => {
            select(sql::<Bool>("EXISTS \
                    (SELECT 1 \
                     FROM sqlite_master \
                     WHERE type = 'table' \
                     AND name = '__diesel_schema_migrations')"))
                .get_result(&conn)
        },
        #[cfg(feature="mysql")]
        InferConnection::Mysql(conn) => {
            select(sql::<Bool>("EXISTS \
                    (SELECT 1 \
                     FROM information_schema.tables \
                     WHERE table_name = '__diesel_schema_migrations'
                     AND table_schema = DATABASE())"))
                .get_result(&conn)
        },
    }.map_err(Into::into)
}

pub fn database_url(matches: &ArgMatches) -> String {
    matches.value_of("DATABASE_URL")
        .map(|s| s.into())
        .or_else(|| env::var("DATABASE_URL").ok())
        .unwrap_or_else(|| {
            handle_error(DatabaseError::DatabaseUrlMissing)
        })
}

#[cfg(any(feature="postgres", feature="mysql"))]
fn change_database_of_url(database_url: &str, default_database: &str) -> (String, String) {
    let mut split: Vec<&str> = database_url.split('/').collect();
    let database = split.pop().unwrap();
    let new_url = format!("{}/{}", split.join("/"), default_database);
    (database.to_owned(), new_url)
}

fn handle_error<E: Error, T>(error: E) -> T {
    println!("{}", error);
    ::std::process::exit(1);
}

#[cfg(all(test, any(feature="postgres", feature="mysql")))]
mod tests {
    use super::change_database_of_url;

    #[test]
    fn split_pg_connection_string_returns_postgres_url_and_database() {
        let database = "database".to_owned();
        let base_url = "postgresql://localhost:5432".to_owned();
        let database_url = format!("{}/{}", base_url, database);
        let postgres_url = format!("{}/{}", base_url, "postgres");
        assert_eq!((database, postgres_url), change_database_of_url(&database_url, "postgres"));
    }

    #[test]
    fn split_pg_connection_string_handles_user_and_password() {
        let database = "database".to_owned();
        let base_url = "postgresql://user:password@localhost:5432".to_owned();
        let database_url = format!("{}/{}", base_url, database);
        let postgres_url = format!("{}/{}", base_url, "postgres");
        assert_eq!((database, postgres_url), change_database_of_url(&database_url, "postgres"));
    }
}
<|endoftext|>


<commit_before>// TODO: Better name?
#[macro_export(local_inner_macros)]
macro_rules! kaze_sugar {
    ($($contents:tt)*) => {
        kaze_sugar_impl!([], [ $($contents)* ])
    };
}

#[doc(hidden)]
#[macro_export]
macro_rules! kaze_sugar_impl {
    // [selector], [token stream]

    // No selector cases
    ([], [ $name:ident = $value:expr; $($rest:tt)* ]) => {
        $name = $value;
        kaze_sugar_impl!([], [ $($rest)* ]);
    };
    ([], [ if ($sel:expr) { $($inner:tt)* } $($rest:tt)* ]) => {
        kaze_sugar_impl!([ $sel ], [ $($inner)* ]);
        kaze_sugar_impl!([], [ $($rest)* ]);
    };
    ([], []) => {};

    // Selector cases
    ([ $sel:expr ], [ $name:ident = $value:expr; $($rest:tt)* ]) => {
        let prev = $name;
        kaze_sugar_impl!([ $sel ], [ $($rest)* ]);
        $name = $sel.mux($value, prev)
    };
    ([ $prev_sel:expr ], [ if ($sel:expr) { $($inner:tt)* } $($rest:tt)* ]) => {
        kaze_sugar_impl!([ $prev_sel & $sel ], [ $($inner)* ]);
        kaze_sugar_impl!([ $prev_sel ], [ $($rest)* ]);
    };
    ([ $_:expr ], []) => {};
}
<commit_msg>Add kaze_sugar impl note<commit_after>#[macro_export(local_inner_macros)]
macro_rules! kaze_sugar {
    ($($contents:tt)*) => {
        kaze_sugar_impl!([], [ $($contents)* ])
    };
}

#[doc(hidden)]
#[macro_export]
// TODO: This formulation can generate a lot of extra mux's with the same or similar conditions, and should be revisited!
macro_rules! kaze_sugar_impl {
    // [selector], [token stream]

    // No selector cases
    ([], [ $name:ident = $value:expr; $($rest:tt)* ]) => {
        $name = $value;
        kaze_sugar_impl!([], [ $($rest)* ]);
    };
    ([], [ if ($sel:expr) { $($inner:tt)* } $($rest:tt)* ]) => {
        kaze_sugar_impl!([ $sel ], [ $($inner)* ]);
        kaze_sugar_impl!([], [ $($rest)* ]);
    };
    ([], []) => {};

    // Selector cases
    ([ $sel:expr ], [ $name:ident = $value:expr; $($rest:tt)* ]) => {
        let prev = $name;
        kaze_sugar_impl!([ $sel ], [ $($rest)* ]);
        $name = $sel.mux($value, prev)
    };
    ([ $prev_sel:expr ], [ if ($sel:expr) { $($inner:tt)* } $($rest:tt)* ]) => {
        kaze_sugar_impl!([ $prev_sel & $sel ], [ $($inner)* ]);
        kaze_sugar_impl!([ $prev_sel ], [ $($rest)* ]);
    };
    ([ $_:expr ], []) => {};
}
<|endoftext|>


<commit_before>static mut seed: u64 = 19940046431; //259261034506304368955239; //1706322144714608529217229883707268827757977089;

/// Generate pseudo random number
pub fn rand() -> usize {
    seed ^= seed << 12;
    seed ^= seed << 25;
    seed ^= seed << 27;
    seed = seed * 82724793451 + 12345;
    seed as usize % ::core::usize::MAX
}

/// Generate pseudo random number via seed
pub fn srand(s: usize) {
    seed = s as u64;
}
<commit_msg>Wrap in unsafe<commit_after>static mut seed: u64 = 19940046431; //259261034506304368955239; //1706322144714608529217229883707268827757977089;

/// Generate pseudo random number
pub fn rand() -> usize {
    unsafe {
        seed ^= seed << 12;
        seed ^= seed << 25;
        seed ^= seed << 27;
        seed = seed * 82724793451 + 12345;
        seed as usize % ::core::usize::MAX
    }
}

/// Set the seed
pub fn srand(s: usize) {
    unsafe {
        seed = s as u64;
    }
}
<|endoftext|>