<commit_before>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

pub enum ReadError {
    Input(std::io::Error),
    Decode(prost::DecodeError),
    Magic(u32),
}

pub enum WriteError {
    Output(std::io::Error),
    Encode(prost::EncodeError),
}

impl From<std::io::Error> for ReadError {

    fn from(err: std::io::Error) -> ReadError {
        ReadError::Input(err)
    }
}

impl From<prost::DecodeError> for ReadError {

    fn from(err: prost::DecodeError) -> ReadError {
        ReadError::Decode(err)
    }
}

impl From<std::io::Error> for WriteError {

    fn from(err: std::io::Error) -> WriteError {
        WriteError::Output(err)
    }
}

impl From<prost::EncodeError> for WriteError {

    fn from(err: prost::EncodeError) -> WriteError {
        WriteError::Encode(err)
    }
}
<commit_msg>Implement display traits for the client errors.<commit_after>// Copyright 2020 Google LLC
//
// Use of this source code is governed by an MIT-style license that can be found
// in the LICENSE file or at https://opensource.org/licenses/MIT.

use std::fmt::{Display, Formatter};

#[derive(Debug)]
pub enum ReadError {
    Input(std::io::Error),
    Decode(prost::DecodeError),
    Magic(u32),
}

#[derive(Debug)]
pub enum WriteError {
    Output(std::io::Error),
    Encode(prost::EncodeError),
}

impl Display for ReadError {

    fn fmt(&self, fmt: &mut Formatter) -> std::fmt::Result {
        use ReadError::*;

        match *self {
            Input(ref err) => write!(fmt, "input error: {}", err),
            Decode(ref err) => write!(fmt, "proto decoding error: {}", err),
            Magic(magic) => write!(fmt, "invalid magic: {}", magic),
        }
    }
}

impl Display for WriteError {

    fn fmt(&self, fmt: &mut Formatter) -> std::fmt::Result {
        use WriteError::*;

        match *self {
            Output(ref err) => write!(fmt, "output error: {}", err),
            Encode(ref err) => write!(fmt, "proto encoding error: {}", err),
        }
    }
}

impl From<std::io::Error> for ReadError {

    fn from(err: std::io::Error) -> ReadError {
        ReadError::Input(err)
    }
}

impl From<prost::DecodeError> for ReadError {

    fn from(err: prost::DecodeError) -> ReadError {
        ReadError::Decode(err)
    }
}

impl From<std::io::Error> for WriteError {

    fn from(err: std::io::Error) -> WriteError {
        WriteError::Output(err)
    }
}

impl From<prost::EncodeError> for WriteError {

    fn from(err: prost::EncodeError) -> WriteError {
        WriteError::Encode(err)
    }
}
<|endoftext|>


<commit_before>/// Lookup hir elements using positions in the source code. This is a lossy
/// transformation: in general, a single source might correspond to several
/// modules, functions, etc, due to macros, cfgs and `#[path=]` attributes on
/// modules.
///
/// So, this modules should not be used during hir construction, it exists
/// purely for "IDE needs".
use std::sync::Arc;

use ra_db::{FileId, FilePosition};
use ra_syntax::{
    SyntaxNode, AstPtr, TextUnit,
    ast::{self, AstNode, NameOwner},
    algo::find_node_at_offset,
};

use crate::{
    HirDatabase, Function, Struct, Enum, Const, Static, Either, DefWithBody,
    AsName, Module, HirFileId, Crate, Trait, Resolver,
    expr::scope::{ReferenceDescriptor, ScopeEntryWithSyntax},
    ids::LocationCtx,
    expr, AstId
};

/// Locates the module by `FileId`. Picks topmost module in the file.
pub fn module_from_file_id(db: &impl HirDatabase, file_id: FileId) -> Option<Module> {
    module_from_source(db, file_id.into(), None)
}

/// Locates the child module by `mod child;` declaration.
pub fn module_from_declaration(
    db: &impl HirDatabase,
    file_id: FileId,
    decl: &ast::Module,
) -> Option<Module> {
    let parent_module = module_from_file_id(db, file_id);
    let child_name = decl.name();
    match (parent_module, child_name) {
        (Some(parent_module), Some(child_name)) => parent_module.child(db, &child_name.as_name()),
        _ => None,
    }
}

/// Locates the module by position in the source code.
pub fn module_from_position(db: &impl HirDatabase, position: FilePosition) -> Option<Module> {
    let file = db.parse(position.file_id);
    match find_node_at_offset::<ast::Module>(file.syntax(), position.offset) {
        Some(m) if !m.has_semi() => module_from_inline(db, position.file_id.into(), m),
        _ => module_from_file_id(db, position.file_id.into()),
    }
}

fn module_from_inline(
    db: &impl HirDatabase,
    file_id: FileId,
    module: &ast::Module,
) -> Option<Module> {
    assert!(!module.has_semi());
    let file_id = file_id.into();
    let ast_id_map = db.ast_id_map(file_id);
    let item_id = ast_id_map.ast_id(module).with_file_id(file_id);
    module_from_source(db, file_id, Some(item_id))
}

/// Locates the module by child syntax element within the module
pub fn module_from_child_node(
    db: &impl HirDatabase,
    file_id: FileId,
    child: &SyntaxNode,
) -> Option<Module> {
    if let Some(m) = child.ancestors().filter_map(ast::Module::cast).find(|it| !it.has_semi()) {
        module_from_inline(db, file_id.into(), m)
    } else {
        module_from_file_id(db, file_id.into())
    }
}

fn module_from_source(
    db: &impl HirDatabase,
    file_id: HirFileId,
    decl_id: Option<AstId<ast::Module>>,
) -> Option<Module> {
    let source_root_id = db.file_source_root(file_id.as_original_file());
    db.source_root_crates(source_root_id).iter().map(|&crate_id| Crate { crate_id }).find_map(
        |krate| {
            let def_map = db.crate_def_map(krate);
            let module_id = def_map.find_module_by_source(file_id, decl_id)?;
            Some(Module { krate, module_id })
        },
    )
}

fn function_from_source(
    db: &impl HirDatabase,
    file_id: FileId,
    fn_def: &ast::FnDef,
) -> Option<Function> {
    let module = module_from_child_node(db, file_id, fn_def.syntax())?;
    let file_id = file_id.into();
    let ctx = LocationCtx::new(db, module, file_id);
    Some(Function { id: ctx.to_def(fn_def) })
}

pub fn struct_from_module(
    db: &impl HirDatabase,
    module: Module,
    struct_def: &ast::StructDef,
) -> Struct {
    let (file_id, _) = module.definition_source(db);
    let file_id = file_id.into();
    let ctx = LocationCtx::new(db, module, file_id);
    Struct { id: ctx.to_def(struct_def) }
}

pub fn enum_from_module(db: &impl HirDatabase, module: Module, enum_def: &ast::EnumDef) -> Enum {
    let (file_id, _) = module.definition_source(db);
    let file_id = file_id.into();
    let ctx = LocationCtx::new(db, module, file_id);
    Enum { id: ctx.to_def(enum_def) }
}

pub fn trait_from_module(
    db: &impl HirDatabase,
    module: Module,
    trait_def: &ast::TraitDef,
) -> Trait {
    let (file_id, _) = module.definition_source(db);
    let file_id = file_id.into();
    let ctx = LocationCtx::new(db, module, file_id);
    Trait { id: ctx.to_def(trait_def) }
}

fn resolver_for_node(
    db: &impl HirDatabase,
    file_id: FileId,
    node: &SyntaxNode,
    offset: Option<TextUnit>,
) -> Resolver {
    node.ancestors()
        .find_map(|node| {
            if ast::Expr::cast(node).is_some() || ast::Block::cast(node).is_some() {
                if let Some(func) = node
                    .ancestors()
                    .find_map(ast::FnDef::cast)
                    .and_then(|it| function_from_source(db, file_id, it))
                {
                    let scopes = func.scopes(db);
                    let scope = match offset {
                        None => scopes.scope_for(&node),
                        Some(offset) => scopes.scope_for_offset(offset),
                    };
                    Some(expr::resolver_for_scope(func.body(db), db, scope))
                } else {
                    // FIXME const/static/array length
                    None
                }
            } else {
                try_get_resolver_for_node(db, file_id, node)
            }
        })
        .unwrap_or_default()
}

fn try_get_resolver_for_node(
    db: &impl HirDatabase,
    file_id: FileId,
    node: &SyntaxNode,
) -> Option<Resolver> {
    if let Some(module) = ast::Module::cast(node) {
        Some(module_from_declaration(db, file_id, module)?.resolver(db))
    } else if let Some(_) = ast::SourceFile::cast(node) {
        Some(module_from_source(db, file_id.into(), None)?.resolver(db))
    } else if let Some(s) = ast::StructDef::cast(node) {
        let module = module_from_child_node(db, file_id, s.syntax())?;
        Some(struct_from_module(db, module, s).resolver(db))
    } else if let Some(e) = ast::EnumDef::cast(node) {
        let module = module_from_child_node(db, file_id, e.syntax())?;
        Some(enum_from_module(db, module, e).resolver(db))
    } else if let Some(f) = ast::FnDef::cast(node) {
        function_from_source(db, file_id, f).map(|f| f.resolver(db))
    } else {
        // FIXME add missing cases
        None
    }
}

pub fn def_with_body_from_child_node(
    db: &impl HirDatabase,
    file_id: FileId,
    node: &SyntaxNode,
) -> Option<DefWithBody> {
    let module = module_from_child_node(db, file_id, node)?;
    let ctx = LocationCtx::new(db, module, file_id.into());
    node.ancestors().find_map(|node| {
        if let Some(def) = ast::FnDef::cast(node) {
            return Some(Function { id: ctx.to_def(def) }.into());
        }
        if let Some(def) = ast::ConstDef::cast(node) {
            return Some(Const { id: ctx.to_def(def) }.into());
        }
        if let Some(def) = ast::StaticDef::cast(node) {
            return Some(Static { id: ctx.to_def(def) }.into());
        }
        None
    })
}

/// `SourceAnalyzer` is a convenience wrapper which exposes HIR API in terms of
/// original source files. It should not be used inside the HIR itself.
#[derive(Debug)]
pub struct SourceAnalyzer {
    resolver: Resolver,
    body_source_map: Option<Arc<crate::expr::BodySourceMap>>,
    infer: Option<Arc<crate::ty::InferenceResult>>,
    scopes: Option<crate::expr::ScopesWithSourceMap>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PathResolution {
    /// An item
    Def(crate::ModuleDef),
    /// A local binding (only value namespace)
    LocalBinding(Either<AstPtr<ast::Pat>, AstPtr<ast::SelfParam>>),
    /// A generic parameter
    GenericParam(u32),
    SelfType(crate::ImplBlock),
    AssocItem(crate::ImplItem),
}

impl SourceAnalyzer {
    pub fn new(
        db: &impl HirDatabase,
        file_id: FileId,
        node: &SyntaxNode,
        offset: Option<TextUnit>,
    ) -> SourceAnalyzer {
        let def_with_body = def_with_body_from_child_node(db, file_id, node);
        SourceAnalyzer {
            resolver: resolver_for_node(db, file_id, node, offset),
            body_source_map: def_with_body.map(|it| it.body_source_map(db)),
            infer: def_with_body.map(|it| it.infer(db)),
            scopes: def_with_body.map(|it| it.scopes(db)),
        }
    }

    pub fn resolver(&self) -> &Resolver {
        &self.resolver
    }

    pub fn type_of(&self, _db: &impl HirDatabase, expr: &ast::Expr) -> Option<crate::Ty> {
        let expr_id = self.body_source_map.as_ref()?.node_expr(expr)?;
        Some(self.infer.as_ref()?[expr_id].clone())
    }

    pub fn type_of_pat(&self, _db: &impl HirDatabase, pat: &ast::Pat) -> Option<crate::Ty> {
        let pat_id = self.body_source_map.as_ref()?.node_pat(pat)?;
        Some(self.infer.as_ref()?[pat_id].clone())
    }

    pub fn resolve_method_call(&self, call: &ast::MethodCallExpr) -> Option<Function> {
        let expr_id = self.body_source_map.as_ref()?.node_expr(call.into())?;
        self.infer.as_ref()?.method_resolution(expr_id)
    }

    pub fn resolve_field(&self, field: &ast::FieldExpr) -> Option<crate::StructField> {
        let expr_id = self.body_source_map.as_ref()?.node_expr(field.into())?;
        self.infer.as_ref()?.field_resolution(expr_id)
    }

    pub fn resolve_path(&self, db: &impl HirDatabase, path: &ast::Path) -> Option<PathResolution> {
        if let Some(path_expr) = path.syntax().parent().and_then(ast::PathExpr::cast) {
            let expr_id = self.body_source_map.as_ref()?.node_expr(path_expr.into())?;
            if let Some(assoc) = self.infer.as_ref()?.assoc_resolutions_for_expr(expr_id) {
                return Some(PathResolution::AssocItem(assoc));
            }
        }
        if let Some(path_pat) = path.syntax().parent().and_then(ast::PathPat::cast) {
            let pat_id = self.body_source_map.as_ref()?.node_pat(path_pat.into())?;
            if let Some(assoc) = self.infer.as_ref()?.assoc_resolutions_for_pat(pat_id) {
                return Some(PathResolution::AssocItem(assoc));
            }
        }
        let hir_path = crate::Path::from_ast(path)?;
        let res = self.resolver.resolve_path(db, &hir_path);
        let res = res.clone().take_types().or_else(|| res.take_values())?;
        let res = match res {
            crate::Resolution::Def(it) => PathResolution::Def(it),
            crate::Resolution::LocalBinding(it) => {
                PathResolution::LocalBinding(self.body_source_map.as_ref()?.pat_syntax(it)?)
            }
            crate::Resolution::GenericParam(it) => PathResolution::GenericParam(it),
            crate::Resolution::SelfType(it) => PathResolution::SelfType(it),
        };
        Some(res)
    }

    pub fn find_all_refs(&self, pat: &ast::BindPat) -> Option<Vec<ReferenceDescriptor>> {
        self.scopes.as_ref().map(|it| it.find_all_refs(pat))
    }

    pub fn resolve_local_name(&self, name_ref: &ast::NameRef) -> Option<ScopeEntryWithSyntax> {
        self.scopes.as_ref()?.resolve_local_name(name_ref)
    }

    #[cfg(test)]
    pub(crate) fn body_source_map(&self) -> Arc<crate::expr::BodySourceMap> {
        self.body_source_map.clone().unwrap()
    }

    #[cfg(test)]
    pub(crate) fn inference_result(&self) -> Arc<crate::ty::InferenceResult> {
        self.infer.clone().unwrap()
    }
}
<commit_msg>generalize<commit_after>/// Lookup hir elements using positions in the source code. This is a lossy
/// transformation: in general, a single source might correspond to several
/// modules, functions, etc, due to macros, cfgs and `#[path=]` attributes on
/// modules.
///
/// So, this modules should not be used during hir construction, it exists
/// purely for "IDE needs".
use std::sync::Arc;

use ra_db::{FileId, FilePosition};
use ra_syntax::{
    SyntaxNode, AstPtr, TextUnit,
    ast::{self, AstNode, NameOwner},
    algo::find_node_at_offset,
};

use crate::{
    HirDatabase, Function, Struct, Enum, Const, Static, Either, DefWithBody,
    AsName, Module, HirFileId, Crate, Trait, Resolver,
    expr::scope::{ReferenceDescriptor, ScopeEntryWithSyntax},
    ids::LocationCtx,
    expr, AstId
};

/// Locates the module by `FileId`. Picks topmost module in the file.
pub fn module_from_file_id(db: &impl HirDatabase, file_id: FileId) -> Option<Module> {
    module_from_source(db, file_id.into(), None)
}

/// Locates the child module by `mod child;` declaration.
pub fn module_from_declaration(
    db: &impl HirDatabase,
    file_id: FileId,
    decl: &ast::Module,
) -> Option<Module> {
    let parent_module = module_from_file_id(db, file_id);
    let child_name = decl.name();
    match (parent_module, child_name) {
        (Some(parent_module), Some(child_name)) => parent_module.child(db, &child_name.as_name()),
        _ => None,
    }
}

/// Locates the module by position in the source code.
pub fn module_from_position(db: &impl HirDatabase, position: FilePosition) -> Option<Module> {
    let file = db.parse(position.file_id);
    match find_node_at_offset::<ast::Module>(file.syntax(), position.offset) {
        Some(m) if !m.has_semi() => module_from_inline(db, position.file_id.into(), m),
        _ => module_from_file_id(db, position.file_id.into()),
    }
}

fn module_from_inline(
    db: &impl HirDatabase,
    file_id: FileId,
    module: &ast::Module,
) -> Option<Module> {
    assert!(!module.has_semi());
    let file_id = file_id.into();
    let ast_id_map = db.ast_id_map(file_id);
    let item_id = ast_id_map.ast_id(module).with_file_id(file_id);
    module_from_source(db, file_id, Some(item_id))
}

/// Locates the module by child syntax element within the module
pub fn module_from_child_node(
    db: &impl HirDatabase,
    file_id: FileId,
    child: &SyntaxNode,
) -> Option<Module> {
    if let Some(m) = child.ancestors().filter_map(ast::Module::cast).find(|it| !it.has_semi()) {
        module_from_inline(db, file_id.into(), m)
    } else {
        module_from_file_id(db, file_id.into())
    }
}

fn module_from_source(
    db: &impl HirDatabase,
    file_id: HirFileId,
    decl_id: Option<AstId<ast::Module>>,
) -> Option<Module> {
    let source_root_id = db.file_source_root(file_id.as_original_file());
    db.source_root_crates(source_root_id).iter().map(|&crate_id| Crate { crate_id }).find_map(
        |krate| {
            let def_map = db.crate_def_map(krate);
            let module_id = def_map.find_module_by_source(file_id, decl_id)?;
            Some(Module { krate, module_id })
        },
    )
}

fn function_from_source(
    db: &impl HirDatabase,
    file_id: FileId,
    fn_def: &ast::FnDef,
) -> Option<Function> {
    let module = module_from_child_node(db, file_id, fn_def.syntax())?;
    let file_id = file_id.into();
    let ctx = LocationCtx::new(db, module, file_id);
    Some(Function { id: ctx.to_def(fn_def) })
}

pub fn struct_from_module(
    db: &impl HirDatabase,
    module: Module,
    struct_def: &ast::StructDef,
) -> Struct {
    let (file_id, _) = module.definition_source(db);
    let file_id = file_id.into();
    let ctx = LocationCtx::new(db, module, file_id);
    Struct { id: ctx.to_def(struct_def) }
}

pub fn enum_from_module(db: &impl HirDatabase, module: Module, enum_def: &ast::EnumDef) -> Enum {
    let (file_id, _) = module.definition_source(db);
    let file_id = file_id.into();
    let ctx = LocationCtx::new(db, module, file_id);
    Enum { id: ctx.to_def(enum_def) }
}

pub fn trait_from_module(
    db: &impl HirDatabase,
    module: Module,
    trait_def: &ast::TraitDef,
) -> Trait {
    let (file_id, _) = module.definition_source(db);
    let file_id = file_id.into();
    let ctx = LocationCtx::new(db, module, file_id);
    Trait { id: ctx.to_def(trait_def) }
}

fn resolver_for_node(
    db: &impl HirDatabase,
    file_id: FileId,
    node: &SyntaxNode,
    offset: Option<TextUnit>,
) -> Resolver {
    node.ancestors()
        .find_map(|node| {
            if ast::Expr::cast(node).is_some() || ast::Block::cast(node).is_some() {
                let def = def_with_body_from_child_node(db, file_id, node)?;
                let scopes = def.scopes(db);
                let scope = match offset {
                    None => scopes.scope_for(&node),
                    Some(offset) => scopes.scope_for_offset(offset),
                };
                Some(expr::resolver_for_scope(def.body(db), db, scope))
            } else {
                try_get_resolver_for_node(db, file_id, node)
            }
        })
        .unwrap_or_default()
}

fn try_get_resolver_for_node(
    db: &impl HirDatabase,
    file_id: FileId,
    node: &SyntaxNode,
) -> Option<Resolver> {
    if let Some(module) = ast::Module::cast(node) {
        Some(module_from_declaration(db, file_id, module)?.resolver(db))
    } else if let Some(_) = ast::SourceFile::cast(node) {
        Some(module_from_source(db, file_id.into(), None)?.resolver(db))
    } else if let Some(s) = ast::StructDef::cast(node) {
        let module = module_from_child_node(db, file_id, s.syntax())?;
        Some(struct_from_module(db, module, s).resolver(db))
    } else if let Some(e) = ast::EnumDef::cast(node) {
        let module = module_from_child_node(db, file_id, e.syntax())?;
        Some(enum_from_module(db, module, e).resolver(db))
    } else if let Some(f) = ast::FnDef::cast(node) {
        function_from_source(db, file_id, f).map(|f| f.resolver(db))
    } else {
        // FIXME add missing cases
        None
    }
}

pub fn def_with_body_from_child_node(
    db: &impl HirDatabase,
    file_id: FileId,
    node: &SyntaxNode,
) -> Option<DefWithBody> {
    let module = module_from_child_node(db, file_id, node)?;
    let ctx = LocationCtx::new(db, module, file_id.into());
    node.ancestors().find_map(|node| {
        if let Some(def) = ast::FnDef::cast(node) {
            return Some(Function { id: ctx.to_def(def) }.into());
        }
        if let Some(def) = ast::ConstDef::cast(node) {
            return Some(Const { id: ctx.to_def(def) }.into());
        }
        if let Some(def) = ast::StaticDef::cast(node) {
            return Some(Static { id: ctx.to_def(def) }.into());
        }
        None
    })
}

/// `SourceAnalyzer` is a convenience wrapper which exposes HIR API in terms of
/// original source files. It should not be used inside the HIR itself.
#[derive(Debug)]
pub struct SourceAnalyzer {
    resolver: Resolver,
    body_source_map: Option<Arc<crate::expr::BodySourceMap>>,
    infer: Option<Arc<crate::ty::InferenceResult>>,
    scopes: Option<crate::expr::ScopesWithSourceMap>,
}

#[derive(Debug, Clone, PartialEq, Eq)]
pub enum PathResolution {
    /// An item
    Def(crate::ModuleDef),
    /// A local binding (only value namespace)
    LocalBinding(Either<AstPtr<ast::Pat>, AstPtr<ast::SelfParam>>),
    /// A generic parameter
    GenericParam(u32),
    SelfType(crate::ImplBlock),
    AssocItem(crate::ImplItem),
}

impl SourceAnalyzer {
    pub fn new(
        db: &impl HirDatabase,
        file_id: FileId,
        node: &SyntaxNode,
        offset: Option<TextUnit>,
    ) -> SourceAnalyzer {
        let def_with_body = def_with_body_from_child_node(db, file_id, node);
        SourceAnalyzer {
            resolver: resolver_for_node(db, file_id, node, offset),
            body_source_map: def_with_body.map(|it| it.body_source_map(db)),
            infer: def_with_body.map(|it| it.infer(db)),
            scopes: def_with_body.map(|it| it.scopes(db)),
        }
    }

    pub fn resolver(&self) -> &Resolver {
        &self.resolver
    }

    pub fn type_of(&self, _db: &impl HirDatabase, expr: &ast::Expr) -> Option<crate::Ty> {
        let expr_id = self.body_source_map.as_ref()?.node_expr(expr)?;
        Some(self.infer.as_ref()?[expr_id].clone())
    }

    pub fn type_of_pat(&self, _db: &impl HirDatabase, pat: &ast::Pat) -> Option<crate::Ty> {
        let pat_id = self.body_source_map.as_ref()?.node_pat(pat)?;
        Some(self.infer.as_ref()?[pat_id].clone())
    }

    pub fn resolve_method_call(&self, call: &ast::MethodCallExpr) -> Option<Function> {
        let expr_id = self.body_source_map.as_ref()?.node_expr(call.into())?;
        self.infer.as_ref()?.method_resolution(expr_id)
    }

    pub fn resolve_field(&self, field: &ast::FieldExpr) -> Option<crate::StructField> {
        let expr_id = self.body_source_map.as_ref()?.node_expr(field.into())?;
        self.infer.as_ref()?.field_resolution(expr_id)
    }

    pub fn resolve_path(&self, db: &impl HirDatabase, path: &ast::Path) -> Option<PathResolution> {
        if let Some(path_expr) = path.syntax().parent().and_then(ast::PathExpr::cast) {
            let expr_id = self.body_source_map.as_ref()?.node_expr(path_expr.into())?;
            if let Some(assoc) = self.infer.as_ref()?.assoc_resolutions_for_expr(expr_id) {
                return Some(PathResolution::AssocItem(assoc));
            }
        }
        if let Some(path_pat) = path.syntax().parent().and_then(ast::PathPat::cast) {
            let pat_id = self.body_source_map.as_ref()?.node_pat(path_pat.into())?;
            if let Some(assoc) = self.infer.as_ref()?.assoc_resolutions_for_pat(pat_id) {
                return Some(PathResolution::AssocItem(assoc));
            }
        }
        let hir_path = crate::Path::from_ast(path)?;
        let res = self.resolver.resolve_path(db, &hir_path);
        let res = res.clone().take_types().or_else(|| res.take_values())?;
        let res = match res {
            crate::Resolution::Def(it) => PathResolution::Def(it),
            crate::Resolution::LocalBinding(it) => {
                PathResolution::LocalBinding(self.body_source_map.as_ref()?.pat_syntax(it)?)
            }
            crate::Resolution::GenericParam(it) => PathResolution::GenericParam(it),
            crate::Resolution::SelfType(it) => PathResolution::SelfType(it),
        };
        Some(res)
    }

    pub fn find_all_refs(&self, pat: &ast::BindPat) -> Option<Vec<ReferenceDescriptor>> {
        self.scopes.as_ref().map(|it| it.find_all_refs(pat))
    }

    pub fn resolve_local_name(&self, name_ref: &ast::NameRef) -> Option<ScopeEntryWithSyntax> {
        self.scopes.as_ref()?.resolve_local_name(name_ref)
    }

    #[cfg(test)]
    pub(crate) fn body_source_map(&self) -> Arc<crate::expr::BodySourceMap> {
        self.body_source_map.clone().unwrap()
    }

    #[cfg(test)]
    pub(crate) fn inference_result(&self) -> Arc<crate::ty::InferenceResult> {
        self.infer.clone().unwrap()
    }
}
<|endoftext|>


<commit_before>//! Module for built-in filter functions
//!
//! Contains all the built-in filter functions for use in templates.
//! You can define your own filters; for more information,
//! see the top-level crate documentation.

#[cfg(feature = "serde-json")]
mod json;

#[cfg(feature = "serde-json")]
pub use self::json::json;

use num_traits::Signed;
use std::fmt;

use super::Result;
use escaping::{self, MarkupDisplay};

// This is used by the code generator to decide whether a named filter is part of
// Askama or should refer to a local `filters` module. It should contain all the
// filters shipped with Askama, even the optional ones (since optional inclusion
// in the const vector based on features seems impossible right now).
pub const BUILT_IN_FILTERS: [&str; 15] = [
    "abs",
    "e",
    "escape",
    "format",
    "join",
    "linebreaks",
    "linebreaksbr",
    "lower",
    "lowercase",
    "safe",
    "trim",
    "truncate",
    "upper",
    "uppercase",
    "json", // Optional feature; reserve the name anyway
];

/// Marks a string (or other `Display` type) as safe
///
/// Use this is you want to allow markup in an expression, or if you know
/// that the expression's contents don't need to be escaped.
pub fn safe<D, I>(v: I) -> Result<MarkupDisplay<D>>
where
    D: fmt::Display,
    MarkupDisplay<D>: From,
{
    let res: MarkupDisplay<D> = v.into();
    Ok(res.mark_safe())
}

/// Escapes `&`, `<` and `>` in strings
pub fn escape<D, I>(i: I) -> Result<MarkupDisplay<String>>
where
    D: fmt::Display,
    MarkupDisplay<D>: From,
{
    let md: MarkupDisplay<D> = i.into();
    Ok(MarkupDisplay::Safe(escaping::escape(md.unsafe_string())))
}

/// Alias for the `escape()` filter
pub fn e<D, I>(i: I) -> Result<MarkupDisplay<String>>
where
    D: fmt::Display,
    MarkupDisplay<D>: From,
{
    escape(i)
}

/// Formats arguments according to the specified format
///
/// The first argument to this filter must be a string literal (as in normal
/// Rust). All arguments are passed through to the `format!()`
/// [macro](https://doc.rust-lang.org/stable/std/macro.format.html) by
/// the Askama code generator.
pub fn format() {}

/// Replaces line breaks in plain text with appropriate HTML
///
/// A single newline becomes an HTML line break `<br>` and a new line
/// followed by a blank line becomes a paragraph break `<p>`.
pub fn linebreaks(s: &fmt::Display) -> Result<String> {
    let s = format!("{}", s);
    let linebroken = s.replace("\n\n", "

").replace("\n", "
");

    Ok(format!("
{}
", linebroken))
}

/// Converts all newlines in a piece of plain text to HTML line breaks
pub fn linebreaksbr(s: &fmt::Display) -> Result<String> {
    let s = format!("{}", s);
    Ok(s.replace("\n", "
"))
}

/// Converts to lowercase
pub fn lower(s: &fmt::Display) -> Result<String> {
    let s = format!("{}", s);
    Ok(s.to_lowercase())
}

/// Alias for the `lower()` filter
pub fn lowercase(s: &fmt::Display) -> Result<String> {
    lower(s)
}

/// Converts to uppercase
pub fn upper(s: &fmt::Display) -> Result<String> {
    let s = format!("{}", s);
    Ok(s.to_uppercase())
}

/// Alias for the `upper()` filter
pub fn uppercase(s: &fmt::Display) -> Result<String> {
    upper(s)
}

/// Strip leading and trailing whitespace
pub fn trim(s: &fmt::Display) -> Result<String> {
    let s = format!("{}", s);
    Ok(s.trim().to_owned())
}

/// Limit string length, appends '...' if truncated
pub fn truncate(s: &fmt::Display, len: &usize) -> Result<String> {
    let mut s = format!("{}", s);
    if s.len() < *len {
        Ok(s)
    } else {
        s.truncate(*len);
        s.push_str("...");
        Ok(s)
    }
}

/// Joins iterable into a string separated by provided argument
pub fn join<T, I, S>(input: I, separator: S) -> Result<String>
where
    T: fmt::Display,
    I: Iterator<Item = T>,
    S: AsRef<str>,
{
    let separator: &str = separator.as_ref();

    let mut rv = String::new();

    for (num, item) in input.enumerate() {
        if num > 0 {
            rv.push_str(separator);
        }

        rv.push_str(&format!("{}", item));
    }

    Ok(rv)
}

/// Absolute value
pub fn abs<T>(number: T) -> Result<T>
where
    T: Signed,
{
    Ok(number.abs())
}

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

    #[test]
    fn test_linebreaks() {
        assert_eq!(
            linebreaks(&"Foo\nBar Baz").unwrap(),
            "
Foo
Bar Baz
"
        );
        assert_eq!(
            linebreaks(&"Foo\nBar\n\nBaz").unwrap(),
            "
Foo
Bar

Baz
"
        );
    }

    #[test]
    fn test_linebreaksbr() {
        assert_eq!(linebreaksbr(&"Foo\nBar").unwrap(), "Foo
Bar");
        assert_eq!(
            linebreaksbr(&"Foo\nBar\n\nBaz").unwrap(),
            "Foo
Bar

Baz"
        );
    }

    #[test]
    fn test_lower() {
        assert_eq!(lower(&"Foo").unwrap(), "foo");
        assert_eq!(lower(&"FOO").unwrap(), "foo");
        assert_eq!(lower(&"FooBar").unwrap(), "foobar");
        assert_eq!(lower(&"foo").unwrap(), "foo");
    }

    #[test]
    fn test_upper() {
        assert_eq!(upper(&"Foo").unwrap(), "FOO");
        assert_eq!(upper(&"FOO").unwrap(), "FOO");
        assert_eq!(upper(&"FooBar").unwrap(), "FOOBAR");
        assert_eq!(upper(&"foo").unwrap(), "FOO");
    }

    #[test]
    fn test_trim() {
        assert_eq!(trim(&" Hello\tworld\t").unwrap(), "Hello\tworld");
    }

    #[test]
    fn test_join() {
        assert_eq!(
            join((&["hello", "world"]).into_iter(), ", ").unwrap(),
            "hello, world"
        );
        assert_eq!(join((&["hello"]).into_iter(), ", ").unwrap(), "hello");

        let empty: &[&str] = &[];
        assert_eq!(join(empty.into_iter(), ", ").unwrap(), "");

        let input: Vec<String> = vec!["foo".into(), "bar".into(), "bazz".into()];
        assert_eq!(
            join((&input).into_iter(), ":".to_string()).unwrap(),
            "foo:bar:bazz"
        );
        assert_eq!(
            join(input.clone().into_iter(), ":").unwrap(),
            "foo:bar:bazz"
        );
        assert_eq!(
            join(input.clone().into_iter(), ":".to_string()).unwrap(),
            "foo:bar:bazz"
        );

        let input: &[String] = &["foo".into(), "bar".into()];
        assert_eq!(join(input.into_iter(), ":").unwrap(), "foo:bar");
        assert_eq!(join(input.into_iter(), ":".to_string()).unwrap(), "foo:bar");

        let real: String = "blah".into();
        let input: Vec<&str> = vec![&real];
        assert_eq!(join(input.into_iter(), ";").unwrap(), "blah");

        assert_eq!(
            join((&&&&&["foo", "bar"]).into_iter(), ", ").unwrap(),
            "foo, bar"
        );
    }

    #[test]
    fn test_abs() {
        assert_eq!(abs(1).unwrap(), 1);
        assert_eq!(abs(-1).unwrap(), 1);
        assert_eq!(abs(1.0).unwrap(), 1.0);
        assert_eq!(abs(-1.0).unwrap(), 1.0);
        assert_eq!(abs(1.0 as f64).unwrap(), 1.0 as f64);
        assert_eq!(abs(-1.0 as f64).unwrap(), 1.0 as f64);
    }
}
<commit_msg>Add capitalize filter<commit_after>//! Module for built-in filter functions
//!
//! Contains all the built-in filter functions for use in templates.
//! You can define your own filters; for more information,
//! see the top-level crate documentation.

#[cfg(feature = "serde-json")]
mod json;

#[cfg(feature = "serde-json")]
pub use self::json::json;

use num_traits::Signed;
use std::fmt;

use super::Result;
use escaping::{self, MarkupDisplay};

// This is used by the code generator to decide whether a named filter is part of
// Askama or should refer to a local `filters` module. It should contain all the
// filters shipped with Askama, even the optional ones (since optional inclusion
// in the const vector based on features seems impossible right now).
pub const BUILT_IN_FILTERS: [&str; 16] = [
    "abs",
    "capitalize",
    "e",
    "escape",
    "format",
    "join",
    "linebreaks",
    "linebreaksbr",
    "lower",
    "lowercase",
    "safe",
    "trim",
    "truncate",
    "upper",
    "uppercase",
    "json", // Optional feature; reserve the name anyway
];

/// Marks a string (or other `Display` type) as safe
///
/// Use this is you want to allow markup in an expression, or if you know
/// that the expression's contents don't need to be escaped.
pub fn safe<D, I>(v: I) -> Result<MarkupDisplay<D>>
where
    D: fmt::Display,
    MarkupDisplay<D>: From,
{
    let res: MarkupDisplay<D> = v.into();
    Ok(res.mark_safe())
}

/// Escapes `&`, `<` and `>` in strings
pub fn escape<D, I>(i: I) -> Result<MarkupDisplay<String>>
where
    D: fmt::Display,
    MarkupDisplay<D>: From,
{
    let md: MarkupDisplay<D> = i.into();
    Ok(MarkupDisplay::Safe(escaping::escape(md.unsafe_string())))
}

/// Alias for the `escape()` filter
pub fn e<D, I>(i: I) -> Result<MarkupDisplay<String>>
where
    D: fmt::Display,
    MarkupDisplay<D>: From,
{
    escape(i)
}

/// Formats arguments according to the specified format
///
/// The first argument to this filter must be a string literal (as in normal
/// Rust). All arguments are passed through to the `format!()`
/// [macro](https://doc.rust-lang.org/stable/std/macro.format.html) by
/// the Askama code generator.
pub fn format() {}

/// Replaces line breaks in plain text with appropriate HTML
///
/// A single newline becomes an HTML line break `<br>` and a new line
/// followed by a blank line becomes a paragraph break `<p>`.
pub fn linebreaks(s: &fmt::Display) -> Result<String> {
    let s = format!("{}", s);
    let linebroken = s.replace("\n\n", "

").replace("\n", "
");

    Ok(format!("
{}
", linebroken))
}

/// Converts all newlines in a piece of plain text to HTML line breaks
pub fn linebreaksbr(s: &fmt::Display) -> Result<String> {
    let s = format!("{}", s);
    Ok(s.replace("\n", "
"))
}

/// Converts to lowercase
pub fn lower(s: &fmt::Display) -> Result<String> {
    let s = format!("{}", s);
    Ok(s.to_lowercase())
}

/// Alias for the `lower()` filter
pub fn lowercase(s: &fmt::Display) -> Result<String> {
    lower(s)
}

/// Converts to uppercase
pub fn upper(s: &fmt::Display) -> Result<String> {
    let s = format!("{}", s);
    Ok(s.to_uppercase())
}

/// Alias for the `upper()` filter
pub fn uppercase(s: &fmt::Display) -> Result<String> {
    upper(s)
}

/// Strip leading and trailing whitespace
pub fn trim(s: &fmt::Display) -> Result<String> {
    let s = format!("{}", s);
    Ok(s.trim().to_owned())
}

/// Limit string length, appends '...' if truncated
pub fn truncate(s: &fmt::Display, len: &usize) -> Result<String> {
    let mut s = format!("{}", s);
    if s.len() < *len {
        Ok(s)
    } else {
        s.truncate(*len);
        s.push_str("...");
        Ok(s)
    }
}

/// Joins iterable into a string separated by provided argument
pub fn join<T, I, S>(input: I, separator: S) -> Result<String>
where
    T: fmt::Display,
    I: Iterator<Item = T>,
    S: AsRef<str>,
{
    let separator: &str = separator.as_ref();

    let mut rv = String::new();

    for (num, item) in input.enumerate() {
        if num > 0 {
            rv.push_str(separator);
        }

        rv.push_str(&format!("{}", item));
    }

    Ok(rv)
}

/// Absolute value
pub fn abs<T>(number: T) -> Result<T>
where
    T: Signed,
{
    Ok(number.abs())
}

/// Capitalize a value. The first character will be uppercase, all others lowercase.
pub fn capitalize(s: &fmt::Display) -> Result<String> {
    let mut s = format!("{}", s);

    match s.get_mut(0..1).map(|s| {
        s.make_ascii_uppercase();
        &*s
    }) {
        None => Ok(s),
        _ => {
            let l = s.len();
            match s.get_mut(1..l).map(|s| {
                s.make_ascii_lowercase();
                &*s
            }) {
                _ => Ok(s),
            }
        }
    }
}

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

    #[test]
    fn test_linebreaks() {
        assert_eq!(
            linebreaks(&"Foo\nBar Baz").unwrap(),
            "
Foo
Bar Baz
"
        );
        assert_eq!(
            linebreaks(&"Foo\nBar\n\nBaz").unwrap(),
            "
Foo
Bar

Baz
"
        );
    }

    #[test]
    fn test_linebreaksbr() {
        assert_eq!(linebreaksbr(&"Foo\nBar").unwrap(), "Foo
Bar");
        assert_eq!(
            linebreaksbr(&"Foo\nBar\n\nBaz").unwrap(),
            "Foo
Bar

Baz"
        );
    }

    #[test]
    fn test_lower() {
        assert_eq!(lower(&"Foo").unwrap(), "foo");
        assert_eq!(lower(&"FOO").unwrap(), "foo");
        assert_eq!(lower(&"FooBar").unwrap(), "foobar");
        assert_eq!(lower(&"foo").unwrap(), "foo");
    }

    #[test]
    fn test_upper() {
        assert_eq!(upper(&"Foo").unwrap(), "FOO");
        assert_eq!(upper(&"FOO").unwrap(), "FOO");
        assert_eq!(upper(&"FooBar").unwrap(), "FOOBAR");
        assert_eq!(upper(&"foo").unwrap(), "FOO");
    }

    #[test]
    fn test_trim() {
        assert_eq!(trim(&" Hello\tworld\t").unwrap(), "Hello\tworld");
    }

    #[test]
    fn test_join() {
        assert_eq!(
            join((&["hello", "world"]).into_iter(), ", ").unwrap(),
            "hello, world"
        );
        assert_eq!(join((&["hello"]).into_iter(), ", ").unwrap(), "hello");

        let empty: &[&str] = &[];
        assert_eq!(join(empty.into_iter(), ", ").unwrap(), "");

        let input: Vec<String> = vec!["foo".into(), "bar".into(), "bazz".into()];
        assert_eq!(
            join((&input).into_iter(), ":".to_string()).unwrap(),
            "foo:bar:bazz"
        );
        assert_eq!(
            join(input.clone().into_iter(), ":").unwrap(),
            "foo:bar:bazz"
        );
        assert_eq!(
            join(input.clone().into_iter(), ":".to_string()).unwrap(),
            "foo:bar:bazz"
        );

        let input: &[String] = &["foo".into(), "bar".into()];
        assert_eq!(join(input.into_iter(), ":").unwrap(), "foo:bar");
        assert_eq!(join(input.into_iter(), ":".to_string()).unwrap(), "foo:bar");

        let real: String = "blah".into();
        let input: Vec<&str> = vec![&real];
        assert_eq!(join(input.into_iter(), ";").unwrap(), "blah");

        assert_eq!(
            join((&&&&&["foo", "bar"]).into_iter(), ", ").unwrap(),
            "foo, bar"
        );
    }

    #[test]
    fn test_abs() {
        assert_eq!(abs(1).unwrap(), 1);
        assert_eq!(abs(-1).unwrap(), 1);
        assert_eq!(abs(1.0).unwrap(), 1.0);
        assert_eq!(abs(-1.0).unwrap(), 1.0);
        assert_eq!(abs(1.0 as f64).unwrap(), 1.0 as f64);
        assert_eq!(abs(-1.0 as f64).unwrap(), 1.0 as f64);
    }

    #[test]
    fn test_capitalize() {
        assert_eq!(capitalize(&"foo").unwrap(), "Foo".to_string());
        assert_eq!(capitalize(&"f").unwrap(), "F".to_string());
        assert_eq!(capitalize(&"fO").unwrap(), "Fo".to_string());
        assert_eq!(capitalize(&"").unwrap(), "".to_string());
        assert_eq!(capitalize(&"FoO").unwrap(), "Foo".to_string());
        assert_eq!(capitalize(&"foO BAR").unwrap(), "Foo bar".to_string());
    }
}
<|endoftext|>


<commit_before>//! Create new `Streams` connected to external inputs.

use std::rc::Rc;
use std::cell::RefCell;
use std::default::Default;

use progress::frontier::Antichain;
use progress::{Operate, Timestamp};
use progress::nested::subgraph::Source;
use progress::ChangeBatch;
use progress::timestamp::RootTimestamp;
use progress::nested::product::Product;

use timely_communication::Allocate;
use {Data, Push};
use dataflow::channels::Content;
use dataflow::channels::pushers::{Tee, Counter};

use dataflow::{Stream, Scope};
use dataflow::scopes::{Child, Root};

// TODO : This is an exogenous input, but it would be nice to wrap a Subgraph in something
// TODO : more like a harness, with direct access to its inputs.

// NOTE : This only takes a &self, not a &mut self, which works but is a bit weird.
// NOTE : Experiments with &mut indicate that the borrow of 'a lives for too long.
// NOTE : Might be able to fix with another lifetime parameter, say 'c: 'a.

/// Create a new `Stream` and `Handle` through which to supply input.
pub trait Input<'a, A: Allocate, T: Timestamp> {
    /// Create a new `Stream` and `Handle` through which to supply input.
    ///
    /// The `new_input` method returns a pair `(Handle, Stream)` where the `Stream` can be used
    /// immediately for timely dataflow construction, and the `Handle` is later used to introduce
    /// data into the timely dataflow computation.
    ///
    /// The `Handle` also provides a means to indicate
    /// to timely dataflow that the input has advanced beyond certain timestamps, allowing timely
    /// to issue progress notifications.
    ///
    /// #Examples
    /// ```
    /// use timely::*;
    /// use timely::dataflow::operators::{Input, Inspect};
    ///
    /// // construct and execute a timely dataflow
    /// timely::execute(Configuration::Thread, |worker| {
    ///
    ///     // add an input and base computation off of it
    ///     let mut input = worker.dataflow(|scope| {
    ///         let (input, stream) = scope.new_input();
    ///         stream.inspect(|x| println!("hello {:?}", x));
    ///         input
    ///     });
    ///
    ///     // introduce input, advance computation
    ///     for round in 0..10 {
    ///         input.send(round);
    ///         input.advance_to(round + 1);
    ///         worker.step();
    ///     }
    /// });
    /// ```
    fn new_input<D: Data>(&mut self) -> (Handle<T, D>, Stream<Child<'a, Root<A>, T>, D>);

    /// Create a new stream from a supplied interactive handle.
    ///
    /// This method creates a new timely stream whose data are supplied interactively through the `handle`
    /// argument. Each handle may be used multiple times (or not at all), and will clone data as appropriate
    /// if it as attached to more than one stream.
    ///
    /// #Examples
    /// ```
    /// use timely::*;
    /// use timely::dataflow::operators::{Input, Inspect};
    /// use timely::dataflow::operators::input::Handle;
    ///
    /// // construct and execute a timely dataflow
    /// timely::execute(Configuration::Thread, |worker| {
    ///
    ///     // add an input and base computation off of it
    ///     let mut input = Handle::new();
    ///     worker.dataflow(|scope| {
    ///         scope.input_from(&mut input)
    ///              .inspect(|x| println!("hello {:?}", x));
    ///     });
    ///
    ///     // introduce input, advance computation
    ///     for round in 0..10 {
    ///         input.send(round);
    ///         input.advance_to(round + 1);
    ///         worker.step();
    ///     }
    /// });
    /// ```
    fn input_from<D: Data>(&mut self, handle: &mut Handle<T, D>) -> Stream<Child<'a, Root<A>, T>, D>;
}

impl<'a, A: Allocate, T: Timestamp> Input<'a, A, T> for Child<'a, Root<A>, T> {
    fn new_input<D: Data>(&mut self) -> (Handle<T, D>, Stream<Child<'a, Root<A>, T>, D>) {
        let mut handle = Handle::new();
        let stream = self.input_from(&mut handle);
        (handle, stream)
    }

    fn input_from<D: Data>(&mut self, handle: &mut Handle<T, D>) -> Stream<Child<'a, Root<A>, T>, D> {

        let (output, registrar) = Tee::<Product<RootTimestamp, T>, D>::new();
        let counter = Counter::new(output);
        let produced = counter.produced().clone();

        let progress = Rc::new(RefCell::new(ChangeBatch::new()));

        handle.register(counter, progress.clone());

        let copies = self.peers();

        let index = self.add_operator(Operator {
            progress: progress,
            messages: produced,
            copies:   copies,
        });

        Stream::new(Source { index: index, port: 0 }, registrar, self.clone())
    }
}

struct Operator<T:Timestamp> {
    progress:   Rc<RefCell<ChangeBatch<Product<RootTimestamp, T>>>>,           // times closed since last asked
    messages:   Rc<RefCell<ChangeBatch<Product<RootTimestamp, T>>>>,           // messages sent since last asked
    copies:     usize,
}

impl<T:Timestamp> Operate<Product<RootTimestamp, T>> for Operator<T> {
    fn name(&self) -> String { "Input".to_owned() }
    fn inputs(&self) -> usize { 0 }
    fn outputs(&self) -> usize { 1 }

    fn get_internal_summary(&mut self) -> (Vec<Vec<Antichain<<Product<RootTimestamp, T> as Timestamp>::Summary>>>,
                                           Vec<ChangeBatch<Product<RootTimestamp, T>>>) {
        let mut map = ChangeBatch::new();
        map.update(Default::default(), self.copies as i64);
        (Vec::new(), vec![map])
    }

    fn pull_internal_progress(&mut self,_messages_consumed: &mut [ChangeBatch<Product<RootTimestamp, T>>],
                                         frontier_progress: &mut [ChangeBatch<Product<RootTimestamp, T>>],
                                         messages_produced: &mut [ChangeBatch<Product<RootTimestamp, T>>]) -> bool
    {
        self.messages.borrow_mut().drain_into(&mut messages_produced[0]);
        self.progress.borrow_mut().drain_into(&mut frontier_progress[0]);
        false
    }

    fn notify_me(&self) -> bool { false }
}

/// A handle to an input `Stream`, used to introduce data to a timely dataflow computation.
pub struct Handle<T: Timestamp, D: Data> {
    progress: Vec<Rc<RefCell<ChangeBatch<Product<RootTimestamp, T>>>>>,
    pushers: Vec<Counter<Product<RootTimestamp, T>, D, Tee<Product<RootTimestamp, T>, D>>>,
    buffer1: Vec<D>,
    buffer2: Vec<D>,
    now_at: Product<RootTimestamp, T>,
}

impl<T:Timestamp, D: Data> Handle<T, D> {

    /// Allocates a new input handle, from which one can create timely streams.
    ///
    /// #Examples
    /// ```
    /// use timely::*;
    /// use timely::dataflow::operators::{Input, Inspect};
    /// use timely::dataflow::operators::input::Handle;
    ///
    /// // construct and execute a timely dataflow
    /// timely::execute(Configuration::Thread, |worker| {
    ///
    ///     // add an input and base computation off of it
    ///     let mut input = Handle::new();
    ///     worker.dataflow(|scope| {
    ///         scope.input_from(&mut input)
    ///              .inspect(|x| println!("hello {:?}", x));
    ///     });
    ///
    ///     // introduce input, advance computation
    ///     for round in 0..10 {
    ///         input.send(round);
    ///         input.advance_to(round + 1);
    ///         worker.step();
    ///     }
    /// });
    /// ```
    pub fn new() -> Self {
        Handle {
            progress: Vec::new(),
            pushers: Vec::new(),
            buffer1: Vec::with_capacity(Content::<D>::default_length()),
            buffer2: Vec::with_capacity(Content::<D>::default_length()),
            now_at: Default::default(),
        }
    }

    /// Creates an input stream from the handle in the supplied scope.
    ///
    /// #Examples
    /// ```
    /// use timely::*;
    /// use timely::dataflow::operators::{Input, Inspect};
    /// use timely::dataflow::operators::input::Handle;
    ///
    /// // construct and execute a timely dataflow
    /// timely::execute(Configuration::Thread, |worker| {
    ///
    ///     // add an input and base computation off of it
    ///     let mut input = Handle::new();
    ///     worker.dataflow(|scope| {
    ///         input.to_stream(scope)
    ///              .inspect(|x| println!("hello {:?}", x));
    ///     });
    ///
    ///     // introduce input, advance computation
    ///     for round in 0..10 {
    ///         input.send(round);
    ///         input.advance_to(round + 1);
    ///         worker.step();
    ///     }
    /// });
    /// ```
    pub fn to_stream<'a, A: Allocate>(&mut self, scope: &mut Child<'a, Root<A>, T>) -> Stream<Child<'a, Root<A>, T>, D>
    where T: Ord {
        scope.input_from(self)
    }

    fn register(
        &mut self,
        pusher: Counter<Product<RootTimestamp, T>, D, Tee<Product<RootTimestamp, T>, D>>,
        progress: Rc<RefCell<ChangeBatch<Product<RootTimestamp, T>>>>
    ) {
        // flush current contents, so new registrant does not see existing data.
        if !self.buffer1.is_empty() { self.flush(); }

        // we need to produce an appropriate update to the capabilities for `progress`, in case a
        // user has decided to drive the handle around a bit before registering it.
        progress.borrow_mut().update(Default::default(), -1);
        progress.borrow_mut().update(self.now_at.clone(), 1);

        self.progress.push(progress);
        self.pushers.push(pusher);
    }

    // flushes our buffer at each of the destinations. there can be more than one; clone if needed.
    fn flush(&mut self) {
        for index in 0 .. self.pushers.len() {
            if index < self.pushers.len() - 1 {
                self.buffer2.extend_from_slice(&self.buffer1[..]);
                Content::push_at(&mut self.buffer2, self.now_at.clone(), &mut self.pushers[index]);
                assert!(self.buffer2.is_empty());
            }
            else {
                Content::push_at(&mut self.buffer1, self.now_at.clone(), &mut self.pushers[index]);
                assert!(self.buffer1.is_empty());
            }
        }
        self.buffer1.clear();
    }

    // closes the current epoch, flushing if needed, shutting if needed, and updating the frontier.
    fn close_epoch(&mut self) {
        if !self.buffer1.is_empty() { self.flush(); }
        for pusher in self.pushers.iter_mut() {
            pusher.done();
        }
        for progress in self.progress.iter() {
            progress.borrow_mut().update(self.now_at.clone(), -1);
        }
    }

    #[inline(always)]
    /// Sends one record into the corresponding timely dataflow `Stream`, at the current epoch.
    pub fn send(&mut self, data: D) {
        // assert!(self.buffer.capacity() == Content::<D>::default_length());
        self.buffer1.push(data);
        if self.buffer1.len() == self.buffer1.capacity() {
            self.flush();
        }
    }

    /// Sends a batch of records into the corresponding timely dataflow `Stream`, at the current epoch.
    ///
    /// This method flushes single elements previously sent with `send`, to keep the insertion order.
    pub fn send_batch(&mut self, buffer: &mut Vec<D>) {

        // flush buffered elements to ensure local fifo.
        if !self.buffer1.is_empty() { self.flush(); }

        // push buffer (or clone of buffer) at each destination.
        for index in 0 .. self.pushers.len() {
            if index < self.pushers.len() - 1 {
                self.buffer2.extend_from_slice(&buffer[..]);
                Content::push_at(&mut self.buffer2, self.now_at.clone(), &mut self.pushers[index]);
                assert!(self.buffer2.is_empty());
            }
            else {
                Content::push_at(buffer, self.now_at.clone(), &mut self.pushers[index]);
                assert!(buffer.is_empty());
            }
        }
        buffer.clear();
    }

    /// Advances the current epoch to `next`.
    ///
    /// This method allows timely dataflow to issue progress notifications as it can now determine
    /// that this input can no longer produce data at earlier timestamps.
    pub fn advance_to(&mut self, next: T) {
        assert!(self.now_at.inner.less_equal(&next));
        self.close_epoch();
        self.now_at = RootTimestamp::new(next);
        for progress in self.progress.iter() {
            progress.borrow_mut().update(self.now_at.clone(), 1);
        }
    }

    /// Closes the input.
    ///
    /// This method allows timely dataflow to issue all progress notifications blocked by this input
    /// and to begin to shut down operators, as this input can no longer produce data.
    pub fn close(self) { }

    /// Reports the current epoch.
    pub fn epoch(&self) -> &T {
        &self.now_at.inner
    }

    /// Reports the current timestamp.
    pub fn time(&self) -> &Product<RootTimestamp, T> {
        &self.now_at
    }
}

impl<T:Timestamp, D: Data> Drop for Handle<T, D> {
    fn drop(&mut self) {
        self.close_epoch();
    }
}
<commit_msg>don't send empty batches<commit_after>//! Create new `Streams` connected to external inputs.

use std::rc::Rc;
use std::cell::RefCell;
use std::default::Default;

use progress::frontier::Antichain;
use progress::{Operate, Timestamp};
use progress::nested::subgraph::Source;
use progress::ChangeBatch;
use progress::timestamp::RootTimestamp;
use progress::nested::product::Product;

use timely_communication::Allocate;
use {Data, Push};
use dataflow::channels::Content;
use dataflow::channels::pushers::{Tee, Counter};

use dataflow::{Stream, Scope};
use dataflow::scopes::{Child, Root};

// TODO : This is an exogenous input, but it would be nice to wrap a Subgraph in something
// TODO : more like a harness, with direct access to its inputs.

// NOTE : This only takes a &self, not a &mut self, which works but is a bit weird.
// NOTE : Experiments with &mut indicate that the borrow of 'a lives for too long.
// NOTE : Might be able to fix with another lifetime parameter, say 'c: 'a.

/// Create a new `Stream` and `Handle` through which to supply input.
pub trait Input<'a, A: Allocate, T: Timestamp> {
    /// Create a new `Stream` and `Handle` through which to supply input.
    ///
    /// The `new_input` method returns a pair `(Handle, Stream)` where the `Stream` can be used
    /// immediately for timely dataflow construction, and the `Handle` is later used to introduce
    /// data into the timely dataflow computation.
    ///
    /// The `Handle` also provides a means to indicate
    /// to timely dataflow that the input has advanced beyond certain timestamps, allowing timely
    /// to issue progress notifications.
    ///
    /// #Examples
    /// ```
    /// use timely::*;
    /// use timely::dataflow::operators::{Input, Inspect};
    ///
    /// // construct and execute a timely dataflow
    /// timely::execute(Configuration::Thread, |worker| {
    ///
    ///     // add an input and base computation off of it
    ///     let mut input = worker.dataflow(|scope| {
    ///         let (input, stream) = scope.new_input();
    ///         stream.inspect(|x| println!("hello {:?}", x));
    ///         input
    ///     });
    ///
    ///     // introduce input, advance computation
    ///     for round in 0..10 {
    ///         input.send(round);
    ///         input.advance_to(round + 1);
    ///         worker.step();
    ///     }
    /// });
    /// ```
    fn new_input<D: Data>(&mut self) -> (Handle<T, D>, Stream<Child<'a, Root<A>, T>, D>);

    /// Create a new stream from a supplied interactive handle.
    ///
    /// This method creates a new timely stream whose data are supplied interactively through the `handle`
    /// argument. Each handle may be used multiple times (or not at all), and will clone data as appropriate
    /// if it as attached to more than one stream.
    ///
    /// #Examples
    /// ```
    /// use timely::*;
    /// use timely::dataflow::operators::{Input, Inspect};
    /// use timely::dataflow::operators::input::Handle;
    ///
    /// // construct and execute a timely dataflow
    /// timely::execute(Configuration::Thread, |worker| {
    ///
    ///     // add an input and base computation off of it
    ///     let mut input = Handle::new();
    ///     worker.dataflow(|scope| {
    ///         scope.input_from(&mut input)
    ///              .inspect(|x| println!("hello {:?}", x));
    ///     });
    ///
    ///     // introduce input, advance computation
    ///     for round in 0..10 {
    ///         input.send(round);
    ///         input.advance_to(round + 1);
    ///         worker.step();
    ///     }
    /// });
    /// ```
    fn input_from<D: Data>(&mut self, handle: &mut Handle<T, D>) -> Stream<Child<'a, Root<A>, T>, D>;
}

impl<'a, A: Allocate, T: Timestamp> Input<'a, A, T> for Child<'a, Root<A>, T> {
    fn new_input<D: Data>(&mut self) -> (Handle<T, D>, Stream<Child<'a, Root<A>, T>, D>) {
        let mut handle = Handle::new();
        let stream = self.input_from(&mut handle);
        (handle, stream)
    }

    fn input_from<D: Data>(&mut self, handle: &mut Handle<T, D>) -> Stream<Child<'a, Root<A>, T>, D> {

        let (output, registrar) = Tee::<Product<RootTimestamp, T>, D>::new();
        let counter = Counter::new(output);
        let produced = counter.produced().clone();

        let progress = Rc::new(RefCell::new(ChangeBatch::new()));

        handle.register(counter, progress.clone());

        let copies = self.peers();

        let index = self.add_operator(Operator {
            progress: progress,
            messages: produced,
            copies:   copies,
        });

        Stream::new(Source { index: index, port: 0 }, registrar, self.clone())
    }
}

struct Operator<T:Timestamp> {
    progress:   Rc<RefCell<ChangeBatch<Product<RootTimestamp, T>>>>,           // times closed since last asked
    messages:   Rc<RefCell<ChangeBatch<Product<RootTimestamp, T>>>>,           // messages sent since last asked
    copies:     usize,
}

impl<T:Timestamp> Operate<Product<RootTimestamp, T>> for Operator<T> {
    fn name(&self) -> String { "Input".to_owned() }
    fn inputs(&self) -> usize { 0 }
    fn outputs(&self) -> usize { 1 }

    fn get_internal_summary(&mut self) -> (Vec<Vec<Antichain<<Product<RootTimestamp, T> as Timestamp>::Summary>>>,
                                           Vec<ChangeBatch<Product<RootTimestamp, T>>>) {
        let mut map = ChangeBatch::new();
        map.update(Default::default(), self.copies as i64);
        (Vec::new(), vec![map])
    }

    fn pull_internal_progress(&mut self,_messages_consumed: &mut [ChangeBatch<Product<RootTimestamp, T>>],
                                         frontier_progress: &mut [ChangeBatch<Product<RootTimestamp, T>>],
                                         messages_produced: &mut [ChangeBatch<Product<RootTimestamp, T>>]) -> bool
    {
        self.messages.borrow_mut().drain_into(&mut messages_produced[0]);
        self.progress.borrow_mut().drain_into(&mut frontier_progress[0]);
        false
    }

    fn notify_me(&self) -> bool { false }
}

/// A handle to an input `Stream`, used to introduce data to a timely dataflow computation.
pub struct Handle<T: Timestamp, D: Data> {
    progress: Vec<Rc<RefCell<ChangeBatch<Product<RootTimestamp, T>>>>>,
    pushers: Vec<Counter<Product<RootTimestamp, T>, D, Tee<Product<RootTimestamp, T>, D>>>,
    buffer1: Vec<D>,
    buffer2: Vec<D>,
    now_at: Product<RootTimestamp, T>,
}

impl<T:Timestamp, D: Data> Handle<T, D> {

    /// Allocates a new input handle, from which one can create timely streams.
    ///
    /// #Examples
    /// ```
    /// use timely::*;
    /// use timely::dataflow::operators::{Input, Inspect};
    /// use timely::dataflow::operators::input::Handle;
    ///
    /// // construct and execute a timely dataflow
    /// timely::execute(Configuration::Thread, |worker| {
    ///
    ///     // add an input and base computation off of it
    ///     let mut input = Handle::new();
    ///     worker.dataflow(|scope| {
    ///         scope.input_from(&mut input)
    ///              .inspect(|x| println!("hello {:?}", x));
    ///     });
    ///
    ///     // introduce input, advance computation
    ///     for round in 0..10 {
    ///         input.send(round);
    ///         input.advance_to(round + 1);
    ///         worker.step();
    ///     }
    /// });
    /// ```
    pub fn new() -> Self {
        Handle {
            progress: Vec::new(),
            pushers: Vec::new(),
            buffer1: Vec::with_capacity(Content::<D>::default_length()),
            buffer2: Vec::with_capacity(Content::<D>::default_length()),
            now_at: Default::default(),
        }
    }

    /// Creates an input stream from the handle in the supplied scope.
    ///
    /// #Examples
    /// ```
    /// use timely::*;
    /// use timely::dataflow::operators::{Input, Inspect};
    /// use timely::dataflow::operators::input::Handle;
    ///
    /// // construct and execute a timely dataflow
    /// timely::execute(Configuration::Thread, |worker| {
    ///
    ///     // add an input and base computation off of it
    ///     let mut input = Handle::new();
    ///     worker.dataflow(|scope| {
    ///         input.to_stream(scope)
    ///              .inspect(|x| println!("hello {:?}", x));
    ///     });
    ///
    ///     // introduce input, advance computation
    ///     for round in 0..10 {
    ///         input.send(round);
    ///         input.advance_to(round + 1);
    ///         worker.step();
    ///     }
    /// });
    /// ```
    pub fn to_stream<'a, A: Allocate>(&mut self, scope: &mut Child<'a, Root<A>, T>) -> Stream<Child<'a, Root<A>, T>, D>
    where T: Ord {
        scope.input_from(self)
    }

    fn register(
        &mut self,
        pusher: Counter<Product<RootTimestamp, T>, D, Tee<Product<RootTimestamp, T>, D>>,
        progress: Rc<RefCell<ChangeBatch<Product<RootTimestamp, T>>>>
    ) {
        // flush current contents, so new registrant does not see existing data.
        if !self.buffer1.is_empty() { self.flush(); }

        // we need to produce an appropriate update to the capabilities for `progress`, in case a
        // user has decided to drive the handle around a bit before registering it.
        progress.borrow_mut().update(Default::default(), -1);
        progress.borrow_mut().update(self.now_at.clone(), 1);

        self.progress.push(progress);
        self.pushers.push(pusher);
    }

    // flushes our buffer at each of the destinations. there can be more than one; clone if needed.
    fn flush(&mut self) {
        for index in 0 .. self.pushers.len() {
            if index < self.pushers.len() - 1 {
                self.buffer2.extend_from_slice(&self.buffer1[..]);
                Content::push_at(&mut self.buffer2, self.now_at.clone(), &mut self.pushers[index]);
                assert!(self.buffer2.is_empty());
            }
            else {
                Content::push_at(&mut self.buffer1, self.now_at.clone(), &mut self.pushers[index]);
                assert!(self.buffer1.is_empty());
            }
        }
        self.buffer1.clear();
    }

    // closes the current epoch, flushing if needed, shutting if needed, and updating the frontier.
    fn close_epoch(&mut self) {
        if !self.buffer1.is_empty() { self.flush(); }
        for pusher in self.pushers.iter_mut() {
            pusher.done();
        }
        for progress in self.progress.iter() {
            progress.borrow_mut().update(self.now_at.clone(), -1);
        }
    }

    #[inline(always)]
    /// Sends one record into the corresponding timely dataflow `Stream`, at the current epoch.
    pub fn send(&mut self, data: D) {
        // assert!(self.buffer.capacity() == Content::<D>::default_length());
        self.buffer1.push(data);
        if self.buffer1.len() == self.buffer1.capacity() {
            self.flush();
        }
    }

    /// Sends a batch of records into the corresponding timely dataflow `Stream`, at the current epoch.
    ///
    /// This method flushes single elements previously sent with `send`, to keep the insertion order.
    pub fn send_batch(&mut self, buffer: &mut Vec<D>) {

        if !buffer.is_empty() {
            // flush buffered elements to ensure local fifo.
            if !self.buffer1.is_empty() { self.flush(); }

            // push buffer (or clone of buffer) at each destination.
            for index in 0 .. self.pushers.len() {
                if index < self.pushers.len() - 1 {
                    self.buffer2.extend_from_slice(&buffer[..]);
                    Content::push_at(&mut self.buffer2, self.now_at.clone(), &mut self.pushers[index]);
                    assert!(self.buffer2.is_empty());
                }
                else {
                    Content::push_at(buffer, self.now_at.clone(), &mut self.pushers[index]);
                    assert!(buffer.is_empty());
                }
            }
            buffer.clear();
        }
    }

    /// Advances the current epoch to `next`.
    ///
    /// This method allows timely dataflow to issue progress notifications as it can now determine
    /// that this input can no longer produce data at earlier timestamps.
    pub fn advance_to(&mut self, next: T) {
        assert!(self.now_at.inner.less_equal(&next));
        self.close_epoch();
        self.now_at = RootTimestamp::new(next);
        for progress in self.progress.iter() {
            progress.borrow_mut().update(self.now_at.clone(), 1);
        }
    }

    /// Closes the input.
    ///
    /// This method allows timely dataflow to issue all progress notifications blocked by this input
    /// and to begin to shut down operators, as this input can no longer produce data.
    pub fn close(self) { }

    /// Reports the current epoch.
    pub fn epoch(&self) -> &T {
        &self.now_at.inner
    }

    /// Reports the current timestamp.
    pub fn time(&self) -> &Product<RootTimestamp, T> {
        &self.now_at
    }
}

impl<T:Timestamp, D: Data> Drop for Handle<T, D> {
    fn drop(&mut self) {
        self.close_epoch();
    }
}
<|endoftext|>


<commit_before>//! Intel Instruction Set Architectures.

pub mod settings;
mod abi;
mod binemit;
mod enc_tables;
mod registers;

use binemit::{CodeSink, MemoryCodeSink, emit_function};
use super::super::settings as shared_settings;
use isa::enc_tables::{self as shared_enc_tables, lookup_enclist, Encodings};
use isa::Builder as IsaBuilder;
use isa::{TargetIsa, RegInfo, RegClass, EncInfo, RegUnit};
use self::registers::RU;
use ir;
use regalloc;
use result;
use ir::InstBuilder;
use ir::immediates::Imm64;
use stack_layout::layout_stack;
use cursor::{Cursor, EncCursor};

#[allow(dead_code)]
struct Isa {
    shared_flags: shared_settings::Flags,
    isa_flags: settings::Flags,
    cpumode: &'static [shared_enc_tables::Level1Entry<u16>],
}

/// Get an ISA builder for creating Intel targets.
pub fn isa_builder() -> IsaBuilder {
    IsaBuilder {
        setup: settings::builder(),
        constructor: isa_constructor,
    }
}

fn isa_constructor(
    shared_flags: shared_settings::Flags,
    builder: &shared_settings::Builder,
) -> Box<TargetIsa> {
    let level1 = if shared_flags.is_64bit() {
        &enc_tables::LEVEL1_I64[..]
    } else {
        &enc_tables::LEVEL1_I32[..]
    };
    Box::new(Isa {
        isa_flags: settings::Flags::new(&shared_flags, builder),
        shared_flags,
        cpumode: level1,
    })
}

impl TargetIsa for Isa {
    fn name(&self) -> &'static str {
        "intel"
    }

    fn flags(&self) -> &shared_settings::Flags {
        &self.shared_flags
    }

    fn register_info(&self) -> RegInfo {
        registers::INFO.clone()
    }

    fn encoding_info(&self) -> EncInfo {
        enc_tables::INFO.clone()
    }

    fn legal_encodings<'a>(
        &'a self,
        dfg: &'a ir::DataFlowGraph,
        inst: &'a ir::InstructionData,
        ctrl_typevar: ir::Type,
    ) -> Encodings<'a> {
        lookup_enclist(
            ctrl_typevar,
            inst,
            dfg,
            self.cpumode,
            &enc_tables::LEVEL2[..],
            &enc_tables::ENCLISTS[..],
            &enc_tables::LEGALIZE_ACTIONS[..],
            &enc_tables::RECIPE_PREDICATES[..],
            &enc_tables::INST_PREDICATES[..],
            self.isa_flags.predicate_view(),
        )
    }

    fn legalize_signature(&self, sig: &mut ir::Signature, current: bool) {
        abi::legalize_signature(sig, &self.shared_flags, current)
    }

    fn regclass_for_abi_type(&self, ty: ir::Type) -> RegClass {
        abi::regclass_for_abi_type(ty)
    }

    fn allocatable_registers(&self, func: &ir::Function) -> regalloc::AllocatableSet {
        abi::allocatable_registers(func, &self.shared_flags)
    }

    fn emit_inst(
        &self,
        func: &ir::Function,
        inst: ir::Inst,
        divert: &mut regalloc::RegDiversions,
        sink: &mut CodeSink,
    ) {
        binemit::emit_inst(func, inst, divert, sink)
    }

    fn emit_function(&self, func: &ir::Function, sink: &mut MemoryCodeSink) {
        emit_function(func, binemit::emit_inst, sink)
    }

    fn reloc_names(&self) -> &'static [&'static str] {
        &binemit::RELOC_NAMES
    }

    fn prologue_epilogue(&self, func: &mut ir::Function) -> result::CtonResult {
        let word_size = if self.flags().is_64bit() { 8 } else { 4 };
        let csr_type = if self.flags().is_64bit() {
            ir::types::I64
        } else {
            ir::types::I32
        };
        let csrs = abi::callee_saved_registers(&self.shared_flags);
        let csr_stack_size = ((csrs.len() + 1) * word_size as usize) as i32;

        func.create_stack_slot(ir::StackSlotData {
            kind: ir::StackSlotKind::IncomingArg,
            size: csr_stack_size as u32,
            offset: -csr_stack_size,
        });

        let total_stack_size = layout_stack(&mut func.stack_slots, word_size)? as i32;
        let local_stack_size = (total_stack_size - csr_stack_size) as i64;

        // Add CSRs to function signature
        let fp_arg = ir::AbiParam::special_reg(
            csr_type,
            ir::ArgumentPurpose::FramePointer,
            RU::rbp as RegUnit,
        );
        func.signature.params.push(fp_arg);
        func.signature.returns.push(fp_arg);

        for csr in csrs.iter() {
            let csr_arg = ir::AbiParam::special_reg(
                csr_type,
                ir::ArgumentPurpose::CalleeSaved,
                *csr as RegUnit,
            );
            func.signature.params.push(csr_arg);
            func.signature.returns.push(csr_arg);
        }

        let entry_ebb = func.layout.entry_block().expect("missing entry block");
        let mut pos = EncCursor::new(func, self).at_first_insertion_point(entry_ebb);

        self.insert_prologue(&mut pos, local_stack_size, csr_type);
        self.insert_epilogues(&mut pos, local_stack_size, csr_type);

        Ok(())
    }
}

impl Isa {
    fn insert_prologue(&self, pos: &mut EncCursor, stack_size: i64, csr_type: ir::types::Type) {
        // Append param to entry EBB
        let ebb = pos.current_ebb().expect("missing ebb under cursor");
        let fp = pos.func.dfg.append_ebb_param(ebb, csr_type);
        pos.func.locations[fp] = ir::ValueLoc::Reg(RU::rbp as RegUnit);

        pos.ins().x86_push(fp);
        pos.ins().copy_special(
            RU::rsp as RegUnit,
            RU::rbp as RegUnit,
        );

        if stack_size > 0 {
            pos.ins().adjust_sp_imm(Imm64::new(-stack_size));
        }

        let csrs = abi::callee_saved_registers(&self.shared_flags);
        for reg in csrs.iter() {
            // Append param to entry EBB
            let csr_arg = pos.func.dfg.append_ebb_param(ebb, csr_type);

            // Assign it a location
            pos.func.locations[csr_arg] = ir::ValueLoc::Reg(*reg as RegUnit);

            // Remember it so we can push it momentarily
            pos.ins().x86_push(csr_arg);
        }
    }

    fn insert_epilogues(&self, pos: &mut EncCursor, stack_size: i64, csr_type: ir::types::Type) {
        while let Some(ebb) = pos.next_ebb() {
            pos.goto_last_inst(ebb);
            if let Some(inst) = pos.current_inst() {
                if pos.func.dfg[inst].opcode().is_return() {
                    self.insert_epilogue(inst, stack_size, pos, csr_type);
                }
            }
        }

    }

    fn insert_epilogue(
        &self,
        inst: ir::Inst,
        stack_size: i64,
        pos: &mut EncCursor,
        csr_type: ir::types::Type,
    ) {
        if stack_size > 0 {
            pos.ins().adjust_sp_imm(Imm64::new(stack_size));
        }

        let csrs = abi::callee_saved_registers(&self.shared_flags);
        for reg in csrs.iter().rev() {
            let csr_ret = pos.ins().x86_pop(csr_type);
            pos.func.locations[csr_ret] = ir::ValueLoc::Reg(*reg as RegUnit);
            pos.func.dfg.append_inst_arg(inst, csr_ret);
        }

        let fp_ret = pos.ins().x86_pop(csr_type);
        pos.func.locations[fp_ret] = ir::ValueLoc::Reg(RU::rbp as RegUnit);
        pos.func.dfg.append_inst_arg(inst, fp_ret);

    }
}
<commit_msg>Fix the ordering of return values.<commit_after>//! Intel Instruction Set Architectures.

pub mod settings;
mod abi;
mod binemit;
mod enc_tables;
mod registers;

use binemit::{CodeSink, MemoryCodeSink, emit_function};
use super::super::settings as shared_settings;
use isa::enc_tables::{self as shared_enc_tables, lookup_enclist, Encodings};
use isa::Builder as IsaBuilder;
use isa::{TargetIsa, RegInfo, RegClass, EncInfo, RegUnit};
use self::registers::RU;
use ir;
use regalloc;
use result;
use ir::InstBuilder;
use ir::immediates::Imm64;
use stack_layout::layout_stack;
use cursor::{Cursor, EncCursor};

#[allow(dead_code)]
struct Isa {
    shared_flags: shared_settings::Flags,
    isa_flags: settings::Flags,
    cpumode: &'static [shared_enc_tables::Level1Entry<u16>],
}

/// Get an ISA builder for creating Intel targets.
pub fn isa_builder() -> IsaBuilder {
    IsaBuilder {
        setup: settings::builder(),
        constructor: isa_constructor,
    }
}

fn isa_constructor(
    shared_flags: shared_settings::Flags,
    builder: &shared_settings::Builder,
) -> Box<TargetIsa> {
    let level1 = if shared_flags.is_64bit() {
        &enc_tables::LEVEL1_I64[..]
    } else {
        &enc_tables::LEVEL1_I32[..]
    };
    Box::new(Isa {
        isa_flags: settings::Flags::new(&shared_flags, builder),
        shared_flags,
        cpumode: level1,
    })
}

impl TargetIsa for Isa {
    fn name(&self) -> &'static str {
        "intel"
    }

    fn flags(&self) -> &shared_settings::Flags {
        &self.shared_flags
    }

    fn register_info(&self) -> RegInfo {
        registers::INFO.clone()
    }

    fn encoding_info(&self) -> EncInfo {
        enc_tables::INFO.clone()
    }

    fn legal_encodings<'a>(
        &'a self,
        dfg: &'a ir::DataFlowGraph,
        inst: &'a ir::InstructionData,
        ctrl_typevar: ir::Type,
    ) -> Encodings<'a> {
        lookup_enclist(
            ctrl_typevar,
            inst,
            dfg,
            self.cpumode,
            &enc_tables::LEVEL2[..],
            &enc_tables::ENCLISTS[..],
            &enc_tables::LEGALIZE_ACTIONS[..],
            &enc_tables::RECIPE_PREDICATES[..],
            &enc_tables::INST_PREDICATES[..],
            self.isa_flags.predicate_view(),
        )
    }

    fn legalize_signature(&self, sig: &mut ir::Signature, current: bool) {
        abi::legalize_signature(sig, &self.shared_flags, current)
    }

    fn regclass_for_abi_type(&self, ty: ir::Type) -> RegClass {
        abi::regclass_for_abi_type(ty)
    }

    fn allocatable_registers(&self, func: &ir::Function) -> regalloc::AllocatableSet {
        abi::allocatable_registers(func, &self.shared_flags)
    }

    fn emit_inst(
        &self,
        func: &ir::Function,
        inst: ir::Inst,
        divert: &mut regalloc::RegDiversions,
        sink: &mut CodeSink,
    ) {
        binemit::emit_inst(func, inst, divert, sink)
    }

    fn emit_function(&self, func: &ir::Function, sink: &mut MemoryCodeSink) {
        emit_function(func, binemit::emit_inst, sink)
    }

    fn reloc_names(&self) -> &'static [&'static str] {
        &binemit::RELOC_NAMES
    }

    fn prologue_epilogue(&self, func: &mut ir::Function) -> result::CtonResult {
        let word_size = if self.flags().is_64bit() { 8 } else { 4 };
        let csr_type = if self.flags().is_64bit() {
            ir::types::I64
        } else {
            ir::types::I32
        };
        let csrs = abi::callee_saved_registers(&self.shared_flags);
        let csr_stack_size = ((csrs.len() + 1) * word_size as usize) as i32;

        func.create_stack_slot(ir::StackSlotData {
            kind: ir::StackSlotKind::IncomingArg,
            size: csr_stack_size as u32,
            offset: -csr_stack_size,
        });

        let total_stack_size = layout_stack(&mut func.stack_slots, word_size)? as i32;
        let local_stack_size = (total_stack_size - csr_stack_size) as i64;

        // Add CSRs to function signature
        let fp_arg = ir::AbiParam::special_reg(
            csr_type,
            ir::ArgumentPurpose::FramePointer,
            RU::rbp as RegUnit,
        );
        func.signature.params.push(fp_arg);
        func.signature.returns.push(fp_arg);

        for csr in csrs.iter() {
            let csr_arg = ir::AbiParam::special_reg(
                csr_type,
                ir::ArgumentPurpose::CalleeSaved,
                *csr as RegUnit,
            );
            func.signature.params.push(csr_arg);
            func.signature.returns.push(csr_arg);
        }

        let entry_ebb = func.layout.entry_block().expect("missing entry block");
        let mut pos = EncCursor::new(func, self).at_first_insertion_point(entry_ebb);

        self.insert_prologue(&mut pos, local_stack_size, csr_type);
        self.insert_epilogues(&mut pos, local_stack_size, csr_type);

        Ok(())
    }
}

impl Isa {
    fn insert_prologue(&self, pos: &mut EncCursor, stack_size: i64, csr_type: ir::types::Type) {
        // Append param to entry EBB
        let ebb = pos.current_ebb().expect("missing ebb under cursor");
        let fp = pos.func.dfg.append_ebb_param(ebb, csr_type);
        pos.func.locations[fp] = ir::ValueLoc::Reg(RU::rbp as RegUnit);

        pos.ins().x86_push(fp);
        pos.ins().copy_special(
            RU::rsp as RegUnit,
            RU::rbp as RegUnit,
        );

        if stack_size > 0 {
            pos.ins().adjust_sp_imm(Imm64::new(-stack_size));
        }

        let csrs = abi::callee_saved_registers(&self.shared_flags);
        for reg in csrs.iter() {
            // Append param to entry EBB
            let csr_arg = pos.func.dfg.append_ebb_param(ebb, csr_type);

            // Assign it a location
            pos.func.locations[csr_arg] = ir::ValueLoc::Reg(*reg as RegUnit);

            // Remember it so we can push it momentarily
            pos.ins().x86_push(csr_arg);
        }
    }

    fn insert_epilogues(&self, pos: &mut EncCursor, stack_size: i64, csr_type: ir::types::Type) {
        while let Some(ebb) = pos.next_ebb() {
            pos.goto_last_inst(ebb);
            if let Some(inst) = pos.current_inst() {
                if pos.func.dfg[inst].opcode().is_return() {
                    self.insert_epilogue(inst, stack_size, pos, csr_type);
                }
            }
        }

    }

    fn insert_epilogue(
        &self,
        inst: ir::Inst,
        stack_size: i64,
        pos: &mut EncCursor,
        csr_type: ir::types::Type,
    ) {
        if stack_size > 0 {
            pos.ins().adjust_sp_imm(Imm64::new(stack_size));
        }

        let fp_ret = pos.ins().x86_pop(csr_type);
        pos.prev_inst();

        pos.func.locations[fp_ret] = ir::ValueLoc::Reg(RU::rbp as RegUnit);
        pos.func.dfg.append_inst_arg(inst, fp_ret);

        let csrs = abi::callee_saved_registers(&self.shared_flags);
        for reg in csrs.iter() {
            let csr_ret = pos.ins().x86_pop(csr_type);
            pos.prev_inst();

            pos.func.locations[csr_ret] = ir::ValueLoc::Reg(*reg as RegUnit);
            pos.func.dfg.append_inst_arg(inst, csr_ret);
        }
    }
}
<|endoftext|>


<commit_before>use flowgger::config::Config;
use flowgger::decoder::Decoder;
use flowgger::encoder::Encoder;
use openssl::ssl::*;
use openssl::ssl::SslMethod::*;
use openssl::x509::X509FileType;
use std::io::{stderr, Write, BufRead, BufReader};
use std::net::{TcpListener, TcpStream};
use std::path::Path;
use std::str;
use std::sync::mpsc::SyncSender;
use std::thread;
use super::Input;

const DEFAULT_CERT: &'static str = "flowgger.pem";
const DEFAULT_CIPHERS: &'static str = "ECDHE-RSA-CHACHA20-POLY1305:ECDHE-ECDSA-CHACHA20-POLY1305:ECDHE-RSA-AES128-GCM-SHA256:ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES256-GCM-SHA384:ECDHE-ECDSA-AES256-GCM-SHA384:DHE-RSA-AES128-GCM-SHA256:DHE-DSS-AES128-GCM-SHA256:kEDH+AESGCM:ECDHE-RSA-AES128-SHA256:ECDHE-ECDSA-AES128-SHA256:ECDHE-RSA-AES128-SHA:ECDHE-ECDSA-AES128-SHA:ECDHE-RSA-AES256-SHA384:ECDHE-ECDSA-AES256-SHA384:ECDHE-RSA-AES256-SHA:ECDHE-ECDSA-AES256-SHA:DHE-RSA-AES128-SHA256:DHE-RSA-AES128-SHA:DHE-DSS-AES128-SHA256:DHE-RSA-AES256-SHA256:DHE-DSS-AES256-SHA:DHE-RSA-AES256-SHA:AES128-GCM-SHA256:AES256-GCM-SHA384:AES128-SHA256:AES256-SHA256:AES128-SHA:AES256-SHA:AES:CAMELLIA:DES-CBC3-SHA:!aNULL:!eNULL:!EXPORT:!DES:!RC4:!MD5:!PSK:!aECDH:!EDH-DSS-DES-CBC3-SHA:!EDH-RSA-DES-CBC3-SHA:!KRB5-DES-CBC3-SH";
const DEFAULT_FRAMED: bool = false;
const DEFAULT_KEY: &'static str = "flowgger.pem";
const DEFAULT_LISTEN: &'static str = "0.0.0.0:6514";
const DEFAULT_TLS_METHOD: &'static str = "TLSv1.2";
const DEFAULT_VERIFY_PEER: bool = false;
const DEFAULT_COMPRESSION: bool = false;
const TLS_VERIFY_DEPTH: u32 = 6;

#[derive(Clone)]
struct TlsConfig {
    cert: String,
    key: String,
    ciphers: String,
    framed: bool,
    tls_method: SslMethod,
    verify_peer: bool,
    compression: bool
}

pub struct TlsInput {
    listen: String,
    tls_config: TlsConfig
}

impl TlsInput {
    pub fn new(config: &Config) -> TlsInput {
        let listen = config.lookup("input.listen").map_or(DEFAULT_LISTEN, |x| x.as_str().
            expect("input.listen must be an ip:port string")).to_owned();
        let cert = config.lookup("input.tls_cert").map_or(DEFAULT_CERT, |x| x.as_str().
            expect("input.tls_cert must be a path to a .pem file")).to_owned();
        let key = config.lookup("input.tls_key").map_or(DEFAULT_KEY, |x| x.as_str().
            expect("input.tls_key must be a path to a .pem file")).to_owned();
        let ciphers = config.lookup("input.tls_ciphers").map_or(DEFAULT_CIPHERS, |x| x.as_str().
            expect("input.tls_ciphers must be a string with a cipher suite")).to_owned();
        let tls_method = match config.lookup("input.tls_method").map_or(DEFAULT_TLS_METHOD, |x| x.as_str().
            expect("input.tls_method must be a string with the TLS method")).to_lowercase().as_ref() {
                "tlsv1" | "tlsv1.0" => SslMethod::Tlsv1,
                "tlsv1.1" => SslMethod::Tlsv1_1,
                "tlsv1.2" => SslMethod::Tlsv1_2,
                _ => panic!(r#"TLS method must be "TLSv1.0", "TLSv1.1" or "TLSv1.2""#)
        };
        let verify_peer = config.lookup("input.tls_verify_peer").map_or(DEFAULT_VERIFY_PEER, |x| x.as_bool().
            expect("input.tls_verify_peer must be a boolean"));
        let compression = config.lookup("input.tls_compression").map_or(DEFAULT_COMPRESSION, |x| x.as_bool().
            expect("input.tls_compression must be a boolean"));
        let framed = config.lookup("input.framed").map_or(DEFAULT_FRAMED, |x| x.as_bool().
            expect("input.framed must be a boolean"));

        let tls_config = TlsConfig {
            cert: cert,
            key: key,
            ciphers: ciphers,
            framed: framed,
            tls_method: tls_method,
            verify_peer: verify_peer,
            compression: compression
        };
        TlsInput {
            listen: listen,
            tls_config: tls_config
        }
    }
}

impl Input for TlsInput {
    fn accept(&self, tx: SyncSender<Vec<u8>>, decoder: Box<Decoder + Send>, encoder: Box<Encoder + Send>) {
        let listener = TcpListener::bind(&self.listen as &str).unwrap();
        for client in listener.incoming() {
            match client {
                Ok(client) => {
                    let tx = tx.clone();
                    let (decoder, encoder) = (decoder.clone_boxed(), encoder.clone_boxed());
                    let tls_config = self.tls_config.clone();
                    thread::spawn(move|| {
                        handle_client(client, tx, decoder, encoder, tls_config);
                    });
                }
                Err(_) => { }
            }
        }
    }
}

fn read_msglen(reader: &mut BufRead) -> Result<usize, &'static str> {
    let mut nbytes_v = Vec::with_capacity(16);
    let nbytes_vl = match reader.read_until(b' ', &mut nbytes_v) {
        Err(_) | Ok(0) | Ok(1) => return Err("EOF"),
        Ok(nbytes_vl) => nbytes_vl
    };
    let nbytes_s = match str::from_utf8(&nbytes_v[..nbytes_vl - 1]) {
        Err(_) => return Err("Invalid or missing message length. Disable framing, maybe?"),
        Ok(nbytes_s) => nbytes_s
    };
    let nbytes: usize = match nbytes_s.parse() {
        Err(_) => return Err("Invalid or missing message length. Disable framing, maybe?"),
        Ok(nbytes) => nbytes
    };
    Ok(nbytes)
}

fn handle_client(client: TcpStream, tx: SyncSender<Vec<u8>>, decoder: Box<Decoder>, encoder: Box<Encoder>, tls_config: TlsConfig) {
    let mut ctx = SslContext::new(Tlsv1_2).unwrap();
    if tls_config.verify_peer == true {
        ctx.set_verify_depth(TLS_VERIFY_DEPTH);
        ctx.set_verify(SSL_VERIFY_PEER | SSL_VERIFY_FAIL_IF_NO_PEER_CERT, None);
    }
    let mut opts = SSL_OP_CIPHER_SERVER_PREFERENCE | SSL_OP_NO_SESSION_RESUMPTION_ON_RENEGOTIATION;
    if tls_config.compression == false {
        opts = opts | SSL_OP_NO_COMPRESSION;
    }
    ctx.set_options(opts);
    ctx.set_certificate_file(&Path::new(&tls_config.cert), X509FileType::PEM).unwrap();
    ctx.set_private_key_file(&Path::new(&tls_config.key), X509FileType::PEM).unwrap();
    ctx.set_cipher_list(&tls_config.ciphers).unwrap();
    if let Ok(peer_addr) = client.peer_addr() {
        println!("Connection over TLS from [{}]", peer_addr);
    }
    let sslclient = match SslStream::accept(&ctx, client) {
        Err(_) => {
            let _ = writeln!(stderr(), "SSL handshake aborted by the client");
            return
        }
        Ok(sslclient) => sslclient
    };
    let mut reader = BufReader::new(sslclient);
    if tls_config.framed == false {
        for line in reader.lines() {
            let line = match line {
                Err(_) => {
                    let _ = writeln!(stderr(), "Invalid UTF-8 input");
                    continue;
                }
                Ok(line) => line
            };
            if let Err(e) = handle_line(&line, &tx, &decoder, &encoder) {
                let _ = writeln!(stderr(), "{}: [{}]", e, line.trim());
            }
        }
    } else {
        loop {
            if let Err(e) = read_msglen(&mut reader) {
                let _ = writeln!(stderr(), "{}", e);
                return;
            }
            let mut line = String::new();
            if reader.read_line(&mut line).is_err() {
                println!("err");
                return;
            }
            if let Err(e) = handle_line(&line, &tx, &decoder, &encoder) {
                let _ = writeln!(stderr(), "{}: [{}]", e, line.trim());
            }
        }
    }
}

fn handle_line(line: &String, tx: &SyncSender<Vec<u8>>, decoder: &Box<Decoder>, encoder: &Box<Encoder>) -> Result<(), &'static str> {
    let decoded = try!(decoder.decode(&line));
    let reencoded = try!(encoder.encode(decoded));
    tx.send(reencoded).unwrap();
    Ok(())
}
<commit_msg>Explicitly use SSL_VERIFY_NONE when this is really what we want<commit_after>use flowgger::config::Config;
use flowgger::decoder::Decoder;
use flowgger::encoder::Encoder;
use openssl::ssl::*;
use openssl::ssl::SslMethod::*;
use openssl::x509::X509FileType;
use std::io::{stderr, Write, BufRead, BufReader};
use std::net::{TcpListener, TcpStream};
use std::path::Path;
use std::str;
use std::sync::mpsc::SyncSender;
use std::thread;
use super::Input;

const DEFAULT_CERT: &'static str = "flowgger.pem";
const DEFAULT_CIPHERS: &'static str = "ECDHE-RSA-CHACHA20-POLY1305:ECDHE-ECDSA-CHACHA20-POLY1305:ECDHE-RSA-AES128-GCM-SHA256:ECDHE-ECDSA-AES128-GCM-SHA256:ECDHE-RSA-AES256-GCM-SHA384:ECDHE-ECDSA-AES256-GCM-SHA384:DHE-RSA-AES128-GCM-SHA256:DHE-DSS-AES128-GCM-SHA256:kEDH+AESGCM:ECDHE-RSA-AES128-SHA256:ECDHE-ECDSA-AES128-SHA256:ECDHE-RSA-AES128-SHA:ECDHE-ECDSA-AES128-SHA:ECDHE-RSA-AES256-SHA384:ECDHE-ECDSA-AES256-SHA384:ECDHE-RSA-AES256-SHA:ECDHE-ECDSA-AES256-SHA:DHE-RSA-AES128-SHA256:DHE-RSA-AES128-SHA:DHE-DSS-AES128-SHA256:DHE-RSA-AES256-SHA256:DHE-DSS-AES256-SHA:DHE-RSA-AES256-SHA:AES128-GCM-SHA256:AES256-GCM-SHA384:AES128-SHA256:AES256-SHA256:AES128-SHA:AES256-SHA:AES:CAMELLIA:DES-CBC3-SHA:!aNULL:!eNULL:!EXPORT:!DES:!RC4:!MD5:!PSK:!aECDH:!EDH-DSS-DES-CBC3-SHA:!EDH-RSA-DES-CBC3-SHA:!KRB5-DES-CBC3-SH";
const DEFAULT_FRAMED: bool = false;
const DEFAULT_KEY: &'static str = "flowgger.pem";
const DEFAULT_LISTEN: &'static str = "0.0.0.0:6514";
const DEFAULT_TLS_METHOD: &'static str = "TLSv1.2";
const DEFAULT_VERIFY_PEER: bool = false;
const DEFAULT_COMPRESSION: bool = false;
const TLS_VERIFY_DEPTH: u32 = 6;

#[derive(Clone)]
struct TlsConfig {
    cert: String,
    key: String,
    ciphers: String,
    framed: bool,
    tls_method: SslMethod,
    verify_peer: bool,
    compression: bool
}

pub struct TlsInput {
    listen: String,
    tls_config: TlsConfig
}

impl TlsInput {
    pub fn new(config: &Config) -> TlsInput {
        let listen = config.lookup("input.listen").map_or(DEFAULT_LISTEN, |x| x.as_str().
            expect("input.listen must be an ip:port string")).to_owned();
        let cert = config.lookup("input.tls_cert").map_or(DEFAULT_CERT, |x| x.as_str().
            expect("input.tls_cert must be a path to a .pem file")).to_owned();
        let key = config.lookup("input.tls_key").map_or(DEFAULT_KEY, |x| x.as_str().
            expect("input.tls_key must be a path to a .pem file")).to_owned();
        let ciphers = config.lookup("input.tls_ciphers").map_or(DEFAULT_CIPHERS, |x| x.as_str().
            expect("input.tls_ciphers must be a string with a cipher suite")).to_owned();
        let tls_method = match config.lookup("input.tls_method").map_or(DEFAULT_TLS_METHOD, |x| x.as_str().
            expect("input.tls_method must be a string with the TLS method")).to_lowercase().as_ref() {
                "tlsv1" | "tlsv1.0" => SslMethod::Tlsv1,
                "tlsv1.1" => SslMethod::Tlsv1_1,
                "tlsv1.2" => SslMethod::Tlsv1_2,
                _ => panic!(r#"TLS method must be "TLSv1.0", "TLSv1.1" or "TLSv1.2""#)
        };
        let verify_peer = config.lookup("input.tls_verify_peer").map_or(DEFAULT_VERIFY_PEER, |x| x.as_bool().
            expect("input.tls_verify_peer must be a boolean"));
        let compression = config.lookup("input.tls_compression").map_or(DEFAULT_COMPRESSION, |x| x.as_bool().
            expect("input.tls_compression must be a boolean"));
        let framed = config.lookup("input.framed").map_or(DEFAULT_FRAMED, |x| x.as_bool().
            expect("input.framed must be a boolean"));

        let tls_config = TlsConfig {
            cert: cert,
            key: key,
            ciphers: ciphers,
            framed: framed,
            tls_method: tls_method,
            verify_peer: verify_peer,
            compression: compression
        };
        TlsInput {
            listen: listen,
            tls_config: tls_config
        }
    }
}

impl Input for TlsInput {
    fn accept(&self, tx: SyncSender<Vec<u8>>, decoder: Box<Decoder + Send>, encoder: Box<Encoder + Send>) {
        let listener = TcpListener::bind(&self.listen as &str).unwrap();
        for client in listener.incoming() {
            match client {
                Ok(client) => {
                    let tx = tx.clone();
                    let (decoder, encoder) = (decoder.clone_boxed(), encoder.clone_boxed());
                    let tls_config = self.tls_config.clone();
                    thread::spawn(move|| {
                        handle_client(client, tx, decoder, encoder, tls_config);
                    });
                }
                Err(_) => { }
            }
        }
    }
}

fn read_msglen(reader: &mut BufRead) -> Result<usize, &'static str> {
    let mut nbytes_v = Vec::with_capacity(16);
    let nbytes_vl = match reader.read_until(b' ', &mut nbytes_v) {
        Err(_) | Ok(0) | Ok(1) => return Err("EOF"),
        Ok(nbytes_vl) => nbytes_vl
    };
    let nbytes_s = match str::from_utf8(&nbytes_v[..nbytes_vl - 1]) {
        Err(_) => return Err("Invalid or missing message length. Disable framing, maybe?"),
        Ok(nbytes_s) => nbytes_s
    };
    let nbytes: usize = match nbytes_s.parse() {
        Err(_) => return Err("Invalid or missing message length. Disable framing, maybe?"),
        Ok(nbytes) => nbytes
    };
    Ok(nbytes)
}

fn handle_client(client: TcpStream, tx: SyncSender<Vec<u8>>, decoder: Box<Decoder>, encoder: Box<Encoder>, tls_config: TlsConfig) {
    let mut ctx = SslContext::new(Tlsv1_2).unwrap();
    if tls_config.verify_peer == false {
        ctx.set_verify(SSL_VERIFY_NONE, None);
    } else {
        ctx.set_verify_depth(TLS_VERIFY_DEPTH);
        ctx.set_verify(SSL_VERIFY_PEER | SSL_VERIFY_FAIL_IF_NO_PEER_CERT, None);
    }
    let mut opts = SSL_OP_CIPHER_SERVER_PREFERENCE | SSL_OP_NO_SESSION_RESUMPTION_ON_RENEGOTIATION;
    if tls_config.compression == false {
        opts = opts | SSL_OP_NO_COMPRESSION;
    }
    ctx.set_options(opts);
    ctx.set_certificate_file(&Path::new(&tls_config.cert), X509FileType::PEM).unwrap();
    ctx.set_private_key_file(&Path::new(&tls_config.key), X509FileType::PEM).unwrap();
    ctx.set_cipher_list(&tls_config.ciphers).unwrap();
    if let Ok(peer_addr) = client.peer_addr() {
        println!("Connection over TLS from [{}]", peer_addr);
    }
    let sslclient = match SslStream::accept(&ctx, client) {
        Err(_) => {
            let _ = writeln!(stderr(), "SSL handshake aborted by the client");
            return
        }
        Ok(sslclient) => sslclient
    };
    let mut reader = BufReader::new(sslclient);
    if tls_config.framed == false {
        for line in reader.lines() {
            let line = match line {
                Err(_) => {
                    let _ = writeln!(stderr(), "Invalid UTF-8 input");
                    continue;
                }
                Ok(line) => line
            };
            if let Err(e) = handle_line(&line, &tx, &decoder, &encoder) {
                let _ = writeln!(stderr(), "{}: [{}]", e, line.trim());
            }
        }
    } else {
        loop {
            if let Err(e) = read_msglen(&mut reader) {
                let _ = writeln!(stderr(), "{}", e);
                return;
            }
            let mut line = String::new();
            if reader.read_line(&mut line).is_err() {
                println!("err");
                return;
            }
            if let Err(e) = handle_line(&line, &tx, &decoder, &encoder) {
                let _ = writeln!(stderr(), "{}: [{}]", e, line.trim());
            }
        }
    }
}

fn handle_line(line: &String, tx: &SyncSender<Vec<u8>>, decoder: &Box<Decoder>, encoder: &Box<Encoder>) -> Result<(), &'static str> {
    let decoded = try!(decoder.decode(&line));
    let reencoded = try!(encoder.encode(decoded));
    tx.send(reencoded).unwrap();
    Ok(())
}
<|endoftext|>