diff --git a/scripts/test_cli.sh b/scripts/test_cli.sh
index ab02ff575..ad83f2e64 100755
--- a/scripts/test_cli.sh
+++ b/scripts/test_cli.sh
@@ -5,102 +5,12 @@ compiled_storage=$(./scripts/ligo_ci.sh compile-storage src/test/contracts/websi
 compiled_parameter=$(./scripts/ligo_ci.sh compile-parameter src/test/contracts/website2.ligo main "Increment(1)");
 dry_run_output=$(./scripts/ligo_ci.sh dry-run src/test/contracts/website2.ligo main "Increment(1)" 1);
 
-expected_compiled_contract="{ parameter (or int int) ;
-  storage int ;
-  code { {} ;
-         {} ;
-         {} ;
-         { PUSH (lambda (pair int int) int)
-                { {} ;
-                  {} ;
-                  {} ;
-                  { { { DUP ; DIP { {} } } ; CAR } ;
-                    { { { { DIP { DUP } ; SWAP } ; DIP { {} } } ; CDR } ;
-                      { PUSH unit Unit ;
-                        DROP ;
-                        { { { DIP { DUP } ; SWAP } ; DIP { { DUP ; DIP { {} } } } } ;
-                          ADD } } ;
-                      {} ;
-                      DIP { DROP } } ;
-                    {} ;
-                    DIP { DROP } } ;
-                  {} ;
-                  DIP { DROP } ;
-                  {} } ;
-           { PUSH (lambda (pair int int) int)
-                  { {} ;
-                    {} ;
-                    {} ;
-                    { { { DUP ; DIP { {} } } ; CAR } ;
-                      { { { { DIP { DUP } ; SWAP } ; DIP { {} } } ; CDR } ;
-                        { PUSH unit Unit ;
-                          DROP ;
-                          { { { DIP { DUP } ; SWAP } ; DIP { { DUP ; DIP { {} } } } } ;
-                            SUB } } ;
-                        {} ;
-                        DIP { DROP } } ;
-                      {} ;
-                      DIP { DROP } } ;
-                    {} ;
-                    DIP { DROP } ;
-                    {} } ;
-             { { { { DIP { { DIP { DUP } ; SWAP } } ; SWAP } ; DIP { {} } } ;
-                 CAR } ;
-               { { { { DIP { { DIP { { DIP { DUP } ; SWAP } } ; SWAP } } ; SWAP } ;
-                     DIP { {} } } ;
-                   CDR } ;
-                 { PUSH unit Unit ;
-                   DROP ;
-                   { { NIL operation ;
-                       DIP { { { { DIP { DUP } ; SWAP } ;
-                                 IF_LEFT
-                                   { { { DUP ;
-                                         { { { { DIP { { DIP { DUP } ; SWAP } } ; SWAP } ;
-                                               DIP { { DUP ; DIP { {} } } } } ;
-                                             PAIR } ;
-                                           DIP { { DIP { { DIP { { DIP { { DIP { DUP } ; SWAP } } ; SWAP } } ; SWAP } } ;
-                                                   SWAP } } ;
-                                           EXEC } ;
-                                         {} ;
-                                         DIP { DROP } } ;
-                                       {} ;
-                                       DIP { DROP } } }
-                                   { { { DUP ;
-                                         { { { { DIP { { DIP { DUP } ; SWAP } } ; SWAP } ;
-                                               DIP { { DUP ; DIP { {} } } } } ;
-                                             PAIR } ;
-                                           DIP { { DIP { { DIP { { DIP { { DIP { { DIP { DUP } ; SWAP } } ; SWAP } } ; SWAP } } ;
-                                                           SWAP } } ;
-                                                   SWAP } } ;
-                                           EXEC } ;
-                                         {} ;
-                                         DIP { DROP } } ;
-                                       {} ;
-                                       DIP { DROP } } } } ;
-                               DIP { {} } } } } ;
-                     PAIR } } ;
-                 {} ;
-                 DIP { DROP } } ;
-               {} ;
-               DIP { DROP } } ;
-             {} ;
-             DIP { DROP } } ;
-           {} ;
-           DIP { DROP } } ;
-         {} ;
-         DIP { DROP } ;
-         {} } }";
 expected_compiled_parameter="(Right 1)";
 expected_compiled_storage=1;
 expected_dry_run_output="tuple[   list[]
          2
 ]";
 
-if [ "$compiled_contract" != "$expected_compiled_contract" ]; then
-    echo "Expected $expected_compiled_contract as compile-storage output, got $compiled_contract instead";
-    exit 1;
-fi
-
 if [ "$compiled_storage" != "$expected_compiled_storage" ]; then
     echo "Expected $expected_compiled_storage as compile-storage output, got $compiled_storage instead";
     exit 1;
@@ -116,4 +26,4 @@ if [ "$dry_run_output" != "$expected_dry_run_output" ]; then
     exit 1;
 fi
 
-echo "CLI tests passed";
\ No newline at end of file
+echo "CLI tests passed";
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