<commit_before>//! Validates the MIR to ensure that invariants are upheld.

use super::{MirPass, MirSource};
use rustc_middle::mir::visit::Visitor;
use rustc_middle::{
    mir::{
        AggregateKind, BasicBlock, Body, BorrowKind, Location, MirPhase, Operand, Rvalue,
        Statement, StatementKind, Terminator, TerminatorKind,
    },
    ty::{
        self,
        relate::{Relate, RelateResult, TypeRelation},
        ParamEnv, Ty, TyCtxt,
    },
};

#[derive(Copy, Clone, Debug)]
enum EdgeKind {
    Unwind,
    Normal,
}

pub struct Validator {
    /// Describes at which point in the pipeline this validation is happening.
    pub when: String,
    /// The phase for which we are upholding the dialect. If the given phase forbids a specific
    /// element, this validator will now emit errors if that specific element is encountered.
    /// Note that phases that change the dialect cause all *following* phases to check the
    /// invariants of the new dialect. A phase that changes dialects never checks the new invariants
    /// itself.
    pub mir_phase: MirPhase,
}

impl<'tcx> MirPass<'tcx> for Validator {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut Body<'tcx>) {
        let param_env = tcx.param_env(source.def_id());
        let mir_phase = self.mir_phase;
        TypeChecker { when: &self.when, source, body, tcx, param_env, mir_phase }.visit_body(body);
    }
}

/// Returns whether the two types are equal up to lifetimes.
/// All lifetimes, including higher-ranked ones, get ignored for this comparison.
/// (This is unlike the `erasing_regions` methods, which keep higher-ranked lifetimes for soundness reasons.)
///
/// The point of this function is to approximate "equal up to subtyping".  However,
/// the approximation is incorrect as variance is ignored.
pub fn equal_up_to_regions(
    tcx: TyCtxt<'tcx>,
    param_env: ParamEnv<'tcx>,
    src: Ty<'tcx>,
    dest: Ty<'tcx>,
) -> bool {
    // Fast path.
    if src == dest {
        return true;
    }

    struct LifetimeIgnoreRelation<'tcx> {
        tcx: TyCtxt<'tcx>,
        param_env: ty::ParamEnv<'tcx>,
    }

    impl TypeRelation<'tcx> for LifetimeIgnoreRelation<'tcx> {
        fn tcx(&self) -> TyCtxt<'tcx> {
            self.tcx
        }

        fn param_env(&self) -> ty::ParamEnv<'tcx> {
            self.param_env
        }

        fn tag(&self) -> &'static str {
            "librustc_mir::transform::validate"
        }

        fn a_is_expected(&self) -> bool {
            true
        }

        fn relate_with_variance<T: Relate<'tcx>>(
            &mut self,
            _: ty::Variance,
            a: T,
            b: T,
        ) -> RelateResult<'tcx, T> {
            // Ignore variance, require types to be exactly the same.
            self.relate(a, b)
        }

        fn tys(&mut self, a: Ty<'tcx>, b: Ty<'tcx>) -> RelateResult<'tcx, Ty<'tcx>> {
            if a == b {
                // Short-circuit.
                return Ok(a);
            }
            ty::relate::super_relate_tys(self, a, b)
        }

        fn regions(
            &mut self,
            a: ty::Region<'tcx>,
            _b: ty::Region<'tcx>,
        ) -> RelateResult<'tcx, ty::Region<'tcx>> {
            // Ignore regions.
            Ok(a)
        }

        fn consts(
            &mut self,
            a: &'tcx ty::Const<'tcx>,
            b: &'tcx ty::Const<'tcx>,
        ) -> RelateResult<'tcx, &'tcx ty::Const<'tcx>> {
            ty::relate::super_relate_consts(self, a, b)
        }

        fn binders<T>(
            &mut self,
            a: ty::Binder<T>,
            b: ty::Binder<T>,
        ) -> RelateResult<'tcx, ty::Binder<T>>
        where
            T: Relate<'tcx>,
        {
            self.relate(a.skip_binder(), b.skip_binder())?;
            Ok(a)
        }
    }

    // Instantiate and run relation.
    let mut relator: LifetimeIgnoreRelation<'tcx> = LifetimeIgnoreRelation { tcx: tcx, param_env };
    relator.relate(src, dest).is_ok()
}

struct TypeChecker<'a, 'tcx> {
    when: &'a str,
    source: MirSource<'tcx>,
    body: &'a Body<'tcx>,
    tcx: TyCtxt<'tcx>,
    param_env: ParamEnv<'tcx>,
    mir_phase: MirPhase,
}

impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
    fn fail(&self, location: Location, msg: impl AsRef<str>) {
        let span = self.body.source_info(location).span;
        // We use `delay_span_bug` as we might see broken MIR when other errors have already
        // occurred.
        self.tcx.sess.diagnostic().delay_span_bug(
            span,
            &format!(
                "broken MIR in {:?} ({}) at {:?}:\n{}",
                self.source.instance,
                self.when,
                location,
                msg.as_ref()
            ),
        );
    }

    fn check_edge(&self, location: Location, bb: BasicBlock, edge_kind: EdgeKind) {
        if let Some(bb) = self.body.basic_blocks().get(bb) {
            let src = self.body.basic_blocks().get(location.block).unwrap();
            match (src.is_cleanup, bb.is_cleanup, edge_kind) {
                // Non-cleanup blocks can jump to non-cleanup blocks along non-unwind edges
                (false, false, EdgeKind::Normal)
                // Non-cleanup blocks can jump to cleanup blocks along unwind edges
                | (false, true, EdgeKind::Unwind)
                // Cleanup blocks can jump to cleanup blocks along non-unwind edges
                | (true, true, EdgeKind::Normal) => {}
                // All other jumps are invalid
                _ => {
                    self.fail(
                        location,
                        format!(
                            "{:?} edge to {:?} violates unwind invariants (cleanup {:?} -> {:?})",
                            edge_kind,
                            bb,
                            src.is_cleanup,
                            bb.is_cleanup,
                        )
                    )
                }
            }
        } else {
            self.fail(location, format!("encountered jump to invalid basic block {:?}", bb))
        }
    }

    /// Check if src can be assigned into dest.
    /// This is not precise, it will accept some incorrect assignments.
    fn mir_assign_valid_types(&self, src: Ty<'tcx>, dest: Ty<'tcx>) -> bool {
        // Fast path before we normalize.
        if src == dest {
            // Equal types, all is good.
            return true;
        }
        // Normalize projections and things like that.
        // FIXME: We need to reveal_all, as some optimizations change types in ways
        // that require unfolding opaque types.
        let param_env = self.param_env.with_reveal_all_normalized(self.tcx);
        let src = self.tcx.normalize_erasing_regions(param_env, src);
        let dest = self.tcx.normalize_erasing_regions(param_env, dest);

        // Type-changing assignments can happen when subtyping is used. While
        // all normal lifetimes are erased, higher-ranked types with their
        // late-bound lifetimes are still around and can lead to type
        // differences. So we compare ignoring lifetimes.
        equal_up_to_regions(self.tcx, param_env, src, dest)
    }
}

impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
    fn visit_operand(&mut self, operand: &Operand<'tcx>, location: Location) {
        // `Operand::Copy` is only supposed to be used with `Copy` types.
        if let Operand::Copy(place) = operand {
            let ty = place.ty(&self.body.local_decls, self.tcx).ty;
            let span = self.body.source_info(location).span;

            if !ty.is_copy_modulo_regions(self.tcx.at(span), self.param_env) {
                self.fail(location, format!("`Operand::Copy` with non-`Copy` type {}", ty));
            }
        }

        self.super_operand(operand, location);
    }

    fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) {
        match &statement.kind {
            StatementKind::Assign(box (dest, rvalue)) => {
                // LHS and RHS of the assignment must have the same type.
                let left_ty = dest.ty(&self.body.local_decls, self.tcx).ty;
                let right_ty = rvalue.ty(&self.body.local_decls, self.tcx);
                if !self.mir_assign_valid_types(right_ty, left_ty) {
                    self.fail(
                        location,
                        format!(
                            "encountered `{:?}` with incompatible types:\n\
                            left-hand side has type: {}\n\
                            right-hand side has type: {}",
                            statement.kind, left_ty, right_ty,
                        ),
                    );
                }
                match rvalue {
                    // The sides of an assignment must not alias. Currently this just checks whether the places
                    // are identical.
                    Rvalue::Use(Operand::Copy(src) | Operand::Move(src)) => {
                        if dest == src {
                            self.fail(
                                location,
                                "encountered `Assign` statement with overlapping memory",
                            );
                        }
                    }
                    // The deaggregator currently does not deaggreagate arrays.
                    // So for now, we ignore them here.
                    Rvalue::Aggregate(box AggregateKind::Array { .. }, _) => {}
                    // All other aggregates must be gone after some phases.
                    Rvalue::Aggregate(box kind, _) => {
                        if self.mir_phase > MirPhase::DropLowering
                            && !matches!(kind, AggregateKind::Generator(..))
                        {
                            // Generators persist until the state machine transformation, but all
                            // other aggregates must have been lowered.
                            self.fail(
                                location,
                                format!("{:?} have been lowered to field assignments", rvalue),
                            )
                        } else if self.mir_phase > MirPhase::GeneratorLowering {
                            // No more aggregates after drop and generator lowering.
                            self.fail(
                                location,
                                format!("{:?} have been lowered to field assignments", rvalue),
                            )
                        }
                    }
                    Rvalue::Ref(_, BorrowKind::Shallow, _) => {
                        if self.mir_phase > MirPhase::DropLowering {
                            self.fail(
                                location,
                                "`Assign` statement with a `Shallow` borrow should have been removed after drop lowering phase",
                            );
                        }
                    }
                    _ => {}
                }
            }
            StatementKind::AscribeUserType(..) => {
                if self.mir_phase > MirPhase::DropLowering {
                    self.fail(
                        location,
                        "`AscribeUserType` should have been removed after drop lowering phase",
                    );
                }
            }
            StatementKind::FakeRead(..) => {
                if self.mir_phase > MirPhase::DropLowering {
                    self.fail(
                        location,
                        "`FakeRead` should have been removed after drop lowering phase",
                    );
                }
            }
            _ => {}
        }
    }

    fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) {
        match &terminator.kind {
            TerminatorKind::Goto { target } => {
                self.check_edge(location, *target, EdgeKind::Normal);
            }
            TerminatorKind::SwitchInt { targets, values, switch_ty, discr } => {
                let ty = discr.ty(&self.body.local_decls, self.tcx);
                if ty != *switch_ty {
                    self.fail(
                        location,
                        format!(
                            "encountered `SwitchInt` terminator with type mismatch: {:?} != {:?}",
                            ty, switch_ty,
                        ),
                    );
                }
                if targets.len() != values.len() + 1 {
                    self.fail(
                        location,
                        format!(
                            "encountered `SwitchInt` terminator with {} values, but {} targets (should be values+1)",
                            values.len(),
                            targets.len(),
                        ),
                    );
                }
                for target in targets {
                    self.check_edge(location, *target, EdgeKind::Normal);
                }
            }
            TerminatorKind::Drop { target, unwind, .. } => {
                self.check_edge(location, *target, EdgeKind::Normal);
                if let Some(unwind) = unwind {
                    self.check_edge(location, *unwind, EdgeKind::Unwind);
                }
            }
            TerminatorKind::DropAndReplace { target, unwind, .. } => {
                if self.mir_phase > MirPhase::DropLowering {
                    self.fail(
                        location,
                        "`DropAndReplace` is not permitted to exist after drop elaboration",
                    );
                }
                self.check_edge(location, *target, EdgeKind::Normal);
                if let Some(unwind) = unwind {
                    self.check_edge(location, *unwind, EdgeKind::Unwind);
                }
            }
            TerminatorKind::Call { func, destination, cleanup, .. } => {
                let func_ty = func.ty(&self.body.local_decls, self.tcx);
                match func_ty.kind() {
                    ty::FnPtr(..) | ty::FnDef(..) => {}
                    _ => self.fail(
                        location,
                        format!("encountered non-callable type {} in `Call` terminator", func_ty),
                    ),
                }
                if let Some((_, target)) = destination {
                    self.check_edge(location, *target, EdgeKind::Normal);
                }
                if let Some(cleanup) = cleanup {
                    self.check_edge(location, *cleanup, EdgeKind::Unwind);
                }
            }
            TerminatorKind::Assert { cond, target, cleanup, .. } => {
                let cond_ty = cond.ty(&self.body.local_decls, self.tcx);
                if cond_ty != self.tcx.types.bool {
                    self.fail(
                        location,
                        format!(
                            "encountered non-boolean condition of type {} in `Assert` terminator",
                            cond_ty
                        ),
                    );
                }
                self.check_edge(location, *target, EdgeKind::Normal);
                if let Some(cleanup) = cleanup {
                    self.check_edge(location, *cleanup, EdgeKind::Unwind);
                }
            }
            TerminatorKind::Yield { resume, drop, .. } => {
                if self.mir_phase > MirPhase::GeneratorLowering {
                    self.fail(location, "`Yield` should have been replaced by generator lowering");
                }
                self.check_edge(location, *resume, EdgeKind::Normal);
                if let Some(drop) = drop {
                    self.check_edge(location, *drop, EdgeKind::Normal);
                }
            }
            TerminatorKind::FalseEdge { real_target, imaginary_target } => {
                self.check_edge(location, *real_target, EdgeKind::Normal);
                self.check_edge(location, *imaginary_target, EdgeKind::Normal);
            }
            TerminatorKind::FalseUnwind { real_target, unwind } => {
                self.check_edge(location, *real_target, EdgeKind::Normal);
                if let Some(unwind) = unwind {
                    self.check_edge(location, *unwind, EdgeKind::Unwind);
                }
            }
            TerminatorKind::InlineAsm { destination, .. } => {
                if let Some(destination) = destination {
                    self.check_edge(location, *destination, EdgeKind::Normal);
                }
            }
            // Nothing to validate for these.
            TerminatorKind::Resume
            | TerminatorKind::Abort
            | TerminatorKind::Return
            | TerminatorKind::Unreachable
            | TerminatorKind::GeneratorDrop => {}
        }
    }
}
<commit_msg>validate: storage must be allocated on local use<commit_after>//! Validates the MIR to ensure that invariants are upheld.

use crate::{
    dataflow::impls::MaybeStorageLive, dataflow::Analysis, dataflow::ResultsCursor,
    util::storage::AlwaysLiveLocals,
};

use super::{MirPass, MirSource};
use rustc_middle::mir::{visit::PlaceContext, visit::Visitor, Local};
use rustc_middle::{
    mir::{
        AggregateKind, BasicBlock, Body, BorrowKind, Location, MirPhase, Operand, Rvalue,
        Statement, StatementKind, Terminator, TerminatorKind,
    },
    ty::{
        self,
        relate::{Relate, RelateResult, TypeRelation},
        ParamEnv, Ty, TyCtxt,
    },
};

#[derive(Copy, Clone, Debug)]
enum EdgeKind {
    Unwind,
    Normal,
}

pub struct Validator {
    /// Describes at which point in the pipeline this validation is happening.
    pub when: String,
    /// The phase for which we are upholding the dialect. If the given phase forbids a specific
    /// element, this validator will now emit errors if that specific element is encountered.
    /// Note that phases that change the dialect cause all *following* phases to check the
    /// invariants of the new dialect. A phase that changes dialects never checks the new invariants
    /// itself.
    pub mir_phase: MirPhase,
}

impl<'tcx> MirPass<'tcx> for Validator {
    fn run_pass(&self, tcx: TyCtxt<'tcx>, source: MirSource<'tcx>, body: &mut Body<'tcx>) {
        let def_id = source.def_id();
        let param_env = tcx.param_env(def_id);
        let mir_phase = self.mir_phase;

        let always_live_locals = AlwaysLiveLocals::new(body);
        let storage_liveness = MaybeStorageLive::new(always_live_locals)
            .into_engine(tcx, body, def_id)
            .iterate_to_fixpoint()
            .into_results_cursor(body);

        TypeChecker { when: &self.when, source, body, tcx, param_env, mir_phase, storage_liveness }
            .visit_body(body);
    }
}

/// Returns whether the two types are equal up to lifetimes.
/// All lifetimes, including higher-ranked ones, get ignored for this comparison.
/// (This is unlike the `erasing_regions` methods, which keep higher-ranked lifetimes for soundness reasons.)
///
/// The point of this function is to approximate "equal up to subtyping".  However,
/// the approximation is incorrect as variance is ignored.
pub fn equal_up_to_regions(
    tcx: TyCtxt<'tcx>,
    param_env: ParamEnv<'tcx>,
    src: Ty<'tcx>,
    dest: Ty<'tcx>,
) -> bool {
    // Fast path.
    if src == dest {
        return true;
    }

    struct LifetimeIgnoreRelation<'tcx> {
        tcx: TyCtxt<'tcx>,
        param_env: ty::ParamEnv<'tcx>,
    }

    impl TypeRelation<'tcx> for LifetimeIgnoreRelation<'tcx> {
        fn tcx(&self) -> TyCtxt<'tcx> {
            self.tcx
        }

        fn param_env(&self) -> ty::ParamEnv<'tcx> {
            self.param_env
        }

        fn tag(&self) -> &'static str {
            "librustc_mir::transform::validate"
        }

        fn a_is_expected(&self) -> bool {
            true
        }

        fn relate_with_variance<T: Relate<'tcx>>(
            &mut self,
            _: ty::Variance,
            a: T,
            b: T,
        ) -> RelateResult<'tcx, T> {
            // Ignore variance, require types to be exactly the same.
            self.relate(a, b)
        }

        fn tys(&mut self, a: Ty<'tcx>, b: Ty<'tcx>) -> RelateResult<'tcx, Ty<'tcx>> {
            if a == b {
                // Short-circuit.
                return Ok(a);
            }
            ty::relate::super_relate_tys(self, a, b)
        }

        fn regions(
            &mut self,
            a: ty::Region<'tcx>,
            _b: ty::Region<'tcx>,
        ) -> RelateResult<'tcx, ty::Region<'tcx>> {
            // Ignore regions.
            Ok(a)
        }

        fn consts(
            &mut self,
            a: &'tcx ty::Const<'tcx>,
            b: &'tcx ty::Const<'tcx>,
        ) -> RelateResult<'tcx, &'tcx ty::Const<'tcx>> {
            ty::relate::super_relate_consts(self, a, b)
        }

        fn binders<T>(
            &mut self,
            a: ty::Binder<T>,
            b: ty::Binder<T>,
        ) -> RelateResult<'tcx, ty::Binder<T>>
        where
            T: Relate<'tcx>,
        {
            self.relate(a.skip_binder(), b.skip_binder())?;
            Ok(a)
        }
    }

    // Instantiate and run relation.
    let mut relator: LifetimeIgnoreRelation<'tcx> = LifetimeIgnoreRelation { tcx: tcx, param_env };
    relator.relate(src, dest).is_ok()
}

struct TypeChecker<'a, 'tcx> {
    when: &'a str,
    source: MirSource<'tcx>,
    body: &'a Body<'tcx>,
    tcx: TyCtxt<'tcx>,
    param_env: ParamEnv<'tcx>,
    mir_phase: MirPhase,
    storage_liveness: ResultsCursor<'a, 'tcx, MaybeStorageLive>,
}

impl<'a, 'tcx> TypeChecker<'a, 'tcx> {
    fn fail(&self, location: Location, msg: impl AsRef<str>) {
        let span = self.body.source_info(location).span;
        // We use `delay_span_bug` as we might see broken MIR when other errors have already
        // occurred.
        self.tcx.sess.diagnostic().delay_span_bug(
            span,
            &format!(
                "broken MIR in {:?} ({}) at {:?}:\n{}",
                self.source.instance,
                self.when,
                location,
                msg.as_ref()
            ),
        );
    }

    fn check_edge(&self, location: Location, bb: BasicBlock, edge_kind: EdgeKind) {
        if let Some(bb) = self.body.basic_blocks().get(bb) {
            let src = self.body.basic_blocks().get(location.block).unwrap();
            match (src.is_cleanup, bb.is_cleanup, edge_kind) {
                // Non-cleanup blocks can jump to non-cleanup blocks along non-unwind edges
                (false, false, EdgeKind::Normal)
                // Non-cleanup blocks can jump to cleanup blocks along unwind edges
                | (false, true, EdgeKind::Unwind)
                // Cleanup blocks can jump to cleanup blocks along non-unwind edges
                | (true, true, EdgeKind::Normal) => {}
                // All other jumps are invalid
                _ => {
                    self.fail(
                        location,
                        format!(
                            "{:?} edge to {:?} violates unwind invariants (cleanup {:?} -> {:?})",
                            edge_kind,
                            bb,
                            src.is_cleanup,
                            bb.is_cleanup,
                        )
                    )
                }
            }
        } else {
            self.fail(location, format!("encountered jump to invalid basic block {:?}", bb))
        }
    }

    /// Check if src can be assigned into dest.
    /// This is not precise, it will accept some incorrect assignments.
    fn mir_assign_valid_types(&self, src: Ty<'tcx>, dest: Ty<'tcx>) -> bool {
        // Fast path before we normalize.
        if src == dest {
            // Equal types, all is good.
            return true;
        }
        // Normalize projections and things like that.
        // FIXME: We need to reveal_all, as some optimizations change types in ways
        // that require unfolding opaque types.
        let param_env = self.param_env.with_reveal_all_normalized(self.tcx);
        let src = self.tcx.normalize_erasing_regions(param_env, src);
        let dest = self.tcx.normalize_erasing_regions(param_env, dest);

        // Type-changing assignments can happen when subtyping is used. While
        // all normal lifetimes are erased, higher-ranked types with their
        // late-bound lifetimes are still around and can lead to type
        // differences. So we compare ignoring lifetimes.
        equal_up_to_regions(self.tcx, param_env, src, dest)
    }
}

impl<'a, 'tcx> Visitor<'tcx> for TypeChecker<'a, 'tcx> {
    fn visit_local(&mut self, local: &Local, context: PlaceContext, location: Location) {
        if context.is_use() {
            // Uses of locals must occur while the local's storage is allocated.
            self.storage_liveness.seek_after_primary_effect(location);
            let locals_with_storage = self.storage_liveness.get();
            if !locals_with_storage.contains(*local) {
                self.fail(location, format!("use of local {:?}, which has no storage here", local));
            }
        }
    }

    fn visit_operand(&mut self, operand: &Operand<'tcx>, location: Location) {
        // `Operand::Copy` is only supposed to be used with `Copy` types.
        if let Operand::Copy(place) = operand {
            let ty = place.ty(&self.body.local_decls, self.tcx).ty;
            let span = self.body.source_info(location).span;

            if !ty.is_copy_modulo_regions(self.tcx.at(span), self.param_env) {
                self.fail(location, format!("`Operand::Copy` with non-`Copy` type {}", ty));
            }
        }

        self.super_operand(operand, location);
    }

    fn visit_statement(&mut self, statement: &Statement<'tcx>, location: Location) {
        match &statement.kind {
            StatementKind::Assign(box (dest, rvalue)) => {
                // LHS and RHS of the assignment must have the same type.
                let left_ty = dest.ty(&self.body.local_decls, self.tcx).ty;
                let right_ty = rvalue.ty(&self.body.local_decls, self.tcx);
                if !self.mir_assign_valid_types(right_ty, left_ty) {
                    self.fail(
                        location,
                        format!(
                            "encountered `{:?}` with incompatible types:\n\
                            left-hand side has type: {}\n\
                            right-hand side has type: {}",
                            statement.kind, left_ty, right_ty,
                        ),
                    );
                }
                match rvalue {
                    // The sides of an assignment must not alias. Currently this just checks whether the places
                    // are identical.
                    Rvalue::Use(Operand::Copy(src) | Operand::Move(src)) => {
                        if dest == src {
                            self.fail(
                                location,
                                "encountered `Assign` statement with overlapping memory",
                            );
                        }
                    }
                    // The deaggregator currently does not deaggreagate arrays.
                    // So for now, we ignore them here.
                    Rvalue::Aggregate(box AggregateKind::Array { .. }, _) => {}
                    // All other aggregates must be gone after some phases.
                    Rvalue::Aggregate(box kind, _) => {
                        if self.mir_phase > MirPhase::DropLowering
                            && !matches!(kind, AggregateKind::Generator(..))
                        {
                            // Generators persist until the state machine transformation, but all
                            // other aggregates must have been lowered.
                            self.fail(
                                location,
                                format!("{:?} have been lowered to field assignments", rvalue),
                            )
                        } else if self.mir_phase > MirPhase::GeneratorLowering {
                            // No more aggregates after drop and generator lowering.
                            self.fail(
                                location,
                                format!("{:?} have been lowered to field assignments", rvalue),
                            )
                        }
                    }
                    Rvalue::Ref(_, BorrowKind::Shallow, _) => {
                        if self.mir_phase > MirPhase::DropLowering {
                            self.fail(
                                location,
                                "`Assign` statement with a `Shallow` borrow should have been removed after drop lowering phase",
                            );
                        }
                    }
                    _ => {}
                }
            }
            StatementKind::AscribeUserType(..) => {
                if self.mir_phase > MirPhase::DropLowering {
                    self.fail(
                        location,
                        "`AscribeUserType` should have been removed after drop lowering phase",
                    );
                }
            }
            StatementKind::FakeRead(..) => {
                if self.mir_phase > MirPhase::DropLowering {
                    self.fail(
                        location,
                        "`FakeRead` should have been removed after drop lowering phase",
                    );
                }
            }
            _ => {}
        }
    }

    fn visit_terminator(&mut self, terminator: &Terminator<'tcx>, location: Location) {
        match &terminator.kind {
            TerminatorKind::Goto { target } => {
                self.check_edge(location, *target, EdgeKind::Normal);
            }
            TerminatorKind::SwitchInt { targets, values, switch_ty, discr } => {
                let ty = discr.ty(&self.body.local_decls, self.tcx);
                if ty != *switch_ty {
                    self.fail(
                        location,
                        format!(
                            "encountered `SwitchInt` terminator with type mismatch: {:?} != {:?}",
                            ty, switch_ty,
                        ),
                    );
                }
                if targets.len() != values.len() + 1 {
                    self.fail(
                        location,
                        format!(
                            "encountered `SwitchInt` terminator with {} values, but {} targets (should be values+1)",
                            values.len(),
                            targets.len(),
                        ),
                    );
                }
                for target in targets {
                    self.check_edge(location, *target, EdgeKind::Normal);
                }
            }
            TerminatorKind::Drop { target, unwind, .. } => {
                self.check_edge(location, *target, EdgeKind::Normal);
                if let Some(unwind) = unwind {
                    self.check_edge(location, *unwind, EdgeKind::Unwind);
                }
            }
            TerminatorKind::DropAndReplace { target, unwind, .. } => {
                if self.mir_phase > MirPhase::DropLowering {
                    self.fail(
                        location,
                        "`DropAndReplace` is not permitted to exist after drop elaboration",
                    );
                }
                self.check_edge(location, *target, EdgeKind::Normal);
                if let Some(unwind) = unwind {
                    self.check_edge(location, *unwind, EdgeKind::Unwind);
                }
            }
            TerminatorKind::Call { func, destination, cleanup, .. } => {
                let func_ty = func.ty(&self.body.local_decls, self.tcx);
                match func_ty.kind() {
                    ty::FnPtr(..) | ty::FnDef(..) => {}
                    _ => self.fail(
                        location,
                        format!("encountered non-callable type {} in `Call` terminator", func_ty),
                    ),
                }
                if let Some((_, target)) = destination {
                    self.check_edge(location, *target, EdgeKind::Normal);
                }
                if let Some(cleanup) = cleanup {
                    self.check_edge(location, *cleanup, EdgeKind::Unwind);
                }
            }
            TerminatorKind::Assert { cond, target, cleanup, .. } => {
                let cond_ty = cond.ty(&self.body.local_decls, self.tcx);
                if cond_ty != self.tcx.types.bool {
                    self.fail(
                        location,
                        format!(
                            "encountered non-boolean condition of type {} in `Assert` terminator",
                            cond_ty
                        ),
                    );
                }
                self.check_edge(location, *target, EdgeKind::Normal);
                if let Some(cleanup) = cleanup {
                    self.check_edge(location, *cleanup, EdgeKind::Unwind);
                }
            }
            TerminatorKind::Yield { resume, drop, .. } => {
                if self.mir_phase > MirPhase::GeneratorLowering {
                    self.fail(location, "`Yield` should have been replaced by generator lowering");
                }
                self.check_edge(location, *resume, EdgeKind::Normal);
                if let Some(drop) = drop {
                    self.check_edge(location, *drop, EdgeKind::Normal);
                }
            }
            TerminatorKind::FalseEdge { real_target, imaginary_target } => {
                self.check_edge(location, *real_target, EdgeKind::Normal);
                self.check_edge(location, *imaginary_target, EdgeKind::Normal);
            }
            TerminatorKind::FalseUnwind { real_target, unwind } => {
                self.check_edge(location, *real_target, EdgeKind::Normal);
                if let Some(unwind) = unwind {
                    self.check_edge(location, *unwind, EdgeKind::Unwind);
                }
            }
            TerminatorKind::InlineAsm { destination, .. } => {
                if let Some(destination) = destination {
                    self.check_edge(location, *destination, EdgeKind::Normal);
                }
            }
            // Nothing to validate for these.
            TerminatorKind::Resume
            | TerminatorKind::Abort
            | TerminatorKind::Return
            | TerminatorKind::Unreachable
            | TerminatorKind::GeneratorDrop => {}
        }
    }
}
<|endoftext|>


<commit_before>extern crate desync;
extern crate futures;

use desync::*;
use futures::stream;
use futures::executor;
use futures::sync::mpsc;

use std::sync::*;
use std::thread;
use std::time::Duration;

#[test]
fn pipe_in_simple_stream() {
    // Create a stream
    let stream  = vec![1, 2, 3];
    let stream  = stream::iter_ok(stream);

    // Create an object for the stream to be piped into
    let obj     = Arc::new(Desync::new(vec![]));

    // Pipe the stream into the object
    pipe_in(Arc::clone(&obj), stream, |core: &mut Vec<Result<i32, ()>>, item| core.push(item));

    // Once the stream is drained, the core should contain Ok(1), Ok(2), Ok(3)
    assert!(obj.sync(|core| core.clone()) == vec![Ok(1), Ok(2), Ok(3)])
}

#[test]
fn pipe_in_mpsc_receiver() {
    // Create a channel to send to the object
    let (sender, receiver) = mpsc::channel(10);

    // Create an object
    let obj = Arc::new(Desync::new(vec![]));

    // Add anything received to the vector via a pipe
    pipe_in(Arc::clone(&obj), receiver, |core, item| core.push(item.unwrap()));

    // Initially empty
    assert!(obj.sync(|core| core.clone()) == vec![]);

    // Send some values
    let mut sender = executor::spawn(sender);
    sender.wait_send(1).unwrap();
    sender.wait_send(2).unwrap();

    // Should be available on the core
    assert!(obj.sync(|core| core.clone()) == vec![1, 2]);
}

#[test]
fn pipe_through() {
    // Create a channel we'll use to send data to the pipe
    let (sender, receiver) = mpsc::channel(10);

    // Create an object to pipe through
    let obj = Arc::new(Desync::new(1));

    // Create a pipe that adds values from the stream to the value in the object
    let pipe_out = pipe(Arc::clone(&obj), receiver, |core, item: Result<i32, ()>| item.map(|item| item + *core));

    // Start things running
    let mut sender      = executor::spawn(sender);
    let mut pipe_out    = executor::spawn(pipe_out);

    // Sending a value should add whatever is in the desync object
    sender.wait_send(2).unwrap();
    assert!(pipe_out.wait_stream() == Some(Ok(3)));

    sender.wait_send(42).unwrap();
    assert!(pipe_out.wait_stream() == Some(Ok(43)));

    // Changing the value should change the output
    obj.async(|core| *core = 2);

    sender.wait_send(42).unwrap();
    assert!(pipe_out.wait_stream() == Some(Ok(44)));
}

#[test]
fn pipe_through_stream_closes() {
    let mut pipe_out_with_closed_stream = {
        // Create a channel we'll use to send data to the pipe
        let (sender, receiver) = mpsc::channel(10);

        // Create an object to pipe through
        let obj = Arc::new(Desync::new(1));

        // Create a pipe that adds values from the stream to the value in the object
        let pipe_out = pipe(Arc::clone(&obj), receiver, |core, item: Result<i32, ()>| item.map(|item| item + *core));

        // Start things running
        let mut sender      = executor::spawn(sender);
        let mut pipe_out    = executor::spawn(pipe_out);

        // Sending a value should add whatever is in the desync object
        sender.wait_send(2).unwrap();
        assert!(pipe_out.wait_stream() == Some(Ok(3)));

        pipe_out
    };

    // The sender is now closed (the sender and receiver are dropped after the block above), so the pipe should close too
    assert!(pipe_out_with_closed_stream.wait_stream() == None);
}

#[test]
fn pipe_through_produces_backpressure() {
    // Create a channel we'll use to send data to the pipe
    let (mut sender, receiver) = mpsc::channel(0);

    // Create an object to pipe through
    let obj = Arc::new(Desync::new(1));

    // Create a pipe that adds values from the stream to the value in the object
    let mut pipe_out    = pipe(Arc::clone(&obj), receiver, |core, item: Result<i32, ()>| item.map(|item| item + *core));

    // Set the backpressure depth to 3
    pipe_out.set_backpressure_depth(3);

    // Start things running. We never read from this pipe here
    let _pipe_out       = executor::spawn(pipe_out);

    // Send 3 events to the pipe. Wait a bit between them to allow for processing time
    for _x in 0..3 {
        assert!(sender.try_send(1) == Ok(()));

        // The wait here allows the message to flow through to the pipe (if we call try_send again before the pipe has a chance to accept the input)
        thread::sleep(Duration::from_millis(5));
    }

    // This will stick in the channel (pipe should not be accepting more input)
    assert!(sender.try_send(2) == Ok(()));
    thread::sleep(Duration::from_millis(5));

    // Channel will push back on this one
    let channel_full = sender.try_send(3);
    assert!(channel_full.is_err());
    assert!(channel_full.unwrap_err().is_full());
}<commit_msg>As the processing is now in the background, these tests need to wait for a while<commit_after>extern crate desync;
extern crate futures;

use desync::*;
use futures::stream;
use futures::executor;
use futures::sync::mpsc;

use std::sync::*;
use std::thread;
use std::time::Duration;

#[test]
fn pipe_in_simple_stream() {
    // Create a stream
    let stream  = vec![1, 2, 3];
    let stream  = stream::iter_ok(stream);

    // Create an object for the stream to be piped into
    let obj     = Arc::new(Desync::new(vec![]));

    // Pipe the stream into the object
    pipe_in(Arc::clone(&obj), stream, |core: &mut Vec<Result<i32, ()>>, item| core.push(item));

    // Delay to allow the messages to be processed on the stream
    thread::sleep(Duration::from_millis(10));

    // Once the stream is drained, the core should contain Ok(1), Ok(2), Ok(3)
    assert!(obj.sync(|core| core.clone()) == vec![Ok(1), Ok(2), Ok(3)])
}

#[test]
fn pipe_in_mpsc_receiver() {
    // Create a channel to send to the object
    let (sender, receiver) = mpsc::channel(0);

    // Create an object
    let obj = Arc::new(Desync::new(vec![]));

    // Add anything received to the vector via a pipe
    pipe_in(Arc::clone(&obj), receiver, |core, item| core.push(item.unwrap()));

    // Initially empty
    assert!(obj.sync(|core| core.clone()) == vec![]);

    // Send some values
    let mut sender = executor::spawn(sender);
    sender.wait_send(1).unwrap();
    sender.wait_send(2).unwrap();

    // Delay to allow the messages to be processed on the stream
    // TODO: fix so this isn't needed. This happens because there's a race between when 'poll'
    // is called in the pipe and the 'async' call
    thread::sleep(Duration::from_millis(10));

    // Should be available on the core
    assert!(obj.sync(|core| core.clone()) == vec![1, 2]);
}

#[test]
fn pipe_through() {
    // Create a channel we'll use to send data to the pipe
    let (sender, receiver) = mpsc::channel(10);

    // Create an object to pipe through
    let obj = Arc::new(Desync::new(1));

    // Create a pipe that adds values from the stream to the value in the object
    let pipe_out = pipe(Arc::clone(&obj), receiver, |core, item: Result<i32, ()>| item.map(|item| item + *core));

    // Start things running
    let mut sender      = executor::spawn(sender);
    let mut pipe_out    = executor::spawn(pipe_out);

    // Sending a value should add whatever is in the desync object
    sender.wait_send(2).unwrap();
    assert!(pipe_out.wait_stream() == Some(Ok(3)));

    sender.wait_send(42).unwrap();
    assert!(pipe_out.wait_stream() == Some(Ok(43)));

    // Changing the value should change the output
    obj.async(|core| *core = 2);

    sender.wait_send(42).unwrap();
    assert!(pipe_out.wait_stream() == Some(Ok(44)));
}

#[test]
fn pipe_through_stream_closes() {
    let mut pipe_out_with_closed_stream = {
        // Create a channel we'll use to send data to the pipe
        let (sender, receiver) = mpsc::channel(10);

        // Create an object to pipe through
        let obj = Arc::new(Desync::new(1));

        // Create a pipe that adds values from the stream to the value in the object
        let pipe_out = pipe(Arc::clone(&obj), receiver, |core, item: Result<i32, ()>| item.map(|item| item + *core));

        // Start things running
        let mut sender      = executor::spawn(sender);
        let mut pipe_out    = executor::spawn(pipe_out);

        // Sending a value should add whatever is in the desync object
        sender.wait_send(2).unwrap();
        assert!(pipe_out.wait_stream() == Some(Ok(3)));

        pipe_out
    };

    // The sender is now closed (the sender and receiver are dropped after the block above), so the pipe should close too
    assert!(pipe_out_with_closed_stream.wait_stream() == None);
}

#[test]
fn pipe_through_produces_backpressure() {
    // Create a channel we'll use to send data to the pipe
    let (mut sender, receiver) = mpsc::channel(0);

    // Create an object to pipe through
    let obj = Arc::new(Desync::new(1));

    // Create a pipe that adds values from the stream to the value in the object
    let mut pipe_out    = pipe(Arc::clone(&obj), receiver, |core, item: Result<i32, ()>| item.map(|item| item + *core));

    // Set the backpressure depth to 3
    pipe_out.set_backpressure_depth(3);

    // Start things running. We never read from this pipe here
    let _pipe_out       = executor::spawn(pipe_out);

    // Send 3 events to the pipe. Wait a bit between them to allow for processing time
    for _x in 0..3 {
        assert!(sender.try_send(1) == Ok(()));

        // The wait here allows the message to flow through to the pipe (if we call try_send again before the pipe has a chance to accept the input)
        thread::sleep(Duration::from_millis(5));
    }

    // This will stick in the channel (pipe should not be accepting more input)
    assert!(sender.try_send(2) == Ok(()));
    thread::sleep(Duration::from_millis(5));

    // Channel will push back on this one
    let channel_full = sender.try_send(3);
    assert!(channel_full.is_err());
    assert!(channel_full.unwrap_err().is_full());
}<|endoftext|>


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

pub use crate::proto::schema::tests::TimestampTx;

use exonum_merkledb::{impl_binary_value_for_message, BinaryValue};
use protobuf::Message as PbMessage;
use rand::{RngCore, SeedableRng};
use rand_xorshift::XorShiftRng;
use semver::Version;

use std::borrow::Cow;

use crate::{
    blockchain::ExecutionResult,
    crypto::{gen_keypair, PublicKey, SecretKey},
    messages::{AnyTx, Message, ServiceInstanceId, Signed},
    runtime::{
        dispatcher::BuiltinService,
        rust::{RustArtifactSpec, Service, ServiceFactory, TransactionContext},
    },
};

pub const DATA_SIZE: usize = 64;

#[service_interface(exonum(crate = "crate"))]
pub trait TimestampingInterface {
    fn timestamp(&self, context: TransactionContext, arg: TimestampTx) -> ExecutionResult;
}

#[derive(Debug)]
pub struct TimestampingService;

impl_service_dispatcher!(TimestampingService, TimestampingInterface);

impl TimestampingInterface for TimestampingService {
    fn timestamp(&self, _context: TransactionContext, _arg: TimestampTx) -> ExecutionResult {
        Ok(())
    }
}

impl Service for TimestampingService {}

impl ServiceFactory for TimestampingService {
    fn artifact(&self) -> RustArtifactSpec {
        RustArtifactSpec {
            name: "timestamping".into(),
            version: Version::new(0, 1, 0),
        }
    }

    fn new_instance(&self) -> Box<dyn Service> {
        Box::new(Self)
    }
}

impl TimestampingService {
    pub const ID: ServiceInstanceId = 1;
}

impl From<TimestampingService> for BuiltinService {
    fn from(factory: TimestampingService) -> Self {
        Self {
            factory: factory.into(),
            instance_id: TimestampingService::ID,
            instance_name: "timestamping".into(),
        }
    }
}

impl_binary_value_for_message! { TimestampTx }

pub struct TimestampingTxGenerator {
    rand: XorShiftRng,
    data_size: usize,
    public_key: PublicKey,
    secret_key: SecretKey,
}

impl TimestampingTxGenerator {
    pub fn new(data_size: usize) -> TimestampingTxGenerator {
        let keypair = gen_keypair();
        TimestampingTxGenerator::with_keypair(data_size, keypair)
    }

    pub fn with_keypair(
        data_size: usize,
        keypair: (PublicKey, SecretKey),
    ) -> TimestampingTxGenerator {
        let rand = XorShiftRng::from_seed([9; 16]);

        TimestampingTxGenerator {
            rand,
            data_size,
            public_key: keypair.0,
            secret_key: keypair.1,
        }
    }
}

impl Iterator for TimestampingTxGenerator {
    type Item = Signed<AnyTx>;

    fn next(&mut self) -> Option<Signed<AnyTx>> {
        let mut data = vec![0; self.data_size];
        self.rand.fill_bytes(&mut data);
        let mut buf = TimestampTx::new();
        buf.set_data(data);
        Some(Message::sign_transaction(
            buf,
            TimestampingService::ID,
            self.public_key,
            &self.secret_key,
        ))
    }
}
<commit_msg>Fix last sandbox test<commit_after>// Copyright 2019 The Exonum Team
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

pub use crate::proto::schema::tests::TimestampTx;

use exonum_merkledb::{impl_binary_value_for_message, BinaryValue, Snapshot};
use protobuf::Message as PbMessage;
use rand::{RngCore, SeedableRng};
use rand_xorshift::XorShiftRng;
use semver::Version;

use std::borrow::Cow;

use crate::{
    blockchain::ExecutionResult,
    crypto::{gen_keypair, Hash, PublicKey, SecretKey, HASH_SIZE},
    messages::{AnyTx, Message, ServiceInstanceId, Signed},
    runtime::{
        dispatcher::BuiltinService,
        rust::{RustArtifactSpec, Service, ServiceFactory, TransactionContext},
    },
};

pub const DATA_SIZE: usize = 64;

#[service_interface(exonum(crate = "crate"))]
pub trait TimestampingInterface {
    fn timestamp(&self, context: TransactionContext, arg: TimestampTx) -> ExecutionResult;
}

#[derive(Debug)]
pub struct TimestampingService;

impl_service_dispatcher!(TimestampingService, TimestampingInterface);

impl TimestampingInterface for TimestampingService {
    fn timestamp(&self, _context: TransactionContext, _arg: TimestampTx) -> ExecutionResult {
        Ok(())
    }
}

impl Service for TimestampingService {
    fn state_hash(&self, _: &dyn Snapshot) -> Vec<Hash> {
        vec![Hash::new([127; HASH_SIZE]), Hash::new([128; HASH_SIZE])]
    }
}

impl ServiceFactory for TimestampingService {
    fn artifact(&self) -> RustArtifactSpec {
        RustArtifactSpec {
            name: "timestamping".into(),
            version: Version::new(0, 1, 0),
        }
    }

    fn new_instance(&self) -> Box<dyn Service> {
        Box::new(Self)
    }
}

impl TimestampingService {
    pub const ID: ServiceInstanceId = 1;
}

impl From<TimestampingService> for BuiltinService {
    fn from(factory: TimestampingService) -> Self {
        Self {
            factory: factory.into(),
            instance_id: TimestampingService::ID,
            instance_name: "timestamping".into(),
        }
    }
}

impl_binary_value_for_message! { TimestampTx }

pub struct TimestampingTxGenerator {
    rand: XorShiftRng,
    data_size: usize,
    public_key: PublicKey,
    secret_key: SecretKey,
}

impl TimestampingTxGenerator {
    pub fn new(data_size: usize) -> TimestampingTxGenerator {
        let keypair = gen_keypair();
        TimestampingTxGenerator::with_keypair(data_size, keypair)
    }

    pub fn with_keypair(
        data_size: usize,
        keypair: (PublicKey, SecretKey),
    ) -> TimestampingTxGenerator {
        let rand = XorShiftRng::from_seed([9; 16]);

        TimestampingTxGenerator {
            rand,
            data_size,
            public_key: keypair.0,
            secret_key: keypair.1,
        }
    }
}

impl Iterator for TimestampingTxGenerator {
    type Item = Signed<AnyTx>;

    fn next(&mut self) -> Option<Signed<AnyTx>> {
        let mut data = vec![0; self.data_size];
        self.rand.fill_bytes(&mut data);
        let mut buf = TimestampTx::new();
        buf.set_data(data);
        Some(Message::sign_transaction(
            buf,
            TimestampingService::ID,
            self.public_key,
            &self.secret_key,
        ))
    }
}
<|endoftext|>


<commit_before>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2016-2017 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Initial 32-bit bootloader for x86_64
#![crate_name = "boot"]
#![crate_type = "staticlib"]
#![feature(asm)]
#![feature(lang_items)]
#![feature(naked_functions)]
#![feature(static_recursion)]
#![feature(linkage)]
#![no_std]

const TABLE_LENGTH: usize = 512;

type Table = [u64; TABLE_LENGTH];

use core::ptr;
//
// #[lang = "panic_fmt"]
// #[cold]
// unsafe extern fn panic_fmt() -> ! {
//     boot_write(b"boot panic occurred!");
//     die();
// }

extern {
    static mut pml4_table: Table;
    static mut pdp_table: Table;
    static mut pd_table: Table;
    // static mut page_table: Table;

}

#[repr(C, packed)]
pub struct Gdt { _null: u64
               , code: u64
           , ptr: GdtPointer
           }

#[repr(C, packed)]
pub struct GdtPointer { /// the length of the descriptor table
                     pub limit: u16
                   , /// pointer to the region in memory
                     /// containing the descriptor table.
                     pub base: &'static Gdt
                   }

#[repr(C, packed)]
pub struct SegmentDescriptor {
}

// TODO: this sucks please fix
#[link_name = ".gdt64"]
#[link_section = ".gdt"]
pub static GDT: Gdt
    = Gdt { _null: 0
          , code: (1<<44) | (1<<47) | (1<<41) | (1<<43) | (1<<53)
          , ptr: GdtPointer { limit: 23
                            , base: &GDT }
     };

#[inline(always)]
#[naked]
fn boot_write(s: &[u8]) {
	// use core::ptr;
    let mut offset = 0;
    unsafe {
        let vga_buf: *mut u16 = 0xb8000 as *mut u16;
    	for c in s {
    		ptr::write_volatile( vga_buf.offset(offset)
                               , 0x0200 + *c as u16);
    		offset += 1;
    	}

    }

}
macro_rules! write_char {
    ($c:expr) => {
        ptr::write_volatile(0xb8000 as *mut u16, 0x0200 + $c as u16)
    }
}

#[naked]
#[inline(always)]
pub unsafe fn create_page_tables() {
    const HUGE_PAGE_SIZE: u64 = 2 * 1024 * 1024; // 2 MiB

    //-- map the PML4 and PDP tables -----------------------------------------
    // recursive map last PML4 entry
    pml4_table[511] = (&pml4_table as *const Table as u64) | 3;
    // map first PML4 entry to PDP table
    pml4_table[0] = (&pdp_table as *const Table as u64) | 3;
    // map first PDPT entry to PD table
    pdp_table[0] = (&pd_table as *const Table as u64) | 3;

    //-- map the PD table ----------------------------------------------------
    for (number, entry) in pd_table.iter_mut().enumerate() {
        // set each PD entry equal to the start address of the page (the page
        // number times the page's size)
        let addr = number as u64 * HUGE_PAGE_SIZE;
        // with the appropriate flags (present + writable + huge)
        // TODO: do we want to do this using bitflags, or is that too
        //       heavyweight for the boot module?
        //          - eliza, 1/23/2017
        *entry = addr | 0b10000011;
    }
}

#[naked]
#[inline(always)]
pub unsafe fn set_long_mode() {
    // load PML4 addr to cr3
    asm!( "mov   cr3, $0"
        : "=r"(&pml4_table as *const Table)
        ::: "intel");
    // enable PAE flag in cr4
    asm!( "mov   eax, cr4
           or    eax, 1 << 5
           mov   cr4, eax"
        :::: "intel");
    // set the long mode bit in EFER MSR (model specific register)
    asm!( "mov   ecx, 0xC0000080
           rdmsr
           or    eax, 1 << 8
           wrmsr"
        :::: "intel");
    // enable paging in cr0
    asm!( "mov  eax, cr0
           or   eax, 1 << 31
           or   eax, 1 << 16
           mov  cr0, eax"
        :::: "intel");

    asm!("ret");
}

#[naked]
#[inline(always)] #[cold]
unsafe fn die() -> ! {
    loop { asm!("hlt"); }
}

#[cold]
#[no_mangle]
#[naked]
pub unsafe extern "C" fn _start() {
    // 0. Move the stack pointer to the top of the stack.
    asm!( "mov esp, stack_top"
        :::: "intel");
    boot_write(b"0");

    // 1. Move Multiboot info pointer to edi
    asm!("mov edi, edx" :::: "intel");
    boot_write(b"1");

    // 2. make sure the system supports SOS
    // TODO: port this from boot.asm
    boot_write(b"2");

    // 3. if everything is okay, create the page tables and start long mode
    const HUGE_PAGE_SIZE: u64 = 2 * 1024 * 1024; // 2 MiB

    //-- map the PML4 and PDP tables -----------------------------------------
    // recursive map last PML4 entry
    pml4_table[511] = (&pml4_table as *const Table as u64) | 3;
    // map first PML4 entry to PDP table
    pml4_table[0] = (&pdp_table as *const Table as u64) | 3;
    // map first PDPT entry to PD table
    pdp_table[0] = (&pd_table as *const Table as u64) | 3;

    boot_write(b"3.1");

    //-- map the PD table ----------------------------------------------------
    for (number, entry) in pd_table.iter_mut().enumerate() {
        // set each PD entry equal to the start address of the page (the page
        // number times the page's size)
        let addr = number as u64 * HUGE_PAGE_SIZE;
        // with the appropriate flags (present + writable + huge)
        // TODO: do we want to do this using bitflags, or is that too
        //       heavyweight for the boot module?
        //          - eliza, 1/23/2017
        *entry = addr | 0b10000011;
    }
    // load PML4 addr to cr3
    asm!( "mov   cr3, $0"
        :: "r"(&pml4_table as *const Table)
        :: "intel");
    boot_write(b"3.2");

    // enable PAE flag in cr4
    asm!( "mov   eax, cr4
           or    eax, 1 << 5
           mov   cr4, eax"
        :::: "intel");
    boot_write(b"3.3");

    // set the long mode bit in EFER MSR (model specific register)
    asm!( "mov   ecx, 0xC0000080
           rdmsr
           or    eax, 1 << 8
           wrmsr"
        :::: "intel");
    boot_write(b"3.4");

    // enable paging in cr0
    asm!( "mov  eax, cr0
           or   eax, 1 << 31
           or   eax, 1 << 16
           mov  cr0, eax"
        :::: "intel");
    boot_write(b"3.5");

    // 4. load the 64-bit GDT
    asm!( "lgdt ($0)"
        :: "r"(&GDT.ptr)
        :  "memory"
        );
    boot_write(b"4");

    // 5. update selectors
    // asm!("mov ax, 16
    //       mov ss, ax
    //       mov ds, ax
    //       mov es, ax"
    //      :::: "intel");
    // update selectors
    asm!("mov ax, 0x10" :::: "intel");
    boot_write(b"5.1");
    // stack selector
    asm!("mov ss, ax" :::: "intel");
    boot_write(b"5.2");
    // data selector
    asm!("mov ds, ax" :::: "intel");
     boot_write(b"5.3");
    // extra selector
    asm!("mov es, ax" :::: "intel");
    boot_write(b"5.4");

    // 6. jump to the 64-bit boot subroutine.
    // asm!("jmp $0:$1" :: "X"(&GDT.descriptors[1]), "X"(arch_init as unsafe extern "C" fn() -> !) :: "intel");
    // asm!("jmp $0:arch_init" :: "r"(&GDT.descriptors[1] as *const u64) :: "intel");
    // arch_init();
    asm!(  "jmp $0:arch_init"
        :: "r"(&GDT.code as *const _ as usize -
               &GDT as *const _ as usize)
        :: "intel");

    loop {
        boot_write(b"kernel returned unexpectedly!");
    }

}
<commit_msg>[boot] re-add functions for page table setup<commit_after>//
//  SOS: the Stupid Operating System
//  by Eliza Weisman (hi@hawkweisman.me)
//
//  Copyright (c) 2016-2017 Eliza Weisman
//  Released under the terms of the MIT license. See `LICENSE` in the root
//  directory of this repository for more information.
//
//! Initial 32-bit bootloader for x86_64
#![crate_name = "boot"]
#![crate_type = "staticlib"]
#![feature(asm)]
#![feature(lang_items)]
#![feature(naked_functions)]
#![feature(static_recursion)]
#![feature(linkage)]
#![no_std]

const TABLE_LENGTH: usize = 512;

type Table = [u64; TABLE_LENGTH];

use core::ptr;
//
// #[lang = "panic_fmt"]
// #[cold]
// unsafe extern fn panic_fmt() -> ! {
//     boot_write(b"boot panic occurred!");
//     die();
// }

#[repr(C, packed)]
pub struct Gdt { _null: u64
               , code: u64
           , ptr: GdtPointer
           }

#[repr(C, packed)]
pub struct GdtPointer { /// the length of the descriptor table
                     pub limit: u16
                   , /// pointer to the region in memory
                     /// containing the descriptor table.
                     pub base: &'static Gdt
                   }

#[repr(C, packed)]
pub struct SegmentDescriptor {
}

// TODO: this sucks please fix
#[link_name = ".gdt64"]
#[link_section = ".gdt"]
pub static GDT: Gdt
    = Gdt { _null: 0
          , code: (1<<44) | (1<<47) | (1<<41) | (1<<43) | (1<<53)
          , ptr: GdtPointer { limit: 23
                            , base: &GDT }
     };

#[inline(always)]
#[naked]
fn boot_write(s: &[u8]) {
	// use core::ptr;
    let mut offset = 0;
    unsafe {
        let vga_buf: *mut u16 = 0xb8000 as *mut u16;
    	for c in s {
    		ptr::write_volatile( vga_buf.offset(offset)
                               , 0x0200 + *c as u16);
    		offset += 1;
    	}

    }

}
macro_rules! write_char {
    ($c:expr) => {
        ptr::write_volatile(0xb8000 as *mut u16, 0x0200 + $c as u16)
    }
}

extern "C" {
    static mut pml4_table: Table;
    static mut pdp_table: Table;
    static mut pd_table: Table;    // static mut page_table: Table;
}

#[naked]
#[inline(always)]
pub unsafe fn create_page_tables() {
    // 3. if everything is okay, create the page tables and start long mode
    const HUGE_PAGE_SIZE: u64 = 2 * 1024 * 1024; // 2 MiB

    //-- map the PML4 and PDP tables -----------------------------------------
    // recursive map last PML4 entry
    pml4_table[511] = (&pml4_table as *const Table as u64) | 3;
    // map first PML4 entry to PDP table
    pml4_table[0] = (&pdp_table as *const Table as u64) | 3;
    // map first PDPT entry to PD table
    pdp_table[0] = (&pd_table as *const Table as u64) | 3;

    boot_write(b"3.1");

    //-- map the PD table ----------------------------------------------------
    for (number, entry) in pd_table.iter_mut().enumerate() {
        // set each PD entry equal to the start address of the page (the page
        // number times the page's size)
        let addr = number as u64 * HUGE_PAGE_SIZE;
        // with the appropriate flags (present + writable + huge)
        // TODO: do we want to do this using bitflags, or is that too
        //       heavyweight for the boot module?
        //          - eliza, 1/23/2017
        *entry = addr | 0b10000011;
    }
}

#[naked]
#[inline(always)]
pub unsafe fn set_long_mode() {
    // load PML4 addr to cr3
    asm!( "mov   cr3, $0"
        :: "r"(&pml4_table)
        :: "intel");
    boot_write(b"3.2");

    // enable PAE flag in cr4
    asm!( "mov   eax, cr4
           or    eax, 1 << 5
           mov   cr4, eax"
        :::: "intel");
    boot_write(b"3.3");

    // set the long mode bit in EFER MSR (model specific register)
    asm!( "mov   ecx, 0xC0000080
           rdmsr
           or    eax, 1 << 8
           wrmsr"
        :::: "intel");
    boot_write(b"3.4");

    // enable paging in cr0
    asm!( "mov  eax, cr0
           or   eax, 1 << 31
           or   eax, 1 << 16
           mov  cr0, eax"
        :::: "intel");
    boot_write(b"3.5");

}

#[cold]
#[no_mangle]
#[naked]
pub unsafe extern "C" fn _start() {
    // 0. Move the stack pointer to the top of the stack.
    asm!( "mov esp, stack_top"
        :::: "intel");
    boot_write(b"0");

    // 1. Move Multiboot info pointer to edi
    asm!("mov edi, edx" :::: "intel");
    boot_write(b"1");

    // 2. make sure the system supports SOS
    // TODO: port this from boot.asm
    boot_write(b"2");

    create_page_tables();
    set_long_mode();

    // 4. load the 64-bit GDT
    asm!( "lgdt ($0)"
        :: "r"(&GDT.ptr)
        :  "memory"
        );
    boot_write(b"4");

    // 5. update selectors
    asm!("mov ax, 0x10" :::: "intel");
    boot_write(b"5.1");
    // stack selector
    asm!("mov ss, ax" :::: "intel");
    boot_write(b"5.2");
    // data selector
    asm!("mov ds, ax" :::: "intel");
     boot_write(b"5.3");
    // extra selector
    asm!("mov es, ax" :::: "intel");
    boot_write(b"5.4");

    // 6. jump to the 64-bit boot subroutine.
    // asm!("jmp $0:$1" :: "X"(&GDT.descriptors[1]), "X"(arch_init as unsafe extern "C" fn() -> !) :: "intel");
    // asm!("jmp $0:arch_init" :: "r"(&GDT.descriptors[1] as *const u64) :: "intel");
    // arch_init();
    asm!(  "jmp $0:arch_init"
        :: "r"(&GDT.code as *const _ as usize -
               &GDT as *const _ as usize)
        :: "intel");

    loop {
        boot_write(b"kernel returned unexpectedly!");
    }

}
<|endoftext|>


<commit_before>#![cfg(feature = "yaml")]

use clap::{load_yaml, App, ValueHint};

#[test]
fn create_app_from_yaml() {
    let yaml = load_yaml!("fixtures/app.yaml");
    App::from(yaml);
}

// TODO: Uncomment to test yaml with 2 spaces https://github.com/chyh1990/yaml-rust/issues/101
// #[test]
// fn create_app_from_yaml_2spaces() {
//     let yaml = load_yaml!("fixtures/app_2space.yaml");
//     App::from(yaml);
// }

#[test]
fn help_message() {
    let yaml = load_yaml!("fixtures/app.yaml");
    let mut app = App::from(yaml);
    // Generate the full help message!
    let _ = app.try_get_matches_from_mut(Vec::<String>::new());

    let mut help_buffer = Vec::new();
    app.write_help(&mut help_buffer).unwrap();
    let help_string = String::from_utf8(help_buffer).unwrap();
    assert!(help_string
        .contains("-h, --help                prints help with a nonstandard description\n"));
}

#[test]
fn author() {
    let yaml = load_yaml!("fixtures/app.yaml");
    let mut app = App::from(yaml);
    // Generate the full help message!
    let _ = app.try_get_matches_from_mut(Vec::<String>::new());

    let mut help_buffer = Vec::new();
    app.write_help(&mut help_buffer).unwrap();
    let help_string = String::from_utf8(help_buffer).unwrap();
    assert!(help_string.contains("Kevin K. <kbknapp@gmail.com>"));
}

// ValueHint must be parsed correctly from Yaml
#[test]
fn value_hint() {
    let yml = load_yaml!("fixtures/app.yaml");
    let app = App::from(yml);

    let arg = app
        .get_arguments()
        .find(|a| a.get_name() == "value_hint")
        .unwrap();
    assert_eq!(arg.get_value_hint(), ValueHint::FilePath);
}
<commit_msg>Provide tests for Arg::default_value_if<commit_after>#![cfg(feature = "yaml")]

use clap::{load_yaml, App, ValueHint};

#[test]
fn create_app_from_yaml() {
    let yaml = load_yaml!("fixtures/app.yaml");
    App::from(yaml);
}

// TODO: Uncomment to test yaml with 2 spaces https://github.com/chyh1990/yaml-rust/issues/101
// #[test]
// fn create_app_from_yaml_2spaces() {
//     let yaml = load_yaml!("fixtures/app_2space.yaml");
//     App::from(yaml);
// }

#[test]
fn help_message() {
    let yaml = load_yaml!("fixtures/app.yaml");
    let mut app = App::from(yaml);
    // Generate the full help message!
    let _ = app.try_get_matches_from_mut(Vec::<String>::new());

    let mut help_buffer = Vec::new();
    app.write_help(&mut help_buffer).unwrap();
    let help_string = String::from_utf8(help_buffer).unwrap();
    assert!(help_string
        .contains("-h, --help                prints help with a nonstandard description\n"));
}

#[test]
fn author() {
    let yaml = load_yaml!("fixtures/app.yaml");
    let mut app = App::from(yaml);
    // Generate the full help message!
    let _ = app.try_get_matches_from_mut(Vec::<String>::new());

    let mut help_buffer = Vec::new();
    app.write_help(&mut help_buffer).unwrap();
    let help_string = String::from_utf8(help_buffer).unwrap();
    assert!(help_string.contains("Kevin K. <kbknapp@gmail.com>"));
}

// ValueHint must be parsed correctly from Yaml
#[test]
fn value_hint() {
    let yml = load_yaml!("fixtures/app.yaml");
    let app = App::from(yml);

    let arg = app
        .get_arguments()
        .find(|a| a.get_name() == "value_hint")
        .unwrap();
    assert_eq!(arg.get_value_hint(), ValueHint::FilePath);
}

#[test]
fn default_value_if_not_triggered_by_argument() {
    let yml = load_yaml!("fixtures/app.yaml");
    let app = App::from(yml);

    // Fixtures use "other" as value
    let matches = app.try_get_matches_from(vec!["prog", "wrong"]).unwrap();

    assert!(matches.value_of("positional2").is_none());
}

#[test]
fn default_value_if_triggered_by_matching_argument() {
    let yml = load_yaml!("fixtures/app.yaml");
    let app = App::from(yml);

    let matches = app.try_get_matches_from(vec!["prog", "other"]).unwrap();
    assert_eq!(matches.value_of("positional2").unwrap(), "something");
}

#[test]
fn default_value_if_triggered_by_flag() {
    let yml = load_yaml!("fixtures/app.yaml");
    let app = App::from(yml);

    let matches = app
        .try_get_matches_from(vec!["prog", "--flag", "flagvalue"])
        .unwrap();

    assert_eq!(matches.value_of("positional2").unwrap(), "some");
}

#[test]
fn default_value_if_triggered_by_flag_and_argument() {
    let yml = load_yaml!("fixtures/app.yaml");
    let app = App::from(yml);

    let matches = app
        .try_get_matches_from(vec!["prog", "--flag", "flagvalue", "other"])
        .unwrap();

    // First condition triggers, therefore "some"
    assert_eq!(matches.value_of("positional2").unwrap(), "some");
}
<|endoftext|>


<commit_before>use linked_hash_map::LinkedHashMap;
use std::collections::HashSet;

use crate::options::{ImportStyle, Options, StringTransform};
use crate::shape::{self, Shape};
use crate::to_singular::to_singular;
use crate::util::{kebab_case, lower_camel_case, snake_case, type_case};

#[derive(PartialEq, PartialOrd, Ord, Eq, Hash, Clone, Copy)]
enum Import {
    Any,
    Optional,
    BaseModel,
    Field,
}

impl Import {
    fn pair(&self) -> (&'static str, &'static str) {
        match self {
            Import::Any => ("typing", "Any"),
            Import::Optional => ("typing", "Optional"),
            Import::BaseModel => ("pydantic", "BaseModel"),
            Import::Field => ("pydantic", "Field"),
        }
    }
    fn identifier(&self) -> &'static str {
        self.pair().1
    }
    fn qualified(&self) -> String {
        let (module, identifier) = self.pair();
        format!("{}.{}", module, identifier)
    }
}

struct Ctxt {
    options: Options,
    type_names: HashSet<String>,
    imports: HashSet<Import>,
    created_classes: Vec<(Shape, Ident)>,
}

pub type Ident = String;
pub type Code = String;

pub fn python_types(name: &str, shape: &Shape, options: Options) -> Code {
    let mut ctxt = Ctxt {
        options,
        type_names: HashSet::new(),
        imports: HashSet::new(),
        created_classes: Vec::new(),
    };

    if !matches!(shape, Shape::Struct { .. }) {
        // reserve the requested name
        ctxt.type_names.insert(name.to_string());
    }

    let (ident, code) = type_from_shape(&mut ctxt, name, shape);
    let mut code = code.unwrap_or_default();

    if !ctxt.imports.is_empty() {
        let mut imports: Vec<_> = ctxt.imports.drain().collect();
        imports.sort();
        let mut import_code = String::new();
        for import in imports {
            match ctxt.options.import_style {
                ImportStyle::AssumeExisting => {}
                ImportStyle::AddImports => {
                    let (module, identifier) = import.pair();
                    import_code += &format!("from {} import {}\n", module, identifier);
                }
                ImportStyle::QualifiedPaths => {
                    let (module, identifier) = import.pair();
                    import_code += &format!("import {}.{}\n", module, identifier);
                }
            }
        }
        import_code += "\n\n";
        code = import_code + &code;
    }

    if ident != name {
        code += &format!("\n\n{} = {}", name, ident);
    }
    code
}

fn type_from_shape(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    use crate::shape::Shape::*;
    match shape {
        Null | Any | Bottom => (import(ctxt, Import::Any), None),
        Bool => ("bool".into(), None),
        StringT => ("str".into(), None),
        Integer => ("int".into(), None),
        Floating => ("float".into(), None),
        Tuple(shapes, _n) => {
            let folded = shape::fold_shapes(shapes.clone());
            if shapes.len() <= 3 && folded == Any && shapes.iter().any(|s| s != &Any) {
                generate_tuple_type(ctxt, path, shapes)
            } else {
                generate_vec_type(ctxt, path, &folded)
            }
        }
        VecT { elem_type: e } => generate_vec_type(ctxt, path, e),
        Struct { fields } => generate_data_class(ctxt, path, fields, shape),
        MapT { val_type: v } => generate_map_type(ctxt, path, v),
        Opaque(t) => (t.clone(), None),
        Optional(e) => {
            let (inner, defs) = type_from_shape(ctxt, path, e);
            if ctxt.options.use_default_for_missing_fields {
                (inner, defs)
            } else {
                let optional = import(ctxt, Import::Optional);
                (format!("{}[{}]", optional, inner), defs)
            }
        }
    }
}

fn generate_vec_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = to_singular(path);
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (format!("list[{}]", inner), defs)
}

fn generate_map_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = to_singular(path);
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (format!("dict[str, {}]", inner), defs)
}

fn generate_tuple_type(ctxt: &mut Ctxt, path: &str, shapes: &[Shape]) -> (Ident, Option<Code>) {
    let mut types = Vec::new();
    let mut defs = Vec::new();

    for shape in shapes {
        let (typ, def) = type_from_shape(ctxt, path, shape);
        types.push(typ);
        if let Some(code) = def {
            defs.push(code)
        }
    }

    (
        format!("tuple[{}]", types.join(", ")),
        Some(defs.join("\n\n")),
    )
}

fn field_name(name: &str, used_names: &HashSet<String>) -> Ident {
    type_or_field_name(name, used_names, "field", snake_case)
}

fn type_name(name: &str, used_names: &HashSet<String>) -> Ident {
    type_or_field_name(name, used_names, "GeneratedType", type_case)
}

// https://docs.python.org/3/reference/lexical_analysis.html#keywords
#[rustfmt::skip]
const PYTHON_KEYWORDS: &[&str] = &[
    "False", "None", "True",
    "and", "as", "assert", "async", "await", "break", "class", "continue",
    "def", "del", "elif", "else", "except", "finally", "for", "from", "global",
    "if", "import", "in", "is", "lambda", "nonlocal", "not", "or", "pass",
    "raise", "return", "try", "while", "with", "yield",
];

fn type_or_field_name(
    name: &str,
    used_names: &HashSet<String>,
    default_name: &str,
    case_fn: fn(&str) -> String,
) -> Ident {
    let name = name.trim();
    let mut output_name = case_fn(name);
    if PYTHON_KEYWORDS.contains(&&*output_name) {
        output_name.push_str("_field");
    }
    if output_name.is_empty() {
        output_name.push_str(default_name);
    }
    if let Some(c) = output_name.chars().next() {
        if c.is_ascii_digit() {
            output_name = String::from("n") + &output_name;
        }
    }
    if !used_names.contains(&output_name) {
        return output_name;
    }
    for n in 2.. {
        let temp = format!("{}{}", output_name, n);
        if !used_names.contains(&temp) {
            return temp;
        }
    }
    unreachable!()
}

fn import(ctxt: &mut Ctxt, import: Import) -> String {
    ctxt.imports.insert(import);
    match ctxt.options.import_style {
        ImportStyle::QualifiedPaths => import.qualified(),
        _ => import.identifier().into(),
    }
}

fn generate_data_class(
    ctxt: &mut Ctxt,
    path: &str,
    field_shapes: &LinkedHashMap<String, Shape>,
    containing_shape: &Shape,
) -> (Ident, Option<Code>) {
    for (created_for_shape, ident) in ctxt.created_classes.iter() {
        if created_for_shape.is_acceptable_substitution_for(containing_shape) {
            return (ident.into(), None);
        }
    }

    let type_name = type_name(path, &ctxt.type_names);
    ctxt.type_names.insert(type_name.clone());
    ctxt.created_classes
        .push((containing_shape.clone(), type_name.clone()));

    let mut field_names = HashSet::new();
    let mut defs = Vec::new();

    let fields: Vec<Code> = field_shapes
        .iter()
        .map(|(name, typ)| {
            let field_name = field_name(name, &field_names);
            field_names.insert(field_name.clone());

            let (field_type, child_defs) = type_from_shape(ctxt, name, typ);

            if let Some(code) = child_defs {
                defs.push(code);
            }

            let mut field_code = String::new();
            let transformed = apply_transform(ctxt, &field_name, &name);
            if transformed != field_name {
                field_code += &format!(" = {}(alias = \"{}\")", import(ctxt, Import::Field), transformed)
            }

            format!("    {}: {}{}", field_name, field_type, field_code)
        })
        .collect();

    let mut code = String::new();

    code += &format!(
        "class {}({}):\n",
        type_name,
        import(ctxt, Import::BaseModel)
    );

    if fields.is_empty() {
        code += "    pass\n";
    } else {
        code += &fields.join("\n");
        code += "\n";
    }

    if !defs.is_empty() {
        let mut d = defs.join("\n\n");
        d += "\n\n";
        d += &code;
        code = d;
    }

    (type_name, Some(code))
}

fn apply_transform(ctxt: &Ctxt, field_name: &str, name: &str) -> String {
    match ctxt.options.property_name_format {
        Some(StringTransform::LowerCase) => field_name.to_ascii_lowercase(),
        Some(StringTransform::PascalCase) => type_case(field_name),
        Some(StringTransform::SnakeCase) => snake_case(field_name),
        Some(StringTransform::KebabCase) => kebab_case(field_name),
        Some(StringTransform::UpperCase) => field_name.to_ascii_uppercase(),
        Some(StringTransform::CamelCase) => lower_camel_case(field_name),
        Some(StringTransform::ScreamingSnakeCase) => snake_case(field_name).to_ascii_uppercase(),
        Some(StringTransform::ScreamingKebabCase) => kebab_case(field_name).to_ascii_uppercase(),
        None => name.to_string(),
    }
}

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

    #[test]
    fn field_names_test() {
        fn field_name_test(from: &str, to: &str) {
            assert_eq!(
                field_name(from, &HashSet::new()),
                to.to_string(),
                r#"From "{}" to "{}""#,
                from,
                to
            );
        }

        field_name_test("valid", "valid");
        field_name_test("1", "n1");
        field_name_test("+1", "n1");
        field_name_test("", "field");
        field_name_test("def", "def_field");
    }
}
<commit_msg>shared: Do not limit tuple size for python.<commit_after>use linked_hash_map::LinkedHashMap;
use std::collections::HashSet;

use crate::options::{ImportStyle, Options, StringTransform};
use crate::shape::{self, Shape};
use crate::to_singular::to_singular;
use crate::util::{kebab_case, lower_camel_case, snake_case, type_case};

#[derive(PartialEq, PartialOrd, Ord, Eq, Hash, Clone, Copy)]
enum Import {
    Any,
    Optional,
    BaseModel,
    Field,
}

impl Import {
    fn pair(&self) -> (&'static str, &'static str) {
        match self {
            Import::Any => ("typing", "Any"),
            Import::Optional => ("typing", "Optional"),
            Import::BaseModel => ("pydantic", "BaseModel"),
            Import::Field => ("pydantic", "Field"),
        }
    }
    fn identifier(&self) -> &'static str {
        self.pair().1
    }
    fn qualified(&self) -> String {
        let (module, identifier) = self.pair();
        format!("{}.{}", module, identifier)
    }
}

struct Ctxt {
    options: Options,
    type_names: HashSet<String>,
    imports: HashSet<Import>,
    created_classes: Vec<(Shape, Ident)>,
}

pub type Ident = String;
pub type Code = String;

pub fn python_types(name: &str, shape: &Shape, options: Options) -> Code {
    let mut ctxt = Ctxt {
        options,
        type_names: HashSet::new(),
        imports: HashSet::new(),
        created_classes: Vec::new(),
    };

    if !matches!(shape, Shape::Struct { .. }) {
        // reserve the requested name
        ctxt.type_names.insert(name.to_string());
    }

    let (ident, code) = type_from_shape(&mut ctxt, name, shape);
    let mut code = code.unwrap_or_default();

    if !ctxt.imports.is_empty() {
        let mut imports: Vec<_> = ctxt.imports.drain().collect();
        imports.sort();
        let mut import_code = String::new();
        for import in imports {
            match ctxt.options.import_style {
                ImportStyle::AssumeExisting => {}
                ImportStyle::AddImports => {
                    let (module, identifier) = import.pair();
                    import_code += &format!("from {} import {}\n", module, identifier);
                }
                ImportStyle::QualifiedPaths => {
                    let (module, identifier) = import.pair();
                    import_code += &format!("import {}.{}\n", module, identifier);
                }
            }
        }
        import_code += "\n\n";
        code = import_code + &code;
    }

    if ident != name {
        code += &format!("\n\n{} = {}", name, ident);
    }
    code
}

fn type_from_shape(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    use crate::shape::Shape::*;
    match shape {
        Null | Any | Bottom => (import(ctxt, Import::Any), None),
        Bool => ("bool".into(), None),
        StringT => ("str".into(), None),
        Integer => ("int".into(), None),
        Floating => ("float".into(), None),
        Tuple(shapes, _n) => {
            let folded = shape::fold_shapes(shapes.clone());
            if folded == Any && shapes.iter().any(|s| s != &Any) {
                generate_tuple_type(ctxt, path, shapes)
            } else {
                generate_vec_type(ctxt, path, &folded)
            }
        }
        VecT { elem_type: e } => generate_vec_type(ctxt, path, e),
        Struct { fields } => generate_data_class(ctxt, path, fields, shape),
        MapT { val_type: v } => generate_map_type(ctxt, path, v),
        Opaque(t) => (t.clone(), None),
        Optional(e) => {
            let (inner, defs) = type_from_shape(ctxt, path, e);
            if ctxt.options.use_default_for_missing_fields {
                (inner, defs)
            } else {
                let optional = import(ctxt, Import::Optional);
                (format!("{}[{}]", optional, inner), defs)
            }
        }
    }
}

fn generate_vec_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = to_singular(path);
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (format!("list[{}]", inner), defs)
}

fn generate_map_type(ctxt: &mut Ctxt, path: &str, shape: &Shape) -> (Ident, Option<Code>) {
    let singular = to_singular(path);
    let (inner, defs) = type_from_shape(ctxt, &singular, shape);
    (format!("dict[str, {}]", inner), defs)
}

fn generate_tuple_type(ctxt: &mut Ctxt, path: &str, shapes: &[Shape]) -> (Ident, Option<Code>) {
    let mut types = Vec::new();
    let mut defs = Vec::new();

    for shape in shapes {
        let (typ, def) = type_from_shape(ctxt, path, shape);
        types.push(typ);
        if let Some(code) = def {
            defs.push(code)
        }
    }

    (
        format!("tuple[{}]", types.join(", ")),
        Some(defs.join("\n\n")),
    )
}

fn field_name(name: &str, used_names: &HashSet<String>) -> Ident {
    type_or_field_name(name, used_names, "field", snake_case)
}

fn type_name(name: &str, used_names: &HashSet<String>) -> Ident {
    type_or_field_name(name, used_names, "GeneratedType", type_case)
}

// https://docs.python.org/3/reference/lexical_analysis.html#keywords
#[rustfmt::skip]
const PYTHON_KEYWORDS: &[&str] = &[
    "False", "None", "True",
    "and", "as", "assert", "async", "await", "break", "class", "continue",
    "def", "del", "elif", "else", "except", "finally", "for", "from", "global",
    "if", "import", "in", "is", "lambda", "nonlocal", "not", "or", "pass",
    "raise", "return", "try", "while", "with", "yield",
];

fn type_or_field_name(
    name: &str,
    used_names: &HashSet<String>,
    default_name: &str,
    case_fn: fn(&str) -> String,
) -> Ident {
    let name = name.trim();
    let mut output_name = case_fn(name);
    if PYTHON_KEYWORDS.contains(&&*output_name) {
        output_name.push_str("_field");
    }
    if output_name.is_empty() {
        output_name.push_str(default_name);
    }
    if let Some(c) = output_name.chars().next() {
        if c.is_ascii_digit() {
            output_name = String::from("n") + &output_name;
        }
    }
    if !used_names.contains(&output_name) {
        return output_name;
    }
    for n in 2.. {
        let temp = format!("{}{}", output_name, n);
        if !used_names.contains(&temp) {
            return temp;
        }
    }
    unreachable!()
}

fn import(ctxt: &mut Ctxt, import: Import) -> String {
    ctxt.imports.insert(import);
    match ctxt.options.import_style {
        ImportStyle::QualifiedPaths => import.qualified(),
        _ => import.identifier().into(),
    }
}

fn generate_data_class(
    ctxt: &mut Ctxt,
    path: &str,
    field_shapes: &LinkedHashMap<String, Shape>,
    containing_shape: &Shape,
) -> (Ident, Option<Code>) {
    for (created_for_shape, ident) in ctxt.created_classes.iter() {
        if created_for_shape.is_acceptable_substitution_for(containing_shape) {
            return (ident.into(), None);
        }
    }

    let type_name = type_name(path, &ctxt.type_names);
    ctxt.type_names.insert(type_name.clone());
    ctxt.created_classes
        .push((containing_shape.clone(), type_name.clone()));

    let mut field_names = HashSet::new();
    let mut defs = Vec::new();

    let fields: Vec<Code> = field_shapes
        .iter()
        .map(|(name, typ)| {
            let field_name = field_name(name, &field_names);
            field_names.insert(field_name.clone());

            let (field_type, child_defs) = type_from_shape(ctxt, name, typ);

            if let Some(code) = child_defs {
                defs.push(code);
            }

            let mut field_code = String::new();
            let transformed = apply_transform(ctxt, &field_name, &name);
            if transformed != field_name {
                field_code += &format!(" = {}(alias = \"{}\")", import(ctxt, Import::Field), transformed)
            }

            format!("    {}: {}{}", field_name, field_type, field_code)
        })
        .collect();

    let mut code = String::new();

    code += &format!(
        "class {}({}):\n",
        type_name,
        import(ctxt, Import::BaseModel)
    );

    if fields.is_empty() {
        code += "    pass\n";
    } else {
        code += &fields.join("\n");
        code += "\n";
    }

    if !defs.is_empty() {
        let mut d = defs.join("\n\n");
        d += "\n\n";
        d += &code;
        code = d;
    }

    (type_name, Some(code))
}

fn apply_transform(ctxt: &Ctxt, field_name: &str, name: &str) -> String {
    match ctxt.options.property_name_format {
        Some(StringTransform::LowerCase) => field_name.to_ascii_lowercase(),
        Some(StringTransform::PascalCase) => type_case(field_name),
        Some(StringTransform::SnakeCase) => snake_case(field_name),
        Some(StringTransform::KebabCase) => kebab_case(field_name),
        Some(StringTransform::UpperCase) => field_name.to_ascii_uppercase(),
        Some(StringTransform::CamelCase) => lower_camel_case(field_name),
        Some(StringTransform::ScreamingSnakeCase) => snake_case(field_name).to_ascii_uppercase(),
        Some(StringTransform::ScreamingKebabCase) => kebab_case(field_name).to_ascii_uppercase(),
        None => name.to_string(),
    }
}

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

    #[test]
    fn field_names_test() {
        fn field_name_test(from: &str, to: &str) {
            assert_eq!(
                field_name(from, &HashSet::new()),
                to.to_string(),
                r#"From "{}" to "{}""#,
                from,
                to
            );
        }

        field_name_test("valid", "valid");
        field_name_test("1", "n1");
        field_name_test("+1", "n1");
        field_name_test("", "field");
        field_name_test("def", "def_field");
    }
}
<|endoftext|>


<commit_before>//! An unbounded set of streams

use std::fmt::{self, Debug};
use std::pin::Pin;
use std::marker::Unpin;

use futures_core::{Poll, Stream, FusedStream};
use futures_core::task::LocalWaker;

use crate::stream::{StreamExt, StreamFuture, FuturesUnordered};

/// An unbounded set of streams
///
/// This "combinator" provides the ability to maintain a set of streams
/// and drive them all to completion.
///
/// Streams are pushed into this set and their realized values are
/// yielded as they become ready. Streams will only be polled when they
/// generate notifications. This allows to coordinate a large number of streams.
///
/// Note that you can create a ready-made `SelectAll` via the
/// `select_all` function in the `stream` module, or you can start with an
/// empty set with the `SelectAll::new` constructor.
#[must_use = "streams do nothing unless polled"]
pub struct SelectAll<St> {
    inner: FuturesUnordered<StreamFuture<St>>,
}

impl<St: Debug> Debug for SelectAll<St> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "SelectAll {{ ... }}")
    }
}

impl<St: Stream + Unpin> SelectAll<St> {
    /// Constructs a new, empty `SelectAll`
    ///
    /// The returned `SelectAll` does not contain any streams and, in this
    /// state, `SelectAll::poll` will return `Poll::Ready(None)`.
    pub fn new() -> SelectAll<St> {
        SelectAll { inner: FuturesUnordered::new() }
    }

    /// Returns the number of streams contained in the set.
    ///
    /// This represents the total number of in-flight streams.
    pub fn len(&self) -> usize {
        self.inner.len()
    }

    /// Returns `true` if the set contains no streams
    pub fn is_empty(&self) -> bool {
        self.inner.is_empty()
    }

    /// Push a stream into the set.
    ///
    /// This function submits the given stream to the set for managing. This
    /// function will not call `poll` on the submitted stream. The caller must
    /// ensure that `SelectAll::poll` is called in order to receive task
    /// notifications.
    pub fn push(&mut self, stream: St) {
        self.inner.push(stream.into_future());
    }
}

impl<St: Stream + Unpin> Default for SelectAll<St> {
    fn default() -> SelectAll<St> {
        SelectAll::new()
    }
}

impl<St: Stream + Unpin> Stream for SelectAll<St> {
    type Item = St::Item;

    fn poll_next(
        mut self: Pin<&mut Self>,
        lw: &LocalWaker,
    ) -> Poll<Option<Self::Item>> {
        match self.inner.poll_next_unpin(lw) {
            Poll::Pending => Poll::Pending,
            Poll::Ready(Some((Some(item), remaining))) => {
                self.push(remaining);
                Poll::Ready(Some(item))
            }
            Poll::Ready(Some((None, _))) => {
                // FuturesUnordered thinks it isn't terminated
                // because it yielded a Some. Here we poll it
                // so it can realize it is terminated.
                self.inner.poll_next_unpin(lw);
                Poll::Ready(None)
            }
            Poll::Ready(_) => Poll::Ready(None),
        }
    }
}

impl<St: Stream + Unpin> FusedStream for SelectAll<St> {
    fn is_terminated(&self) -> bool {
        self.inner.is_terminated()
    }
}

/// Convert a list of streams into a `Stream` of results from the streams.
///
/// This essentially takes a list of streams (e.g. a vector, an iterator, etc.)
/// and bundles them together into a single stream.
/// The stream will yield items as they become available on the underlying
/// streams internally, in the order they become available.
///
/// Note that the returned set can also be used to dynamically push more
/// futures into the set as they become available.
pub fn select_all<I>(streams: I) -> SelectAll<I::Item>
    where I: IntoIterator,
          I::Item: Stream + Unpin
{
    let mut set = SelectAll::new();

    for stream in streams {
        set.push(stream);
    }

    set
}
<commit_msg>Revert "impl Default for SelectAll"<commit_after>//! An unbounded set of streams

use std::fmt::{self, Debug};
use std::pin::Pin;
use std::marker::Unpin;

use futures_core::{Poll, Stream, FusedStream};
use futures_core::task::LocalWaker;

use crate::stream::{StreamExt, StreamFuture, FuturesUnordered};

/// An unbounded set of streams
///
/// This "combinator" provides the ability to maintain a set of streams
/// and drive them all to completion.
///
/// Streams are pushed into this set and their realized values are
/// yielded as they become ready. Streams will only be polled when they
/// generate notifications. This allows to coordinate a large number of streams.
///
/// Note that you can create a ready-made `SelectAll` via the
/// `select_all` function in the `stream` module, or you can start with an
/// empty set with the `SelectAll::new` constructor.
#[must_use = "streams do nothing unless polled"]
pub struct SelectAll<St> {
    inner: FuturesUnordered<StreamFuture<St>>,
}

impl<St: Debug> Debug for SelectAll<St> {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "SelectAll {{ ... }}")
    }
}

impl<St: Stream + Unpin> SelectAll<St> {
    /// Constructs a new, empty `SelectAll`
    ///
    /// The returned `SelectAll` does not contain any streams and, in this
    /// state, `SelectAll::poll` will return `Poll::Ready(None)`.
    pub fn new() -> SelectAll<St> {
        SelectAll { inner: FuturesUnordered::new() }
    }

    /// Returns the number of streams contained in the set.
    ///
    /// This represents the total number of in-flight streams.
    pub fn len(&self) -> usize {
        self.inner.len()
    }

    /// Returns `true` if the set contains no streams
    pub fn is_empty(&self) -> bool {
        self.inner.is_empty()
    }

    /// Push a stream into the set.
    ///
    /// This function submits the given stream to the set for managing. This
    /// function will not call `poll` on the submitted stream. The caller must
    /// ensure that `SelectAll::poll` is called in order to receive task
    /// notifications.
    pub fn push(&mut self, stream: St) {
        self.inner.push(stream.into_future());
    }
}

impl<St: Stream + Unpin> Stream for SelectAll<St> {
    type Item = St::Item;

    fn poll_next(
        mut self: Pin<&mut Self>,
        lw: &LocalWaker,
    ) -> Poll<Option<Self::Item>> {
        match self.inner.poll_next_unpin(lw) {
            Poll::Pending => Poll::Pending,
            Poll::Ready(Some((Some(item), remaining))) => {
                self.push(remaining);
                Poll::Ready(Some(item))
            }
            Poll::Ready(Some((None, _))) => {
                // FuturesUnordered thinks it isn't terminated
                // because it yielded a Some. Here we poll it
                // so it can realize it is terminated.
                self.inner.poll_next_unpin(lw);
                Poll::Ready(None)
            }
            Poll::Ready(_) => Poll::Ready(None),
        }
    }
}

impl<St: Stream + Unpin> FusedStream for SelectAll<St> {
    fn is_terminated(&self) -> bool {
        self.inner.is_terminated()
    }
}

/// Convert a list of streams into a `Stream` of results from the streams.
///
/// This essentially takes a list of streams (e.g. a vector, an iterator, etc.)
/// and bundles them together into a single stream.
/// The stream will yield items as they become available on the underlying
/// streams internally, in the order they become available.
///
/// Note that the returned set can also be used to dynamically push more
/// futures into the set as they become available.
pub fn select_all<I>(streams: I) -> SelectAll<I::Item>
    where I: IntoIterator,
          I::Item: Stream + Unpin
{
    let mut set = SelectAll::new();

    for stream in streams {
        set.push(stream);
    }

    set
}
<|endoftext|>