Change get_state and set_state fcode primitives to take a register

src/aeso_ast_to_fcode.erl

@@ -24,6 +24,8 @@
 -type var_name() :: string().
 -type sophia_name() :: [string()].
 
+-type state_reg() :: pos_integer().
+
 -type builtin() :: atom().
 
 -type op() :: '+' | '-' | '*' | '/' | mod | '^' | '++' | '::' |
@@ -61,6 +63,8 @@
                | {funcall, fexpr(), [fexpr()]}  %% Call to unknown function
                | {closure, fun_name(), fexpr()}
                | {switch, fsplit()}
+               | {set_state, state_reg(), fexpr()}
+               | {get_state, state_reg()}
                %% The following (unapplied top-level functions/builtins and
                %% lambdas) are generated by the fcode compiler, but translated
                %% to closures by the lambda lifter.
@@ -910,6 +914,10 @@ op_builtins() ->
      crypto_ecverify_secp256k1, crypto_ecrecover_secp256k1
     ].
 
+builtin_to_fcode(set_state, [Val]) ->
+    {set_state, 1, Val};
+builtin_to_fcode(get_state, []) ->
+    {get_state, 1};
 builtin_to_fcode(require, [Cond, Msg]) ->
     make_if(Cond, {tuple, []}, {builtin, abort, [Msg]});
 builtin_to_fcode(chain_event, [Event]) ->
@@ -942,7 +950,7 @@ add_init_function(Env, Main, StateType, Funs0) ->
             InitName = {entrypoint, <<"init">>},
             InitFun  = #{ body := InitBody} = maps:get(InitName, Funs),
             Funs#{ InitName => InitFun#{ return => {tuple, []},
-                                         body   => {builtin, set_state, [InitBody]} } }
+                                         body   => builtin_to_fcode(set_state, [InitBody]) } }
     end.
 
 add_default_init_function(_Env, Main, StateType, Funs) ->
@@ -1072,6 +1080,8 @@ lambda_lift_expr(Expr) ->
         {op, Op, As}           -> {op, Op, lambda_lift_exprs(As)};
         {'let', X, A, B}       -> {'let', X, lambda_lift_expr(A), lambda_lift_expr(B)};
         {funcall, A, Bs}       -> {funcall, lambda_lift_expr(A), lambda_lift_exprs(Bs)};
+        {set_state, R, A}      -> {set_state, R, lambda_lift_expr(A)};
+        {get_state, _}         -> Expr;
         {switch, S}            -> {switch, lambda_lift_expr(S)};
         {split, Type, X, Alts} -> {split, Type, X, lambda_lift_exprs(Alts)};
         {nosplit, A}           -> {nosplit, lambda_lift_expr(A)};
@@ -1236,7 +1246,7 @@ resolve_var(Env, Q) -> resolve_fun(Env, Q).
 resolve_fun(#{ fun_env := Funs, builtins := Builtin }, Q) ->
     case {maps:get(Q, Funs, not_found), maps:get(Q, Builtin, not_found)} of
         {not_found, not_found} -> internal_error({unbound_variable, Q});
-        {_, {B, none}}         -> {builtin, B, []};
+        {_, {B, none}}         -> builtin_to_fcode(B, []);
         {_, {B, Ar}}           -> {builtin_u, B, Ar};
         {{Fun, Ar}, _}         -> {def_u, Fun, Ar}
     end.
@@ -1301,6 +1311,8 @@ free_vars(Expr) ->
         {op, _, As}          -> free_vars(As);
         {'let', X, A, B}     -> free_vars([A, {lam, [X], B}]);
         {funcall, A, Bs}     -> free_vars([A | Bs]);
+        {set_state, _, A}    -> free_vars(A);
+        {get_state, _}       -> [];
         {lam, Xs, B}         -> free_vars(B) -- lists:sort(Xs);
         {closure, _, A}      -> free_vars(A);
         {switch, A}          -> free_vars(A);
@@ -1330,6 +1342,8 @@ used_defs(Expr) ->
         {op, _, As}          -> used_defs(As);
         {'let', _, A, B}     -> used_defs([A, B]);
         {funcall, A, Bs}     -> used_defs([A | Bs]);
+        {set_state, _, A}    -> used_defs(A);
+        {get_state, _}       -> [];
         {lam, _, B}          -> used_defs(B);
         {closure, F, A}      -> lists:umerge([F], used_defs(A));
         {switch, A}          -> used_defs(A);
@@ -1372,6 +1386,8 @@ rename(Ren, Expr) ->
         {set_proj, R, I, E}   -> {set_proj, rename(Ren, R), I, rename(Ren, E)};
         {op, Op, Es}          -> {op, Op, [rename(Ren, E) || E <- Es]};
         {funcall, Fun, Es}    -> {funcall, rename(Ren, Fun), [rename(Ren, E) || E <- Es]};
+        {set_state, R, E}     -> {set_state, R, rename(Ren, E)};
+        {get_state, _}        -> Expr;
         {closure, F, Env}     -> {closure, F, rename(Ren, Env)};
         {switch, Split}       -> {switch, rename_split(Ren, Split)};
         {lam, Xs, B} ->
@@ -1595,6 +1611,10 @@ pp_fexpr({remote, ArgsT, RetT, Ct, Fun, As}) ->
     pp_call(pp_parens(pp_beside([pp_fexpr(Ct), pp_text("."), pp_fun_name(Fun), pp_text(" : "), pp_ftype({function, ArgsT, RetT})])), As);
 pp_fexpr({funcall, Fun, As}) ->
     pp_call(pp_fexpr(Fun), As);
+pp_fexpr({set_state, R, A}) ->
+    pp_call(pp_text("set_state"), [{lit, {int, R}}, A]);
+pp_fexpr({get_state, R}) ->
+    pp_call(pp_text("get_state"), [{lit, {int, R}}]);
 pp_fexpr({switch, Split}) -> pp_split(Split).
 
 pp_call(Fun, Args) ->

src/aeso_fcode_to_fate.erl

@@ -298,6 +298,12 @@ to_scode(Env, {remote, ArgsT, RetT, Ct, Fun, [Gas, Value | Args]}) ->
             call_to_scode(Env, Call, [Ct, Value, Gas | Args])
     end;
 
+to_scode(_Env, {get_state, Reg}) ->
+    [push(?s(Reg))];
+to_scode(Env, {set_state, Reg, Val}) ->
+    call_to_scode(Env, [aeb_fate_ops:store(?s(Reg), ?a),
+                        tuple(0)], [Val]);
+
 to_scode(Env, {closure, Fun, FVs}) ->
     to_scode(Env, {tuple, [{lit, {string, make_function_id(Fun)}}, FVs]});
 
@@ -420,11 +426,6 @@ call_to_scode(Env, CallCode, Args) ->
     [[to_scode(notail(Env), A) || A <- lists:reverse(Args)],
      CallCode].
 
-builtin_to_scode(_Env, get_state, []) ->
-    [push(?s(1))];
-builtin_to_scode(Env, set_state, [_] = Args) ->
-    call_to_scode(Env, [aeb_fate_ops:store(?s(1), ?a),
-                        tuple(0)], Args);
 builtin_to_scode(Env, chain_event, Args) ->
     call_to_scode(Env, [erlang:apply(aeb_fate_ops, log, lists:duplicate(length(Args), ?a)),
                         tuple(0)], Args);