From 8a04ee8097536830a8a570e66440800e916f95e1 Mon Sep 17 00:00:00 2001
From: Tom Jack <tom@tomjack.co>
Date: Wed, 25 Sep 2019 13:29:00 -0500
Subject: [PATCH] Fix bug for compile-parameter and compile-storage

---
 src/main/run/of_simplified.ml |  7 +------
 src/main/run/of_source.ml     | 16 ++++++++++++----
 src/stages/ast_typed/misc.ml  |  5 +++++
 3 files changed, 18 insertions(+), 10 deletions(-)

diff --git a/src/main/run/of_simplified.ml b/src/main/run/of_simplified.ml
index 523ceed23..4bc7729b8 100644
--- a/src/main/run/of_simplified.ml
+++ b/src/main/run/of_simplified.ml
@@ -1,11 +1,6 @@
 open Trace
 open Ast_simplified
 
-let get_final_environment program =
-  let last_declaration = Location.unwrap List.(hd @@ rev program) in
-  let (Ast_typed.Declaration_constant (_ , (_ , post_env))) = last_declaration in
-  post_env
-
 let compile_expression ?(value = false) ?env expr =
   if value
   then (
@@ -22,7 +17,7 @@ let run_typed_program
     (input : expression) : expression result =
   let%bind code = Compile.Of_typed.compile_function_entry program entry in
   let%bind input =
-    let env = get_final_environment program in
+    let env = Ast_typed.program_environment program in
     compile_expression ?value:input_to_value ~env input
   in
   let%bind ex_ty_value = Of_michelson.run ?options code input in
diff --git a/src/main/run/of_source.ml b/src/main/run/of_source.ml
index 5fed28c09..f9a8e776c 100644
--- a/src/main/run/of_source.ml
+++ b/src/main/run/of_source.ml
@@ -50,15 +50,19 @@ end
 
 let compile_file_contract_parameter : string -> string -> string -> Compile.Helpers.s_syntax -> Michelson.t result =
   fun source_filename _entry_point expression syntax ->
+  let%bind program = Compile.Of_source.type_file syntax source_filename in
+  let env = Ast_typed.program_environment program in
   let%bind syntax = Compile.Helpers.syntax_to_variant syntax (Some source_filename) in
   let%bind simplified = Compile.Helpers.parsify_expression syntax expression in
-  Of_simplified.compile_expression simplified
+  Of_simplified.compile_expression simplified ~env
 
 let compile_file_expression : string -> string -> string -> Compile.Helpers.s_syntax -> Michelson.t result =
   fun source_filename _entry_point expression syntax ->
+  let%bind program = Compile.Of_source.type_file syntax source_filename in
+  let env = Ast_typed.program_environment program in
   let%bind syntax = Compile.Helpers.syntax_to_variant syntax (Some source_filename) in
   let%bind simplified = Compile.Helpers.parsify_expression syntax expression in
-  Of_simplified.compile_expression simplified
+  Of_simplified.compile_expression simplified ~env
 
 let compile_expression : string -> Compile.Helpers.s_syntax -> Michelson.t result =
   fun expression syntax ->
@@ -68,17 +72,21 @@ let compile_expression : string -> Compile.Helpers.s_syntax -> Michelson.t resul
 
 let compile_file_contract_storage ~value : string -> string -> string -> Compile.Helpers.s_syntax -> Michelson.t result =
   fun source_filename _entry_point expression syntax ->
+  let%bind program = Compile.Of_source.type_file syntax source_filename in
+  let env = Ast_typed.program_environment program in
   let%bind syntax = Compile.Helpers.syntax_to_variant syntax (Some source_filename) in
   let%bind simplified = Compile.Helpers.parsify_expression syntax expression in
-  Of_simplified.compile_expression ~value simplified
+  Of_simplified.compile_expression ~value simplified ~env
 
 let compile_file_contract_args =
   fun ?value source_filename _entry_point storage parameter syntax ->
+  let%bind program = Compile.Of_source.type_file syntax source_filename in
+  let env = Ast_typed.program_environment program in
   let%bind syntax = Compile.Helpers.syntax_to_variant syntax (Some source_filename) in
   let%bind storage_simplified = Compile.Helpers.parsify_expression syntax storage in
   let%bind parameter_simplified = Compile.Helpers.parsify_expression syntax parameter in
   let args = Ast_simplified.e_pair storage_simplified parameter_simplified in
-  Of_simplified.compile_expression ?value args
+  Of_simplified.compile_expression ?value args ~env
 
 
 let run_contract ?amount ?storage_value source_filename entry_point storage parameter syntax =
diff --git a/src/stages/ast_typed/misc.ml b/src/stages/ast_typed/misc.ml
index 5aaf28550..db33f6062 100644
--- a/src/stages/ast_typed/misc.ml
+++ b/src/stages/ast_typed/misc.ml
@@ -500,3 +500,8 @@ let get_entry (lst : program) (name : string) : annotated_expression result =
     else None
   in
   List.find_map aux lst
+
+let program_environment (program : program) : full_environment =
+  let last_declaration = Location.unwrap List.(hd @@ rev program) in
+  match last_declaration with
+  | Declaration_constant (_ , (_ , post_env)) -> post_env