<commit_before>use std::fmt;

use base::ast::{self, DisplayEnv, Expr, MutVisitor, Pattern, SpannedExpr, Typed, TypedIdent};
use base::error::Errors;
use base::fnv::FnvMap;
use base::kind::{ArcKind, Kind, KindEnv};
use base::pos::{self, BytePos, Span, Spanned};
use base::scoped_map::ScopedMap;
use base::symbol::{Symbol, SymbolModule, SymbolRef};
use base::types::{self, Alias, ArcType, RecordSelector, Type, TypeEnv};
use unify_type::{State, TypeError};
use unify::{Error as UnifyError, Unifiable, Unifier, UnifierState};

pub type Error = Errors<Spanned<RenameError, BytePos>>;

#[derive(Clone, Debug, PartialEq)]
pub enum RenameError {
    NoMatchingType {
        symbol: String,
        expected: ArcType,
        possible_types: Vec<(Option<Span<BytePos>>, ArcType)>,
    },
}

impl fmt::Display for RenameError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            RenameError::NoMatchingType {
                ref symbol,
                ref expected,
                ref possible_types,
            } => {
                writeln!(
                    f,
                    "Could not resolve a binding for `{}` with type `{}`",
                    symbol,
                    expected
                )?;
                writeln!(f, "Possibilities:")?;
                for &(ref span, ref typ) in possible_types {
                    match *span {
                        Some(ref span) => writeln!(f, "{} at {}", typ, span.start)?,
                        None => writeln!(f, "{} at 'global'", typ)?,
                    }
                }
                Ok(())
            }
        }
    }
}

struct Environment<'b> {
    env: &'b TypeEnv,
    stack: ScopedMap<Symbol, (Symbol, Span<BytePos>, ArcType)>,
    stack_types: ScopedMap<Symbol, Alias<Symbol, ArcType>>,
}

impl<'a> KindEnv for Environment<'a> {
    fn find_kind(&self, _type_name: &SymbolRef) -> Option<ArcKind> {
        None
    }
}

impl<'a> TypeEnv for Environment<'a> {
    fn find_type(&self, id: &SymbolRef) -> Option<&ArcType> {
        self.stack
            .get(id)
            .map(|t| &t.2)
            .or_else(|| self.env.find_type(id))
    }

    fn find_type_info(&self, id: &SymbolRef) -> Option<&Alias<Symbol, ArcType>> {
        self.stack_types
            .get(id)
            .or_else(|| self.env.find_type_info(id))
    }

    fn find_record(
        &self,
        _fields: &[Symbol],
        _selector: RecordSelector,
    ) -> Option<(ArcType, ArcType)> {
        None
    }
}

pub fn rename(
    symbols: &mut SymbolModule,
    env: &TypeEnv,
    expr: &mut SpannedExpr<Symbol>,
) -> Result<(), Error> {
    use base::resolve;

    struct RenameVisitor<'a: 'b, 'b> {
        symbols: &'b mut SymbolModule<'a>,
        env: Environment<'b>,
        errors: Error,
    }

    impl<'a, 'b> RenameVisitor<'a, 'b> {
        fn find_fields(&self, typ: &ArcType) -> Vec<types::Field<Symbol, ArcType>> {
            // Walk through all type aliases
            let record = resolve::remove_aliases(&self.env, typ.remove_forall().clone());
            record.row_iter().cloned().collect()
        }

        fn new_pattern(&mut self, typ: &ArcType, pattern: &mut ast::SpannedPattern<Symbol>) {
            match pattern.value {
                ast::Pattern::Record {
                    ref mut fields,
                    ref types,
                    ..
                } => {
                    let field_types = self.find_fields(typ);
                    for field in fields {
                        let field_type = &field_types
                            .iter()
                            .find(|field_type| field_type.name.name_eq(&field.name.value))
                            .expect("ICE: Existing field")
                            .typ;
                        match field.value {
                            Some(ref mut pat) => self.new_pattern(field_type, pat),
                            None => {
                                let id = field.name.value.clone();
                                let pat = Pattern::Ident(TypedIdent {
                                    name: self.stack_var(id, pattern.span, field_type.clone()),
                                    typ: field_type.clone(),
                                });
                                field.value = Some(pos::spanned(pattern.span, pat));
                            }
                        }
                    }

                    let record_type = resolve::remove_aliases(&self.env, typ.clone()).clone();
                    for ast_field in types {
                        let field_type = record_type
                            .remove_forall()
                            .type_field_iter()
                            .find(|field| field.name.name_eq(&ast_field.name.value))
                            .unwrap_or_else(|| {
                                panic!(
                                    "ICE: Type `{}` does not have type field `{}`",
                                    record_type,
                                    ast_field.name.value
                                )
                            });
                        self.stack_type(
                            ast_field.name.value.clone(),
                            pattern.span,
                            &field_type.typ,
                        );
                    }
                }
                ast::Pattern::Ident(ref mut id) => {
                    let new_name = self.stack_var(id.name.clone(), pattern.span, id.typ.clone());
                    id.name = new_name;
                }
                ast::Pattern::Tuple {
                    ref typ,
                    ref mut elems,
                } => for (field, elem) in typ.row_iter().zip(elems) {
                    self.new_pattern(&field.typ, elem);
                },
                ast::Pattern::Constructor(ref mut id, ref mut args) => {
                    let typ = self.env
                        .find_type(&id.name)
                        .unwrap_or_else(|| panic!("ICE: Expected constructor: {}", id.name))
                        .clone();
                    for (arg_type, arg) in types::arg_iter(&typ).zip(args) {
                        self.new_pattern(arg_type, arg);
                    }
                }
                ast::Pattern::Error => (),
            }
        }

        fn stack_var(&mut self, id: Symbol, span: Span<BytePos>, typ: ArcType) -> Symbol {
            let old_id = id.clone();
            let name = self.symbols.string(&id).to_owned();
            let new_id = self.symbols.symbol(format!("{}:{}", name, span.start));
            debug!(
                "Rename binding `{}` = `{}` `{}`",
                self.symbols.string(&old_id),
                self.symbols.string(&new_id),
                typ
            );
            self.env.stack.insert(old_id, (new_id.clone(), span, typ));
            new_id
        }

        fn stack_type(&mut self, id: Symbol, span: Span<BytePos>, alias: &Alias<Symbol, ArcType>) {
            // Insert variant constructors into the local scope
            let aliased_type = alias.typ();
            if let Type::Variant(ref row) = **aliased_type.remove_forall() {
                for field in row.row_iter().cloned() {
                    self.env
                        .stack
                        .insert(field.name.clone(), (field.name, span, field.typ));
                }
            }

            // FIXME: Workaround so that both the types name in this module and its global
            // name are imported. Without this aliases may not be traversed properly
            self.env
                .stack_types
                .insert(alias.name.clone(), alias.clone());
            self.env.stack_types.insert(id, alias.clone());
        }

        /// Renames `id` to the unique identifier which have the type `expected`
        /// Returns `Some(new_id)` if renaming was necessary or `None` if no renaming was necessary
        /// as `id` was currently unique (#Int+, #Float*, etc)
        fn rename(&self, id: &Symbol, expected: &ArcType) -> Result<Option<Symbol>, RenameError> {
            let locals = self.env.stack.get_all(id);
            let global = self.env.env.find_type(id).map(|typ| (id, None, typ));
            let candidates = || {
                locals
                    .iter()
                    .flat_map(|bindings| {
                        bindings
                            .iter()
                            .rev()
                            .map(|bind| (&bind.0, Some(&bind.1), &bind.2))
                    })
                    .chain(global)
            };
            // If there is a single binding (or no binding in case of primitives such as #Int+)
            // there is no need to check for equivalency as typechecker couldnt have infered a
            // different binding
            if candidates().count() <= 1 {
                return Ok(candidates().next().map(|tup| tup.0.clone()));
            }
            candidates()
                .find(|tup| {
                    equivalent(&self.env, tup.2.remove_forall(), expected.remove_forall())
                })
                .map(|tup| Some(tup.0.clone()))
                .ok_or_else(|| {
                    RenameError::NoMatchingType {
                        symbol: String::from(self.symbols.string(id)),
                        expected: expected.clone(),
                        possible_types: candidates()
                            .map(|tup| (tup.1.cloned(), tup.2.clone()))
                            .collect(),
                    }
                })
        }

        fn rename_expr(&mut self, expr: &mut SpannedExpr<Symbol>) -> Result<(), RenameError> {
            match expr.value {
                Expr::Ident(ref mut id) => if let Some(new_id) = self.rename(&id.name, &id.typ)? {
                    debug!("Rename identifier {} = {}", id.name, new_id);
                    id.name = new_id;
                },
                Expr::Record {
                    ref mut typ,
                    ref mut exprs,
                    ref mut base,
                    ..
                } => {
                    let field_types = self.find_fields(typ);
                    for (field, expr_field) in field_types.iter().zip(exprs) {
                        match expr_field.value {
                            Some(ref mut expr) => self.visit_expr(expr),
                            None => if let Some(new_id) =
                                self.rename(&expr_field.name.value, &field.typ)?
                            {
                                debug!("Rename record field {} = {}", expr_field.name, new_id);
                                expr_field.value = Some(pos::spanned(
                                    expr.span,
                                    Expr::Ident(TypedIdent {
                                        name: new_id,
                                        typ: field.typ.clone(),
                                    }),
                                ));
                            },
                        }
                    }

                    if let Some(ref mut base) = *base {
                        self.visit_expr(base);
                    }
                }
                Expr::Infix(ref mut lhs, ref mut op, ref mut rhs) => {
                    if let Some(new_id) = self.rename(&op.value.name, &op.value.typ)? {
                        debug!(
                            "Rename {} = {}",
                            self.symbols.string(&op.value.name),
                            self.symbols.string(&new_id)
                        );
                        op.value.name = new_id;
                    }
                    self.visit_expr(lhs);
                    self.visit_expr(rhs);
                }
                Expr::Match(ref mut expr, ref mut alts) => {
                    self.visit_expr(expr);
                    for alt in alts {
                        self.env.stack_types.enter_scope();
                        self.env.stack.enter_scope();
                        let typ = expr.env_type_of(&self.env);
                        self.new_pattern(&typ, &mut alt.pattern);
                        self.visit_expr(&mut alt.expr);
                        self.env.stack.exit_scope();
                        self.env.stack_types.exit_scope();
                    }
                }
                Expr::LetBindings(ref mut bindings, ref mut expr) => {
                    self.env.stack_types.enter_scope();
                    self.env.stack.enter_scope();
                    let is_recursive = bindings.iter().all(|bind| !bind.args.is_empty());
                    for bind in bindings.iter_mut() {
                        if !is_recursive {
                            self.visit_expr(&mut bind.expr);
                        }
                        self.new_pattern(&bind.resolved_type, &mut bind.name);
                    }
                    if is_recursive {
                        for bind in bindings {
                            self.env.stack.enter_scope();
                            for (typ, arg) in types::arg_iter(bind.resolved_type.remove_forall())
                                .zip(&mut bind.args)
                            {
                                arg.value.name =
                                    self.stack_var(arg.value.name.clone(), expr.span, typ.clone());
                            }
                            self.visit_expr(&mut bind.expr);
                            self.env.stack.exit_scope();
                        }
                    }
                    self.visit_expr(expr);
                    self.env.stack.exit_scope();
                    self.env.stack_types.exit_scope();
                }
                Expr::Lambda(ref mut lambda) => {
                    self.env.stack.enter_scope();
                    for (typ, arg) in types::arg_iter(&lambda.id.typ).zip(&mut lambda.args) {
                        arg.value.name =
                            self.stack_var(arg.value.name.clone(), expr.span, typ.clone());
                    }
                    self.visit_expr(&mut lambda.body);
                    self.env.stack.exit_scope();
                }
                Expr::TypeBindings(ref bindings, ref mut body) => {
                    self.env.stack_types.enter_scope();
                    for bind in bindings {
                        self.stack_type(
                            bind.name.value.clone(),
                            expr.span,
                            bind.finalized_alias.as_ref().expect(
                                "ICE: Alias should have been finalized \
                                 before renaming",
                            ),
                        );
                    }
                    self.visit_expr(body);
                    self.env.stack_types.exit_scope();
                }
                _ => ast::walk_mut_expr(self, expr),
            }
            Ok(())
        }
    }

    impl<'a, 'b> MutVisitor for RenameVisitor<'a, 'b> {
        type Ident = Symbol;

        fn visit_expr(&mut self, expr: &mut SpannedExpr<Self::Ident>) {
            if let Err(err) = self.rename_expr(expr) {
                self.errors.push(Spanned {
                    span: expr.span,
                    value: err,
                });
            }
        }
    }

    let mut visitor = RenameVisitor {
        symbols: symbols,
        errors: Errors::new(),
        env: Environment {
            env: env,
            stack: ScopedMap::new(),
            stack_types: ScopedMap::new(),
        },
    };
    visitor.visit_expr(expr);
    if visitor.errors.has_errors() {
        Err(visitor.errors)
    } else {
        Ok(())
    }
}

pub fn equivalent(env: &TypeEnv, actual: &ArcType, inferred: &ArcType) -> bool {
    use substitution::Substitution;
    // FIXME This Substitution is unnecessary for equivalence unification
    let subs = Substitution::new(Kind::typ());
    let mut unifier = UnifierState {
        state: State::new(env, &subs),
        unifier: Equivalent {
            map: FnvMap::default(),
            equiv: true,
        },
    };
    unifier.try_match(actual, inferred);
    unifier.unifier.equiv
}

struct Equivalent {
    map: FnvMap<Symbol, ArcType>,
    equiv: bool,
}

impl<'a> Unifier<State<'a>, ArcType> for Equivalent {
    fn report_error(
        unifier: &mut UnifierState<State<'a>, Self>,
        _error: UnifyError<ArcType, TypeError<Symbol>>,
    ) {
        unifier.unifier.equiv = false;
    }

    fn try_match_res(
        unifier: &mut UnifierState<State<'a>, Self>,
        l: &ArcType,
        r: &ArcType,
    ) -> Result<Option<ArcType>, UnifyError<ArcType, TypeError<Symbol>>> {
        debug!("{} ====> {}", l, r);
        match (&**l, &**r) {
            (&Type::Generic(ref gl), &Type::Generic(ref gr)) if gl == gr => Ok(None),
            (&Type::Skolem(ref gl), &Type::Skolem(ref gr)) if gl == gr => Ok(None),
            (&Type::Generic(ref gl), _) => match unifier.unifier.map.get(&gl.id).cloned() {
                Some(ref typ) => unifier.try_match_res(typ, r),
                None => {
                    unifier.unifier.map.insert(gl.id.clone(), r.clone());
                    Ok(None)
                }
            },
            (&Type::Skolem(ref gl), _) => match unifier.unifier.map.get(&gl.name).cloned() {
                Some(ref typ) => unifier.try_match_res(typ, r),
                None => {
                    unifier.unifier.map.insert(gl.name.clone(), r.clone());
                    Ok(None)
                }
            },
            _ => l.zip_match(r, unifier),
        }
    }

    fn error_type(_unifier: &mut UnifierState<State<'a>, Self>) -> Option<ArcType> {
        None
    }
}
<commit_msg>fix(check): Don't consider global variables for overload resolution<commit_after>use std::fmt;

use base::ast::{self, DisplayEnv, Expr, MutVisitor, Pattern, SpannedExpr, Typed, TypedIdent};
use base::error::Errors;
use base::fnv::FnvMap;
use base::kind::{ArcKind, Kind, KindEnv};
use base::pos::{self, BytePos, Span, Spanned};
use base::scoped_map::ScopedMap;
use base::symbol::{Symbol, SymbolModule, SymbolRef};
use base::types::{self, Alias, ArcType, RecordSelector, Type, TypeEnv};
use unify_type::{State, TypeError};
use unify::{Error as UnifyError, Unifiable, Unifier, UnifierState};

pub type Error = Errors<Spanned<RenameError, BytePos>>;

#[derive(Clone, Debug, PartialEq)]
pub enum RenameError {
    NoMatchingType {
        symbol: String,
        expected: ArcType,
        possible_types: Vec<(Option<Span<BytePos>>, ArcType)>,
    },
}

impl fmt::Display for RenameError {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        match *self {
            RenameError::NoMatchingType {
                ref symbol,
                ref expected,
                ref possible_types,
            } => {
                writeln!(
                    f,
                    "Could not resolve a binding for `{}` with type `{}`",
                    symbol,
                    expected
                )?;
                writeln!(f, "Possibilities:")?;
                for &(ref span, ref typ) in possible_types {
                    match *span {
                        Some(ref span) => writeln!(f, "{} at {}", typ, span.start)?,
                        None => writeln!(f, "{} at 'global'", typ)?,
                    }
                }
                Ok(())
            }
        }
    }
}

struct Environment<'b> {
    env: &'b TypeEnv,
    stack: ScopedMap<Symbol, (Symbol, Span<BytePos>, ArcType)>,
    stack_types: ScopedMap<Symbol, Alias<Symbol, ArcType>>,
}

impl<'a> KindEnv for Environment<'a> {
    fn find_kind(&self, _type_name: &SymbolRef) -> Option<ArcKind> {
        None
    }
}

impl<'a> TypeEnv for Environment<'a> {
    fn find_type(&self, id: &SymbolRef) -> Option<&ArcType> {
        self.stack
            .get(id)
            .map(|t| &t.2)
            .or_else(|| self.env.find_type(id))
    }

    fn find_type_info(&self, id: &SymbolRef) -> Option<&Alias<Symbol, ArcType>> {
        self.stack_types
            .get(id)
            .or_else(|| self.env.find_type_info(id))
    }

    fn find_record(
        &self,
        _fields: &[Symbol],
        _selector: RecordSelector,
    ) -> Option<(ArcType, ArcType)> {
        None
    }
}

pub fn rename(
    symbols: &mut SymbolModule,
    env: &TypeEnv,
    expr: &mut SpannedExpr<Symbol>,
) -> Result<(), Error> {
    use base::resolve;

    struct RenameVisitor<'a: 'b, 'b> {
        symbols: &'b mut SymbolModule<'a>,
        env: Environment<'b>,
        errors: Error,
    }

    impl<'a, 'b> RenameVisitor<'a, 'b> {
        fn find_fields(&self, typ: &ArcType) -> Vec<types::Field<Symbol, ArcType>> {
            // Walk through all type aliases
            let record = resolve::remove_aliases(&self.env, typ.remove_forall().clone());
            record.row_iter().cloned().collect()
        }

        fn new_pattern(&mut self, typ: &ArcType, pattern: &mut ast::SpannedPattern<Symbol>) {
            match pattern.value {
                ast::Pattern::Record {
                    ref mut fields,
                    ref types,
                    ..
                } => {
                    let field_types = self.find_fields(typ);
                    for field in fields {
                        let field_type = &field_types
                            .iter()
                            .find(|field_type| field_type.name.name_eq(&field.name.value))
                            .expect("ICE: Existing field")
                            .typ;
                        match field.value {
                            Some(ref mut pat) => self.new_pattern(field_type, pat),
                            None => {
                                let id = field.name.value.clone();
                                let pat = Pattern::Ident(TypedIdent {
                                    name: self.stack_var(id, pattern.span, field_type.clone()),
                                    typ: field_type.clone(),
                                });
                                field.value = Some(pos::spanned(pattern.span, pat));
                            }
                        }
                    }

                    let record_type = resolve::remove_aliases(&self.env, typ.clone()).clone();
                    for ast_field in types {
                        let field_type = record_type
                            .remove_forall()
                            .type_field_iter()
                            .find(|field| field.name.name_eq(&ast_field.name.value))
                            .unwrap_or_else(|| {
                                panic!(
                                    "ICE: Type `{}` does not have type field `{}`",
                                    record_type,
                                    ast_field.name.value
                                )
                            });
                        self.stack_type(
                            ast_field.name.value.clone(),
                            pattern.span,
                            &field_type.typ,
                        );
                    }
                }
                ast::Pattern::Ident(ref mut id) => {
                    let new_name = self.stack_var(id.name.clone(), pattern.span, id.typ.clone());
                    id.name = new_name;
                }
                ast::Pattern::Tuple {
                    ref typ,
                    ref mut elems,
                } => for (field, elem) in typ.row_iter().zip(elems) {
                    self.new_pattern(&field.typ, elem);
                },
                ast::Pattern::Constructor(ref mut id, ref mut args) => {
                    let typ = self.env
                        .find_type(&id.name)
                        .unwrap_or_else(|| panic!("ICE: Expected constructor: {}", id.name))
                        .clone();
                    for (arg_type, arg) in types::arg_iter(&typ).zip(args) {
                        self.new_pattern(arg_type, arg);
                    }
                }
                ast::Pattern::Error => (),
            }
        }

        fn stack_var(&mut self, id: Symbol, span: Span<BytePos>, typ: ArcType) -> Symbol {
            let old_id = id.clone();
            let name = self.symbols.string(&id).to_owned();
            let new_id = self.symbols.symbol(format!("{}:{}", name, span.start));
            debug!(
                "Rename binding `{}` = `{}` `{}`",
                self.symbols.string(&old_id),
                self.symbols.string(&new_id),
                typ
            );
            self.env.stack.insert(old_id, (new_id.clone(), span, typ));
            new_id
        }

        fn stack_type(&mut self, id: Symbol, span: Span<BytePos>, alias: &Alias<Symbol, ArcType>) {
            // Insert variant constructors into the local scope
            let aliased_type = alias.typ();
            if let Type::Variant(ref row) = **aliased_type.remove_forall() {
                for field in row.row_iter().cloned() {
                    self.env
                        .stack
                        .insert(field.name.clone(), (field.name, span, field.typ));
                }
            }

            // FIXME: Workaround so that both the types name in this module and its global
            // name are imported. Without this aliases may not be traversed properly
            self.env
                .stack_types
                .insert(alias.name.clone(), alias.clone());
            self.env.stack_types.insert(id, alias.clone());
        }

        /// Renames `id` to the unique identifier which have the type `expected`
        /// Returns `Some(new_id)` if renaming was necessary or `None` if no renaming was necessary
        /// as `id` was currently unique (#Int+, #Float*, etc)
        fn rename(&self, id: &Symbol, expected: &ArcType) -> Result<Option<Symbol>, RenameError> {
            let locals = self.env.stack.get_all(id);
            let candidates = || {
                locals
                    .iter()
                    .flat_map(|bindings| {
                        bindings
                            .iter()
                            .rev()
                            .map(|bind| (&bind.0, Some(&bind.1), &bind.2))
                    })
            };
            // If there is a single binding (or no binding in case of primitives such as #Int+)
            // there is no need to check for equivalency as typechecker couldnt have infered a
            // different binding
            if candidates().count() <= 1 {
                return Ok(candidates().next().map(|tup| tup.0.clone()));
            }
            candidates()
                .find(|tup| {
                    equivalent(&self.env, tup.2.remove_forall(), expected.remove_forall())
                })
                .map(|tup| Some(tup.0.clone()))
                .ok_or_else(|| {
                    RenameError::NoMatchingType {
                        symbol: String::from(self.symbols.string(id)),
                        expected: expected.clone(),
                        possible_types: candidates()
                            .map(|tup| (tup.1.cloned(), tup.2.clone()))
                            .collect(),
                    }
                })
        }

        fn rename_expr(&mut self, expr: &mut SpannedExpr<Symbol>) -> Result<(), RenameError> {
            match expr.value {
                Expr::Ident(ref mut id) => if let Some(new_id) = self.rename(&id.name, &id.typ)? {
                    debug!("Rename identifier {} = {}", id.name, new_id);
                    id.name = new_id;
                },
                Expr::Record {
                    ref mut typ,
                    ref mut exprs,
                    ref mut base,
                    ..
                } => {
                    let field_types = self.find_fields(typ);
                    for (field, expr_field) in field_types.iter().zip(exprs) {
                        match expr_field.value {
                            Some(ref mut expr) => self.visit_expr(expr),
                            None => if let Some(new_id) =
                                self.rename(&expr_field.name.value, &field.typ)?
                            {
                                debug!("Rename record field {} = {}", expr_field.name, new_id);
                                expr_field.value = Some(pos::spanned(
                                    expr.span,
                                    Expr::Ident(TypedIdent {
                                        name: new_id,
                                        typ: field.typ.clone(),
                                    }),
                                ));
                            },
                        }
                    }

                    if let Some(ref mut base) = *base {
                        self.visit_expr(base);
                    }
                }
                Expr::Infix(ref mut lhs, ref mut op, ref mut rhs) => {
                    if let Some(new_id) = self.rename(&op.value.name, &op.value.typ)? {
                        debug!(
                            "Rename {} = {}",
                            self.symbols.string(&op.value.name),
                            self.symbols.string(&new_id)
                        );
                        op.value.name = new_id;
                    }
                    self.visit_expr(lhs);
                    self.visit_expr(rhs);
                }
                Expr::Match(ref mut expr, ref mut alts) => {
                    self.visit_expr(expr);
                    for alt in alts {
                        self.env.stack_types.enter_scope();
                        self.env.stack.enter_scope();
                        let typ = expr.env_type_of(&self.env);
                        self.new_pattern(&typ, &mut alt.pattern);
                        self.visit_expr(&mut alt.expr);
                        self.env.stack.exit_scope();
                        self.env.stack_types.exit_scope();
                    }
                }
                Expr::LetBindings(ref mut bindings, ref mut expr) => {
                    self.env.stack_types.enter_scope();
                    self.env.stack.enter_scope();
                    let is_recursive = bindings.iter().all(|bind| !bind.args.is_empty());
                    for bind in bindings.iter_mut() {
                        if !is_recursive {
                            self.visit_expr(&mut bind.expr);
                        }
                        self.new_pattern(&bind.resolved_type, &mut bind.name);
                    }
                    if is_recursive {
                        for bind in bindings {
                            self.env.stack.enter_scope();
                            for (typ, arg) in types::arg_iter(bind.resolved_type.remove_forall())
                                .zip(&mut bind.args)
                            {
                                arg.value.name =
                                    self.stack_var(arg.value.name.clone(), expr.span, typ.clone());
                            }
                            self.visit_expr(&mut bind.expr);
                            self.env.stack.exit_scope();
                        }
                    }
                    self.visit_expr(expr);
                    self.env.stack.exit_scope();
                    self.env.stack_types.exit_scope();
                }
                Expr::Lambda(ref mut lambda) => {
                    self.env.stack.enter_scope();
                    for (typ, arg) in types::arg_iter(&lambda.id.typ).zip(&mut lambda.args) {
                        arg.value.name =
                            self.stack_var(arg.value.name.clone(), expr.span, typ.clone());
                    }
                    self.visit_expr(&mut lambda.body);
                    self.env.stack.exit_scope();
                }
                Expr::TypeBindings(ref bindings, ref mut body) => {
                    self.env.stack_types.enter_scope();
                    for bind in bindings {
                        self.stack_type(
                            bind.name.value.clone(),
                            expr.span,
                            bind.finalized_alias.as_ref().expect(
                                "ICE: Alias should have been finalized \
                                 before renaming",
                            ),
                        );
                    }
                    self.visit_expr(body);
                    self.env.stack_types.exit_scope();
                }
                _ => ast::walk_mut_expr(self, expr),
            }
            Ok(())
        }
    }

    impl<'a, 'b> MutVisitor for RenameVisitor<'a, 'b> {
        type Ident = Symbol;

        fn visit_expr(&mut self, expr: &mut SpannedExpr<Self::Ident>) {
            if let Err(err) = self.rename_expr(expr) {
                self.errors.push(Spanned {
                    span: expr.span,
                    value: err,
                });
            }
        }
    }

    let mut visitor = RenameVisitor {
        symbols: symbols,
        errors: Errors::new(),
        env: Environment {
            env: env,
            stack: ScopedMap::new(),
            stack_types: ScopedMap::new(),
        },
    };
    visitor.visit_expr(expr);
    if visitor.errors.has_errors() {
        Err(visitor.errors)
    } else {
        Ok(())
    }
}

pub fn equivalent(env: &TypeEnv, actual: &ArcType, inferred: &ArcType) -> bool {
    use substitution::Substitution;
    // FIXME This Substitution is unnecessary for equivalence unification
    let subs = Substitution::new(Kind::typ());
    let mut unifier = UnifierState {
        state: State::new(env, &subs),
        unifier: Equivalent {
            map: FnvMap::default(),
            equiv: true,
        },
    };
    unifier.try_match(actual, inferred);
    unifier.unifier.equiv
}

struct Equivalent {
    map: FnvMap<Symbol, ArcType>,
    equiv: bool,
}

impl<'a> Unifier<State<'a>, ArcType> for Equivalent {
    fn report_error(
        unifier: &mut UnifierState<State<'a>, Self>,
        _error: UnifyError<ArcType, TypeError<Symbol>>,
    ) {
        unifier.unifier.equiv = false;
    }

    fn try_match_res(
        unifier: &mut UnifierState<State<'a>, Self>,
        l: &ArcType,
        r: &ArcType,
    ) -> Result<Option<ArcType>, UnifyError<ArcType, TypeError<Symbol>>> {
        debug!("{} ====> {}", l, r);
        match (&**l, &**r) {
            (&Type::Generic(ref gl), &Type::Generic(ref gr)) if gl == gr => Ok(None),
            (&Type::Skolem(ref gl), &Type::Skolem(ref gr)) if gl == gr => Ok(None),
            (&Type::Generic(ref gl), _) => match unifier.unifier.map.get(&gl.id).cloned() {
                Some(ref typ) => unifier.try_match_res(typ, r),
                None => {
                    unifier.unifier.map.insert(gl.id.clone(), r.clone());
                    Ok(None)
                }
            },
            (&Type::Skolem(ref gl), _) => match unifier.unifier.map.get(&gl.name).cloned() {
                Some(ref typ) => unifier.try_match_res(typ, r),
                None => {
                    unifier.unifier.map.insert(gl.name.clone(), r.clone());
                    Ok(None)
                }
            },
            _ => l.zip_match(r, unifier),
        }
    }

    fn error_type(_unifier: &mut UnifierState<State<'a>, Self>) -> Option<ArcType> {
        None
    }
}
<|endoftext|>


<commit_before>#![doc(html_favicon_url = "https://www.ruma.io/favicon.ico")]
#![doc(html_logo_url = "https://www.ruma.io/images/logo.png")]
//! Types for [Matrix](https://matrix.org/) identifiers for devices, events, keys, rooms, servers,
//! users and URIs.

#![warn(missing_docs)]
// FIXME: Remove once lint doesn't trigger on std::convert::TryFrom in macros.rs anymore
#![allow(unused_qualifications)]
#![cfg_attr(docsrs, feature(doc_cfg))]

// Renamed in `Cargo.toml` so we can have a serde feature.
// Rename it back here because `serde1` is ugly.
#[cfg(feature = "serde")]
extern crate serde1 as serde;

use std::convert::TryFrom;

#[cfg(feature = "serde")]
use serde::de::{self, Deserializer, Unexpected};

#[doc(inline)]
pub use crate::{
    client_secret::{ClientSecret, ClientSecretBox},
    crypto_algorithms::{DeviceKeyAlgorithm, EventEncryptionAlgorithm, SigningKeyAlgorithm},
    device_key_id::DeviceKeyId,
    event_id::EventId,
    key_id::{DeviceSigningKeyId, KeyId, ServerSigningKeyId, SigningKeyId},
    mxc_uri::MxcUri,
    opaque_ids::{DeviceId, DeviceIdBox, KeyName, KeyNameBox},
    room_alias_id::RoomAliasId,
    room_id::RoomId,
    room_id_or_room_alias_id::RoomIdOrAliasId,
    room_version_id::RoomVersionId,
    server_name::{ServerName, ServerNameBox},
    session_id::{SessionId, SessionIdBox},
    signatures::{DeviceSignatures, EntitySignatures, ServerSignatures, Signatures},
    user_id::UserId,
};
#[doc(inline)]
pub use ruma_identifiers_validation::error::Error;

#[macro_use]
mod macros;

pub mod user_id;

mod client_secret;
mod crypto_algorithms;
mod device_key_id;
mod event_id;
mod key_id;
mod mxc_uri;
mod opaque_ids;
mod room_alias_id;
mod room_id;
mod room_id_or_room_alias_id;
mod room_version_id;
mod server_name;
mod session_id;
mod signatures;

/// Generates a random identifier localpart.
#[cfg(feature = "rand")]
fn generate_localpart(length: usize) -> Box<str> {
    use rand::Rng as _;
    rand::thread_rng()
        .sample_iter(&rand::distributions::Alphanumeric)
        .map(char::from)
        .take(length)
        .collect::<String>()
        .into_boxed_str()
}

/// Deserializes any type of id using the provided TryFrom implementation.
///
/// This is a helper function to reduce the boilerplate of the Deserialize implementations.
#[cfg(feature = "serde")]
fn deserialize_id<'de, D, T>(deserializer: D, expected_str: &str) -> Result<T, D::Error>
where
    D: Deserializer<'de>,
    T: for<'a> TryFrom<&'a str>,
{
    ruma_serde::deserialize_cow_str(deserializer).and_then(|v| {
        T::try_from(&v).map_err(|_| de::Error::invalid_value(Unexpected::Str(&v), &expected_str))
    })
}

/// Shorthand for `Box::<DeviceId>::from`.
#[macro_export]
macro_rules! device_id {
    ($s:tt) => {
        ::std::boxed::Box::<$crate::DeviceId>::from($s)
    };
}

// A plain re-export shows up in rustdoc despite doc(hidden). Use a module instead.
// Bug report: https://github.com/rust-lang/rust/issues/83939
#[doc(hidden)]
pub mod _macros {
    pub use ruma_identifiers_macros::*;
}

/// Compile-time checked `DeviceKeyId` construction.
#[macro_export]
macro_rules! device_key_id {
    ($s:literal) => {
        $crate::_macros::device_key_id!($crate, $s)
    };
}

/// Compile-time checked `EventId` construction.
#[macro_export]
macro_rules! event_id {
    ($s:literal) => {
        $crate::_macros::event_id!($crate, $s)
    };
}

/// Compile-time checked `RoomAliasId` construction.
#[macro_export]
macro_rules! room_alias_id {
    ($s:literal) => {
        $crate::_macros::room_alias_id!($crate, $s)
    };
}

/// Compile-time checked `RoomId` construction.
#[macro_export]
macro_rules! room_id {
    ($s:literal) => {
        $crate::_macros::room_id!($crate, $s)
    };
}

/// Compile-time checked `RoomVersionId` construction.
#[macro_export]
macro_rules! room_version_id {
    ($s:literal) => {
        $crate::_macros::room_version_id!($crate, $s)
    };
}

/// Compile-time checked `ServerSigningKeyId` construction.
#[macro_export]
macro_rules! server_signing_key_id {
    ($s:literal) => {
        $crate::_macros::server_signing_key_id!($crate, $s)
    };
}

/// Compile-time checked `ServerName` construction.
#[macro_export]
macro_rules! server_name {
    ($s:literal) => {
        $crate::_macros::server_name!($crate, $s)
    };
}

/// Compile-time checked `MxcUri` construction.
#[macro_export]
macro_rules! mxc_uri {
    ($s:literal) => {
        $crate::_macros::mxc_uri!($crate, $s)
    };
}

/// Compile-time checked `UserId` construction.
#[macro_export]
macro_rules! user_id {
    ($s:literal) => {
        $crate::_macros::user_id!($crate, $s)
    };
}
<commit_msg>identifiers: Fix unused import without serde feature<commit_after>#![doc(html_favicon_url = "https://www.ruma.io/favicon.ico")]
#![doc(html_logo_url = "https://www.ruma.io/images/logo.png")]
//! Types for [Matrix](https://matrix.org/) identifiers for devices, events, keys, rooms, servers,
//! users and URIs.

#![warn(missing_docs)]
// FIXME: Remove once lint doesn't trigger on std::convert::TryFrom in macros.rs anymore
#![allow(unused_qualifications)]
#![cfg_attr(docsrs, feature(doc_cfg))]

// Renamed in `Cargo.toml` so we can have a serde feature.
// Rename it back here because `serde1` is ugly.
#[cfg(feature = "serde")]
extern crate serde1 as serde;

#[cfg(feature = "serde")]
use std::convert::TryFrom;

#[cfg(feature = "serde")]
use serde::de::{self, Deserializer, Unexpected};

#[doc(inline)]
pub use crate::{
    client_secret::{ClientSecret, ClientSecretBox},
    crypto_algorithms::{DeviceKeyAlgorithm, EventEncryptionAlgorithm, SigningKeyAlgorithm},
    device_key_id::DeviceKeyId,
    event_id::EventId,
    key_id::{DeviceSigningKeyId, KeyId, ServerSigningKeyId, SigningKeyId},
    mxc_uri::MxcUri,
    opaque_ids::{DeviceId, DeviceIdBox, KeyName, KeyNameBox},
    room_alias_id::RoomAliasId,
    room_id::RoomId,
    room_id_or_room_alias_id::RoomIdOrAliasId,
    room_version_id::RoomVersionId,
    server_name::{ServerName, ServerNameBox},
    session_id::{SessionId, SessionIdBox},
    signatures::{DeviceSignatures, EntitySignatures, ServerSignatures, Signatures},
    user_id::UserId,
};
#[doc(inline)]
pub use ruma_identifiers_validation::error::Error;

#[macro_use]
mod macros;

pub mod user_id;

mod client_secret;
mod crypto_algorithms;
mod device_key_id;
mod event_id;
mod key_id;
mod mxc_uri;
mod opaque_ids;
mod room_alias_id;
mod room_id;
mod room_id_or_room_alias_id;
mod room_version_id;
mod server_name;
mod session_id;
mod signatures;

/// Generates a random identifier localpart.
#[cfg(feature = "rand")]
fn generate_localpart(length: usize) -> Box<str> {
    use rand::Rng as _;
    rand::thread_rng()
        .sample_iter(&rand::distributions::Alphanumeric)
        .map(char::from)
        .take(length)
        .collect::<String>()
        .into_boxed_str()
}

/// Deserializes any type of id using the provided TryFrom implementation.
///
/// This is a helper function to reduce the boilerplate of the Deserialize implementations.
#[cfg(feature = "serde")]
fn deserialize_id<'de, D, T>(deserializer: D, expected_str: &str) -> Result<T, D::Error>
where
    D: Deserializer<'de>,
    T: for<'a> TryFrom<&'a str>,
{
    ruma_serde::deserialize_cow_str(deserializer).and_then(|v| {
        T::try_from(&v).map_err(|_| de::Error::invalid_value(Unexpected::Str(&v), &expected_str))
    })
}

/// Shorthand for `Box::<DeviceId>::from`.
#[macro_export]
macro_rules! device_id {
    ($s:tt) => {
        ::std::boxed::Box::<$crate::DeviceId>::from($s)
    };
}

// A plain re-export shows up in rustdoc despite doc(hidden). Use a module instead.
// Bug report: https://github.com/rust-lang/rust/issues/83939
#[doc(hidden)]
pub mod _macros {
    pub use ruma_identifiers_macros::*;
}

/// Compile-time checked `DeviceKeyId` construction.
#[macro_export]
macro_rules! device_key_id {
    ($s:literal) => {
        $crate::_macros::device_key_id!($crate, $s)
    };
}

/// Compile-time checked `EventId` construction.
#[macro_export]
macro_rules! event_id {
    ($s:literal) => {
        $crate::_macros::event_id!($crate, $s)
    };
}

/// Compile-time checked `RoomAliasId` construction.
#[macro_export]
macro_rules! room_alias_id {
    ($s:literal) => {
        $crate::_macros::room_alias_id!($crate, $s)
    };
}

/// Compile-time checked `RoomId` construction.
#[macro_export]
macro_rules! room_id {
    ($s:literal) => {
        $crate::_macros::room_id!($crate, $s)
    };
}

/// Compile-time checked `RoomVersionId` construction.
#[macro_export]
macro_rules! room_version_id {
    ($s:literal) => {
        $crate::_macros::room_version_id!($crate, $s)
    };
}

/// Compile-time checked `ServerSigningKeyId` construction.
#[macro_export]
macro_rules! server_signing_key_id {
    ($s:literal) => {
        $crate::_macros::server_signing_key_id!($crate, $s)
    };
}

/// Compile-time checked `ServerName` construction.
#[macro_export]
macro_rules! server_name {
    ($s:literal) => {
        $crate::_macros::server_name!($crate, $s)
    };
}

/// Compile-time checked `MxcUri` construction.
#[macro_export]
macro_rules! mxc_uri {
    ($s:literal) => {
        $crate::_macros::mxc_uri!($crate, $s)
    };
}

/// Compile-time checked `UserId` construction.
#[macro_export]
macro_rules! user_id {
    ($s:literal) => {
        $crate::_macros::user_id!($crate, $s)
    };
}
<|endoftext|>


<commit_before>///! Helper functions and macros for tests
use std::{fs::read_to_string, path::PathBuf};

/// Get the path of the home directory of this repository
pub fn home() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")) // the crate
        .parent().expect("Should have parent") // the `rust` dir
        .parent().expect("Should have parent") // the `stencila` dir
        .canonicalize()
        .expect("Unable to get repository home directory")
}

/// Get the path of the `fixtures` directory
pub fn fixtures() -> PathBuf {
    home().join("fixtures")
}

/// Generate snapshots from the string content of fixtures matching
/// a glob pattern.
///
/// # Arguments
///
/// - `pattern`: glob pattern _within_ the fixtures folder
/// - `func`: the test function to run on the string content of each
///           file matching the `pattern`.
///
/// `insta`'s `glob` macro seems to have difficulties with our Rust module
/// layout (workspaces and nested modules). This function deals with that
/// by making the pattern relative to the fixtures and adding some other
/// conveniences.
pub fn snapshot_fixtures<F: FnMut(&str, &str)>(pattern: &str, mut func: F) {
    let mut settings = insta::Settings::clone_current();
    settings.set_prepend_module_to_snapshot(false);
    settings.bind(|| {
        insta::_macro_support::glob_exec(&fixtures(), pattern, |path| {
            let content = read_to_string(path).unwrap();
            let path = pathdiff::diff_paths(path, fixtures())
                .unwrap()
                .display()
                .to_string();
            func(&path, &content)
        });
    });
}

/// Should slow tests be skipped?
///
/// Use at the start of slow tests to return early except on CI or when
/// the env var `RUN_SLOW_TESTS` is set.
///
/// Inspired by https://github.com/rust-analyzer/rust-analyzer/pull/2491
pub fn skip_slow_tests() -> bool {
    if std::env::var("CI").is_err() && std::env::var("RUN_SLOW_TESTS").is_err() {
        eprintln!("Skipping slow test");
        true
    } else {
        false
    }
}

/// Assert that an expression is serialized to expected JSON
#[macro_export]
macro_rules! assert_json {
    ($expr:expr, $json:tt) => {
        pretty_assertions::assert_eq!(
            serde_json::to_value(&$expr).unwrap(),
            serde_json::json!($json)
        );
    };
}

/// Assert that the JSON serialization of two expression is equal
#[macro_export]
macro_rules! assert_json_eq {
    ($expr1:expr, $expr2:expr) => {
        pretty_assertions::assert_eq!(
            serde_json::to_value(&$expr1).unwrap(),
            serde_json::to_value(&$expr2).unwrap()
        );
    };
}
<commit_msg>chore(Tests): Formatting<commit_after>///! Helper functions and macros for tests
use std::{fs::read_to_string, path::PathBuf};

/// Get the path of the home directory of this repository
pub fn home() -> PathBuf {
    PathBuf::from(env!("CARGO_MANIFEST_DIR")) // the crate
        .parent()
        .expect("Should have parent") // the `rust` dir
        .parent()
        .expect("Should have parent") // the `stencila` dir
        .canonicalize()
        .expect("Unable to get repository home directory")
}

/// Get the path of the `fixtures` directory
pub fn fixtures() -> PathBuf {
    home().join("fixtures")
}

/// Generate snapshots from the string content of fixtures matching
/// a glob pattern.
///
/// # Arguments
///
/// - `pattern`: glob pattern _within_ the fixtures folder
/// - `func`: the test function to run on the string content of each
///           file matching the `pattern`.
///
/// `insta`'s `glob` macro seems to have difficulties with our Rust module
/// layout (workspaces and nested modules). This function deals with that
/// by making the pattern relative to the fixtures and adding some other
/// conveniences.
pub fn snapshot_fixtures<F: FnMut(&str, &str)>(pattern: &str, mut func: F) {
    let mut settings = insta::Settings::clone_current();
    settings.set_prepend_module_to_snapshot(false);
    settings.bind(|| {
        insta::_macro_support::glob_exec(&fixtures(), pattern, |path| {
            let content = read_to_string(path).unwrap();
            let path = pathdiff::diff_paths(path, fixtures())
                .unwrap()
                .display()
                .to_string();
            func(&path, &content)
        });
    });
}

/// Should slow tests be skipped?
///
/// Use at the start of slow tests to return early except on CI or when
/// the env var `RUN_SLOW_TESTS` is set.
///
/// Inspired by https://github.com/rust-analyzer/rust-analyzer/pull/2491
pub fn skip_slow_tests() -> bool {
    if std::env::var("CI").is_err() && std::env::var("RUN_SLOW_TESTS").is_err() {
        eprintln!("Skipping slow test");
        true
    } else {
        false
    }
}

/// Assert that an expression is serialized to expected JSON
#[macro_export]
macro_rules! assert_json {
    ($expr:expr, $json:tt) => {
        pretty_assertions::assert_eq!(
            serde_json::to_value(&$expr).unwrap(),
            serde_json::json!($json)
        );
    };
}

/// Assert that the JSON serialization of two expression is equal
#[macro_export]
macro_rules! assert_json_eq {
    ($expr1:expr, $expr2:expr) => {
        pretty_assertions::assert_eq!(
            serde_json::to_value(&$expr1).unwrap(),
            serde_json::to_value(&$expr2).unwrap()
        );
    };
}
<|endoftext|>


<commit_before>// aux-build:arc_wake.rs
// edition:2018
// run-pass

#![allow(unused_variables)]

// Test that the drop order for parameters in a fn and async fn matches up. Also test that
// parameters (used or unused) are not dropped until the async fn completes execution.
// See also #54716.

extern crate arc_wake;

use arc_wake::ArcWake;
use std::cell::RefCell;
use std::future::Future;
use std::sync::Arc;
use std::rc::Rc;
use std::task::Context;

struct EmptyWaker;

impl ArcWake for EmptyWaker {
    fn wake(self: Arc<Self>) {}
}

#[derive(Debug, Eq, PartialEq)]
enum DropOrder {
    Function,
    Val(&'static str),
}

type DropOrderListPtr = Rc<RefCell<Vec<DropOrder>>>;

struct D(&'static str, DropOrderListPtr);

impl Drop for D {
    fn drop(&mut self) {
        self.1.borrow_mut().push(DropOrder::Val(self.0));
    }
}

/// Check drop order of temporary "temp" as compared to x, y, and z.
///
/// Expected order:
/// - z
/// - temp
/// - y
/// - x
async fn foo_async(x: D, _y: D) {
    let l = x.1.clone();
    let z = D("z", l.clone());
    l.borrow_mut().push(DropOrder::Function);
    helper_async(&D("temp", l)).await
}

async fn helper_async(v: &D) { }

fn foo_sync(x: D, _y: D) {
    let l = x.1.clone();
    let z = D("z", l.clone());
    l.borrow_mut().push(DropOrder::Function);
    helper_sync(&D("temp", l))
}

fn helper_sync(v: &D) { }

fn assert_drop_order_after_poll<Fut: Future<Output = ()>>(
    f: impl FnOnce(DropOrderListPtr) -> Fut,
    g: impl FnOnce(DropOrderListPtr),
) {
    let empty = Arc::new(EmptyWaker);
    let waker = ArcWake::into_waker(empty);
    let mut cx = Context::from_waker(&waker);

    let actual_order = Rc::new(RefCell::new(Vec::new()));
    let mut fut = Box::pin(f(actual_order.clone()));
    let r = fut.as_mut().poll(&mut cx);

    assert!(match r {
        std::task::Poll::Ready(()) => true,
        _ => false,
    });

    let expected_order = Rc::new(RefCell::new(Vec::new()));
    g(expected_order.clone());

    assert_eq!(*actual_order.borrow(), *expected_order.borrow());
}

fn main() {
    // Free functions (see doc comment on function for what it tests).
    assert_drop_order_after_poll(|l| foo_async(D("x", l.clone()), D("_y", l.clone())),
                                 |l| foo_sync(D("x", l.clone()), D("_y", l.clone())));
}
<commit_msg>apply nits from centril<commit_after>// aux-build:arc_wake.rs
// edition:2018
// run-pass

#![allow(unused_variables)]

// Test the drop order for parameters relative to local variables and
// temporaries created in the tail return expression of the function
// body. In particular, check that this drop order is the same between
// a `async fn` and an ordinary `fn`. See #64512.

extern crate arc_wake;

use arc_wake::ArcWake;
use std::cell::RefCell;
use std::future::Future;
use std::sync::Arc;
use std::rc::Rc;
use std::task::Context;

struct EmptyWaker;

impl ArcWake for EmptyWaker {
    fn wake(self: Arc<Self>) {}
}

#[derive(Debug, Eq, PartialEq)]
enum DropOrder {
    Function,
    Val(&'static str),
}

type DropOrderListPtr = Rc<RefCell<Vec<DropOrder>>>;

struct D(&'static str, DropOrderListPtr);

impl Drop for D {
    fn drop(&mut self) {
        self.1.borrow_mut().push(DropOrder::Val(self.0));
    }
}

/// Check drop order of temporary "temp" as compared to `x`, `y`, and `z`.
///
/// Expected order:
/// - `z`
/// - temp
/// - `y`
/// - `x`
async fn foo_async(x: D, _y: D) {
    let l = x.1.clone();
    let z = D("z", l.clone());
    l.borrow_mut().push(DropOrder::Function);
    helper_async(&D("temp", l)).await
}

async fn helper_async(v: &D) { }

fn foo_sync(x: D, _y: D) {
    let l = x.1.clone();
    let z = D("z", l.clone());
    l.borrow_mut().push(DropOrder::Function);
    helper_sync(&D("temp", l))
}

fn helper_sync(v: &D) { }

fn assert_drop_order_after_poll<Fut: Future<Output = ()>>(
    f: impl FnOnce(DropOrderListPtr) -> Fut,
    g: impl FnOnce(DropOrderListPtr),
) {
    let empty = Arc::new(EmptyWaker);
    let waker = ArcWake::into_waker(empty);
    let mut cx = Context::from_waker(&waker);

    let actual_order = Rc::new(RefCell::new(Vec::new()));
    let mut fut = Box::pin(f(actual_order.clone()));
    let r = fut.as_mut().poll(&mut cx);

    assert!(match r {
        std::task::Poll::Ready(()) => true,
        _ => false,
    });

    let expected_order = Rc::new(RefCell::new(Vec::new()));
    g(expected_order.clone());

    assert_eq!(*actual_order.borrow(), *expected_order.borrow());
}

fn main() {
    // Free functions (see doc comment on function for what it tests).
    assert_drop_order_after_poll(|l| foo_async(D("x", l.clone()), D("_y", l.clone())),
                                 |l| foo_sync(D("x", l.clone()), D("_y", l.clone())));
}
<|endoftext|>


<commit_before>use futures_01::{
    task as task01, Async as Async01, Future as Future01, Poll as Poll01,
    Stream as Stream01,
};
#[cfg(feature = "sink")]
use futures_01::{
    AsyncSink as AsyncSink01, Sink as Sink01, StartSend as StartSend01,
};
use futures_core::{
    task::{RawWaker, RawWakerVTable},
    future::TryFuture as TryFuture03,
    stream::TryStream as TryStream03,
};
#[cfg(feature = "sink")]
use futures_sink::Sink as Sink03;
use crate::task::{
    self as task03,
    ArcWake as ArcWake03,
    WakerRef,
};
#[cfg(feature = "sink")]
use std::marker::PhantomData;
use std::{
    mem,
    pin::Pin,
    sync::Arc,
    task::Context,
};

/// Converts a futures 0.3 [`TryFuture`](futures_core::future::TryFuture),
/// [`TryStream`](futures_core::stream::TryStream) or
/// [`Sink`](futures_sink::Sink) into a futures 0.1
/// [`Future`](futures_01::future::Future),
/// [`Stream`](futures_01::stream::Stream) or
/// [`Sink`](futures_01::sink::Sink).
#[derive(Debug, Clone, Copy)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct Compat<T> {
    pub(crate) inner: T,
}

/// Converts a futures 0.3 Sink object to a futures 0.1-compatible version
#[cfg(feature = "sink")]
#[derive(Debug)]
#[must_use = "sinks do nothing unless polled"]
pub struct CompatSink<T, Item> {
    inner: T,
    _phantom: PhantomData<fn(Item)>,
}

impl<T> Compat<T> {
    /// Creates a new [`Compat`].
    ///
    /// For types which implement appropriate futures `0.3`
    /// traits, the result will be a type which implements
    /// the corresponding futures 0.1 type.
    pub fn new(inner: T) -> Compat<T> {
        Compat { inner }
    }

    /// Get a reference to 0.3 Future, Stream, AsyncRead, or AsyncWrite object
    /// contained within.
    pub fn get_ref(&self) -> &T {
        &self.inner
    }

    /// Get a mutable reference to 0.3 Future, Stream, AsyncRead, or AsyncWrite object
    /// contained within.
    pub fn get_mut(&mut self) -> &mut T {
        &mut self.inner
    }

    /// Returns the inner item.
    pub fn into_inner(self) -> T {
        self.inner
    }
}

#[cfg(feature = "sink")]
impl<T, Item> CompatSink<T, Item> {
    /// Creates a new [`CompatSink`].
    pub fn new(inner: T) -> Self {
        CompatSink {
            inner,
            _phantom: PhantomData,
        }
    }

    /// Get a reference to 0.3 Sink contained within.
    pub fn get_ref(&self) -> &T {
        &self.inner
    }

    /// Get a mutable reference to 0.3 Sink contained within.
    pub fn get_mut(&mut self) -> &mut T {
        &mut self.inner
    }

    /// Returns the inner item.
    pub fn into_inner(self) -> T {
        self.inner
    }
}

fn poll_03_to_01<T, E>(x: task03::Poll<Result<T, E>>)
    -> Result<Async01<T>, E>
{
    match x? {
        task03::Poll::Ready(t) => Ok(Async01::Ready(t)),
        task03::Poll::Pending => Ok(Async01::NotReady),
    }
}

impl<Fut> Future01 for Compat<Fut>
where
    Fut: TryFuture03 + Unpin,
{
    type Item = Fut::Ok;
    type Error = Fut::Error;

    fn poll(&mut self) -> Poll01<Self::Item, Self::Error> {
        with_context(self, |inner, cx| poll_03_to_01(inner.try_poll(cx)))
    }
}

impl<St> Stream01 for Compat<St>
where
    St: TryStream03 + Unpin,
{
    type Item = St::Ok;
    type Error = St::Error;

    fn poll(&mut self) -> Poll01<Option<Self::Item>, Self::Error> {
        with_context(self, |inner, cx| match inner.try_poll_next(cx)? {
            task03::Poll::Ready(None) => Ok(Async01::Ready(None)),
            task03::Poll::Ready(Some(t)) => Ok(Async01::Ready(Some(t))),
            task03::Poll::Pending => Ok(Async01::NotReady),
        })
    }
}

#[cfg(feature = "sink")]
impl<T, Item> Sink01 for CompatSink<T, Item>
where
    T: Sink03<Item> + Unpin,
{
    type SinkItem = Item;
    type SinkError = T::Error;

    fn start_send(
        &mut self,
        item: Self::SinkItem,
    ) -> StartSend01<Self::SinkItem, Self::SinkError> {
        with_sink_context(self, |mut inner, cx| {
            match inner.as_mut().poll_ready(cx)? {
                task03::Poll::Ready(()) => {
                    inner.start_send(item).map(|()| AsyncSink01::Ready)
                }
                task03::Poll::Pending => Ok(AsyncSink01::NotReady(item)),
            }
        })
    }

    fn poll_complete(&mut self) -> Poll01<(), Self::SinkError> {
        with_sink_context(self, |inner, cx| poll_03_to_01(inner.poll_flush(cx)))
    }

    fn close(&mut self) -> Poll01<(), Self::SinkError> {
        with_sink_context(self, |inner, cx| poll_03_to_01(inner.poll_close(cx)))
    }
}

#[derive(Clone)]
struct Current(task01::Task);

impl Current {
    fn new() -> Current {
        Current(task01::current())
    }

    fn as_waker(&self) -> WakerRef<'_> {
        unsafe fn ptr_to_current<'a>(ptr: *const ()) -> &'a Current {
            &*(ptr as *const Current)
        }
        fn current_to_ptr(current: &Current) -> *const () {
            current as *const Current as *const ()
        }

        unsafe fn clone(ptr: *const ()) -> RawWaker {
            // Lazily create the `Arc` only when the waker is actually cloned.
            // FIXME: remove `transmute` when a `Waker` -> `RawWaker` conversion
            // function is landed in `core`.
            mem::transmute::<task03::Waker, RawWaker>(
                task03::waker(Arc::new(ptr_to_current(ptr).clone()))
            )
        }
        unsafe fn drop(_: *const ()) {}
        unsafe fn wake(ptr: *const ()) {
            ptr_to_current(ptr).0.notify()
        }

        let ptr = current_to_ptr(self);
        let vtable = &RawWakerVTable::new(clone, wake, wake, drop);
        WakerRef::new_unowned(std::mem::ManuallyDrop::new(unsafe {
            task03::Waker::from_raw(RawWaker::new(ptr, vtable))
        }))
    }
}

impl ArcWake03 for Current {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        arc_self.0.notify();
    }
}

fn with_context<T, R, F>(compat: &mut Compat<T>, f: F) -> R
where
    T: Unpin,
    F: FnOnce(Pin<&mut T>, &mut Context<'_>) -> R,
{
    let current = Current::new();
    let waker = current.as_waker();
    let mut cx = Context::from_waker(&waker);
    f(Pin::new(&mut compat.inner), &mut cx)
}

#[cfg(feature = "sink")]
fn with_sink_context<T, Item, R, F>(compat: &mut CompatSink<T, Item>, f: F) -> R
where
    T: Unpin,
    F: FnOnce(Pin<&mut T>, &mut Context<'_>) -> R,
{
    let current = Current::new();
    let waker = current.as_waker();
    let mut cx = Context::from_waker(&waker);
    f(Pin::new(&mut compat.inner), &mut cx)
}

#[cfg(feature = "io-compat")]
mod io {
    use super::*;
    use futures_io::{AsyncRead as AsyncRead03, AsyncWrite as AsyncWrite03};
    use tokio_io::{AsyncRead as AsyncRead01, AsyncWrite as AsyncWrite01};

    fn poll_03_to_io<T>(x: task03::Poll<Result<T, std::io::Error>>)
        -> Result<T, std::io::Error>
    {
        match x {
            task03::Poll::Ready(Ok(t)) => Ok(t),
            task03::Poll::Pending => Err(std::io::ErrorKind::WouldBlock.into()),
            task03::Poll::Ready(Err(e)) => Err(e),
        }
    }

    impl<R: AsyncRead03 + Unpin> std::io::Read for Compat<R> {
        fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
            let current = Current::new();
            let waker = current.as_waker();
            let mut cx = Context::from_waker(&waker);
            poll_03_to_io(Pin::new(&mut self.inner).poll_read(&mut cx, buf))
        }
    }

    impl<R: AsyncRead03 + Unpin> AsyncRead01 for Compat<R> {
        #[cfg(feature = "read-initializer")]
        unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
            let initializer = self.inner.initializer();
            let does_init = initializer.should_initialize();
            if does_init {
                initializer.initialize(buf);
            }
            does_init
        }
    }

    impl<W: AsyncWrite03 + Unpin> std::io::Write for Compat<W> {
        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
            let current = Current::new();
            let waker = current.as_waker();
            let mut cx = Context::from_waker(&waker);
            poll_03_to_io(Pin::new(&mut self.inner).poll_write(&mut cx, buf))
        }

        fn flush(&mut self) -> std::io::Result<()> {
            let current = Current::new();
            let waker = current.as_waker();
            let mut cx = Context::from_waker(&waker);
            poll_03_to_io(Pin::new(&mut self.inner).poll_flush(&mut cx))
        }
    }

    impl<W: AsyncWrite03 + Unpin> AsyncWrite01 for Compat<W> {
        fn shutdown(&mut self) -> std::io::Result<Async01<()>> {
            let current = Current::new();
            let waker = current.as_waker();
            let mut cx = Context::from_waker(&waker);
            poll_03_to_01(Pin::new(&mut self.inner).poll_close(&mut cx))
        }
    }
}
<commit_msg>futures-util/src/compat: Fix Compat and CompatSink documentation<commit_after>use futures_01::{
    task as task01, Async as Async01, Future as Future01, Poll as Poll01,
    Stream as Stream01,
};
#[cfg(feature = "sink")]
use futures_01::{
    AsyncSink as AsyncSink01, Sink as Sink01, StartSend as StartSend01,
};
use futures_core::{
    task::{RawWaker, RawWakerVTable},
    future::TryFuture as TryFuture03,
    stream::TryStream as TryStream03,
};
#[cfg(feature = "sink")]
use futures_sink::Sink as Sink03;
use crate::task::{
    self as task03,
    ArcWake as ArcWake03,
    WakerRef,
};
#[cfg(feature = "sink")]
use std::marker::PhantomData;
use std::{
    mem,
    pin::Pin,
    sync::Arc,
    task::Context,
};

/// Converts a futures 0.3 [`TryFuture`](futures_core::future::TryFuture) or
/// [`TryStream`](futures_core::stream::TryStream) into a futures 0.1
/// [`Future`](futures_01::future::Future) or
/// [`Stream`](futures_01::stream::Stream).
#[derive(Debug, Clone, Copy)]
#[must_use = "futures do nothing unless you `.await` or poll them"]
pub struct Compat<T> {
    pub(crate) inner: T,
}

/// Converts a futures 0.3 [`Sink`](futures_sink::Sink) into a futures 0.1
/// [`Sink`](futures_01::sink::Sink).
#[cfg(feature = "sink")]
#[derive(Debug)]
#[must_use = "sinks do nothing unless polled"]
pub struct CompatSink<T, Item> {
    inner: T,
    _phantom: PhantomData<fn(Item)>,
}

impl<T> Compat<T> {
    /// Creates a new [`Compat`].
    ///
    /// For types which implement appropriate futures `0.3`
    /// traits, the result will be a type which implements
    /// the corresponding futures 0.1 type.
    pub fn new(inner: T) -> Compat<T> {
        Compat { inner }
    }

    /// Get a reference to 0.3 Future, Stream, AsyncRead, or AsyncWrite object
    /// contained within.
    pub fn get_ref(&self) -> &T {
        &self.inner
    }

    /// Get a mutable reference to 0.3 Future, Stream, AsyncRead, or AsyncWrite object
    /// contained within.
    pub fn get_mut(&mut self) -> &mut T {
        &mut self.inner
    }

    /// Returns the inner item.
    pub fn into_inner(self) -> T {
        self.inner
    }
}

#[cfg(feature = "sink")]
impl<T, Item> CompatSink<T, Item> {
    /// Creates a new [`CompatSink`].
    pub fn new(inner: T) -> Self {
        CompatSink {
            inner,
            _phantom: PhantomData,
        }
    }

    /// Get a reference to 0.3 Sink contained within.
    pub fn get_ref(&self) -> &T {
        &self.inner
    }

    /// Get a mutable reference to 0.3 Sink contained within.
    pub fn get_mut(&mut self) -> &mut T {
        &mut self.inner
    }

    /// Returns the inner item.
    pub fn into_inner(self) -> T {
        self.inner
    }
}

fn poll_03_to_01<T, E>(x: task03::Poll<Result<T, E>>)
    -> Result<Async01<T>, E>
{
    match x? {
        task03::Poll::Ready(t) => Ok(Async01::Ready(t)),
        task03::Poll::Pending => Ok(Async01::NotReady),
    }
}

impl<Fut> Future01 for Compat<Fut>
where
    Fut: TryFuture03 + Unpin,
{
    type Item = Fut::Ok;
    type Error = Fut::Error;

    fn poll(&mut self) -> Poll01<Self::Item, Self::Error> {
        with_context(self, |inner, cx| poll_03_to_01(inner.try_poll(cx)))
    }
}

impl<St> Stream01 for Compat<St>
where
    St: TryStream03 + Unpin,
{
    type Item = St::Ok;
    type Error = St::Error;

    fn poll(&mut self) -> Poll01<Option<Self::Item>, Self::Error> {
        with_context(self, |inner, cx| match inner.try_poll_next(cx)? {
            task03::Poll::Ready(None) => Ok(Async01::Ready(None)),
            task03::Poll::Ready(Some(t)) => Ok(Async01::Ready(Some(t))),
            task03::Poll::Pending => Ok(Async01::NotReady),
        })
    }
}

#[cfg(feature = "sink")]
impl<T, Item> Sink01 for CompatSink<T, Item>
where
    T: Sink03<Item> + Unpin,
{
    type SinkItem = Item;
    type SinkError = T::Error;

    fn start_send(
        &mut self,
        item: Self::SinkItem,
    ) -> StartSend01<Self::SinkItem, Self::SinkError> {
        with_sink_context(self, |mut inner, cx| {
            match inner.as_mut().poll_ready(cx)? {
                task03::Poll::Ready(()) => {
                    inner.start_send(item).map(|()| AsyncSink01::Ready)
                }
                task03::Poll::Pending => Ok(AsyncSink01::NotReady(item)),
            }
        })
    }

    fn poll_complete(&mut self) -> Poll01<(), Self::SinkError> {
        with_sink_context(self, |inner, cx| poll_03_to_01(inner.poll_flush(cx)))
    }

    fn close(&mut self) -> Poll01<(), Self::SinkError> {
        with_sink_context(self, |inner, cx| poll_03_to_01(inner.poll_close(cx)))
    }
}

#[derive(Clone)]
struct Current(task01::Task);

impl Current {
    fn new() -> Current {
        Current(task01::current())
    }

    fn as_waker(&self) -> WakerRef<'_> {
        unsafe fn ptr_to_current<'a>(ptr: *const ()) -> &'a Current {
            &*(ptr as *const Current)
        }
        fn current_to_ptr(current: &Current) -> *const () {
            current as *const Current as *const ()
        }

        unsafe fn clone(ptr: *const ()) -> RawWaker {
            // Lazily create the `Arc` only when the waker is actually cloned.
            // FIXME: remove `transmute` when a `Waker` -> `RawWaker` conversion
            // function is landed in `core`.
            mem::transmute::<task03::Waker, RawWaker>(
                task03::waker(Arc::new(ptr_to_current(ptr).clone()))
            )
        }
        unsafe fn drop(_: *const ()) {}
        unsafe fn wake(ptr: *const ()) {
            ptr_to_current(ptr).0.notify()
        }

        let ptr = current_to_ptr(self);
        let vtable = &RawWakerVTable::new(clone, wake, wake, drop);
        WakerRef::new_unowned(std::mem::ManuallyDrop::new(unsafe {
            task03::Waker::from_raw(RawWaker::new(ptr, vtable))
        }))
    }
}

impl ArcWake03 for Current {
    fn wake_by_ref(arc_self: &Arc<Self>) {
        arc_self.0.notify();
    }
}

fn with_context<T, R, F>(compat: &mut Compat<T>, f: F) -> R
where
    T: Unpin,
    F: FnOnce(Pin<&mut T>, &mut Context<'_>) -> R,
{
    let current = Current::new();
    let waker = current.as_waker();
    let mut cx = Context::from_waker(&waker);
    f(Pin::new(&mut compat.inner), &mut cx)
}

#[cfg(feature = "sink")]
fn with_sink_context<T, Item, R, F>(compat: &mut CompatSink<T, Item>, f: F) -> R
where
    T: Unpin,
    F: FnOnce(Pin<&mut T>, &mut Context<'_>) -> R,
{
    let current = Current::new();
    let waker = current.as_waker();
    let mut cx = Context::from_waker(&waker);
    f(Pin::new(&mut compat.inner), &mut cx)
}

#[cfg(feature = "io-compat")]
mod io {
    use super::*;
    use futures_io::{AsyncRead as AsyncRead03, AsyncWrite as AsyncWrite03};
    use tokio_io::{AsyncRead as AsyncRead01, AsyncWrite as AsyncWrite01};

    fn poll_03_to_io<T>(x: task03::Poll<Result<T, std::io::Error>>)
        -> Result<T, std::io::Error>
    {
        match x {
            task03::Poll::Ready(Ok(t)) => Ok(t),
            task03::Poll::Pending => Err(std::io::ErrorKind::WouldBlock.into()),
            task03::Poll::Ready(Err(e)) => Err(e),
        }
    }

    impl<R: AsyncRead03 + Unpin> std::io::Read for Compat<R> {
        fn read(&mut self, buf: &mut [u8]) -> std::io::Result<usize> {
            let current = Current::new();
            let waker = current.as_waker();
            let mut cx = Context::from_waker(&waker);
            poll_03_to_io(Pin::new(&mut self.inner).poll_read(&mut cx, buf))
        }
    }

    impl<R: AsyncRead03 + Unpin> AsyncRead01 for Compat<R> {
        #[cfg(feature = "read-initializer")]
        unsafe fn prepare_uninitialized_buffer(&self, buf: &mut [u8]) -> bool {
            let initializer = self.inner.initializer();
            let does_init = initializer.should_initialize();
            if does_init {
                initializer.initialize(buf);
            }
            does_init
        }
    }

    impl<W: AsyncWrite03 + Unpin> std::io::Write for Compat<W> {
        fn write(&mut self, buf: &[u8]) -> std::io::Result<usize> {
            let current = Current::new();
            let waker = current.as_waker();
            let mut cx = Context::from_waker(&waker);
            poll_03_to_io(Pin::new(&mut self.inner).poll_write(&mut cx, buf))
        }

        fn flush(&mut self) -> std::io::Result<()> {
            let current = Current::new();
            let waker = current.as_waker();
            let mut cx = Context::from_waker(&waker);
            poll_03_to_io(Pin::new(&mut self.inner).poll_flush(&mut cx))
        }
    }

    impl<W: AsyncWrite03 + Unpin> AsyncWrite01 for Compat<W> {
        fn shutdown(&mut self) -> std::io::Result<Async01<()>> {
            let current = Current::new();
            let waker = current.as_waker();
            let mut cx = Context::from_waker(&waker);
            poll_03_to_01(Pin::new(&mut self.inner).poll_close(&mut cx))
        }
    }
}
<|endoftext|>


<commit_before>extern crate cursive;

use cursive::Cursive;
use cursive::view::{TextView,Dialog};

use std::fs::File;
use std::io::Read;

fn main() {
    let mut siv = Cursive::new();

    // Read some long text from a file.
    let mut file = File::open("assets/lorem.txt").unwrap();
    let mut content = String::new();
    file.read_to_string(&mut content).unwrap();

    // We can quit by pressing q
    siv.add_global_callback('q' as i32, |s| s.quit());

    // The text is too long to fit on a line, so the view will wrap lines,
    // and will adapt to the terminal size.
    siv.add_layer(TextView::new(&content));
    // Show a popup on top of the view.
    siv.add_layer(Dialog::new(TextView::new("Try resizing the terminal!\n(Press 'q' to quit when you're done.)"))
                  .dismiss_button("Ok"));

    siv.run();
}

<commit_msg>Update lorem example<commit_after>extern crate cursive;

use cursive::Cursive;
use cursive::view::{TextView,Dialog};

use std::fs::File;
use std::io::Read;

fn main() {
    // Read some long text from a file.
    let mut file = File::open("assets/lorem.txt").unwrap();
    let mut content = String::new();
    file.read_to_string(&mut content).unwrap();

    let mut siv = Cursive::new();

    // We can quit by pressing q
    siv.add_global_callback('q' as i32, |s| s.quit());

    // The text is too long to fit on a line, so the view will wrap lines,
    // and will adapt to the terminal size.
    siv.add_layer(TextView::new(&content));
    // Show a popup on top of the view.
    siv.add_layer(Dialog::new(TextView::new("Try resizing the terminal!\n(Press 'q' to quit when you're done.)"))
                  .dismiss_button("Ok"));

    siv.run();
}

<|endoftext|>


<commit_before>#![feature(zero_one)]

extern crate deeplearn;
extern crate matrix;

use std::fs::File;
use std::io::{
    self,
    BufReader,
    Read,
};
use std::num::{Zero, One};
use std::path::Path;

use deeplearn::Graph;
use deeplearn::op::{Add, MatMul, Mse, Relu};
use deeplearn::init;
use matrix::Matrix;

fn main() {
    println!("Reading training labels...");
    let train_labels = read_mnist_labels("data/mnist/train-labels-idx1-ubyte");
    let train_labels: Vec<Matrix<f32>> = train_labels.unwrap().into_iter().map(|x| one_hot(x, 10)).collect();
    println!("Label count: {}", train_labels.len());

    println!("Reading training images...");
    let (rows, columns, train_images) = read_mnist_images("data/mnist/train-images-idx3-ubyte").unwrap();

    let ctx = matrix::Context::new();

    // Setup the graph. It's going to have 2 inputs, 3 hidden nodes in it's 1 hidden layer, and 1
    // output node
    let mut graph = Graph::new();

    //////////////////////////
    // Layer 1

    // Input.  1 batches of rows*columns inputs each.
    let input = graph.add_variable(&ctx, (1, rows*columns), vec![0.0; rows*columns]);
    // Weights for layer 1: [2 inputs x 3 nodes]
    let l1_w = graph.add_variable(&ctx, (rows*columns, 1000), init::Normal(0.001, 0.005));
    // Use matrix multiplication to do a fully connected layer
    let l1_mat_mul = graph.add_node(&ctx,
                                    MatMul::new(&ctx, (1, rows*columns), (rows*columns, 1000)),
                                    vec![input, l1_w],
                                    &[(1, 1000)]); // out shape: [5x2]*[2x3] = [5 batches x 3 outputs]
    // Grab VarIndex for l1_mat_mul's output
    let l1_mat_mul_out = l1_mat_mul.get(&graph).outputs[0];
    // 3 biases; one for each node in layer 1
    let l1_b = graph.add_variable(&ctx, (1, 1000), init::Normal(0.001, 0.005));
    // Here we add the biases to the matrix multiplication output
    let l1_biased = graph.add_node(&ctx,
                               Add::new(0),
                               vec![l1_mat_mul_out, l1_b],
                               &[(1, 1000)]);
    // Grab VarIndex for l1_biased's output
    let l1_biased_out = l1_biased.get(&graph).outputs[0];
    // Run the biased input*weight sums through an ReLU activation
    let l1_relu = graph.add_node(&ctx,
                                 Relu::new(),
                                 vec![l1_biased_out],
                                 &[(1, 1000)]);
    let l1_relu_out = l1_relu.get(&graph).outputs[0];

    //////////////////////////
    // Layer 2

    // Weights for layer 2: [3 inputs x 1 nodes]
    let l2_w = graph.add_variable(&ctx, (1000, 10), init::Normal(0.001, 0.005));
    // Fully connected layer 2. Use layer 1's output as layer 2's input
    let l2_mat_mul = graph.add_node(&ctx,
                                    MatMul::new(&ctx, (1, 1000), (1000, 10)),
                                    vec![l1_relu_out, l2_w],
                                    &[(1, 10)]); // out shape: [5x3]*[3x1] = [5 batches x 1 output]
    // Grab VarIndex for l2_mat_mul's output
    let l2_mat_mul_out = l2_mat_mul.get(&graph).outputs[0];
    // 1 bias for 1 output node
    let l2_b = graph.add_variable(&ctx, (1, 10), init::Normal(0.001, 0.005));
    // Here we add the bias to the matrix multiplication output
    let l2_biased = graph.add_node(&ctx,
                               Add::new(0),
                               vec![l2_mat_mul_out, l2_b],
                               &[(1, 10)]);
    // Grab VarIndex for l2_biased's output
    let l2_biased_out = l2_biased.get(&graph).outputs[0];
    // Run the biased input*weight sums through an ReLU activation
    let l2_relu = graph.add_node(&ctx,
                                 Relu::new(),
                                 vec![l2_biased_out],
                                 &[(1, 10)]);
    // Grab VarIndex for l2_relu's output
    let l2_relu_out = l2_relu.get(&graph).outputs[0];

    //////////////////////////
    // Loss

    // Add a variable for the training outputs: [5 batches x 1 output]
    let train_out = graph.add_variable(&ctx, (1, 10), vec![0.0; 10]);
    // Use mean squared error loss function
    let loss = graph.add_node(&ctx,
                              Mse::new(),
                              vec![l2_relu_out, train_out],
                              &[(1, 10)]);
    // Grab the VarIndex for loss's output
    let loss_out = loss.get(&graph).outputs[0];
    // Create a gradient to apply to the loss function
    let loss_d = graph.add_gradient(&ctx, loss, 0);

    ////////////////////////////////////////////////////////////////////////////////////////////////

    // We apply a gradient of -0.1 to the loss function
    let loss_d_cpu = Matrix::from_vec(1, 10, vec![-0.001; 10]);
    loss_d.get(&graph).set(&ctx, &loss_d_cpu);

    /////////////////////////
    // Run/Train the network
    for epoch in 0..100000 {
        // Upload training data
        let train_sample = epoch%train_images.len();
        input.get(&graph).set(&ctx, &train_images[train_sample]);
        train_out.get(&graph).set(&ctx, &train_labels[train_sample]);

        graph.run(&ctx);
        let out = l2_relu_out.get(&graph).get(&ctx);
        let out_d = l2_relu.get(&graph)
                           .out_grad[0]
                           .gradient()
                           .get(&graph)
                           .get(&ctx);
        let l = loss_out.get(&graph).get(&ctx);
        let l1_w_d = graph.get_input_gradient(l1_w).unwrap().get(&graph);
        let l2_w_d = graph.get_input_gradient(l2_w).unwrap().get(&graph);
        let l1_b_d = graph.get_input_gradient(l1_b).unwrap().get(&graph);
        let l2_b_d = graph.get_input_gradient(l2_b).unwrap().get(&graph);
        l1_w.get(&graph).add(&ctx, -1, &*l1_w_d, &*l1_w.get(&graph));
        l2_w.get(&graph).add(&ctx, -1, &*l2_w_d, &*l2_w.get(&graph));
        l1_b.get(&graph).add(&ctx, -1, &*l1_b_d, &*l1_b.get(&graph));
        l2_b.get(&graph).add(&ctx, -1, &*l2_b_d, &*l2_b.get(&graph));
        if epoch % 1000 == 1 {
            println!("===================");
            println!("Epoch: {}", epoch);
            println!("out = {:?}", out);
            println!("out_d = {:?}", out_d);
            println!("loss = {:?}", l);
        }
    }
}

fn read_mnist_labels<P: AsRef<Path>>(path: P) -> io::Result<Vec<u8>> {
    use std::io::{Error, ErrorKind};

    let ref mut file = BufReader::new(File::open(path).unwrap());

    let magic = u32::from_be(try!(read_u32(file)));
    if magic != 2049 {
        return Err(Error::new(ErrorKind::Other,
                              format!("Invalid magic number. Got expect 2049, got {}",
                                      magic).as_ref()))
    }

    let label_count = u32::from_be(try!(read_u32(file)));

    let mut labels = Vec::with_capacity(label_count as usize);
    for _ in 0..label_count {
        labels.push(try!(read_u8(file)));
    }

    Ok(labels)
}

fn read_mnist_images<P: AsRef<Path>>(path: P) -> io::Result<(usize, usize, Vec<Matrix<f32>>)> {
    use std::io::{Error, ErrorKind};

    let ref mut file = BufReader::new(File::open(path).unwrap());

    let magic = u32::from_be(try!(read_u32(file)));
    if magic != 2051 {
        return Err(Error::new(ErrorKind::Other,
                              format!("Invalid magic number. Got expect 2051, got {}",
                                      magic).as_ref()))
    }

    let image_count = u32::from_be(try!(read_u32(file))) as usize;
    let rows = u32::from_be(try!(read_u32(file))) as usize;
    let columns = u32::from_be(try!(read_u32(file))) as usize;

    let mut images = Vec::with_capacity(image_count);
    for _ in 0..image_count {
        let mut pixel_buf = vec![0u8; rows*columns];
        try!(file.read_exact(pixel_buf.as_mut()));
        let matrix = Matrix::from_vec(rows, columns,
                                      pixel_buf.into_iter().map(|x| (x as f32)/255.0).collect());
        images.push(matrix);
    }

    Ok((rows, columns, images))
}

fn one_hot<N, M>(label: N, classes: N) -> Matrix<M>
    where usize: From<N>,
          M:     matrix::num::Num+Zero+One,
{
    let classes: usize = From::from(classes); // Cast class count to usize
    let label: usize = From::from(label); // Cast label to usize
    let mut buf: Vec<M> = vec![Zero::zero(); classes]; // Create array of zeroes
    buf[label] = One::one(); // Set the one-hot component
    Matrix::from_vec(1, buf.len(), buf) // Construct the matrix
}

fn read_u8<T: Read>(reader: &mut T) -> io::Result<u8> {
    use std::mem;

    let mut buf: [u8; 1] = [0];
    reader.read_exact(&mut buf).map(|_| {
        let data: u8 = unsafe { mem::transmute(buf) };
        data
    })
}

fn read_u32<T: Read>(reader: &mut T) -> io::Result<u32> {
    use std::mem;

    let mut buf: [u8; 4] = [0, 0, 0, 0];
    reader.read_exact(&mut buf).map(|_| {
        let data: u32 = unsafe { mem::transmute(buf) };
        data
    })
}
<commit_msg>Added accuracy score to mnist example output<commit_after>#![feature(zero_one)]

extern crate deeplearn;
extern crate matrix;

use std::fs::File;
use std::io::{
    self,
    BufReader,
    Read,
};
use std::num::{Zero, One};
use std::path::Path;

use deeplearn::Graph;
use deeplearn::op::{Add, MatMul, Mse, Relu};
use deeplearn::init;
use matrix::Matrix;

fn main() {
    println!("Reading training labels...");
    let train_labels = read_mnist_labels("data/mnist/train-labels-idx1-ubyte").unwrap();
    let train_labels_logits: Vec<Matrix<f32>> = train_labels.iter().cloned().map(|x| one_hot(x, 10)).collect();
    println!("Label count: {}", train_labels.len());

    println!("Reading training images...");
    let (rows, columns, train_images) = read_mnist_images("data/mnist/train-images-idx3-ubyte").unwrap();

    let ctx = matrix::Context::new();

    // Setup the graph. It's going to have 2 inputs, 3 hidden nodes in it's 1 hidden layer, and 1
    // output node
    let mut graph = Graph::new();

    //////////////////////////
    // Layer 1

    // Input.  1 batches of rows*columns inputs each.
    let input = graph.add_variable(&ctx, (1, rows*columns), vec![0.0; rows*columns]);
    // Weights for layer 1: [2 inputs x 3 nodes]
    let l1_w = graph.add_variable(&ctx, (rows*columns, 1000), init::Normal(0.001, 0.005));
    // Use matrix multiplication to do a fully connected layer
    let l1_mat_mul = graph.add_node(&ctx,
                                    MatMul::new(&ctx, (1, rows*columns), (rows*columns, 1000)),
                                    vec![input, l1_w],
                                    &[(1, 1000)]); // out shape: [5x2]*[2x3] = [5 batches x 3 outputs]
    // Grab VarIndex for l1_mat_mul's output
    let l1_mat_mul_out = l1_mat_mul.get(&graph).outputs[0];
    // 3 biases; one for each node in layer 1
    let l1_b = graph.add_variable(&ctx, (1, 1000), init::Normal(0.001, 0.005));
    // Here we add the biases to the matrix multiplication output
    let l1_biased = graph.add_node(&ctx,
                               Add::new(0),
                               vec![l1_mat_mul_out, l1_b],
                               &[(1, 1000)]);
    // Grab VarIndex for l1_biased's output
    let l1_biased_out = l1_biased.get(&graph).outputs[0];
    // Run the biased input*weight sums through an ReLU activation
    let l1_relu = graph.add_node(&ctx,
                                 Relu::new(),
                                 vec![l1_biased_out],
                                 &[(1, 1000)]);
    let l1_relu_out = l1_relu.get(&graph).outputs[0];

    //////////////////////////
    // Layer 2

    // Weights for layer 2: [3 inputs x 1 nodes]
    let l2_w = graph.add_variable(&ctx, (1000, 10), init::Normal(0.001, 0.005));
    // Fully connected layer 2. Use layer 1's output as layer 2's input
    let l2_mat_mul = graph.add_node(&ctx,
                                    MatMul::new(&ctx, (1, 1000), (1000, 10)),
                                    vec![l1_relu_out, l2_w],
                                    &[(1, 10)]); // out shape: [5x3]*[3x1] = [5 batches x 1 output]
    // Grab VarIndex for l2_mat_mul's output
    let l2_mat_mul_out = l2_mat_mul.get(&graph).outputs[0];
    // 1 bias for 1 output node
    let l2_b = graph.add_variable(&ctx, (1, 10), init::Normal(0.001, 0.005));
    // Here we add the bias to the matrix multiplication output
    let l2_biased = graph.add_node(&ctx,
                               Add::new(0),
                               vec![l2_mat_mul_out, l2_b],
                               &[(1, 10)]);
    // Grab VarIndex for l2_biased's output
    let l2_biased_out = l2_biased.get(&graph).outputs[0];
    // Run the biased input*weight sums through an ReLU activation
    let l2_relu = graph.add_node(&ctx,
                                 Relu::new(),
                                 vec![l2_biased_out],
                                 &[(1, 10)]);
    // Grab VarIndex for l2_relu's output
    let l2_relu_out = l2_relu.get(&graph).outputs[0];

    //////////////////////////
    // Loss

    // Add a variable for the training outputs: [5 batches x 1 output]
    let train_out = graph.add_variable(&ctx, (1, 10), vec![0.0; 10]);
    // Use mean squared error loss function
    let loss = graph.add_node(&ctx,
                              Mse::new(),
                              vec![l2_relu_out, train_out],
                              &[(1, 10)]);
    // Grab the VarIndex for loss's output
    let loss_out = loss.get(&graph).outputs[0];
    // Create a gradient to apply to the loss function
    let loss_d = graph.add_gradient(&ctx, loss, 0);

    ////////////////////////////////////////////////////////////////////////////////////////////////

    // We apply a gradient of -0.1 to the loss function
    let loss_d_cpu = Matrix::from_vec(1, 10, vec![-0.001; 10]);
    loss_d.get(&graph).set(&ctx, &loss_d_cpu);

    /////////////////////////
    // Run/Train the network
    let mut correct = 0;
    for epoch in 0..100000 {
        // Upload training data
        let train_sample = epoch%train_images.len();
        input.get(&graph).set(&ctx, &train_images[train_sample]);
        train_out.get(&graph).set(&ctx, &train_labels_logits[train_sample]);

        graph.run(&ctx);
        let out = l2_relu_out.get(&graph).get(&ctx);
        let out_d = l2_relu.get(&graph)
                           .out_grad[0]
                           .gradient()
                           .get(&graph)
                           .get(&ctx);
        let l = loss_out.get(&graph).get(&ctx);
        let l1_w_d = graph.get_input_gradient(l1_w).unwrap().get(&graph);
        let l2_w_d = graph.get_input_gradient(l2_w).unwrap().get(&graph);
        let l1_b_d = graph.get_input_gradient(l1_b).unwrap().get(&graph);
        let l2_b_d = graph.get_input_gradient(l2_b).unwrap().get(&graph);
        l1_w.get(&graph).add(&ctx, -1, &*l1_w_d, &*l1_w.get(&graph));
        l2_w.get(&graph).add(&ctx, -1, &*l2_w_d, &*l2_w.get(&graph));
        l1_b.get(&graph).add(&ctx, -1, &*l1_b_d, &*l1_b.get(&graph));
        l2_b.get(&graph).add(&ctx, -1, &*l2_b_d, &*l2_b.get(&graph));

        let (mut max_index, mut max_value) = (0, out[(0, 0)]);
        for (i, val) in out.buffer().iter().enumerate() {
            if *val > max_value {
                max_index = i;
                max_value = *val;
            }
        }

        if max_index == train_labels[train_sample] as usize {
            correct += 1;
        }

        if epoch % 1000 == 999 {
            println!("===================");
            println!("Epoch: {}", epoch);
            println!("out = {:?}", out);
            println!("out_d = {:?}", out_d);
            println!("loss = {:?}", l);
            println!("Accuracy: {}%", (correct as f32)/(1000 as f32) * 100.0);
            correct = 0;
        }
    }
}

fn read_mnist_labels<P: AsRef<Path>>(path: P) -> io::Result<Vec<u8>> {
    use std::io::{Error, ErrorKind};

    let ref mut file = BufReader::new(File::open(path).unwrap());

    let magic = u32::from_be(try!(read_u32(file)));
    if magic != 2049 {
        return Err(Error::new(ErrorKind::Other,
                              format!("Invalid magic number. Got expect 2049, got {}",
                                      magic).as_ref()))
    }

    let label_count = u32::from_be(try!(read_u32(file)));

    let mut labels = Vec::with_capacity(label_count as usize);
    for _ in 0..label_count {
        labels.push(try!(read_u8(file)));
    }

    Ok(labels)
}

fn read_mnist_images<P: AsRef<Path>>(path: P) -> io::Result<(usize, usize, Vec<Matrix<f32>>)> {
    use std::io::{Error, ErrorKind};

    let ref mut file = BufReader::new(File::open(path).unwrap());

    let magic = u32::from_be(try!(read_u32(file)));
    if magic != 2051 {
        return Err(Error::new(ErrorKind::Other,
                              format!("Invalid magic number. Got expect 2051, got {}",
                                      magic).as_ref()))
    }

    let image_count = u32::from_be(try!(read_u32(file))) as usize;
    let rows = u32::from_be(try!(read_u32(file))) as usize;
    let columns = u32::from_be(try!(read_u32(file))) as usize;

    let mut images = Vec::with_capacity(image_count);
    for _ in 0..image_count {
        let mut pixel_buf = vec![0u8; rows*columns];
        try!(file.read_exact(pixel_buf.as_mut()));
        let matrix = Matrix::from_vec(rows, columns,
                                      pixel_buf.into_iter().map(|x| (x as f32)/255.0).collect());
        images.push(matrix);
    }

    Ok((rows, columns, images))
}

fn one_hot<N, M>(label: N, classes: N) -> Matrix<M>
    where usize: From<N>,
          M:     matrix::num::Num+Zero+One,
{
    let classes: usize = From::from(classes); // Cast class count to usize
    let label: usize = From::from(label); // Cast label to usize
    let mut buf: Vec<M> = vec![Zero::zero(); classes]; // Create array of zeroes
    buf[label] = One::one(); // Set the one-hot component
    Matrix::from_vec(1, buf.len(), buf) // Construct the matrix
}

fn read_u8<T: Read>(reader: &mut T) -> io::Result<u8> {
    use std::mem;

    let mut buf: [u8; 1] = [0];
    reader.read_exact(&mut buf).map(|_| {
        let data: u8 = unsafe { mem::transmute(buf) };
        data
    })
}

fn read_u32<T: Read>(reader: &mut T) -> io::Result<u32> {
    use std::mem;

    let mut buf: [u8; 4] = [0, 0, 0, 0];
    reader.read_exact(&mut buf).map(|_| {
        let data: u32 = unsafe { mem::transmute(buf) };
        data
    })
}
<|endoftext|>