From ffae3b39ff232f0f556687707760f42b526bd8db Mon Sep 17 00:00:00 2001
From: Pranav Gaddamadugu <pranavsaig@gmail.com>
Date: Wed, 29 Sep 2021 07:56:28 -0700
Subject: [PATCH] Fixed JSON order bug in AST on removal of values

---
 ast/src/lib.rs                | 24 +++++++++++-------------
 compiler/Cargo.toml           |  1 +
 parser/Cargo.toml             |  1 +
 test-framework/Cargo.toml     |  1 +
 test-framework/src/bin/tgc.rs |  2 ++
 5 files changed, 16 insertions(+), 13 deletions(-)

diff --git a/ast/src/lib.rs b/ast/src/lib.rs
index 275c73153e..4c67a5134d 100644
--- a/ast/src/lib.rs
+++ b/ast/src/lib.rs
@@ -133,7 +133,7 @@ impl Ast {
 
         let mut value = self.to_json_value().unwrap();
         for key in excluded_keys {
-            remove_key_from_json(&mut value, key);
+            value = remove_key_from_json(value, key);
         }
         value = normalize_json_value(value);
 
@@ -161,20 +161,18 @@ impl AsRef<Program> for Ast {
 }
 
 /// Helper function to recursively filter keys from AST JSON
-fn remove_key_from_json(value: &mut serde_json::Value, key: &str) {
+fn remove_key_from_json(value: serde_json::Value, key: &str) -> serde_json::Value {
     match value {
-        serde_json::value::Value::Object(map) => {
-            map.remove(key);
-            for val in map.values_mut() {
-                remove_key_from_json(val, key);
-            }
+        serde_json::Value::Object(map) => serde_json::Value::Object(
+            map.into_iter()
+                .filter(|(k, _)| k != key)
+                .map(|(k, v)| (k, remove_key_from_json(v, key)))
+                .collect(),
+        ),
+        serde_json::Value::Array(values) => {
+            serde_json::Value::Array(values.into_iter().map(|v| remove_key_from_json(v, key)).collect())
         }
-        serde_json::value::Value::Array(values) => {
-            for val in values.iter_mut() {
-                remove_key_from_json(val, key);
-            }
-        }
-        _ => (),
+        _ => value,
     }
 }
 
diff --git a/compiler/Cargo.toml b/compiler/Cargo.toml
index e915fa9674..66ac496d45 100644
--- a/compiler/Cargo.toml
+++ b/compiler/Cargo.toml
@@ -106,6 +106,7 @@ version = "1.0"
 
 [dependencies.serde_json]
 version = "1.0"
+features = [ "preserve_order" ]
 
 [dependencies.sha2]
 version = "0.9"
diff --git a/parser/Cargo.toml b/parser/Cargo.toml
index f32cab30ec..7c97005bd9 100644
--- a/parser/Cargo.toml
+++ b/parser/Cargo.toml
@@ -51,6 +51,7 @@ version = "0.3"
 
 [dev-dependencies.serde_json]
 version = "1.0"
+features = [ "preserve_order" ]
 
 [dev-dependencies.serde_yaml]
 version = "0.8"
diff --git a/test-framework/Cargo.toml b/test-framework/Cargo.toml
index bea2f59c5c..508f8997d9 100644
--- a/test-framework/Cargo.toml
+++ b/test-framework/Cargo.toml
@@ -23,6 +23,7 @@ features = [ "derive" ]
 
 [dependencies.serde_json]
 version = "1.0"
+features = [ "preserve_order" ]
 
 [dependencies.serde_yaml]
 version = "0.8"
diff --git a/test-framework/src/bin/tgc.rs b/test-framework/src/bin/tgc.rs
index e7a68def7e..1f5c456be8 100644
--- a/test-framework/src/bin/tgc.rs
+++ b/test-framework/src/bin/tgc.rs
@@ -128,6 +128,8 @@ fn generate_asts(src_path: PathBuf, target_path: PathBuf, text: &str) -> Result<
 
     let mut ast = leo_parser::parse_ast(src_path.clone().into_os_string().into_string().unwrap(), text)?;
 
+    ast.to_json_file(target_path.clone(), "initial_ast_with_span.json")?;
+
     ast.to_json_file_without_keys(target_path.clone(), "initial_ast.json", &["span"])?;
 
     ast = leo_ast_passes::Importer::do_pass(ast.into_repr(), &mut ImportParser::new(src_path, Default::default()))?;