From ebd72d920352f287a12f82448871c75f23f35902 Mon Sep 17 00:00:00 2001
From: Gaith Hallak <gaithhallak@gmail.com>
Date: Thu, 1 Sep 2022 15:31:23 +0300
Subject: [PATCH] Return mapping from variables to registers

---
 src/aeso_compiler.erl      |  5 +++--
 src/aeso_fcode_to_fate.erl | 32 +++++++++++++++++++++++++-------
 2 files changed, 28 insertions(+), 9 deletions(-)

diff --git a/src/aeso_compiler.erl b/src/aeso_compiler.erl
index d49c47d..0bc1329 100644
--- a/src/aeso_compiler.erl
+++ b/src/aeso_compiler.erl
@@ -115,7 +115,7 @@ from_string1(ContractString, Options) ->
      , fcode_env := #{child_con_env := ChildContracts}
      , folded_typed_ast := FoldedTypedAst
      , warnings := Warnings } = string_to_code(ContractString, Options),
-    FateCode = aeso_fcode_to_fate:compile(ChildContracts, FCode, Options),
+    {FateCode, VarsRegs} = aeso_fcode_to_fate:compile(ChildContracts, FCode, Options),
     pp_assembler(FateCode, Options),
     ByteCode = aeb_fate_code:serialize(FateCode, []),
     {ok, Version} = version(),
@@ -126,6 +126,7 @@ from_string1(ContractString, Options) ->
             fate_code => FateCode,
             abi_version => aeb_fate_abi:abi_version(),
             payable => maps:get(payable, FCode),
+            variables_registers => VarsRegs,
             warnings => Warnings
            },
     {ok, maybe_generate_aci(Res, FoldedTypedAst, Options)}.
@@ -185,7 +186,7 @@ check_call1(ContractString0, FunName, Args, Options) ->
         #{fcode := OrgFcode
         , fcode_env := #{child_con_env := ChildContracts}
         , ast := Ast} = string_to_code(ContractString0, Options),
-        FateCode = aeso_fcode_to_fate:compile(ChildContracts, OrgFcode, []),
+        {FateCode, _VarsRegs} = aeso_fcode_to_fate:compile(ChildContracts, OrgFcode, []),
         %% collect all hashes and compute the first name without hash collision to
         SymbolHashes = maps:keys(aeb_fate_code:symbols(FateCode)),
         CallName = first_none_match(?CALL_NAME, SymbolHashes,
diff --git a/src/aeso_fcode_to_fate.erl b/src/aeso_fcode_to_fate.erl
index cdf0c8f..0712f44 100644
--- a/src/aeso_fcode_to_fate.erl
+++ b/src/aeso_fcode_to_fate.erl
@@ -76,13 +76,18 @@ compile(FCode, Options) ->
 compile(ChildContracts, FCode, Options) ->
     #{ contract_name := ContractName,
        functions     := Functions } = FCode,
-    SFuns  = functions_to_scode(ChildContracts, ContractName, Functions, Options),
-    SFuns1 = optimize_scode(SFuns, Options),
-    FateCode = to_basic_blocks(SFuns1),
-    ?debug(compile, Options, "~s\n", [aeb_fate_asm:pp(FateCode)]),
-    case proplists:get_value(include_child_contract_symbols, Options, false) of
-        false -> FateCode;
-        true -> add_child_symbols(ChildContracts, FateCode)
+    try
+        SFuns  = functions_to_scode(ChildContracts, ContractName, Functions, Options),
+        SFuns1 = optimize_scode(SFuns, Options),
+        FateCode = to_basic_blocks(SFuns1),
+        ?debug(compile, Options, "~s\n", [aeb_fate_asm:pp(FateCode)]),
+        FateCode1 = case proplists:get_value(include_child_contract_symbols, Options, false) of
+                        false -> FateCode;
+                        true -> add_child_symbols(ChildContracts, FateCode)
+                    end,
+        {FateCode1, get_variables_registers()}
+    after
+        put(variables_registers, undefined)
     end.
 
 make_function_id(X) ->
@@ -110,9 +115,21 @@ function_to_scode(ChildContracts, ContractName, Functions, Name, Attrs0, Args, B
     {ArgTypes, ResType1} = typesig_to_scode(Args, ResType),
     Attrs = Attrs0 -- [stateful], %% Only track private and payable from here.
     Env = init_env(ChildContracts, ContractName, Functions, Name, Args, Options),
+    [ add_variables_register(Env, Arg, Register) || {Arg, Register} <- Env#env.vars ],
     SCode = to_scode(Env, Body),
     {Attrs, {ArgTypes, ResType1}, SCode}.
 
+get_variables_registers() ->
+    case get(variables_registers) of
+        undefined -> [];
+        Vs        -> Vs
+    end.
+
+add_variables_register(Env, Name, Register) ->
+    Olds = get_variables_registers(),
+    New = {Env#env.contract, Env#env.current_function, Name, Register},
+    put(variables_registers, [New | Olds]).
+
 -define(tvars, '$tvars').
 
 typesig_to_scode(Args, Res) ->
@@ -170,6 +187,7 @@ next_var(#env{ vars = Vars }) ->
     1 + lists:max([-1 | [J || {_, {var, J}} <- Vars]]).
 
 bind_var(Name, Var, Env = #env{ vars = Vars }) ->
+    add_variables_register(Env, Name, Var),
     Env#env{ vars = [{Name, Var} | Vars] }.
 
 bind_local(Name, Env) ->