From 1377d60f070a10cf259f31b83b79f69eb4f3cb6d Mon Sep 17 00:00:00 2001
From: gluax <16431709+gluax@users.noreply.github.com>
Date: Wed, 16 Feb 2022 10:10:26 -0800
Subject: [PATCH] test canonicalization

---
 compiler/src/lib.rs   | 14 ++++++++++++--
 leo/commands/build.rs |  2 +-
 2 files changed, 13 insertions(+), 3 deletions(-)

diff --git a/compiler/src/lib.rs b/compiler/src/lib.rs
index 50238038d9..fd203591ae 100644
--- a/compiler/src/lib.rs
+++ b/compiler/src/lib.rs
@@ -22,6 +22,7 @@
 #![allow(clippy::upper_case_acronyms)]
 #![doc = include_str!("../README.md")]
 
+use leo_ast::AstPass;
 use leo_errors::emitter::Handler;
 use leo_errors::{CompilerError, Result};
 use leo_span::symbol::create_session_if_not_set_then;
@@ -34,16 +35,18 @@ use std::path::PathBuf;
 pub struct Compiler<'a> {
     handler: &'a Handler,
     main_file_path: PathBuf,
+    output_directory: PathBuf,
 }
 
 impl<'a> Compiler<'a> {
     ///
     /// Returns a new Leo compiler.
     ///
-    pub fn new(handler: &'a Handler, main_file_path: PathBuf) -> Self {
+    pub fn new(handler: &'a Handler, main_file_path: PathBuf, output_directory: PathBuf) -> Self {
         Self {
             handler,
             main_file_path,
+            output_directory,
         }
     }
 
@@ -72,11 +75,18 @@ impl<'a> Compiler<'a> {
             .map_err(|e| CompilerError::file_read_error(self.main_file_path.clone(), e))?;
 
         // Use the parser to construct the abstract syntax tree (ast).
-        let ast: leo_ast::Ast = leo_parser::parse_ast(
+        let mut ast: leo_ast::Ast = leo_parser::parse_ast(
             self.handler,
             self.main_file_path.to_str().unwrap_or_default(),
             program_string,
         )?;
+        // Write the AST snapshot post parsing.
+        ast.to_json_file_without_keys(self.output_directory.clone(), "inital_ast.json", &["span"])?;
+
+        // Canonicalize the AST.
+        ast = leo_ast_passes::Canonicalizer::do_pass(Default::default(), ast.into_repr())?;
+        // Write the AST snapshot post parsing
+        ast.to_json_file_without_keys(self.output_directory.clone(), "canonicalization_ast.json", &["span"])?;
 
         Ok(ast)
     }
diff --git a/leo/commands/build.rs b/leo/commands/build.rs
index 3161d24f7b..053c113d75 100644
--- a/leo/commands/build.rs
+++ b/leo/commands/build.rs
@@ -178,7 +178,7 @@ impl Command for Build {
         // Initialize error handler
         let handler = leo_errors::emitter::Handler::default();
 
-        let program = Compiler::new(&handler, main_file_path);
+        let program = Compiler::new(&handler, main_file_path, output_directory);
 
         // Compute the current program checksum
         let program_checksum = program.checksum()?;