From 3b76af8de5bd49ebf2834ddd59ca57d1948eae9c Mon Sep 17 00:00:00 2001
From: WenyXu <wenymedia@gmail.com>
Date: Wed, 3 Jul 2024 11:55:45 +0000
Subject: [PATCH] chore: apply suggestions from CR

---
 tests-fuzz/src/ir/insert_expr.rs              | 12 ++++++----
 tests-fuzz/targets/fuzz_insert.rs             | 22 +++----------------
 .../targets/fuzz_insert_logical_table.rs      | 22 +++----------------
 3 files changed, 14 insertions(+), 42 deletions(-)

diff --git a/tests-fuzz/src/ir/insert_expr.rs b/tests-fuzz/src/ir/insert_expr.rs
index a315680f4483..da8ca202cabc 100644
--- a/tests-fuzz/src/ir/insert_expr.rs
+++ b/tests-fuzz/src/ir/insert_expr.rs
@@ -36,8 +36,7 @@ impl InsertIntoExpr {
         self.columns
             .iter()
             .enumerate()
-            .find(|(_, c)| c.is_time_index())
-            .map(|(idx, _)| idx)
+            .find_map(|(idx, c)| if c.is_time_index() { Some(idx) } else { None })
     }
 
     /// Returns a vector of columns that are primary keys or time indices.
@@ -54,8 +53,13 @@ impl InsertIntoExpr {
         self.columns
             .iter()
             .enumerate()
-            .filter(|(_, c)| c.is_primary_key() || c.is_time_index())
-            .map(|(i, _)| i)
+            .filter_map(|(i, c)| {
+                if c.is_primary_key() || c.is_time_index() {
+                    Some(i)
+                } else {
+                    None
+                }
+            })
             .collect::<Vec<_>>()
     }
 }
diff --git a/tests-fuzz/targets/fuzz_insert.rs b/tests-fuzz/targets/fuzz_insert.rs
index fdb258775187..ce2971408583 100644
--- a/tests-fuzz/targets/fuzz_insert.rs
+++ b/tests-fuzz/targets/fuzz_insert.rs
@@ -33,8 +33,8 @@ use tests_fuzz::generator::create_expr::CreateTableExprGeneratorBuilder;
 use tests_fuzz::generator::insert_expr::InsertExprGeneratorBuilder;
 use tests_fuzz::generator::Generator;
 use tests_fuzz::ir::{
-    generate_random_timestamp_for_mysql, generate_random_value, replace_default, CreateTableExpr,
-    InsertIntoExpr, MySQLTsColumnTypeGenerator,
+    generate_random_timestamp_for_mysql, generate_random_value, replace_default,
+    sort_by_primary_keys, CreateTableExpr, InsertIntoExpr, MySQLTsColumnTypeGenerator,
 };
 use tests_fuzz::translator::mysql::create_expr::CreateTableExprTranslator;
 use tests_fuzz::translator::mysql::insert_expr::InsertIntoExprTranslator;
@@ -184,23 +184,7 @@ async fn execute_insert(ctx: FuzzContext, input: FuzzInput) -> Result<()> {
     );
     let fetched_rows = validator::row::fetch_values(&ctx.greptime, select_sql.as_str()).await?;
     let mut expected_rows = replace_default(&insert_expr.values_list, &table_ctx, &insert_expr);
-    expected_rows.sort_by(|a, b| {
-        let a_keys: Vec<_> = primary_keys_idxs_in_insert_expr
-            .iter()
-            .map(|&i| &a[i])
-            .collect();
-        let b_keys: Vec<_> = primary_keys_idxs_in_insert_expr
-            .iter()
-            .map(|&i| &b[i])
-            .collect();
-        for (a_key, b_key) in a_keys.iter().zip(b_keys.iter()) {
-            match a_key.cmp(b_key) {
-                Some(std::cmp::Ordering::Equal) => continue,
-                non_eq => return non_eq.unwrap(),
-            }
-        }
-        std::cmp::Ordering::Equal
-    });
+    sort_by_primary_keys(&mut expected_rows, primary_keys_idxs_in_insert_expr);
     validator::row::assert_eq::<MySql>(&insert_expr.columns, &fetched_rows, &expected_rows)?;
 
     // Cleans up
diff --git a/tests-fuzz/targets/fuzz_insert_logical_table.rs b/tests-fuzz/targets/fuzz_insert_logical_table.rs
index de518f284ac0..af7903a9f101 100644
--- a/tests-fuzz/targets/fuzz_insert_logical_table.rs
+++ b/tests-fuzz/targets/fuzz_insert_logical_table.rs
@@ -36,8 +36,8 @@ use tests_fuzz::generator::create_expr::{
 use tests_fuzz::generator::insert_expr::InsertExprGeneratorBuilder;
 use tests_fuzz::generator::Generator;
 use tests_fuzz::ir::{
-    generate_random_timestamp_for_mysql, generate_random_value, replace_default, CreateTableExpr,
-    InsertIntoExpr,
+    generate_random_timestamp_for_mysql, generate_random_value, replace_default,
+    sort_by_primary_keys, CreateTableExpr, InsertIntoExpr,
 };
 use tests_fuzz::translator::mysql::create_expr::CreateTableExprTranslator;
 use tests_fuzz::translator::mysql::insert_expr::InsertIntoExprTranslator;
@@ -181,23 +181,7 @@ async fn validate_values(
     let fetched_rows = validator::row::fetch_values(&ctx.greptime, select_sql.as_str()).await?;
     let mut expected_rows =
         replace_default(&insert_expr.values_list, &logical_table_ctx, insert_expr);
-    expected_rows.sort_by(|a, b| {
-        let a_keys: Vec<_> = primary_keys_idxs_in_insert_expr
-            .iter()
-            .map(|&i| &a[i])
-            .collect();
-        let b_keys: Vec<_> = primary_keys_idxs_in_insert_expr
-            .iter()
-            .map(|&i| &b[i])
-            .collect();
-        for (a_key, b_key) in a_keys.iter().zip(b_keys.iter()) {
-            match a_key.cmp(b_key) {
-                Some(std::cmp::Ordering::Equal) => continue,
-                non_eq => return non_eq.unwrap(),
-            }
-        }
-        std::cmp::Ordering::Equal
-    });
+    sort_by_primary_keys(&mut expected_rows, primary_keys_idxs_in_insert_expr);
     validator::row::assert_eq::<MySql>(&insert_expr.columns, &fetched_rows, &expected_rows)?;
 
     Ok(())