<commit_before>// Copyright (c) 2015, The Radare Project. All rights reserved.
// See the COPYING file at the top-level directory of this distribution.
// Licensed under the BSD 3-Clause License:
// <http://opensource.org/licenses/BSD-3-Clause>
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Implements the `GraphDot` trait for `SSAStorage`

use petgraph::graph;
use petgraph::visit::{IntoEdgeReferences, EdgeRef};
use petgraph::graph::NodeIndex;

use middle::ir::MOpcode;
use middle::dot::{DotAttrBlock, GraphDot};
use super::ssastorage::{EdgeData, NodeData, SSAStorage};
use super::ssa_traits::{SSA, SSAExtra};
use middle::ssa::cfg_traits::CFG;

///////////////////////////////////////////////////////////////////////////////
//// Implementation of GraphDot to emit Dot for SSAStorage.
///////////////////////////////////////////////////////////////////////////////

impl GraphDot for SSAStorage {
    type NodeIndex = graph::NodeIndex;
    type EdgeIndex = graph::EdgeIndex;

    fn configure(&self) -> String {
        "digraph cfg {\nsplines=\"ortho\"\nranksep=1\nnodesep=1;\ngraph [fontsize=12 fontname=\"Verdana\" compound=true \
         rankdir=TB;]\n"
            .to_owned()
    }

    fn nodes(&self) -> Vec<Self::NodeIndex> {
        self.g.node_indices().collect()
    }

    fn edges(&self) -> Vec<Self::EdgeIndex> {
        self.g.edge_references().map(|x| x.id()).collect()
    }

    fn node_count(&self) -> usize {
        self.g.node_count()
    }

    fn edge_count(&self) -> usize {
        self.g.edge_count()
    }

    fn node_index_new(i: usize) -> Self::NodeIndex {
        graph::NodeIndex::new(i)
    }

    fn edge_index_new(i: usize) -> Self::EdgeIndex {
        graph::EdgeIndex::new(i)
    }

    fn node_cluster(&self, i: &Self::NodeIndex) -> Option<usize> {
        let invalid_block = self.invalid_action().expect("Invalid Action is not defined");
        match self.g.node_weight(*i) {
            Some(&NodeData::BasicBlock(_, _)) | Some(&NodeData::DynamicAction) => Some(i.index()),
            _ => Some(self.block_for(*i).unwrap_or(invalid_block).index()),
        }
    }

    fn edge_source(&self, i: &Self::EdgeIndex) -> Self::NodeIndex {
        let edge_opt = &self.g.edge_references().find(|x| x.id() == *i);
        if edge_opt.is_none() {
            radeco_err!("Invalid edge index");
            return NodeIndex::end();
        };
        let edge = edge_opt.unwrap();
        match *edge.weight() {
            EdgeData::Data(_) => edge.target(),
            _ => edge.source(),
        }
    }

    fn edge_target(&self, i: &Self::EdgeIndex) -> Self::NodeIndex {
        let edge_opt = &self.g.edge_references().find(|x| x.id() == *i);
        if edge_opt.is_none() {
            radeco_err!("Invalid edge index");
            return NodeIndex::end();
        };
        let edge = edge_opt.unwrap();
         match *edge.weight() {
            EdgeData::Data(_) => edge.source(),
            _ => edge.target(),
        }
    }

    fn edge_skip(&self, i: &Self::EdgeIndex) -> bool {
        let edge_opt = &self.g.edge_references().find(|x| x.id() == *i);
        if edge_opt.is_none() {
            radeco_err!("Invalid edge index");
            return false;
        };
        let edge = edge_opt.unwrap();
        match *edge.weight() {
            EdgeData::ContainedInBB(_) | EdgeData::RegisterState => true,
            EdgeData::Data(_) => {
                // Skip data edges from registers to registerstates.
                if let NodeData::RegisterState = self.g[edge.source()] {
                    self.comment(edge.target()).is_some()
                } else {
                    false
                }
            },
            _ => false,
        }
    }

    // TODO: Ordering of clusters for ssa is kind of hacky and may not run top to
    // bottom in some
    // cases.
    fn edge_attrs(&self, i: &Self::EdgeIndex) -> DotAttrBlock {
        let edge_opt = &self.g.edge_references().find(|x| x.id() == *i);
        if edge_opt.is_none() {
            radeco_err!("Invalid edge index");
            return DotAttrBlock::Raw(String::new());
        };
        let edge = edge_opt.unwrap();
        let mut prefix = String::new();
        let src = edge.source();
        let target = edge.target();

        prefix.push_str(&format!("n{} -> n{}",
                                 src.index(),
                                 target.index()));
        let target_is_bb = if let NodeData::BasicBlock(_, _) = self.g[edge.target()] {
            true
        } else {
            false
        };
        let attr = match *edge.weight() {
            EdgeData::Control(_) if !target_is_bb => {
                vec![("color".to_string(), "red".to_string())]
            }
            EdgeData::Control(i) => {
                // Determine the source and destination clusters.
                let source_cluster = edge.source().index();
                let dst_cluster = edge.target().index();
                let (color, label) = match i {
                    0 => ("red", "F"),
                    1 => ("green", "T"),
                    2 => ("blue", "U"),
                    _ => unreachable!(),
                };
                vec![("color".to_string(), color.to_string()),
                     ("xlabel".to_string(), label.to_string()),
                     ("ltail".to_string(), format!("cluster_{}", source_cluster)),
                     ("lhead".to_string(), format!("cluster_{}", dst_cluster)),
                     ("minlen".to_string(), "9".to_owned())]
            }
            EdgeData::Data(i) => {
                vec![("dir".to_string(), "back".to_string()),
                     ("xlabel".to_string(), format!("{}", i))]
            }
            EdgeData::ContainedInBB(_) => {
                vec![("color".to_string(), "gray".to_string())]
            }
            EdgeData::Selector => {
                vec![("color".to_string(), "purple".to_string())]
            }
            EdgeData::ReplacedBy => {
                vec![("color".to_string(), "brown".to_string())]
            }
            EdgeData::RegisterInfo => {
                vec![("color".to_string(), "yellow".to_string())]
            }
            EdgeData::RegisterState => unreachable!(),
        };

        DotAttrBlock::Hybrid(prefix, attr)
    }

    fn node_attrs(&self, i: &Self::NodeIndex) -> DotAttrBlock {
        let node = &self.g[*i];
        let mut prefix = String::new();
        prefix.push_str(&format!("n{}", i.index()));

        let attr = match *node {
            NodeData::Op(ref opc, vi) => {
                let w = vi.width().get_width().unwrap_or(64);
                let mut attrs = Vec::new();
                let mut r = String::new();
                let addr = self.addr(i);
                if addr.is_some() {
                    r.push_str(&format!("<{}: ",
                                        addr.as_ref().unwrap()))
                }
                r.push_str(&format!("\"[i{}] {:?}\"", w, opc));
                if addr.is_some() {
                    r.push_str(">");
                }

                attrs.push(("style".to_owned(), "filled".to_owned()));
                if let MOpcode::OpConst(c) = *opc {
                    attrs.push(("color".to_owned(), "black".to_owned()));
                    attrs.push(("fillcolor".to_owned(), "yellow".to_owned()));
                    attrs.push(("label".to_string(), format!("\"0x{:x}\"", c)));
                } else {
                    attrs.push(("color".to_owned(), "black".to_owned()));
                    attrs.push(("label".to_string(), r));

                    if self.is_marked(i) {
                        attrs.push(("fillcolor".to_owned(), "green".to_owned()));
                    } else {
                        attrs.push(("fillcolor".to_owned(), "white".to_owned()));
                    }
                }
                attrs
            }
            NodeData::BasicBlock(addr, _) => {
                let label_str = format!("<Basic Block \
                                         Information
Start Address: {}>",
                                        addr);
                let mut attrs = Vec::new();
                if let Some(e) = self.entry_node() {
                    if *i == e {
                        attrs.push(("rank".to_string(), "min".to_string()));
                    }
                };

                attrs.extend([("style".to_string(), "filled".to_string()),
                              ("fillcolor".to_string(), "white".to_string()),
                              ("label".to_string(), label_str),
                              ("shape".to_string(), "box".to_string()),
                              ("color".to_string(), "\"grey\"".to_string()),]
                                 .iter()
                                 .cloned());
                attrs
            }
            NodeData::Comment(_, ref msg) => {
                vec![("label".to_string(),
                      format!("\"{}\"", msg.replace("\"", "\\\""))),
                     ("shape".to_owned(), "box".to_owned()),
                     ("style".to_owned(), "filled".to_owned()),
                     ("color".to_owned(), "black".to_owned()),
                     ("fillcolor".to_owned(), "greenyellow".to_owned())]
            }
            NodeData::Phi(_, _) => {
                let mut attrs = Vec::new();
                let mut label = format!("{:?}", node);
                label = label.replace("\"", "\\\"");
                label = format!("\"{}\"", label);
                if let Some(addr) = self.addr(i) {
                    label = format!("<{}:  {}>", addr, label);
                }
                attrs.push(("label".to_string(), label));
                attrs.push(("style".to_owned(), "filled".to_owned()));
                attrs.push(("color".to_owned(), "black".to_owned()));
                if self.is_marked(i) {
                    attrs.push(("fillcolor".to_owned(), "green".to_owned()));
                } else {
                    attrs.push(("fillcolor".to_owned(), "orange".to_owned()));
                }
                attrs
            }
            _ => {
                let mut attrs = Vec::new();
                let mut label = format!("{:?}", node);
                label = label.replace("\"", "\\\"");
                label = format!("\"{}\"", label);
                if let Some(addr) = self.addr(i) {
                    label = format!("<{}:  {}>", addr, label);
                }
                attrs.push(("label".to_string(), label));
                attrs.push(("style".to_owned(), "filled".to_owned()));
                attrs.push(("color".to_owned(), "black".to_owned()));
                if self.is_marked(i) {
                    attrs.push(("fillcolor".to_owned(), "green".to_owned()));
                } else {
                    attrs.push(("fillcolor".to_owned(), "grey90".to_owned()));
                }
                attrs
            }
        };
        DotAttrBlock::Hybrid(prefix, attr)
    }
}
<commit_msg>Fix xdot parsing<commit_after>// Copyright (c) 2015, The Radare Project. All rights reserved.
// See the COPYING file at the top-level directory of this distribution.
// Licensed under the BSD 3-Clause License:
// <http://opensource.org/licenses/BSD-3-Clause>
// This file may not be copied, modified, or distributed
// except according to those terms.

//! Implements the `GraphDot` trait for `SSAStorage`

use petgraph::graph;
use petgraph::visit::{IntoEdgeReferences, EdgeRef};
use petgraph::graph::NodeIndex;

use middle::ir::MOpcode;
use middle::dot::{DotAttrBlock, GraphDot};
use super::ssastorage::{EdgeData, NodeData, SSAStorage};
use super::ssa_traits::{SSA, SSAExtra};
use middle::ssa::cfg_traits::CFG;

///////////////////////////////////////////////////////////////////////////////
//// Implementation of GraphDot to emit Dot for SSAStorage.
///////////////////////////////////////////////////////////////////////////////

impl GraphDot for SSAStorage {
    type NodeIndex = graph::NodeIndex;
    type EdgeIndex = graph::EdgeIndex;

    fn configure(&self) -> String {
        "digraph cfg {\nsplines=\"ortho\"\nranksep=1\nnodesep=1;\ngraph [fontsize=12 fontname=\"Verdana\" compound=true \
         rankdir=TB;]\n"
            .to_owned()
    }

    fn nodes(&self) -> Vec<Self::NodeIndex> {
        self.g.node_indices().collect()
    }

    fn edges(&self) -> Vec<Self::EdgeIndex> {
        self.g.edge_references().map(|x| x.id()).collect()
    }

    fn node_count(&self) -> usize {
        self.g.node_count()
    }

    fn edge_count(&self) -> usize {
        self.g.edge_count()
    }

    fn node_index_new(i: usize) -> Self::NodeIndex {
        graph::NodeIndex::new(i)
    }

    fn edge_index_new(i: usize) -> Self::EdgeIndex {
        graph::EdgeIndex::new(i)
    }

    fn node_cluster(&self, i: &Self::NodeIndex) -> Option<usize> {
        let invalid_block = self.invalid_action().expect("Invalid Action is not defined");
        match self.g.node_weight(*i) {
            Some(&NodeData::BasicBlock(_, _)) | Some(&NodeData::DynamicAction) => Some(i.index()),
            _ => Some(self.block_for(*i).unwrap_or(invalid_block).index()),
        }
    }

    fn edge_source(&self, i: &Self::EdgeIndex) -> Self::NodeIndex {
        let edge_opt = &self.g.edge_references().find(|x| x.id() == *i);
        if edge_opt.is_none() {
            radeco_err!("Invalid edge index");
            return NodeIndex::end();
        };
        let edge = edge_opt.unwrap();
        match *edge.weight() {
            EdgeData::Data(_) => edge.target(),
            _ => edge.source(),
        }
    }

    fn edge_target(&self, i: &Self::EdgeIndex) -> Self::NodeIndex {
        let edge_opt = &self.g.edge_references().find(|x| x.id() == *i);
        if edge_opt.is_none() {
            radeco_err!("Invalid edge index");
            return NodeIndex::end();
        };
        let edge = edge_opt.unwrap();
         match *edge.weight() {
            EdgeData::Data(_) => edge.source(),
            _ => edge.target(),
        }
    }

    fn edge_skip(&self, i: &Self::EdgeIndex) -> bool {
        let edge_opt = &self.g.edge_references().find(|x| x.id() == *i);
        if edge_opt.is_none() {
            radeco_err!("Invalid edge index");
            return false;
        };
        let edge = edge_opt.unwrap();
        match *edge.weight() {
            EdgeData::ContainedInBB(_) | EdgeData::RegisterState => true,
            EdgeData::Data(_) => {
                // Skip data edges from registers to registerstates.
                if let NodeData::RegisterState = self.g[edge.source()] {
                    self.comment(edge.target()).is_some()
                } else {
                    false
                }
            },
            _ => false,
        }
    }

    // TODO: Ordering of clusters for ssa is kind of hacky and may not run top to
    // bottom in some
    // cases.
    fn edge_attrs(&self, i: &Self::EdgeIndex) -> DotAttrBlock {
        let edge_opt = &self.g.edge_references().find(|x| x.id() == *i);
        if edge_opt.is_none() {
            radeco_err!("Invalid edge index");
            return DotAttrBlock::Raw(String::new());
        };
        let edge = edge_opt.unwrap();
        let mut prefix = String::new();
        let src = edge.source();
        let target = edge.target();

        prefix.push_str(&format!("n{} -> n{}",
                                 src.index(),
                                 target.index()));
        let target_is_bb = if let NodeData::BasicBlock(_, _) = self.g[edge.target()] {
            true
        } else {
            false
        };
        let attr = match *edge.weight() {
            EdgeData::Control(_) if !target_is_bb => {
                vec![("color".to_string(), "red".to_string())]
            }
            EdgeData::Control(i) => {
                // Determine the source and destination clusters.
                let source_cluster = edge.source().index();
                let dst_cluster = edge.target().index();
                let (color, label) = match i {
                    0 => ("red", "F"),
                    1 => ("green", "T"),
                    2 => ("blue", "U"),
                    _ => unreachable!(),
                };
                vec![("color".to_string(), color.to_string()),
                     ("xlabel".to_string(), label.to_string()),
                     ("ltail".to_string(), format!("cluster_{}", source_cluster)),
                     ("lhead".to_string(), format!("cluster_{}", dst_cluster)),
                     ("minlen".to_string(), "9".to_owned())]
            }
            EdgeData::Data(i) => {
                vec![("dir".to_string(), "back".to_string()),
                     ("xlabel".to_string(), format!("{}", i))]
            }
            EdgeData::ContainedInBB(_) => {
                vec![("color".to_string(), "gray".to_string())]
            }
            EdgeData::Selector => {
                vec![("color".to_string(), "purple".to_string())]
            }
            EdgeData::ReplacedBy => {
                vec![("color".to_string(), "brown".to_string())]
            }
            EdgeData::RegisterInfo => {
                vec![("color".to_string(), "yellow".to_string())]
            }
            EdgeData::RegisterState => unreachable!(),
        };

        DotAttrBlock::Hybrid(prefix, attr)
    }

    fn node_attrs(&self, i: &Self::NodeIndex) -> DotAttrBlock {
        let node = &self.g[*i];
        let mut prefix = String::new();
        prefix.push_str(&format!("n{}", i.index()));

        let attr = match *node {
            NodeData::Op(ref opc, vi) => {
                let w = vi.width().get_width().unwrap_or(64);
                let mut attrs = Vec::new();
                let mut r = String::new();
                let addr = self.addr(i);
                if addr.is_some() {
                    r.push_str(&format!("<{}: ",
                                        addr.as_ref().unwrap()))
                }
                r.push_str(&format!("\"[i{}] {:?}\"", w, opc));
                if addr.is_some() {
                    r.push_str(">");
                }

                attrs.push(("style".to_owned(), "filled".to_owned()));
                if let MOpcode::OpConst(c) = *opc {
                    attrs.push(("color".to_owned(), "black".to_owned()));
                    attrs.push(("fillcolor".to_owned(), "yellow".to_owned()));
                    attrs.push(("label".to_string(), format!("\"0x{:x}\"", c)));
                } else {
                    attrs.push(("color".to_owned(), "black".to_owned()));
                    attrs.push(("label".to_string(), r));

                    if self.is_marked(i) {
                        attrs.push(("fillcolor".to_owned(), "green".to_owned()));
                    } else {
                        attrs.push(("fillcolor".to_owned(), "white".to_owned()));
                    }
                }
                attrs
            }
            NodeData::BasicBlock(addr, _) => {
                let label_str = format!("<Basic Block \
                                         Information
Start Address: {}>",
                                        addr);
                let mut attrs = Vec::new();
                if let Some(e) = self.entry_node() {
                    if *i == e {
                        attrs.push(("rank".to_string(), "min".to_string()));
                    }
                };

                attrs.extend([("style".to_string(), "filled".to_string()),
                              ("fillcolor".to_string(), "white".to_string()),
                              ("label".to_string(), label_str),
                              ("shape".to_string(), "box".to_string()),
                              ("color".to_string(), "\"grey\"".to_string()),]
                                 .iter()
                                 .cloned());
                attrs
            }
            NodeData::Comment(_, ref msg) => {
                vec![("label".to_string(),
                      format!("\"{}\"", msg.replace("\"", "\\\""))),
                     ("shape".to_owned(), "box".to_owned()),
                     ("style".to_owned(), "filled".to_owned()),
                     ("color".to_owned(), "black".to_owned()),
                     ("fillcolor".to_owned(), "greenyellow".to_owned())]
            }
            NodeData::Phi(_, _) => {
                let mut attrs = Vec::new();
                let mut label = format!("{:?}", node);
                label = format!("\"{}\"", label);
                if let Some(addr) = self.addr(i) {
                    label = format!("<{}:  {}>", addr, label);
                }
                attrs.push(("label".to_string(), label));
                attrs.push(("style".to_owned(), "filled".to_owned()));
                attrs.push(("color".to_owned(), "black".to_owned()));
                if self.is_marked(i) {
                    attrs.push(("fillcolor".to_owned(), "green".to_owned()));
                } else {
                    attrs.push(("fillcolor".to_owned(), "orange".to_owned()));
                }
                attrs
            }
            _ => {
                let mut attrs = Vec::new();
                let mut label = format!("{:?}", node);
                label = label.replace("\"", "\\\"");
                label = format!("\"{}\"", label);
                if let Some(addr) = self.addr(i) {
                    label = format!("<{}:  {}>", addr, label);
                }
                attrs.push(("label".to_string(), label));
                attrs.push(("style".to_owned(), "filled".to_owned()));
                attrs.push(("color".to_owned(), "black".to_owned()));
                if self.is_marked(i) {
                    attrs.push(("fillcolor".to_owned(), "green".to_owned()));
                } else {
                    attrs.push(("fillcolor".to_owned(), "grey90".to_owned()));
                }
                attrs
            }
        };
        DotAttrBlock::Hybrid(prefix, attr)
    }
}
<|endoftext|>


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

//! Projection Push Down optimizer rule ensures that only referenced columns are
//! loaded into memory

use crate::error::{ExecutionError, Result};
use crate::logicalplan::LogicalPlan;
use crate::logicalplan::{Expr, LogicalPlanBuilder};
use crate::optimizer::optimizer::OptimizerRule;
use crate::optimizer::utils;
use arrow::datatypes::{Field, Schema};
use std::collections::{HashMap, HashSet};

/// Projection Push Down optimizer rule ensures that only referenced columns are
/// loaded into memory
pub struct ProjectionPushDown {}

impl OptimizerRule for ProjectionPushDown {
    fn optimize(&mut self, plan: &LogicalPlan) -> Result<LogicalPlan> {
        let mut accum: HashSet<usize> = HashSet::new();
        let mut mapping: HashMap<usize, usize> = HashMap::new();
        self.optimize_plan(plan, &mut accum, &mut mapping)
    }
}

impl ProjectionPushDown {
    #[allow(missing_docs)]
    pub fn new() -> Self {
        Self {}
    }

    fn optimize_plan(
        &self,
        plan: &LogicalPlan,
        accum: &mut HashSet<usize>,
        mapping: &mut HashMap<usize, usize>,
    ) -> Result<LogicalPlan> {
        match plan {
            LogicalPlan::Projection { expr, input, .. } => {
                // collect all columns referenced by projection expressions
                utils::exprlist_to_column_indices(&expr, accum)?;

                LogicalPlanBuilder::from(&self.optimize_plan(&input, accum, mapping)?)
                    .project(self.rewrite_expr_list(expr, mapping)?)?
                    .build()
            }
            LogicalPlan::Selection { expr, input } => {
                // collect all columns referenced by filter expression
                utils::expr_to_column_indices(expr, accum)?;

                LogicalPlanBuilder::from(&self.optimize_plan(&input, accum, mapping)?)
                    .filter(self.rewrite_expr(expr, mapping)?)?
                    .build()
            }
            LogicalPlan::Aggregate {
                input,
                group_expr,
                aggr_expr,
                ..
            } => {
                // collect all columns referenced by grouping and aggregate expressions
                utils::exprlist_to_column_indices(&group_expr, accum)?;
                utils::exprlist_to_column_indices(&aggr_expr, accum)?;

                LogicalPlanBuilder::from(&self.optimize_plan(&input, accum, mapping)?)
                    .aggregate(
                        self.rewrite_expr_list(group_expr, mapping)?,
                        self.rewrite_expr_list(aggr_expr, mapping)?,
                    )?
                    .build()
            }
            LogicalPlan::Sort { expr, input, .. } => {
                // collect all columns referenced by sort expressions
                utils::exprlist_to_column_indices(&expr, accum)?;

                LogicalPlanBuilder::from(&self.optimize_plan(&input, accum, mapping)?)
                    .sort(self.rewrite_expr_list(expr, mapping)?)?
                    .build()
            }
            LogicalPlan::EmptyRelation { schema } => Ok(LogicalPlan::EmptyRelation {
                schema: schema.clone(),
            }),
            LogicalPlan::TableScan {
                schema_name,
                table_name,
                table_schema,
                projection,
                ..
            } => {
                let (projection, projected_schema) =
                    get_projected_schema(&table_schema, projection, accum, mapping)?;

                // return the table scan with projection
                Ok(LogicalPlan::TableScan {
                    schema_name: schema_name.to_string(),
                    table_name: table_name.to_string(),
                    table_schema: table_schema.clone(),
                    projected_schema: Box::new(projected_schema),
                    projection: Some(projection),
                })
            }
            LogicalPlan::InMemoryScan {
                data,
                schema,
                projection,
                ..
            } => {
                let (projection, projected_schema) =
                    get_projected_schema(&schema, projection, accum, mapping)?;

                Ok(LogicalPlan::InMemoryScan {
                    data: data.clone(),
                    schema: schema.clone(),
                    projection: Some(projection),
                    projected_schema: Box::new(projected_schema),
                })
            }
            LogicalPlan::CsvScan {
                path,
                has_header,
                delimiter,
                schema,
                projection,
                ..
            } => {
                let (projection, projected_schema) =
                    get_projected_schema(&schema, projection, accum, mapping)?;

                Ok(LogicalPlan::CsvScan {
                    path: path.to_owned(),
                    has_header: *has_header,
                    schema: schema.clone(),
                    delimiter: *delimiter,
                    projection: Some(projection),
                    projected_schema: Box::new(projected_schema),
                })
            }
            LogicalPlan::ParquetScan {
                path,
                schema,
                projection,
                ..
            } => {
                let (projection, projected_schema) =
                    get_projected_schema(&schema, projection, accum, mapping)?;

                Ok(LogicalPlan::ParquetScan {
                    path: path.to_owned(),
                    schema: schema.clone(),
                    projection: Some(projection),
                    projected_schema: Box::new(projected_schema),
                })
            }
            LogicalPlan::Limit { expr, input, .. } => {
                // Note that limit expressions are scalar values so there is no need to
                // rewrite them but we do need to optimize the input to the limit plan
                LogicalPlanBuilder::from(&self.optimize_plan(&input, accum, mapping)?)
                    .limit(expr.clone())?
                    .build()
            }
            LogicalPlan::CreateExternalTable {
                schema,
                name,
                location,
                file_type,
                has_header,
            } => Ok(LogicalPlan::CreateExternalTable {
                schema: schema.clone(),
                name: name.to_string(),
                location: location.to_string(),
                file_type: file_type.clone(),
                has_header: *has_header,
            }),
        }
    }

    fn rewrite_expr_list(
        &self,
        expr: &[Expr],
        mapping: &HashMap<usize, usize>,
    ) -> Result<Vec<Expr>> {
        Ok(expr
            .iter()
            .map(|e| self.rewrite_expr(e, mapping))
            .collect::<Result<Vec<Expr>>>()?)
    }

    fn rewrite_expr(&self, expr: &Expr, mapping: &HashMap<usize, usize>) -> Result<Expr> {
        match expr {
            Expr::Alias(expr, name) => Ok(Expr::Alias(
                Box::new(self.rewrite_expr(expr, mapping)?),
                name.clone(),
            )),
            Expr::Column(i) => Ok(Expr::Column(self.new_index(mapping, i)?)),
            Expr::UnresolvedColumn(_) => Err(ExecutionError::ExecutionError(
                "Columns need to be resolved before projection push down rule can run"
                    .to_owned(),
            )),
            Expr::Literal(_) => Ok(expr.clone()),
            Expr::Not(e) => Ok(Expr::Not(Box::new(self.rewrite_expr(e, mapping)?))),
            Expr::IsNull(e) => Ok(Expr::IsNull(Box::new(self.rewrite_expr(e, mapping)?))),
            Expr::IsNotNull(e) => {
                Ok(Expr::IsNotNull(Box::new(self.rewrite_expr(e, mapping)?)))
            }
            Expr::BinaryExpr { left, op, right } => Ok(Expr::BinaryExpr {
                left: Box::new(self.rewrite_expr(left, mapping)?),
                op: op.clone(),
                right: Box::new(self.rewrite_expr(right, mapping)?),
            }),
            Expr::Cast { expr, data_type } => Ok(Expr::Cast {
                expr: Box::new(self.rewrite_expr(expr, mapping)?),
                data_type: data_type.clone(),
            }),
            Expr::Sort {
                expr,
                asc,
                nulls_first,
            } => Ok(Expr::Sort {
                expr: Box::new(self.rewrite_expr(expr, mapping)?),
                asc: *asc,
                nulls_first: *nulls_first,
            }),
            Expr::AggregateFunction {
                name,
                args,
                return_type,
            } => Ok(Expr::AggregateFunction {
                name: name.to_string(),
                args: self.rewrite_expr_list(args, mapping)?,
                return_type: return_type.clone(),
            }),
            Expr::ScalarFunction {
                name,
                args,
                return_type,
            } => Ok(Expr::ScalarFunction {
                name: name.to_string(),
                args: self.rewrite_expr_list(args, mapping)?,
                return_type: return_type.clone(),
            }),
            Expr::Wildcard => Err(ExecutionError::General(
                "Wildcard expressions are not valid in a logical query plan".to_owned(),
            )),
        }
    }

    fn new_index(&self, mapping: &HashMap<usize, usize>, i: &usize) -> Result<usize> {
        match mapping.get(i) {
            Some(j) => Ok(*j),
            _ => Err(ExecutionError::InternalError(
                "Internal error computing new column index".to_string(),
            )),
        }
    }
}

fn get_projected_schema(
    table_schema: &Schema,
    projection: &Option<Vec<usize>>,
    accum: &HashSet<usize>,
    mapping: &mut HashMap<usize, usize>,
) -> Result<(Vec<usize>, Schema)> {
    if projection.is_some() {
        return Err(ExecutionError::General(
            "Cannot run projection push-down rule more than once".to_string(),
        ));
    }

    // once we reach the table scan, we can use the accumulated set of column
    // indexes as the projection in the table scan
    let mut projection: Vec<usize> = Vec::with_capacity(accum.len());
    accum.iter().for_each(|i| projection.push(*i));

    // Ensure that we are reading at least one column from the table in case the query
    // does not reference any columns directly such as "SELECT COUNT(1) FROM table"
    if projection.is_empty() {
        projection.push(0);
    }

    // sort the projection otherwise we get non-deterministic behavior
    projection.sort();

    // create the projected schema
    let mut projected_fields: Vec<Field> = Vec::with_capacity(projection.len());
    for i in &projection {
        projected_fields.push(table_schema.fields()[*i].clone());
    }

    // now that the table scan is returning a different schema we need to
    // create a mapping from the original column index to the
    // new column index so that we can rewrite expressions as
    // we walk back up the tree

    if mapping.len() != 0 {
        return Err(ExecutionError::InternalError("illegal state".to_string()));
    }

    for i in 0..table_schema.fields().len() {
        if let Some(n) = projection.iter().position(|v| *v == i) {
            mapping.insert(i, n);
        }
    }

    Ok((projection, Schema::new(projected_fields)))
}

#[cfg(test)]
mod tests {

    use super::*;
    use crate::logicalplan::Expr::*;
    use crate::logicalplan::ScalarValue;
    use crate::test::*;
    use arrow::datatypes::DataType;

    #[test]
    fn aggregate_no_group_by() -> Result<()> {
        let table_scan = test_table_scan()?;

        let plan = LogicalPlanBuilder::from(&table_scan)
            .aggregate(vec![], vec![max(Column(1))])?
            .build()?;

        let expected = "Aggregate: groupBy=[[]], aggr=[[MAX(#0)]]\
        \n  TableScan: test projection=Some([1])";

        assert_optimized_plan_eq(&plan, expected);

        Ok(())
    }

    #[test]
    fn aggregate_group_by() -> Result<()> {
        let table_scan = test_table_scan()?;

        let plan = LogicalPlanBuilder::from(&table_scan)
            .aggregate(vec![Column(2)], vec![max(Column(1))])?
            .build()?;

        let expected = "Aggregate: groupBy=[[#1]], aggr=[[MAX(#0)]]\
        \n  TableScan: test projection=Some([1, 2])";

        assert_optimized_plan_eq(&plan, expected);

        Ok(())
    }

    #[test]
    fn aggregate_no_group_by_with_selection() -> Result<()> {
        let table_scan = test_table_scan()?;

        let plan = LogicalPlanBuilder::from(&table_scan)
            .filter(Column(2))?
            .aggregate(vec![], vec![max(Column(1))])?
            .build()?;

        let expected = "Aggregate: groupBy=[[]], aggr=[[MAX(#0)]]\
        \n  Selection: #1\
        \n    TableScan: test projection=Some([1, 2])";

        assert_optimized_plan_eq(&plan, expected);

        Ok(())
    }

    #[test]
    fn cast() -> Result<()> {
        let table_scan = test_table_scan()?;

        let projection = LogicalPlanBuilder::from(&table_scan)
            .project(vec![Cast {
                expr: Box::new(Column(2)),
                data_type: DataType::Float64,
            }])?
            .build()?;

        let expected = "Projection: CAST(#0 AS Float64)\
        \n  TableScan: test projection=Some([2])";

        assert_optimized_plan_eq(&projection, expected);

        Ok(())
    }

    #[test]
    fn table_scan_projected_schema() -> Result<()> {
        let table_scan = test_table_scan()?;
        assert_eq!(3, table_scan.schema().fields().len());
        assert_fields_eq(&table_scan, vec!["a", "b", "c"]);

        let plan = LogicalPlanBuilder::from(&table_scan)
            .project(vec![Column(0), Column(1)])?
            .build()?;

        assert_fields_eq(&plan, vec!["a", "b"]);

        let expected = "Projection: #0, #1\
        \n  TableScan: test projection=Some([0, 1])";

        assert_optimized_plan_eq(&plan, expected);

        Ok(())
    }

    #[test]
    fn table_limit() -> Result<()> {
        let table_scan = test_table_scan()?;
        assert_eq!(3, table_scan.schema().fields().len());
        assert_fields_eq(&table_scan, vec!["a", "b", "c"]);

        let plan = LogicalPlanBuilder::from(&table_scan)
            .project(vec![Column(2), Column(0)])?
            .limit(Expr::Literal(ScalarValue::UInt32(5)))?
            .build()?;

        assert_fields_eq(&plan, vec!["c", "a"]);

        let expected = "Limit: UInt32(5)\
        \n  Projection: #1, #0\
        \n    TableScan: test projection=Some([0, 2])";

        assert_optimized_plan_eq(&plan, expected);

        Ok(())
    }

    fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) {
        let optimized_plan = optimize(plan).expect("failed to optimize plan");
        let formatted_plan = format!("{:?}", optimized_plan);
        assert_eq!(formatted_plan, expected);
    }

    fn optimize(plan: &LogicalPlan) -> Result<LogicalPlan> {
        let mut rule = ProjectionPushDown::new();
        rule.optimize(plan)
    }
}
<commit_msg>ARROW-9184: [Rust][Datafusion] table scan without projection should return all columns<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

//! Projection Push Down optimizer rule ensures that only referenced columns are
//! loaded into memory

use crate::error::{ExecutionError, Result};
use crate::logicalplan::LogicalPlan;
use crate::logicalplan::{Expr, LogicalPlanBuilder};
use crate::optimizer::optimizer::OptimizerRule;
use crate::optimizer::utils;
use arrow::datatypes::{Field, Schema};
use std::collections::{HashMap, HashSet};

/// Projection Push Down optimizer rule ensures that only referenced columns are
/// loaded into memory
pub struct ProjectionPushDown {}

impl OptimizerRule for ProjectionPushDown {
    fn optimize(&mut self, plan: &LogicalPlan) -> Result<LogicalPlan> {
        let mut accum: HashSet<usize> = HashSet::new();
        let mut mapping: HashMap<usize, usize> = HashMap::new();
        self.optimize_plan(plan, &mut accum, &mut mapping, false)
    }
}

impl ProjectionPushDown {
    #[allow(missing_docs)]
    pub fn new() -> Self {
        Self {}
    }

    fn optimize_plan(
        &self,
        plan: &LogicalPlan,
        accum: &mut HashSet<usize>,
        mapping: &mut HashMap<usize, usize>,
        has_projection: bool,
    ) -> Result<LogicalPlan> {
        match plan {
            LogicalPlan::Projection { expr, input, .. } => {
                // collect all columns referenced by projection expressions
                utils::exprlist_to_column_indices(&expr, accum)?;

                LogicalPlanBuilder::from(
                    &self.optimize_plan(&input, accum, mapping, true)?,
                )
                .project(self.rewrite_expr_list(expr, mapping)?)?
                .build()
            }
            LogicalPlan::Selection { expr, input } => {
                // collect all columns referenced by filter expression
                utils::expr_to_column_indices(expr, accum)?;

                LogicalPlanBuilder::from(&self.optimize_plan(
                    &input,
                    accum,
                    mapping,
                    has_projection,
                )?)
                .filter(self.rewrite_expr(expr, mapping)?)?
                .build()
            }
            LogicalPlan::Aggregate {
                input,
                group_expr,
                aggr_expr,
                ..
            } => {
                // collect all columns referenced by grouping and aggregate expressions
                utils::exprlist_to_column_indices(&group_expr, accum)?;
                utils::exprlist_to_column_indices(&aggr_expr, accum)?;

                LogicalPlanBuilder::from(&self.optimize_plan(
                    &input,
                    accum,
                    mapping,
                    has_projection,
                )?)
                .aggregate(
                    self.rewrite_expr_list(group_expr, mapping)?,
                    self.rewrite_expr_list(aggr_expr, mapping)?,
                )?
                .build()
            }
            LogicalPlan::Sort { expr, input, .. } => {
                // collect all columns referenced by sort expressions
                utils::exprlist_to_column_indices(&expr, accum)?;

                LogicalPlanBuilder::from(&self.optimize_plan(
                    &input,
                    accum,
                    mapping,
                    has_projection,
                )?)
                .sort(self.rewrite_expr_list(expr, mapping)?)?
                .build()
            }
            LogicalPlan::EmptyRelation { schema } => Ok(LogicalPlan::EmptyRelation {
                schema: schema.clone(),
            }),
            LogicalPlan::TableScan {
                schema_name,
                table_name,
                table_schema,
                projection,
                ..
            } => {
                let (projection, projected_schema) = get_projected_schema(
                    &table_schema,
                    projection,
                    accum,
                    mapping,
                    has_projection,
                )?;

                // return the table scan with projection
                Ok(LogicalPlan::TableScan {
                    schema_name: schema_name.to_string(),
                    table_name: table_name.to_string(),
                    table_schema: table_schema.clone(),
                    projected_schema: Box::new(projected_schema),
                    projection: Some(projection),
                })
            }
            LogicalPlan::InMemoryScan {
                data,
                schema,
                projection,
                ..
            } => {
                let (projection, projected_schema) = get_projected_schema(
                    &schema,
                    projection,
                    accum,
                    mapping,
                    has_projection,
                )?;
                Ok(LogicalPlan::InMemoryScan {
                    data: data.clone(),
                    schema: schema.clone(),
                    projection: Some(projection),
                    projected_schema: Box::new(projected_schema),
                })
            }
            LogicalPlan::CsvScan {
                path,
                has_header,
                delimiter,
                schema,
                projection,
                ..
            } => {
                let (projection, projected_schema) = get_projected_schema(
                    &schema,
                    projection,
                    accum,
                    mapping,
                    has_projection,
                )?;

                Ok(LogicalPlan::CsvScan {
                    path: path.to_owned(),
                    has_header: *has_header,
                    schema: schema.clone(),
                    delimiter: *delimiter,
                    projection: Some(projection),
                    projected_schema: Box::new(projected_schema),
                })
            }
            LogicalPlan::ParquetScan {
                path,
                schema,
                projection,
                ..
            } => {
                let (projection, projected_schema) = get_projected_schema(
                    &schema,
                    projection,
                    accum,
                    mapping,
                    has_projection,
                )?;

                Ok(LogicalPlan::ParquetScan {
                    path: path.to_owned(),
                    schema: schema.clone(),
                    projection: Some(projection),
                    projected_schema: Box::new(projected_schema),
                })
            }
            LogicalPlan::Limit { expr, input, .. } => {
                // Note that limit expressions are scalar values so there is no need to
                // rewrite them but we do need to optimize the input to the limit plan
                LogicalPlanBuilder::from(&self.optimize_plan(
                    &input,
                    accum,
                    mapping,
                    has_projection,
                )?)
                .limit(expr.clone())?
                .build()
            }
            LogicalPlan::CreateExternalTable {
                schema,
                name,
                location,
                file_type,
                has_header,
            } => Ok(LogicalPlan::CreateExternalTable {
                schema: schema.clone(),
                name: name.to_string(),
                location: location.to_string(),
                file_type: file_type.clone(),
                has_header: *has_header,
            }),
        }
    }

    fn rewrite_expr_list(
        &self,
        expr: &[Expr],
        mapping: &HashMap<usize, usize>,
    ) -> Result<Vec<Expr>> {
        Ok(expr
            .iter()
            .map(|e| self.rewrite_expr(e, mapping))
            .collect::<Result<Vec<Expr>>>()?)
    }

    fn rewrite_expr(&self, expr: &Expr, mapping: &HashMap<usize, usize>) -> Result<Expr> {
        match expr {
            Expr::Alias(expr, name) => Ok(Expr::Alias(
                Box::new(self.rewrite_expr(expr, mapping)?),
                name.clone(),
            )),
            Expr::Column(i) => Ok(Expr::Column(self.new_index(mapping, i)?)),
            Expr::UnresolvedColumn(_) => Err(ExecutionError::ExecutionError(
                "Columns need to be resolved before projection push down rule can run"
                    .to_owned(),
            )),
            Expr::Literal(_) => Ok(expr.clone()),
            Expr::Not(e) => Ok(Expr::Not(Box::new(self.rewrite_expr(e, mapping)?))),
            Expr::IsNull(e) => Ok(Expr::IsNull(Box::new(self.rewrite_expr(e, mapping)?))),
            Expr::IsNotNull(e) => {
                Ok(Expr::IsNotNull(Box::new(self.rewrite_expr(e, mapping)?)))
            }
            Expr::BinaryExpr { left, op, right } => Ok(Expr::BinaryExpr {
                left: Box::new(self.rewrite_expr(left, mapping)?),
                op: op.clone(),
                right: Box::new(self.rewrite_expr(right, mapping)?),
            }),
            Expr::Cast { expr, data_type } => Ok(Expr::Cast {
                expr: Box::new(self.rewrite_expr(expr, mapping)?),
                data_type: data_type.clone(),
            }),
            Expr::Sort {
                expr,
                asc,
                nulls_first,
            } => Ok(Expr::Sort {
                expr: Box::new(self.rewrite_expr(expr, mapping)?),
                asc: *asc,
                nulls_first: *nulls_first,
            }),
            Expr::AggregateFunction {
                name,
                args,
                return_type,
            } => Ok(Expr::AggregateFunction {
                name: name.to_string(),
                args: self.rewrite_expr_list(args, mapping)?,
                return_type: return_type.clone(),
            }),
            Expr::ScalarFunction {
                name,
                args,
                return_type,
            } => Ok(Expr::ScalarFunction {
                name: name.to_string(),
                args: self.rewrite_expr_list(args, mapping)?,
                return_type: return_type.clone(),
            }),
            Expr::Wildcard => Err(ExecutionError::General(
                "Wildcard expressions are not valid in a logical query plan".to_owned(),
            )),
        }
    }

    fn new_index(&self, mapping: &HashMap<usize, usize>, i: &usize) -> Result<usize> {
        match mapping.get(i) {
            Some(j) => Ok(*j),
            _ => Err(ExecutionError::InternalError(
                "Internal error computing new column index".to_string(),
            )),
        }
    }
}

fn get_projected_schema(
    table_schema: &Schema,
    projection: &Option<Vec<usize>>,
    accum: &HashSet<usize>,
    mapping: &mut HashMap<usize, usize>,
    has_projection: bool,
) -> Result<(Vec<usize>, Schema)> {
    if projection.is_some() {
        return Err(ExecutionError::General(
            "Cannot run projection push-down rule more than once".to_string(),
        ));
    }

    // once we reach the table scan, we can use the accumulated set of column
    // indexes as the projection in the table scan
    let mut projection = accum.iter().map(|i| *i).collect::<Vec<usize>>();

    if projection.is_empty() {
        if has_projection {
            // Ensure that we are reading at least one column from the table in case the query
            // does not reference any columns directly such as "SELECT COUNT(1) FROM table"
            projection.push(0);
        } else {
            // for table scan without projection, we default to return all columns
            projection = table_schema
                .fields()
                .iter()
                .enumerate()
                .map(|(i, _)| i)
                .collect::<Vec<usize>>();
        }
    }

    // sort the projection otherwise we get non-deterministic behavior
    projection.sort();

    // create the projected schema
    let mut projected_fields: Vec<Field> = Vec::with_capacity(projection.len());
    for i in &projection {
        projected_fields.push(table_schema.fields()[*i].clone());
    }

    // now that the table scan is returning a different schema we need to
    // create a mapping from the original column index to the
    // new column index so that we can rewrite expressions as
    // we walk back up the tree

    if mapping.len() != 0 {
        return Err(ExecutionError::InternalError("illegal state".to_string()));
    }

    for i in 0..table_schema.fields().len() {
        if let Some(n) = projection.iter().position(|v| *v == i) {
            mapping.insert(i, n);
        }
    }

    Ok((projection, Schema::new(projected_fields)))
}

#[cfg(test)]
mod tests {

    use super::*;
    use crate::logicalplan::Expr::*;
    use crate::logicalplan::ScalarValue;
    use crate::test::*;
    use arrow::datatypes::DataType;

    #[test]
    fn aggregate_no_group_by() -> Result<()> {
        let table_scan = test_table_scan()?;

        let plan = LogicalPlanBuilder::from(&table_scan)
            .aggregate(vec![], vec![max(Column(1))])?
            .build()?;

        let expected = "Aggregate: groupBy=[[]], aggr=[[MAX(#0)]]\
        \n  TableScan: test projection=Some([1])";

        assert_optimized_plan_eq(&plan, expected);

        Ok(())
    }

    #[test]
    fn aggregate_group_by() -> Result<()> {
        let table_scan = test_table_scan()?;

        let plan = LogicalPlanBuilder::from(&table_scan)
            .aggregate(vec![Column(2)], vec![max(Column(1))])?
            .build()?;

        let expected = "Aggregate: groupBy=[[#1]], aggr=[[MAX(#0)]]\
        \n  TableScan: test projection=Some([1, 2])";

        assert_optimized_plan_eq(&plan, expected);

        Ok(())
    }

    #[test]
    fn aggregate_no_group_by_with_selection() -> Result<()> {
        let table_scan = test_table_scan()?;

        let plan = LogicalPlanBuilder::from(&table_scan)
            .filter(Column(2))?
            .aggregate(vec![], vec![max(Column(1))])?
            .build()?;

        let expected = "Aggregate: groupBy=[[]], aggr=[[MAX(#0)]]\
        \n  Selection: #1\
        \n    TableScan: test projection=Some([1, 2])";

        assert_optimized_plan_eq(&plan, expected);

        Ok(())
    }

    #[test]
    fn cast() -> Result<()> {
        let table_scan = test_table_scan()?;

        let projection = LogicalPlanBuilder::from(&table_scan)
            .project(vec![Cast {
                expr: Box::new(Column(2)),
                data_type: DataType::Float64,
            }])?
            .build()?;

        let expected = "Projection: CAST(#0 AS Float64)\
        \n  TableScan: test projection=Some([2])";

        assert_optimized_plan_eq(&projection, expected);

        Ok(())
    }

    #[test]
    fn table_scan_projected_schema() -> Result<()> {
        let table_scan = test_table_scan()?;
        assert_eq!(3, table_scan.schema().fields().len());
        assert_fields_eq(&table_scan, vec!["a", "b", "c"]);

        let plan = LogicalPlanBuilder::from(&table_scan)
            .project(vec![Column(0), Column(1)])?
            .build()?;

        assert_fields_eq(&plan, vec!["a", "b"]);

        let expected = "Projection: #0, #1\
        \n  TableScan: test projection=Some([0, 1])";

        assert_optimized_plan_eq(&plan, expected);

        Ok(())
    }

    #[test]
    fn table_limit() -> Result<()> {
        let table_scan = test_table_scan()?;
        assert_eq!(3, table_scan.schema().fields().len());
        assert_fields_eq(&table_scan, vec!["a", "b", "c"]);

        let plan = LogicalPlanBuilder::from(&table_scan)
            .project(vec![Column(2), Column(0)])?
            .limit(Expr::Literal(ScalarValue::UInt32(5)))?
            .build()?;

        assert_fields_eq(&plan, vec!["c", "a"]);

        let expected = "Limit: UInt32(5)\
        \n  Projection: #1, #0\
        \n    TableScan: test projection=Some([0, 2])";

        assert_optimized_plan_eq(&plan, expected);

        Ok(())
    }

    #[test]
    fn table_scan_without_projection() -> Result<()> {
        let table_scan = test_table_scan()?;
        let plan = LogicalPlanBuilder::from(&table_scan).build()?;
        // should expand projection to all columns without projection
        let expected = "TableScan: test projection=Some([0, 1, 2])";
        assert_optimized_plan_eq(&plan, expected);
        Ok(())
    }

    #[test]
    fn table_scan_with_literal_projection() -> Result<()> {
        let table_scan = test_table_scan()?;
        let plan = LogicalPlanBuilder::from(&table_scan)
            .project(vec![
                Expr::Literal(ScalarValue::Int64(1)),
                Expr::Literal(ScalarValue::Int64(2)),
            ])?
            .build()?;
        let expected = "Projection: Int64(1), Int64(2)\
                      \n  TableScan: test projection=Some([0])";
        assert_optimized_plan_eq(&plan, expected);
        Ok(())
    }

    fn assert_optimized_plan_eq(plan: &LogicalPlan, expected: &str) {
        let optimized_plan = optimize(plan).expect("failed to optimize plan");
        let formatted_plan = format!("{:?}", optimized_plan);
        assert_eq!(formatted_plan, expected);
    }

    fn optimize(plan: &LogicalPlan) -> Result<LogicalPlan> {
        let mut rule = ProjectionPushDown::new();
        rule.optimize(plan)
    }
}
<|endoftext|>


<commit_before>pub mod gc;
pub mod object;

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

//     #[test]
//     fn mark_cons_references() {
//         let mut gc: GcPool<Object> = GcPool::new();
//         let head = gc.allocate(Object::Integer(10));
//         let tail = gc.allocate(Object::Integer(12));
//         let cons = gc.allocate(Object::List(head, tail));
//         let null = gc.allocate(Object::Null);

//         assert_eq!(cons.is_marked(), false);
//         gc.mark(&mut vec![cons]);
//         assert_eq!(cons.is_marked(), true);
//         assert_eq!(head.is_marked(), true);
//         assert_eq!(null.is_marked(), false);
//     }
// }

// fn main() {
//     println!("Hello, world!");

//     let mut gc = Heap::new();
//     let head = gc.allocate_int(10);

//     let tail = gc.allocate_int(12);
//     let _drop = gc.allocate(Object::Integer(2));
//     let cons = gc.allocate(Object::List(head, tail));
//     let _drop2 = gc.allocate(Object::Null);

//     println!("{:?}, {:?}", *cons, cons.is_marked());
//     println!("{:?}", *head);

//     gc.mark(&mut vec![cons]);

//     println!("cons marked? {:?}", cons.is_marked());
//     println!("head marked? {:?}", head.is_marked());

//     gc.sweep();

//     for obj in gc.pool.iter() {
//         println!("obj {:x}", (obj.as_ref()  as *const GcBox<Object>) as usize)
//     };

// }

// #[cfg(test)]
// mod tests {
//     #[test]
//     fn it_works() {
//         assert_eq!(2 + 2, 4);
//     }
// }
<commit_msg>remove old code from lib.rs<commit_after>pub mod gc;
pub mod object;
<|endoftext|>


<commit_before>extern crate libc;

use libc::c_char;

use std::ffi::CStr;
use std::ffi::CString;

#[no_mangle]
pub extern fn hello() {
    println!("hello")
}

#[no_mangle]
pub extern fn add(a: i64, b: i64) -> i64 {
    a + b
}

#[no_mangle]
pub extern fn length(str: *const c_char) -> i64 {
    let c_str = unsafe { CStr::from_ptr(str) };
    c_str.to_bytes().len() as i64
}

#[no_mangle]
pub extern fn duplicate(count: i64, c_msg: *const c_char) -> *mut c_char {
    let msg = unsafe { CStr::from_ptr(c_msg) }.to_str().unwrap();

    let out = std::iter::repeat(msg).take(count as usize)
        .collect::<String>();

    CString::new(out).unwrap().into_raw()
}

#[no_mangle]
pub extern fn release(msg: *mut c_char) {
    unsafe { CString::from_raw(msg) };
}

#[repr(C)]
pub struct Point {
    x : i32,
    y : i32,
}

#[no_mangle]
pub extern fn add_points(c_p1: *const Point, c_p2: *const Point, c_result: *mut Point) {
    let (p1, p2, mut result) = unsafe { (&*c_p1, &*c_p2, &mut *c_result) };

    result.x = p1.x + p2.x;
    result.y = p1.y + p2.y;
}
<commit_msg>Minor style fixes<commit_after>extern crate libc;

use libc::c_char;

use std::ffi::{
    CStr,
    CString,
};

#[no_mangle]
pub extern fn hello() {
    println!("hello")
}

#[no_mangle]
pub extern fn add(a: i64, b: i64) -> i64 {
    a + b
}

#[no_mangle]
pub extern fn length(str: *const c_char) -> i64 {
    let c_str = unsafe { CStr::from_ptr(str) };

    c_str.to_bytes().len() as i64
}

#[no_mangle]
pub extern fn duplicate(count: i64, c_msg: *const c_char) -> *mut c_char {
    let msg = unsafe { CStr::from_ptr(c_msg) }.to_str().unwrap();

    let out = std::iter::repeat(msg)
        .take(count as usize)
        .collect::<String>();

    CString::new(out).unwrap().into_raw()
}

#[no_mangle]
pub extern fn release(msg: *mut c_char) {
    unsafe { CString::from_raw(msg) };
}

#[repr(C)]
pub struct Point {
    x: i32,
    y: i32,
}

#[no_mangle]
pub extern fn add_points(
    c_p1: *const Point,
    c_p2: *const Point,
    c_result: *mut Point,
) {
    let (p1, p2, mut result) = unsafe { (&*c_p1, &*c_p2, &mut *c_result) };

    result.x = p1.x + p2.x;
    result.y = p1.y + p2.y;
}
<|endoftext|>


<commit_before>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

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

use kvproto::metapb;

use tikv::pd::{PdClient, RpcClient, validate_endpoints, Error as PdError};

use super::mock::mocker::*;
use super::mock::Server as MockServer;

#[test]
fn test_rpc_client() {
    let eps = "http://127.0.0.1:52729".to_owned();

    let se = Arc::new(Service::new(vec![eps.clone()]));
    let _server = MockServer::run("127.0.0.1:52729", se.clone(), Some(se.clone()));

    thread::sleep(Duration::from_secs(1));

    let client = RpcClient::new(&eps).unwrap();
    assert_ne!(client.get_cluster_id().unwrap(), 0);

    let store_id = client.alloc_id().unwrap();
    let mut store = metapb::Store::new();
    store.set_id(store_id);
    debug!("bootstrap store {:?}", store);

    let peer_id = client.alloc_id().unwrap();
    let mut peer = metapb::Peer::new();
    peer.set_id(peer_id);
    peer.set_store_id(store_id);

    let region_id = client.alloc_id().unwrap();
    let mut region = metapb::Region::new();
    region.set_id(region_id);
    region.mut_peers().push(peer.clone());
    debug!("bootstrap region {:?}", region);

    client.bootstrap_cluster(store.clone(), region.clone()).unwrap();
    assert_eq!(client.is_cluster_bootstrapped().unwrap(), true);

    let tmp_store = client.get_store(store_id).unwrap();
    assert_eq!(tmp_store.get_id(), store.get_id());

    let tmp_region = client.get_region_by_id(region_id).unwrap().unwrap();
    assert_eq!(tmp_region.get_id(), region.get_id());

    let mut prev_id = 0;
    for _ in 0..100 {
        let client = RpcClient::new(&eps).unwrap();
        let alloc_id = client.alloc_id().unwrap();
        assert!(alloc_id > prev_id);
        prev_id = alloc_id;
    }
}

#[test]
fn test_reboot() {
    let mut eps = vec![
    "http://127.0.0.1:52730".to_owned(),
    ];

    let se = Arc::new(Service::new(eps.clone()));
    let lc = Arc::new(AlreadyBootstrap::new());

    let _server = MockServer::run("127.0.0.1:52730", se.clone(), Some(lc.clone()));
    thread::sleep(Duration::from_secs(1));

    let client = RpcClient::new(&eps.pop().unwrap()).unwrap();

    assert_eq!(client.is_cluster_bootstrapped().unwrap(), false);

    match client.bootstrap_cluster(metapb::Store::new(), metapb::Region::new()) {
        Err(PdError::ClusterBootstrapped(_)) => (),
        _ => {
            panic!("failed, should return ClusterBootstrapped");
        }
    }
}

#[test]
fn test_validate_endpoints() {
    let eps = vec![
        "http://127.0.0.1:13079".to_owned(),
        "http://127.0.0.1:23079".to_owned(),
        "http://127.0.0.1:33079".to_owned(),
    ];

    let se = Arc::new(Service::new(eps.clone()));
    let sp = Arc::new(Split::new(eps.clone()));

    let _server_a = MockServer::run("127.0.0.1:13079", se.clone(), Some(sp.clone()));
    let _server_b = MockServer::run("127.0.0.1:23079", se.clone(), Some(sp.clone()));
    let _server_c = MockServer::run("127.0.0.1:33079", se.clone(), Some(sp.clone()));

    thread::sleep(Duration::from_secs(1));

    assert!(validate_endpoints(&eps).is_err());
}

#[test]
fn test_change_leader() {
    let mut eps = vec![
        "http://127.0.0.1:43079".to_owned(),
        "http://127.0.0.1:53079".to_owned(),
        "http://127.0.0.1:63079".to_owned(),
    ];

    let se = Arc::new(Service::new(eps.clone()));
    let lc = Arc::new(LeaderChange::new(eps.clone()));

    let _server_a = MockServer::run("127.0.0.1:43079", se.clone(), Some(lc.clone()));
    let _server_b = MockServer::run("127.0.0.1:53079", se.clone(), Some(lc.clone()));
    let _server_c = MockServer::run("127.0.0.1:63079", se.clone(), Some(lc.clone()));

    thread::sleep(Duration::from_secs(1));

    let client = RpcClient::new(&eps.pop().unwrap()).unwrap();
    let leader = client.get_leader();

    for _ in 0..5 {
        client.is_cluster_bootstrapped().unwrap();
        let new = client.get_leader();
        if new != leader {
            return;
        }
        thread::sleep(LeaderChange::get_leader_interval());
    }

    panic!("failed, leader should changed");
}
<commit_msg>*: cleanup test. (#1727)<commit_after>// Copyright 2016 PingCAP, Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// See the License for the specific language governing permissions and
// limitations under the License.

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

use kvproto::metapb;

use tikv::pd::{PdClient, RpcClient, validate_endpoints, Error as PdError};

use super::mock::mocker::*;
use super::mock::Server as MockServer;

#[test]
fn test_rpc_client() {
    let eps = "http://127.0.0.1:52729".to_owned();

    let se = Arc::new(Service::new(vec![eps.clone()]));
    let _server = MockServer::run("127.0.0.1:52729", se.clone(), Some(se.clone()));

    thread::sleep(Duration::from_secs(1));

    let client = RpcClient::new(&eps).unwrap();
    assert_ne!(client.get_cluster_id().unwrap(), 0);

    let store_id = client.alloc_id().unwrap();
    let mut store = metapb::Store::new();
    store.set_id(store_id);
    debug!("bootstrap store {:?}", store);

    let peer_id = client.alloc_id().unwrap();
    let mut peer = metapb::Peer::new();
    peer.set_id(peer_id);
    peer.set_store_id(store_id);

    let region_id = client.alloc_id().unwrap();
    let mut region = metapb::Region::new();
    region.set_id(region_id);
    region.mut_peers().push(peer.clone());
    debug!("bootstrap region {:?}", region);

    client.bootstrap_cluster(store.clone(), region.clone()).unwrap();
    assert_eq!(client.is_cluster_bootstrapped().unwrap(), true);

    let tmp_store = client.get_store(store_id).unwrap();
    assert_eq!(tmp_store.get_id(), store.get_id());

    let tmp_region = client.get_region_by_id(region_id).unwrap().unwrap();
    assert_eq!(tmp_region.get_id(), region.get_id());

    let mut prev_id = 0;
    for _ in 0..100 {
        let client = RpcClient::new(&eps).unwrap();
        let alloc_id = client.alloc_id().unwrap();
        assert!(alloc_id > prev_id);
        prev_id = alloc_id;
    }
}

#[test]
fn test_reboot() {
    let mut eps = vec![
        "http://127.0.0.1:52730".to_owned(),
    ];

    let se = Arc::new(Service::new(eps.clone()));
    let lc = Arc::new(AlreadyBootstrap::new());

    let _server = MockServer::run("127.0.0.1:52730", se.clone(), Some(lc.clone()));
    thread::sleep(Duration::from_secs(1));

    let client = RpcClient::new(&eps.pop().unwrap()).unwrap();

    assert!(!client.is_cluster_bootstrapped().unwrap());

    match client.bootstrap_cluster(metapb::Store::new(), metapb::Region::new()) {
        Err(PdError::ClusterBootstrapped(_)) => (),
        _ => {
            panic!("failed, should return ClusterBootstrapped");
        }
    }
}

#[test]
fn test_validate_endpoints() {
    let eps = vec![
        "http://127.0.0.1:13079".to_owned(),
        "http://127.0.0.1:23079".to_owned(),
        "http://127.0.0.1:33079".to_owned(),
    ];

    let se = Arc::new(Service::new(eps.clone()));
    let sp = Arc::new(Split::new(eps.clone()));

    let _server_a = MockServer::run("127.0.0.1:13079", se.clone(), Some(sp.clone()));
    let _server_b = MockServer::run("127.0.0.1:23079", se.clone(), Some(sp.clone()));
    let _server_c = MockServer::run("127.0.0.1:33079", se.clone(), Some(sp.clone()));

    thread::sleep(Duration::from_secs(1));

    assert!(validate_endpoints(&eps).is_err());
}

#[test]
fn test_change_leader() {
    let mut eps = vec![
        "http://127.0.0.1:43079".to_owned(),
        "http://127.0.0.1:53079".to_owned(),
        "http://127.0.0.1:63079".to_owned(),
    ];

    let se = Arc::new(Service::new(eps.clone()));
    let lc = Arc::new(LeaderChange::new(eps.clone()));

    let _server_a = MockServer::run("127.0.0.1:43079", se.clone(), Some(lc.clone()));
    let _server_b = MockServer::run("127.0.0.1:53079", se.clone(), Some(lc.clone()));
    let _server_c = MockServer::run("127.0.0.1:63079", se.clone(), Some(lc.clone()));

    thread::sleep(Duration::from_secs(1));

    let client = RpcClient::new(&eps.pop().unwrap()).unwrap();
    let leader = client.get_leader();

    for _ in 0..5 {
        client.is_cluster_bootstrapped().unwrap();
        let new = client.get_leader();
        if new != leader {
            return;
        }
        thread::sleep(LeaderChange::get_leader_interval());
    }

    panic!("failed, leader should changed");
}
<|endoftext|>


<commit_before>#![allow(unused)]

macro_rules! log {
    (target: $target:expr, $lvl:expr, $($arg:tt)+) => {
        let _ = $target;
        let _ = log!($lvl, $($arg)+);
    }
    ($lvl:expr, $($arg:tt)+) => {
        let _ = $lvl;
        let _ = format_args!($($arg)+);
    }
}
macro_rules! error {
    (target: $target:expr, $($arg:tt)+) => { log!(target: $target, "", $($arg)+); }
    ($($arg:tt)+) => { log!("", $($arg)+); }
}
macro_rules! warn {
    (target: $target:expr, $($arg:tt)*) => { log!(target: $target, "", $($arg)*); }
    ($($arg:tt)*) => { log!("", $($arg)*); }
}
macro_rules! info {
    (target: $target:expr, $($arg:tt)+) => { log!(target: $target, "", $($arg)+); }
    ($($arg:tt)+) => { log!("", $($arg)+); }
}
macro_rules! debug {
    (target: $target:expr, $($arg:tt)+) => { log!(target: $target, "", $($arg)+); }
    ($($arg:tt)+) => { log!("", $($arg)+); }
}
macro_rules! trace {
    (target: $target:expr, $($arg:tt)+) => { log!(target: $target, "", $($arg)+); }
    ($($arg:tt)+) => { log!("", $($arg)+); }
}
<commit_msg>Fix separator and treating expression<commit_after>#![allow(unused)]

macro_rules! log {
    (target: $target:expr, $lvl:expr, $($arg:tt)+) => {{
        let _ = $target;
        let _ = log!($lvl, $($arg)+);
    }};
    ($lvl:expr, $($arg:tt)+) => {{
        let _ = $lvl;
        let _ = format_args!($($arg)+);
    }};
}
macro_rules! error {
    (target: $target:expr, $($arg:tt)+) => { log!(target: $target, "", $($arg)+) };
    ($($arg:tt)+) => { log!("", $($arg)+) };
}
macro_rules! warn {
    (target: $target:expr, $($arg:tt)*) => { log!(target: $target, "", $($arg)*) };
    ($($arg:tt)*) => { log!("", $($arg)*) };
}
macro_rules! info {
    (target: $target:expr, $($arg:tt)+) => { log!(target: $target, "", $($arg)+) };
    ($($arg:tt)+) => { log!("", $($arg)+) };
}
macro_rules! debug {
    (target: $target:expr, $($arg:tt)+) => { log!(target: $target, "", $($arg)+) };
    ($($arg:tt)+) => { log!("", $($arg)+) };
}
macro_rules! trace {
    (target: $target:expr, $($arg:tt)+) => { log!(target: $target, "", $($arg)+) };
    ($($arg:tt)+) => { log!("", $($arg)+) };
}
<|endoftext|>


<commit_before>//! Types defined by Lua thread

use std::fmt::{Debug, Formatter};
use std::fmt::Result as FmtResult;
use std::cmp::{PartialEq, Eq};

use rlua;

use keys::KeyPress;

/// Messages sent to the lua thread
#[allow(dead_code)]
pub enum LuaQuery {
    /// Pings the lua thread
    Ping,
    /// Halt the lua thread
    Terminate,
    // Restart the lua thread
    Restart,

    /// Execute a string
    Execute(String),
    /// Execute a file
    ExecFile(String),
    /// Execute some Rust using the Lua context.
    ExecRust(fn(&mut rlua::Lua) -> rlua::Value<'static>),

    /// Handle the key press for the given key.
    HandleKey(KeyPress),

    /// Update the registry value from Lua's registry cache.
    UpdateRegistryFromCache,
}

impl Debug for LuaQuery {
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        match *self {
            LuaQuery::Ping => write!(f, "LuaQuery::Ping"),
            LuaQuery::Terminate => write!(f, "LuaQuery::Terminate"),
            LuaQuery::Restart => write!(f, "LuaQuery::Restart"),
            LuaQuery::Execute(ref val) =>
                write!(f, "LuaQuery::Execute({:?})", val),
            LuaQuery::ExecFile(ref val) =>
                write!(f, "LuaQuery::ExecFile({:?})", val),
            // This is why there's no #[derive(Debug)],
            // and why we have lua/types.rs
            LuaQuery::ExecRust(_) =>
                write!(f, "LuaQuery::ExecRust()"),
            LuaQuery::HandleKey(ref press) =>
                write!(f, "LuaQuery::HandleKey({:?})", press),
            LuaQuery::UpdateRegistryFromCache =>
                write!(f, "LuaQuery::UpdateRegistryFromCache")
        }
    }
}

unsafe impl Send for LuaQuery { }
unsafe impl Sync for LuaQuery { }

impl PartialEq for LuaQuery {
    fn eq(&self, other: &LuaQuery) -> bool {
        match (self, other) {
            (&LuaQuery::Ping, &LuaQuery::Ping) => true,
            (&LuaQuery::Terminate, &LuaQuery::Terminate) => true,
            (&LuaQuery::Restart, &LuaQuery::Restart) => true,

            (&LuaQuery::Execute(ref s1), &LuaQuery::Execute(ref s2)) =>
                s1 == s2,
            (&LuaQuery::ExecFile(ref s1), &LuaQuery::ExecFile(ref s2)) =>
                s1 == s2,
            (&LuaQuery::ExecRust(_), &LuaQuery::ExecRust(_)) => true,
            (&LuaQuery::HandleKey(ref p1), &LuaQuery::HandleKey(ref p2)) =>
                p1 == p2,
            _ => false
        }
    }
}

impl Eq for LuaQuery { }

/// Messages received from lua thread
#[allow(dead_code)]
pub enum LuaResponse {
    /// If the identifier had length 0
    InvalidName,
    /// Lua variable obtained
    Variable(Option<rlua::Value<'static>>),
    /// Lua error
    Error(rlua::Error),
    /// A function is returned
    Function(rlua::Function<'static>),
    /// Pong response from lua ping
    Pong,
}

#[allow(dead_code)]
impl LuaResponse {
    /// Whether this response is an InvalidName or Error
    pub fn is_err(&self) -> bool {
        match self {
            &LuaResponse::InvalidName | &LuaResponse::Error(_) => true,
            _ => false
        }
    }

    /// If this response is a Variable, Function, or Pong
    #[inline]
    pub fn is_ok(&self) -> bool {
        !self.is_err()
    }
}

impl PartialEq for LuaResponse {
    fn eq(&self, other: &LuaResponse) -> bool {
        match (self, other) {
            (&LuaResponse::InvalidName, &LuaResponse::InvalidName) => true,
            (&LuaResponse::Pong, &LuaResponse::Pong) => true,

            (&LuaResponse::Variable(ref v1), &LuaResponse::Variable(ref v2)) =>
                // TODO FIXME
                false,//v1 == v2,
            (&LuaResponse::Error(ref e1), &LuaResponse::Error(ref e2)) =>
                format!("{:?}", e1) == format!("{:?}", e2),
            (&LuaResponse::Function(_), &LuaResponse::Function(_)) => true,

            _ => false
        }
    }
}

unsafe impl Send for LuaResponse { }
unsafe impl Sync for LuaResponse { }

impl Debug for LuaResponse {
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        match *self {
            LuaResponse::InvalidName =>
                write!(f, "LuaReponse::InvalidName"),
            LuaResponse::Variable(ref var) =>
                write!(f, "LuaResponse::Variable({:?})", var),
            LuaResponse::Error(ref err) =>
                write!(f, "LuaResponse::Error({:?})", err),
            LuaResponse::Function(_) =>
                write!(f, "LuaResponse::Function"),
            LuaResponse::Pong =>
                write!(f, "LuaResponse::Pong")
        }
    }
}
<commit_msg>Removed unused impls for LuaQuery and LuaResponse<commit_after>//! Types defined by Lua thread

use std::fmt::{Debug, Formatter};
use std::fmt::Result as FmtResult;
use std::cmp::{PartialEq, Eq};

use rlua;

use keys::KeyPress;

/// Messages sent to the lua thread
#[allow(dead_code)]
pub enum LuaQuery {
    /// Pings the lua thread
    Ping,
    /// Halt the lua thread
    Terminate,
    // Restart the lua thread
    Restart,

    /// Execute a string
    Execute(String),
    /// Execute a file
    ExecFile(String),
    /// Execute some Rust using the Lua context.
    ExecRust(fn(&mut rlua::Lua) -> rlua::Value<'static>),

    /// Handle the key press for the given key.
    HandleKey(KeyPress),

    /// Update the registry value from Lua's registry cache.
    UpdateRegistryFromCache,
}

impl Debug for LuaQuery {
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        match *self {
            LuaQuery::Ping => write!(f, "LuaQuery::Ping"),
            LuaQuery::Terminate => write!(f, "LuaQuery::Terminate"),
            LuaQuery::Restart => write!(f, "LuaQuery::Restart"),
            LuaQuery::Execute(ref val) =>
                write!(f, "LuaQuery::Execute({:?})", val),
            LuaQuery::ExecFile(ref val) =>
                write!(f, "LuaQuery::ExecFile({:?})", val),
            // This is why there's no #[derive(Debug)],
            // and why we have lua/types.rs
            LuaQuery::ExecRust(_) =>
                write!(f, "LuaQuery::ExecRust()"),
            LuaQuery::HandleKey(ref press) =>
                write!(f, "LuaQuery::HandleKey({:?})", press),
            LuaQuery::UpdateRegistryFromCache =>
                write!(f, "LuaQuery::UpdateRegistryFromCache")
        }
    }
}

/// Messages received from lua thread
#[allow(dead_code)]
pub enum LuaResponse {
    /// If the identifier had length 0
    InvalidName,
    /// Lua variable obtained
    Variable(Option<rlua::Value<'static>>),
    /// Lua error
    Error(rlua::Error),
    /// A function is returned
    Function(rlua::Function<'static>),
    /// Pong response from lua ping
    Pong,
}

#[allow(dead_code)]
impl LuaResponse {
    /// Whether this response is an InvalidName or Error
    pub fn is_err(&self) -> bool {
        match self {
            &LuaResponse::InvalidName | &LuaResponse::Error(_) => true,
            _ => false
        }
    }

    /// If this response is a Variable, Function, or Pong
    #[inline]
    pub fn is_ok(&self) -> bool {
        !self.is_err()
    }
}

impl Debug for LuaResponse {
    fn fmt(&self, f: &mut Formatter) -> FmtResult {
        match *self {
            LuaResponse::InvalidName =>
                write!(f, "LuaReponse::InvalidName"),
            LuaResponse::Variable(ref var) =>
                write!(f, "LuaResponse::Variable({:?})", var),
            LuaResponse::Error(ref err) =>
                write!(f, "LuaResponse::Error({:?})", err),
            LuaResponse::Function(_) =>
                write!(f, "LuaResponse::Function"),
            LuaResponse::Pong =>
                write!(f, "LuaResponse::Pong")
        }
    }
}
<|endoftext|>


<commit_before>//! Information about Intel's performance events.
pub mod events;
// The types need to be in a spearate file so we don't get circular
// dependencies with build.rs include:
mod description;
pub use self::description::{Counter, EventDescription, MSRIndex, PebsType, Tuple};

use crate::cpuid;
use core::fmt::{Error, Result, Write};
use core::str;
use phf;

const MODEL_LEN: usize = 30;

#[derive(Default)]
struct ModelWriter {
    buffer: [u8; MODEL_LEN],
    index: usize,
}

impl ModelWriter {
    fn as_str(&self) -> &str {
        str::from_utf8(&self.buffer[..self.index]).unwrap()
    }
}

impl Write for ModelWriter {
    fn write_str(&mut self, s: &str) -> Result {
        // TODO: There exists probably a more efficient way of doing this:
        for c in s.chars() {
            if self.index >= self.buffer.len() {
                return Err(Error);
            }
            self.buffer[self.index] = c as u8;
            self.index += 1;
        }
        Ok(())
    }
}

// Format must be a string literal
macro_rules! get_events {
    ($format:expr) => {{
        let cpuid = cpuid::CpuId::new();

        cpuid.get_vendor_info().map_or(None, |vf| {
            cpuid.get_feature_info().map_or(None, |fi| {
                let vendor = vf.as_str();
                let (family, extended_model, model) =
                    (fi.base_family_id(), fi.extended_model_id(), fi.base_model_id());

                let mut writer: ModelWriter = Default::default();
                // Should work as long as it fits in MODEL_LEN bytes:
                write!(writer, $format, vendor, family, extended_model, model).unwrap();
                let key = writer.as_str();

                events::COUNTER_MAP.get(key)
            })
        })
    }};
}

/// Return all core performance events for the running micro-architecture.
pub fn events() -> Option<&'static phf::Map<&'static str, EventDescription<'static>>> {
    // Should be something like: GenuineIntel-6-2C
    get_events!("{}-{}-{:X}{:X}")
}

#[test]
fn events_test() {
    // Note: This will silently fail in case the counter is not available.
    events().map(|cc| {
        cc.get("INST_RETIRED.ANY").map(|p| {
            assert!(p.event_name == "INST_RETIRED.ANY");
        });
    });
}
<commit_msg>Formatting.<commit_after>//! Information about Intel's performance events.
pub mod events;
// The types need to be in a spearate file so we don't get circular
// dependencies with build.rs include:
mod description;
pub use self::description::{Counter, EventDescription, MSRIndex, PebsType, Tuple};

use crate::cpuid;
use core::fmt::{Error, Result, Write};
use core::str;
use phf;

const MODEL_LEN: usize = 30;

#[derive(Default)]
struct ModelWriter {
    buffer: [u8; MODEL_LEN],
    index: usize,
}

impl ModelWriter {
    fn as_str(&self) -> &str {
        str::from_utf8(&self.buffer[..self.index]).unwrap()
    }
}

impl Write for ModelWriter {
    fn write_str(&mut self, s: &str) -> Result {
        // TODO: There exists probably a more efficient way of doing this:
        for c in s.chars() {
            if self.index >= self.buffer.len() {
                return Err(Error);
            }
            self.buffer[self.index] = c as u8;
            self.index += 1;
        }
        Ok(())
    }
}

// Format must be a string literal
macro_rules! get_events {
    ($format:expr) => {{
        let cpuid = cpuid::CpuId::new();

        cpuid.get_vendor_info().map_or(None, |vf| {
            cpuid.get_feature_info().map_or(None, |fi| {
                let vendor = vf.as_str();
                let (family, extended_model, model) = (
                    fi.base_family_id(),
                    fi.extended_model_id(),
                    fi.base_model_id(),
                );

                let mut writer: ModelWriter = Default::default();
                // Should work as long as it fits in MODEL_LEN bytes:
                write!(writer, $format, vendor, family, extended_model, model).unwrap();
                let key = writer.as_str();

                events::COUNTER_MAP.get(key)
            })
        })
    }};
}

/// Return all core performance events for the running micro-architecture.
pub fn events() -> Option<&'static phf::Map<&'static str, EventDescription<'static>>> {
    // Should be something like: GenuineIntel-6-2C
    get_events!("{}-{}-{:X}{:X}")
}

#[test]
fn events_test() {
    // Note: This will silently fail in case the counter is not available.
    events().map(|cc| {
        cc.get("INST_RETIRED.ANY").map(|p| {
            assert!(p.event_name == "INST_RETIRED.ANY");
        });
    });
}
<|endoftext|>


<commit_before>// vim: tw=80

use mockall::predicate::*;

mod no_argument {
    mockall::omnimock!{FooExpectation, __foo_priv, u32, [], [], [], [], []}

    #[test]
    fn t() {
        let mut e = FooExpectation::default();
        e.returning(|| 42);
        assert_eq!(42u32, e.call());
    }
}

mod reference_argument {
    use super::*;

    mockall::omnimock!{FooExpectation, __foo_priv, u32,
        [&u32], [&i0], [i0], [p0], [u32]}

    #[test]
    fn t() {
        let mut e = FooExpectation::default();
        e.with(eq(42u32))
            .returning(|x| *x);
        let x = 42u32;
        assert_eq!(42u32, e.call(&x));
    }

    #[test]
    fn match_fn() {
        let mut e = FooExpectation::default();
        e.withf(|x| *x == 42)
            .returning(|x| *x);
        let x = 42u32;
        assert_eq!(42u32, e.call(&x));
    }
}

mod ref_and_nonref_arguments {
    use super::*;

    mockall::omnimock!{FooExpectation, __foo_priv, i32,
        [i32, &u16],
        [&i0, i1],
        [i0, i1],
        [p0, p1],
        [i32, u16]
    }

    #[test]
    fn t() {
        let mut e = FooExpectation::default();
        e.with(eq(42), eq(1))
            .returning(|x, y| x + i32::from(*y));
        let x = 42i32;
        let y = 1u16;
        assert_eq!(43i32, e.call(x, &y));
    }

    #[test]
    fn match_fn() {
        let mut e = FooExpectation::default();
        e.withf(|x, y| *x == i32::from(*y))
            .returning(|x, y| x + i32::from(*y));
        let x = 42i32;
        let y = 42u16;
        assert_eq!(84i32, e.call(x, &y));
    }
}

mod two_reference_arguments {
    use super::*;

    mockall::omnimock!{FooExpectation, __foo_priv, u32,
        [&u32, &u16],
        [i0, i1],
        [i0, i1],
        [p0, p1],
        [u32, u16]
    }

    #[test]
    fn t() {
        let mut e = FooExpectation::default();
        e.with(eq(42), eq(1))
            .returning(|x, y| *x + u32::from(*y));
        let x = 42u32;
        let y = 1u16;
        assert_eq!(43u32, e.call(&x, &y));
    }
}
<commit_msg>Fixup to "s/omnimock/expectation/g"<commit_after>// vim: tw=80

use mockall::predicate::*;

mod no_argument {
    mockall::expectation!{FooExpectation, __foo_priv, u32, [], [], [], [], []}

    #[test]
    fn t() {
        let mut e = FooExpectation::default();
        e.returning(|| 42);
        assert_eq!(42u32, e.call());
    }
}

mod reference_argument {
    use super::*;

    mockall::expectation!{FooExpectation, __foo_priv, u32,
        [&u32], [&i0], [i0], [p0], [u32]}

    #[test]
    fn t() {
        let mut e = FooExpectation::default();
        e.with(eq(42u32))
            .returning(|x| *x);
        let x = 42u32;
        assert_eq!(42u32, e.call(&x));
    }

    #[test]
    fn match_fn() {
        let mut e = FooExpectation::default();
        e.withf(|x| *x == 42)
            .returning(|x| *x);
        let x = 42u32;
        assert_eq!(42u32, e.call(&x));
    }
}

mod ref_and_nonref_arguments {
    use super::*;

    mockall::expectation!{FooExpectation, __foo_priv, i32,
        [i32, &u16],
        [&i0, i1],
        [i0, i1],
        [p0, p1],
        [i32, u16]
    }

    #[test]
    fn t() {
        let mut e = FooExpectation::default();
        e.with(eq(42), eq(1))
            .returning(|x, y| x + i32::from(*y));
        let x = 42i32;
        let y = 1u16;
        assert_eq!(43i32, e.call(x, &y));
    }

    #[test]
    fn match_fn() {
        let mut e = FooExpectation::default();
        e.withf(|x, y| *x == i32::from(*y))
            .returning(|x, y| x + i32::from(*y));
        let x = 42i32;
        let y = 42u16;
        assert_eq!(84i32, e.call(x, &y));
    }
}

mod two_reference_arguments {
    use super::*;

    mockall::expectation!{FooExpectation, __foo_priv, u32,
        [&u32, &u16],
        [i0, i1],
        [i0, i1],
        [p0, p1],
        [u32, u16]
    }

    #[test]
    fn t() {
        let mut e = FooExpectation::default();
        e.with(eq(42), eq(1))
            .returning(|x, y| *x + u32::from(*y));
        let x = 42u32;
        let y = 1u16;
        assert_eq!(43u32, e.call(&x, &y));
    }
}
<|endoftext|>


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

use libc;
use std::mem;
use std::ptr;
use std::slice;

use super::buffer::*;
use super::memory::*;

#[cfg(windows)]
#[link(name = "msvcrt")]
extern "C" {
    fn _aligned_free(prt: *const u8);
}

/// Buffer builder with zero-copy build method
pub struct Builder<T> {
    data: *mut T,
    len: usize,
    capacity: usize,
}

impl<T> Builder<T> {
    /// Creates a builder with a default capacity
    pub fn new() -> Self {
        Builder::with_capacity(64)
    }

    /// Creates a builder with a fixed capacity
    pub fn with_capacity(capacity: usize) -> Self {
        let sz = mem::size_of::<T>();
        let buffer = allocate_aligned((capacity * sz) as i64).unwrap();
        Builder {
            len: 0,
            capacity,
            data: unsafe { mem::transmute::<*const u8, *mut T>(buffer) },
        }
    }

    /// Get the internal byte-aligned memory buffer as a mutable slice
    pub fn slice_mut(&self, start: usize, end: usize) -> &mut [T] {
        assert!(start <= end);
        assert!(start < self.len as usize);
        assert!(end <= self.len as usize);
        unsafe {
            slice::from_raw_parts_mut(self.data.offset(start as isize), (end - start) as usize)
        }
    }

    /// Override the length
    pub fn set_len(&mut self, len: usize) {
        self.len = len;
    }

    /// Push a value into the builder, growing the internal buffer as needed
    pub fn push(&mut self, v: T) {
        assert!(!self.data.is_null());
        if self.len == self.capacity {
            let sz = mem::size_of::<T>();
            let new_capacity = self.capacity * 2;
            unsafe {
                let old_buffer = self.data;
                let new_buffer = allocate_aligned((new_capacity * sz) as i64).unwrap();
                libc::memcpy(
                    mem::transmute::<*const u8, *mut libc::c_void>(new_buffer),
                    mem::transmute::<*const T, *const libc::c_void>(old_buffer),
                    self.len * sz,
                );
                self.capacity = new_capacity;
                self.data = mem::transmute::<*const u8, *mut T>(new_buffer);
                mem::drop(old_buffer);
            }
        }
        assert!(self.len < self.capacity);
        unsafe {
            *self.data.offset(self.len as isize) = v;
        }
        self.len += 1;
    }

    /// Build a Buffer from the existing memory
    pub fn finish(&mut self) -> Buffer<T> {
        assert!(!self.data.is_null());
        let p = unsafe { mem::transmute::<*mut T, *const T>(self.data) };
        self.data = ptr::null_mut(); // ensure builder cannot be re-used
        Buffer::from_raw_parts(p, self.len as i32)
    }
}

impl<T> Drop for Builder<T> {
    #[cfg(windows)]
    fn drop(&mut self) {
        if !self.data.is_null() {
            unsafe {
                let p = mem::transmute::<*const T, *const u8>(self.data);
                _aligned_free(p);
            }
        }
    }

    #[cfg(not(windows))]
    fn drop(&mut self) {
        if !self.data.is_null() {
            unsafe {
                let p = mem::transmute::<*const T, *mut libc::c_void>(self.data);
                libc::free(p);
            }
        }
    }
}

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

    #[test]
    fn test_builder_i32_empty() {
        let mut b: Builder<i32> = Builder::with_capacity(5);
        let a = b.finish();
        assert_eq!(0, a.len());
    }

    #[test]
    fn test_builder_i32() {
        let mut b: Builder<i32> = Builder::with_capacity(5);
        for i in 0..5 {
            b.push(i);
        }
        let a = b.finish();
        assert_eq!(5, a.len());
        for i in 0..5 {
            assert_eq!(&i, a.get(i as usize));
        }
    }

    #[test]
    fn test_builder_i32_grow_buffer() {
        let mut b: Builder<i32> = Builder::with_capacity(2);
        for i in 0..5 {
            b.push(i);
        }
        let a = b.finish();
        assert_eq!(5, a.len());
        for i in 0..5 {
            assert_eq!(&i, a.get(i as usize));
        }
    }

}
<commit_msg>ARROW-2433: [Rust] Add Builder.push_slice(&[T])<commit_after>// Licensed to the Apache Software Foundation (ASF) under one
// or more contributor license agreements.  See the NOTICE file
// distributed with this work for additional information
// regarding copyright ownership.  The ASF licenses this file
// to you under the Apache License, Version 2.0 (the
// "License"); you may not use this file except in compliance
// with the License.  You may obtain a copy of the License at
//
//   http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing,
// software distributed under the License is distributed on an
// "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
// KIND, either express or implied.  See the License for the
// specific language governing permissions and limitations
// under the License.

use libc;
use std::cmp;
use std::mem;
use std::ptr;
use std::slice;

use super::buffer::*;
use super::memory::*;

#[cfg(windows)]
#[link(name = "msvcrt")]
extern "C" {
    fn _aligned_free(prt: *const u8);
}

/// Buffer builder with zero-copy build method
pub struct Builder<T> {
    data: *mut T,
    len: usize,
    capacity: usize,
}

impl<T> Builder<T> {
    /// Creates a builder with a default capacity
    pub fn new() -> Self {
        Builder::with_capacity(64)
    }

    /// Creates a builder with a fixed capacity
    pub fn with_capacity(capacity: usize) -> Self {
        let sz = mem::size_of::<T>();
        let buffer = allocate_aligned((capacity * sz) as i64).unwrap();
        Builder {
            len: 0,
            capacity,
            data: unsafe { mem::transmute::<*const u8, *mut T>(buffer) },
        }
    }

    pub fn len(&self) -> usize {
        self.len
    }

    pub fn capacity(&self) -> usize {
        self.capacity
    }

    /// Get the internal byte-aligned memory buffer as a mutable slice
    pub fn slice_mut(&self, start: usize, end: usize) -> &mut [T] {
        assert!(start <= end);
        assert!(start < self.len as usize);
        assert!(end <= self.len as usize);
        unsafe {
            slice::from_raw_parts_mut(self.data.offset(start as isize), (end - start) as usize)
        }
    }

    /// Override the length
    pub fn set_len(&mut self, len: usize) {
        self.len = len;
    }

    /// Push a value into the builder, growing the internal buffer as needed
    pub fn push(&mut self, v: T) {
        assert!(!self.data.is_null());
        if self.len == self.capacity {
            let new_capacity = self.capacity;
            self.grow(new_capacity * 2);
        }
        assert!(self.len < self.capacity);
        unsafe {
            *self.data.offset(self.len as isize) = v;
        }
        self.len += 1;
    }

    /// push a slice of type T, growing the internal buffer as needed
    pub fn push_slice(&mut self, slice: &[T]) {
        self.reserve(slice.len());
        let sz = mem::size_of::<T>();
        unsafe {
            libc::memcpy(
                mem::transmute::<*mut T, *mut libc::c_void>(self.data.offset(self.len() as isize)),
                mem::transmute::<*const T, *const libc::c_void>(slice.as_ptr()),
                slice.len() * sz,
            );
        }
        self.len += slice.len();
    }

    /// Reserve memory for n elements of type T
    pub fn reserve(&mut self, n: usize) {
        if self.len + n > self.capacity {
            let new_capacity = cmp::max(self.capacity * 2, n);
            self.grow(new_capacity);
        }
    }

    /// Grow the buffer to the new size n (number of elements of type T)
    fn grow(&mut self, new_capacity: usize) {
        let sz = mem::size_of::<T>();
        unsafe {
            let old_buffer = self.data;
            let new_buffer = allocate_aligned((new_capacity * sz) as i64).unwrap();
            libc::memcpy(
                mem::transmute::<*const u8, *mut libc::c_void>(new_buffer),
                mem::transmute::<*const T, *const libc::c_void>(old_buffer),
                self.len * sz,
            );
            self.capacity = new_capacity;
            self.data = mem::transmute::<*const u8, *mut T>(new_buffer);
            libc::free(mem::transmute::<*mut T, *mut libc::c_void>(old_buffer));
        }
    }

    /// Build a Buffer from the existing memory
    pub fn finish(&mut self) -> Buffer<T> {
        assert!(!self.data.is_null());
        let p = unsafe { mem::transmute::<*mut T, *const T>(self.data) };
        self.data = ptr::null_mut(); // ensure builder cannot be re-used
        Buffer::from_raw_parts(p, self.len as i32)
    }
}

impl<T> Drop for Builder<T> {
    #[cfg(windows)]
    fn drop(&mut self) {
        if !self.data.is_null() {
            unsafe {
                let p = mem::transmute::<*const T, *const u8>(self.data);
                _aligned_free(p);
            }
        }
    }

    #[cfg(not(windows))]
    fn drop(&mut self) {
        if !self.data.is_null() {
            unsafe {
                let p = mem::transmute::<*const T, *mut libc::c_void>(self.data);
                libc::free(p);
            }
        }
    }
}

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

    #[test]
    fn test_builder_i32_empty() {
        let mut b: Builder<i32> = Builder::with_capacity(5);
        let a = b.finish();
        assert_eq!(0, a.len());
    }

    #[test]
    fn test_builder_i32() {
        let mut b: Builder<i32> = Builder::with_capacity(5);
        for i in 0..5 {
            b.push(i);
        }
        let a = b.finish();
        assert_eq!(5, a.len());
        for i in 0..5 {
            assert_eq!(&i, a.get(i as usize));
        }
    }

    #[test]
    fn test_builder_i32_grow_buffer() {
        let mut b: Builder<i32> = Builder::with_capacity(2);
        for i in 0..5 {
            b.push(i);
        }
        let a = b.finish();
        assert_eq!(5, a.len());
        for i in 0..5 {
            assert_eq!(&i, a.get(i as usize));
        }
    }

    #[test]
    fn test_reserve() {
        let mut b: Builder<u8> = Builder::with_capacity(2);
        assert_eq!(2, b.capacity());
        b.reserve(2);
        assert_eq!(2, b.capacity());
        b.reserve(3);
        assert_eq!(4, b.capacity());
    }

    #[test]
    fn test_push_slice() {
        let mut b: Builder<u8> = Builder::new();
        b.push_slice("Hello, ".as_bytes());
        b.push_slice("World!".as_bytes());
        let buffer = b.finish();
        assert_eq!(13, buffer.len());

        let s = String::from_utf8(buffer.iter().collect::<Vec<u8>>()).unwrap();
        assert_eq!("Hello, World!", s);
    }

}
<|endoftext|>