From 15f9af937504eb67082dabd787bed0aea98c927b Mon Sep 17 00:00:00 2001
From: Qinxuan Chen <koushiro.cqx@gmail.com>
Date: Sat, 30 Nov 2019 12:14:23 +0800
Subject: [PATCH] copr: vectorize md5 (#6091)

Signed-off-by: koushiro <koushiro.cqx@gmail.com>
---
 .../src/rpn_expr/impl_encryption.rs           | 104 ++++++++++++++----
 components/tidb_query/src/rpn_expr/mod.rs     |   1 +
 2 files changed, 81 insertions(+), 24 deletions(-)

diff --git a/components/tidb_query/src/rpn_expr/impl_encryption.rs b/components/tidb_query/src/rpn_expr/impl_encryption.rs
index 0934317d232..4a8b2ca02b6 100644
--- a/components/tidb_query/src/rpn_expr/impl_encryption.rs
+++ b/components/tidb_query/src/rpn_expr/impl_encryption.rs
@@ -6,18 +6,22 @@ use tidb_query_codegen::rpn_fn;
 use crate::codec::data_type::*;
 use crate::Result;
 
+#[rpn_fn]
+#[inline]
+pub fn md5(arg: &Option<Bytes>) -> Result<Option<Bytes>> {
+    match arg {
+        Some(arg) => hex_digest(MessageDigest::md5(), arg).map(Some),
+        None => Ok(None),
+    }
+}
+
 #[rpn_fn]
 #[inline]
 pub fn sha1(arg: &Option<Bytes>) -> Result<Option<Bytes>> {
-    Ok(match arg {
-        Some(arg) => {
-            match hex_digest(MessageDigest::sha1(), arg) {
-                Ok(s) => return Ok(Some(s)),
-                Err(err) => return Err(err),
-            };
-        }
-        _ => None,
-    })
+    match arg {
+        Some(arg) => hex_digest(MessageDigest::sha1(), arg).map(Some),
+        None => Ok(None),
+    }
 }
 
 #[inline]
@@ -29,36 +33,88 @@ fn hex_digest(hashtype: MessageDigest, input: &[u8]) -> Result<Bytes> {
 
 #[cfg(test)]
 mod tests {
+    use tipb::ScalarFuncSig;
+
     use super::*;
     use crate::rpn_expr::types::test_util::RpnFnScalarEvaluator;
-    use tipb::ScalarFuncSig;
+
+    fn test_unary_func_ok_none<I: Evaluable, O: Evaluable>(sig: ScalarFuncSig)
+    where
+        O: PartialEq,
+        Option<I>: Into<ScalarValue>,
+        Option<O>: From<ScalarValue>,
+    {
+        assert_eq!(
+            None,
+            RpnFnScalarEvaluator::new()
+                .push_param(Option::<I>::None)
+                .evaluate::<O>(sig)
+                .unwrap()
+        );
+    }
 
     #[test]
-    fn test_sha1() {
-        let cases = vec![
-            (Some(""), Some("da39a3ee5e6b4b0d3255bfef95601890afd80709")),
-            (Some("a"), Some("86f7e437faa5a7fce15d1ddcb9eaeaea377667b8")),
-            (Some("ab"), Some("da23614e02469a0d7c7bd1bdab5c9c474b1904dc")),
+    fn test_md5() {
+        let test_cases = vec![
+            (vec![], "d41d8cd98f00b204e9800998ecf8427e"),
+            (b"a".to_vec(), "0cc175b9c0f1b6a831c399e269772661"),
+            (b"ab".to_vec(), "187ef4436122d1cc2f40dc2b92f0eba0"),
+            (b"abc".to_vec(), "900150983cd24fb0d6963f7d28e17f72"),
+            (b"123".to_vec(), "202cb962ac59075b964b07152d234b70"),
             (
-                Some("abc"),
-                Some("a9993e364706816aba3e25717850c26c9cd0d89d"),
+                "你好".as_bytes().to_vec(),
+                "7eca689f0d3389d9dea66ae112e5cfd7",
             ),
             (
-                Some("123"),
-                Some("40bd001563085fc35165329ea1ff5c5ecbdbbeef"),
+                "分布式データベース".as_bytes().to_vec(),
+                "63c0354797bd261e2cbf8581147eeeda",
             ),
-            (None, None),
+            (vec![0xc0, 0x80], "b26555f33aedac7b2684438cc5d4d05e"),
+            (vec![0xED, 0xA0, 0x80], "546d3dc8de10fbf8b448f678a47901e4"),
         ];
+        for (arg, expect_output) in test_cases {
+            let expect_output = Some(Bytes::from(expect_output));
 
-        for (arg, expect_output) in cases {
-            let arg = arg.map(|s| s.as_bytes().to_vec());
-            let expect_output = expect_output.map(|s| Bytes::from(s));
+            let output = RpnFnScalarEvaluator::new()
+                .push_param(arg)
+                .evaluate::<Bytes>(ScalarFuncSig::Md5)
+                .unwrap();
+            assert_eq!(output, expect_output);
+        }
+        test_unary_func_ok_none::<Bytes, Bytes>(ScalarFuncSig::Md5);
+    }
+
+    #[test]
+    fn test_sha1() {
+        let test_cases = vec![
+            (vec![], "da39a3ee5e6b4b0d3255bfef95601890afd80709"),
+            (b"a".to_vec(), "86f7e437faa5a7fce15d1ddcb9eaeaea377667b8"),
+            (b"ab".to_vec(), "da23614e02469a0d7c7bd1bdab5c9c474b1904dc"),
+            (b"abc".to_vec(), "a9993e364706816aba3e25717850c26c9cd0d89d"),
+            (b"123".to_vec(), "40bd001563085fc35165329ea1ff5c5ecbdbbeef"),
+            (
+                "你好".as_bytes().to_vec(),
+                "440ee0853ad1e99f962b63e459ef992d7c211722",
+            ),
+            (
+                "分布式データベース".as_bytes().to_vec(),
+                "82aa64080df2ca37550ddfc3419d75ac1df3e0d0",
+            ),
+            (vec![0xc0, 0x80], "8bf4822782a21d7ac68ece130ac36987548003bd"),
+            (
+                vec![0xED, 0xA0, 0x80],
+                "10db70ec072d000c68dd95879f9b831e43a859fd",
+            ),
+        ];
+        for (arg, expect_output) in test_cases {
+            let expect_output = Some(Bytes::from(expect_output));
 
             let output = RpnFnScalarEvaluator::new()
                 .push_param(arg)
-                .evaluate(ScalarFuncSig::Sha1)
+                .evaluate::<Bytes>(ScalarFuncSig::Sha1)
                 .unwrap();
             assert_eq!(output, expect_output);
         }
+        test_unary_func_ok_none::<Bytes, Bytes>(ScalarFuncSig::Sha1);
     }
 }
diff --git a/components/tidb_query/src/rpn_expr/mod.rs b/components/tidb_query/src/rpn_expr/mod.rs
index d25ce4aaa9f..44919b48490 100644
--- a/components/tidb_query/src/rpn_expr/mod.rs
+++ b/components/tidb_query/src/rpn_expr/mod.rs
@@ -251,6 +251,7 @@ fn map_expr_node_to_rpn_func(expr: &Expr) -> Result<RpnFnMeta> {
         ScalarFuncSig::CaseWhenTime => case_when_fn_meta::<DateTime>(),
         ScalarFuncSig::CaseWhenDuration => case_when_fn_meta::<Duration>(),
         ScalarFuncSig::CaseWhenJson => case_when_fn_meta::<Json>(),
+        ScalarFuncSig::Md5 => md5_fn_meta(),
         ScalarFuncSig::Sha1 => sha1_fn_meta(),
         ScalarFuncSig::DateFormatSig => date_format_fn_meta(),
         ScalarFuncSig::DayOfYear => day_of_year_fn_meta(),