From d5881519cb5b364b42f2b2780cea48752c1ec79c Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= <clement@meilisearch.com>
Date: Wed, 15 Mar 2023 11:01:51 +0100
Subject: [PATCH] Make the json flattener return the original values

---
 flatten-serde-json/src/lib.rs | 64 ++++++++++++++++++++++++++---------
 1 file changed, 48 insertions(+), 16 deletions(-)

diff --git a/flatten-serde-json/src/lib.rs b/flatten-serde-json/src/lib.rs
index e1b2b20c7..b2e36d5b3 100644
--- a/flatten-serde-json/src/lib.rs
+++ b/flatten-serde-json/src/lib.rs
@@ -3,45 +3,45 @@
 use serde_json::{Map, Value};
 
 pub fn flatten(json: &Map<String, Value>) -> Map<String, Value> {
-    let mut obj = Map::new();
-    let mut all_keys = vec![];
-    insert_object(&mut obj, None, json, &mut all_keys);
-    for key in all_keys {
-        obj.entry(key).or_insert(Value::Array(vec![]));
+    let mut obj = json.clone();
+    let mut all_entries = vec![];
+    insert_object(&mut obj, None, json, &mut all_entries);
+    for (key, old_val) in all_entries {
+        obj.entry(key).or_insert(old_val.clone());
     }
     obj
 }
 
-fn insert_object(
+fn insert_object<'a>(
     base_json: &mut Map<String, Value>,
     base_key: Option<&str>,
-    object: &Map<String, Value>,
-    all_keys: &mut Vec<String>,
+    object: &'a Map<String, Value>,
+    all_entries: &mut Vec<(String, &'a Value)>,
 ) {
     for (key, value) in object {
         let new_key = base_key.map_or_else(|| key.clone(), |base_key| format!("{base_key}.{key}"));
-        all_keys.push(new_key.clone());
+        all_entries.push((new_key.clone(), value));
         if let Some(array) = value.as_array() {
-            insert_array(base_json, &new_key, array, all_keys);
+            insert_array(base_json, &new_key, array, all_entries);
         } else if let Some(object) = value.as_object() {
-            insert_object(base_json, Some(&new_key), object, all_keys);
+            insert_object(base_json, Some(&new_key), object, all_entries);
         } else {
             insert_value(base_json, &new_key, value.clone());
         }
     }
 }
 
-fn insert_array(
+fn insert_array<'a>(
     base_json: &mut Map<String, Value>,
     base_key: &str,
-    array: &Vec<Value>,
-    all_keys: &mut Vec<String>,
+    array: &'a Vec<Value>,
+    all_entries: &mut Vec<(String, &'a Value)>,
 ) {
     for value in array {
         if let Some(object) = value.as_object() {
-            insert_object(base_json, Some(base_key), object, all_keys);
+            insert_object(base_json, Some(base_key), object, all_entries);
         } else if let Some(sub_array) = value.as_array() {
-            insert_array(base_json, base_key, sub_array, all_keys);
+            insert_array(base_json, base_key, sub_array, all_entries);
         } else {
             insert_value(base_json, base_key, value.clone());
         }
@@ -302,4 +302,36 @@ mod tests {
             .unwrap()
         );
     }
+
+    #[test]
+    fn flatten_nested_values_keep_original_values() {
+        let mut base: Value = json!({
+            "tags": {
+                "t1": "v1"
+            },
+            "prices": {
+                "p1": [null]
+            }
+        });
+        let json = std::mem::take(base.as_object_mut().unwrap());
+        let flat = flatten(&json);
+
+        println!("{}", serde_json::to_string_pretty(&flat).unwrap());
+
+        assert_eq!(
+            &flat,
+            json!({
+                "tags": {
+                    "t1": "v1"
+                },
+                "tags.t1": "v1",
+                "prices": {
+                    "p1": [null]
+                },
+                "prices.p1": [null]
+            })
+            .as_object()
+            .unwrap()
+        );
+    }
 }