Support transform exp tree into RPN (#4329)

Signed-off-by: Breezewish <breezewish@pingcap.com>

src/coprocessor/dag/mod.rs

@@ -39,6 +39,7 @@ mod builder;
 pub mod executor;
 pub mod expr;
 pub mod handler;
+pub mod rpn_expr;
 
 pub use self::executor::{ScanOn, Scanner};
 pub use self::handler::DAGRequestHandler;

src/coprocessor/dag/rpn_expr/function.rs

@@ -0,0 +1,68 @@
+// Copyright 2019 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.
+
+/// A trait for all RPN functions.
+pub trait RpnFunction: std::fmt::Debug + Send + Sync + 'static {
+    /// The display name of the function.
+    fn name(&self) -> &'static str;
+
+    /// The accepted argument length of this RPN function.
+    ///
+    /// Currently we do not support variable arguments.
+    fn args_len(&self) -> usize;
+}
+
+impl<T: RpnFunction + ?Sized> RpnFunction for Box<T> {
+    #[inline]
+    fn name(&self) -> &'static str {
+        (**self).name()
+    }
+
+    #[inline]
+    fn args_len(&self) -> usize {
+        (**self).args_len()
+    }
+}
+
+/// Implements `RpnFunction` automatically for structure that accepts 0, 1, 2 or 3 arguments.
+///
+/// The structure must have a `call` member function accepting corresponding number of scalar
+/// arguments.
+#[macro_export]
+macro_rules! impl_template_fn {
+    (0 arg @ $name:ident) => {
+        impl_template_fn! { @inner $name, 0 }
+    };
+    (1 arg @ $name:ident) => {
+        impl_template_fn! { @inner $name, 1 }
+    };
+    (2 arg @ $name:ident) => {
+        impl_template_fn! { @inner $name, 2 }
+    };
+    (3 arg @ $name:ident) => {
+        impl_template_fn! { @inner $name, 3 }
+    };
+    (@inner $name:ident, $args:expr) => {
+        impl $crate::coprocessor::dag::rpn_expr::RpnFunction for $name {
+            #[inline]
+            fn name(&self) -> &'static str {
+                stringify!($name)
+            }
+
+            #[inline]
+            fn args_len(&self) -> usize {
+                $args
+            }
+        }
+    };
+}

src/coprocessor/dag/rpn_expr/mod.rs

@@ -0,0 +1,34 @@
+// Copyright 2019 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.
+
+#[macro_use]
+mod function;
+mod types;
+
+pub use self::function::RpnFunction;
+pub use self::types::RpnExpression;
+
+use tipb::expression::ScalarFuncSig;
+
+use crate::coprocessor::Error;
+
+// TODO: We should not expose this function as `pub` in future once all executors are batch
+// executors.
+pub fn map_pb_sig_to_rpn_func(value: ScalarFuncSig) -> Result<Box<dyn RpnFunction>, Error> {
+    match value {
+        v => Err(box_err!(
+            "ScalarFunction {:?} is not supported in batch mode",
+            v
+        )),
+    }
+}

src/coprocessor/dag/rpn_expr/types/expr.rs

@@ -0,0 +1,135 @@
+// Copyright 2019 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 tipb::expression::FieldType;
+
+use super::super::function::RpnFunction;
+use crate::coprocessor::codec::data_type::ScalarValue;
+
+/// A type for each node in the RPN expression list.
+#[derive(Debug)]
+pub enum RpnExpressionNode {
+    /// Represents a function call.
+    FnCall {
+        func: Box<dyn RpnFunction>,
+        field_type: FieldType,
+    },
+
+    /// Represents a scalar constant value.
+    Constant {
+        value: ScalarValue,
+        field_type: FieldType,
+    },
+
+    /// Represents a reference to a column in the columns specified in evaluation.
+    ColumnRef {
+        offset: usize,
+
+        // Although we can know `ColumnInfo` according to `offset` and columns info in scan
+        // executors, its type is `ColumnInfo` instead of `FieldType`.
+        // Maybe we can remove this field in future.
+        field_type: FieldType,
+    },
+}
+
+impl RpnExpressionNode {
+    /// Gets the field type.
+    #[inline]
+    pub fn field_type(&self) -> &FieldType {
+        match self {
+            RpnExpressionNode::FnCall { ref field_type, .. } => field_type,
+            RpnExpressionNode::Constant { ref field_type, .. } => field_type,
+            RpnExpressionNode::ColumnRef { ref field_type, .. } => field_type,
+        }
+    }
+
+    /// Borrows the function instance for `FnCall` variant.
+    #[inline]
+    pub fn fn_call_func(&self) -> Option<&dyn RpnFunction> {
+        match self {
+            RpnExpressionNode::FnCall { ref func, .. } => Some(&*func),
+            _ => None,
+        }
+    }
+
+    /// Borrows the constant value for `Constant` variant.
+    #[inline]
+    pub fn constant_value(&self) -> Option<&ScalarValue> {
+        match self {
+            RpnExpressionNode::Constant { ref value, .. } => Some(value),
+            _ => None,
+        }
+    }
+
+    /// Gets the column offset for `ColumnRef` variant.
+    #[inline]
+    pub fn column_ref_offset(&self) -> Option<usize> {
+        match self {
+            RpnExpressionNode::ColumnRef { ref offset, .. } => Some(*offset),
+            _ => None,
+        }
+    }
+}
+
+/// An expression in Reverse Polish notation, which is simply a list of RPN expression nodes.
+///
+/// You may want to build it using `RpnExpressionBuilder`.
+#[derive(Debug)]
+pub struct RpnExpression(Vec<RpnExpressionNode>);
+
+impl std::ops::Deref for RpnExpression {
+    type Target = Vec<RpnExpressionNode>;
+
+    fn deref(&self) -> &Vec<RpnExpressionNode> {
+        &self.0
+    }
+}
+
+impl std::ops::DerefMut for RpnExpression {
+    fn deref_mut(&mut self) -> &mut Vec<RpnExpressionNode> {
+        &mut self.0
+    }
+}
+
+impl From<Vec<RpnExpressionNode>> for RpnExpression {
+    fn from(v: Vec<RpnExpressionNode>) -> Self {
+        Self(v)
+    }
+}
+
+#[cfg(test)]
+pub mod tests {
+    /// An RPN function for test. It accepts 1 int argument, returns the value in float.
+    #[derive(Debug, Clone, Copy)]
+    pub struct FnA;
+
+    impl_template_fn! { 1 arg @ FnA }
+
+    /// An RPN function for test. It accepts 2 float arguments, returns their sum in int.
+    #[derive(Debug, Clone, Copy)]
+    pub struct FnB;
+
+    impl_template_fn! { 2 arg @ FnB }
+
+    /// An RPN function for test. It accepts 3 int arguments, returns their sum in int.
+    #[derive(Debug, Clone, Copy)]
+    pub struct FnC;
+
+    impl_template_fn! { 3 arg @ FnC }
+
+    /// An RPN function for test. It accepts 3 float arguments, returns their sum in float.
+    #[derive(Debug, Clone, Copy)]
+    pub struct FnD;
+
+    impl_template_fn! { 3 arg @ FnD }
+}