Implemented option to remove from serialized AST

Cargo.lock

@@ -1296,6 +1296,7 @@ dependencies = [
  "rand_core 0.6.3",
  "rand_xorshift",
  "serde",
+ "serde_json",
  "serde_yaml",
  "sha2",
  "snarkvm-algorithms",

ast/src/lib.rs

@@ -103,12 +103,17 @@ impl Ast {
         Ok(serde_json::to_string_pretty(&self.ast).map_err(|e| AstError::failed_to_convert_ast_to_json_string(&e))?)
     }
 
+    // Converts the ast into a JSON value
+    pub fn to_json_value(&self) -> Result<serde_json::Value> {
+        Ok(serde_json::to_value(&self.ast).map_err(|e| AstError::failed_to_convert_ast_to_json_value(&e))?)
+    }
+
     /// Serializes the ast into a JSON file.
-    pub fn to_json_file(&self, mut path: std::path::PathBuf, file_name: &str) -> Result<()> {
+    pub fn to_json_file(value: serde_json::Value, mut path: std::path::PathBuf, file_name: &str) -> Result<()> {
         path.push(file_name);
         let file = std::fs::File::create(&path).map_err(|e| AstError::failed_to_create_ast_json_file(&path, &e))?;
         let writer = std::io::BufWriter::new(file);
-        Ok(serde_json::to_writer_pretty(writer, &self.ast)
+        Ok(serde_json::to_writer_pretty(writer, &value)
             .map_err(|e| AstError::failed_to_write_ast_to_json_file(&path, &e))?)
     }
 

compiler/Cargo.toml

@@ -104,6 +104,9 @@ version = "0.8"
 [dependencies.serde]
 version = "1.0"
 
+[dependencies.serde_json]
+version = "1.0"
+
 [dependencies.sha2]
 version = "0.9"
 

compiler/src/compiler.rs

@@ -246,11 +246,11 @@ impl<'a, F: PrimeField, G: GroupType<F>> Compiler<'a, F, G> {
         let mut ast: leo_ast::Ast = parse_ast(self.main_file_path.to_str().unwrap_or_default(), program_string)?;
 
         if self.ast_snapshot_options.initial {
+            let mut value = ast.to_json_value()?;
             if !self.ast_snapshot_options.spans_enabled {
-                //TODO: Implement
-            } else {
-                ast.to_json_file(self.output_directory.clone(), "initial_ast.json")?;
+                remove_key_from_json(&mut value, "span")
             }
+            leo_ast::Ast::to_json_file(value, self.output_directory.clone(), "initial_ast.json")?;
         }
 
         // Preform import resolution.
@@ -260,22 +260,24 @@ impl<'a, F: PrimeField, G: GroupType<F>> Compiler<'a, F, G> {
         )?;
 
         if self.ast_snapshot_options.imports_resolved {
+            let mut value = ast.to_json_value()?;
             if !self.ast_snapshot_options.spans_enabled {
-                //TODO: Implement
-            } else {
-                ast.to_json_file(self.output_directory.clone(), "imports_resolved_ast.json")?;
+                remove_key_from_json(&mut value, "span");
             }
+
+            leo_ast::Ast::to_json_file(value, self.output_directory.clone(), "imports_resolved_ast.json")?;
         }
 
         // Preform canonicalization of AST always.
         ast = leo_ast_passes::Canonicalizer::do_pass(ast.into_repr())?;
 
         if self.ast_snapshot_options.canonicalized {
+            let mut value = ast.to_json_value()?;
             if !self.ast_snapshot_options.spans_enabled {
-                //TODO: Implement
-            } else {
-                ast.to_json_file(self.output_directory.clone(), "canonicalization_ast.json")?;
+                remove_key_from_json(&mut value, "span");
             }
+
+            leo_ast::Ast::to_json_file(value, self.output_directory.clone(), "canonicalization_ast.json")?;
         }
 
         // Store the main program file.
@@ -291,11 +293,13 @@ impl<'a, F: PrimeField, G: GroupType<F>> Compiler<'a, F, G> {
             let new_ast = TypeInferencePhase::default()
                 .phase_ast(&self.program, &asg.clone().into_repr())
                 .expect("Failed to produce type inference ast.");
+
+            let mut value = new_ast.to_json_value()?;
             if self.ast_snapshot_options.spans_enabled {
-                //TODO: Implement
-            } else {
-                new_ast.to_json_file(self.output_directory.clone(), "type_inferenced_ast.json")?;
+                remove_key_from_json(&mut value, "span");
             }
+
+            leo_ast::Ast::to_json_file(value, self.output_directory.clone(), "type_inferenced_ast.json")?;
         }
 
         tracing::debug!("ASG generation complete");
@@ -405,3 +409,13 @@ impl<'a, F: PrimeField, G: GroupType<F>> ConstraintSynthesizer<F> for Compiler<'
         Ok(())
     }
 }
+
+// Helper function to recursively filter keys from AST JSON
+fn remove_key_from_json(value: &mut serde_json::Value, key: &str) {
+    if let serde_json::value::Value::Object(map) = value {
+        map.remove(key);
+        for val in map.values_mut() {
+            remove_key_from_json(val, key);
+        }
+    }
+}