%%%-------------------------------------------------------------------
%%% @author Ulf Norell
%%% @copyright (C) 2025, QPQ AG
%%% @copyright (C) 2019, Aeternity Anstalt
%%% @doc
%%%     Fate backend for Sophia compiler
%%% @end
%%%-------------------------------------------------------------------
-module(so_fcode_to_fate).
-vsn("9.0.0").

-export([compile/3, compile/4, term_to_fate/1, term_to_fate/2]).

-ifdef(TEST).
-export([optimize_fun/4, to_basic_blocks/1]).
-endif.

%% -- Preamble ---------------------------------------------------------------

-type scode()  :: [sinstr()].
-type sinstr() :: {switch, arg(), stype(), [maybe_scode()], maybe_scode()}  %% last arg is catch-all
                | switch_body
                | loop
                | tuple() | atom().    %% FATE instruction

-type arg() :: tuple(). %% Not exported: gmb_fate_ops:fate_arg().

%% Annotated scode
-type scode_a()  :: [sinstr_a()].
-type sinstr_a() :: {switch, arg(), stype(), [maybe_scode_a()], maybe_scode_a()}  %% last arg is catch-all
                  | {i, ann(), tuple()}.    %% FATE instruction

-type ann() :: #{ live_in := vars(), live_out := vars() }.
-type var() :: {var, integer()}.
-type vars() :: ordsets:ordset(var()).

-type stype()         :: tuple | boolean | {variant, [non_neg_integer()]}.
-type maybe_scode()   :: missing | scode().
-type maybe_scode_a() :: missing | scode_a().

-define(TODO(What), error({todo, ?FILE, ?LINE, ?FUNCTION_NAME, What})).

-define(i(X), {immediate, X}).
-define(a,    {stack, 0}).
-define(s(N), {store, N}).
-define(void, {var, 9999}).

-record(env, { contract,
               vars              = [],
               locals            = [],
               current_function,
               tailpos           = true,
               child_contracts   = #{},
               saved_fresh_names = #{},
               options           = [],
               debug_info        = false }).

%% -- Debugging --------------------------------------------------------------

is_debug(Tag, Options) ->
    Tags = proplists:get_value(debug, Options, []),
    Tags == all orelse lists:member(Tag, Tags).

-define(debug(Tag, Options, Fmt, Args),
        debug(Tag, Options, fun() -> io:format(Fmt, Args) end)).

debug(Tag, Options, Fun) ->
    case is_debug(Tag, Options) of
        true  -> Fun();
        false -> ok
    end.

-dialyzer({nowarn_function, [code_error/1]}).
code_error(Err) ->
    Pos = so_errors:pos(0, 0),
    Msg = lists:flatten(io_lib:format("Unknown error: ~p\n", [Err])),
    so_errors:throw(so_errors:new(code_error, Pos, Msg)).

%% -- Main -------------------------------------------------------------------

%% @doc Main entry point.
compile(FCode, SavedFreshNames, Options) ->
    compile(#{}, FCode, SavedFreshNames, Options).
compile(ChildContracts, FCode, SavedFreshNames, Options) ->
    #{ contract_name := ContractName,
       functions     := Functions } = FCode,
    SFuns  = functions_to_scode(ChildContracts, ContractName, Functions, SavedFreshNames, Options),
    SFuns1 = optimize_scode(SFuns, Options),
    FateCode = to_basic_blocks(SFuns1),
    ?debug(compile, Options, "~s\n", [gmb_fate_asm:pp(FateCode)]),
    case proplists:get_value(include_child_contract_symbols, Options, false) of
        false -> FateCode;
        true  -> add_child_symbols(ChildContracts, FateCode)
    end.

make_function_id(X) ->
    gmb_fate_code:symbol_identifier(make_function_name(X)).

make_function_name(event)              -> <<"Chain.event">>;
make_function_name({entrypoint, Name}) -> Name;
make_function_name({local_fun, Xs})    -> list_to_binary("." ++ string:join(Xs, ".")).

add_child_symbols(ChildContracts, FateCode) ->
    Funs = lists:flatten([ maps:keys(ChildFuns) || {_, #{functions := ChildFuns}} <- maps:to_list(ChildContracts) ]),
    Symbols = maps:from_list([ {make_function_id(FName), make_function_name(FName)} || FName <- Funs ]),
    gmb_fate_code:update_symbols(FateCode, Symbols).

functions_to_scode(ChildContracts, ContractName, Functions, SavedFreshNames, Options) ->
    FunNames = maps:keys(Functions),
    maps:from_list(
        [ {make_function_name(Name), function_to_scode(ChildContracts, ContractName, FunNames, Name, Attrs, Args, Body, Type, SavedFreshNames, Options)}
        || {Name, #{args   := Args,
                    body   := Body,
                    attrs  := Attrs,
                    return := Type}} <- maps:to_list(Functions)]).

function_to_scode(ChildContracts, ContractName, Functions, Name, Attrs0, Args, Body, ResType, SavedFreshNames, Options) ->
    {ArgTypes, ResType1} = typesig_to_scode(Args, ResType),
    Attrs = [ A || A <- Attrs0, A == private orelse A == payable ],
    Env = init_env(ChildContracts, ContractName, Functions, Name, Args, SavedFreshNames, Options),
    ArgsNames = [ X || {X, _} <- lists:reverse(Env#env.vars) ],

    %% DBG_LOC is added before the function body to make it possible to break
    %% at the function signature
    SCode = to_scode(Env, Body),
    DbgSCode = dbg_contract(Env) ++ dbg_loc(Env, Attrs0) ++ dbg_scoped_vars(Env, ArgsNames, SCode),
    {Attrs, {ArgTypes, ResType1}, DbgSCode}.

-define(tvars, '$tvars').

typesig_to_scode(Args, Res) ->
    put(?tvars, {0, #{}}),
    R = {[type_to_scode(T) || {_, T} <- Args], type_to_scode(Res)},
    erase(?tvars),
    R.

type_to_scode(integer)         -> integer;
type_to_scode(boolean)         -> boolean;
type_to_scode(string)          -> string;
type_to_scode(address)         -> address;
type_to_scode({bytes, N})      -> {bytes, N};
type_to_scode(contract)        -> contract;
type_to_scode(name)            -> name;
type_to_scode(channel)         -> channel;
type_to_scode(bits)            -> bits;
type_to_scode(any)             -> any;
type_to_scode({variant, Cons}) -> {variant, [{tuple, types_to_scode(Con)} || Con <- Cons]};
type_to_scode({list, Type})    -> {list, type_to_scode(Type)};
type_to_scode({tuple, [Type]}) -> type_to_scode(Type);
type_to_scode({tuple, Types})  -> {tuple, types_to_scode(Types)};
type_to_scode({map, Key, Val}) -> {map, type_to_scode(Key), type_to_scode(Val)};
type_to_scode({function, _Args, _Res}) -> {tuple, [string, any]};
type_to_scode({tvar, X}) ->
    {I, Vars} = get(?tvars),
    case maps:get(X, Vars, false) of
        false ->
            put(?tvars, {I + 1, Vars#{ X => I }}),
            {tvar, I};
        J -> {tvar, J}
    end;
type_to_scode(L) when is_list(L) -> {tuple, types_to_scode(L)}.


types_to_scode(Ts) -> lists:map(fun type_to_scode/1, Ts).

%% -- Phase I ----------------------------------------------------------------
%%  Icode to structured assembly

%% -- Environment functions --

init_env(ChildContracts, ContractName, FunNames, Name, Args, SavedFreshNames, Options) ->
    #env{ vars              = [ {X, {arg, I}} || {I, {X, _}} <- with_ixs(Args) ],
          contract          = ContractName,
          child_contracts   = ChildContracts,
          locals            = FunNames,
          current_function  = Name,
          options           = Options,
          tailpos           = true,
          saved_fresh_names = SavedFreshNames,
          debug_info        = proplists:get_value(debug_info, Options, false) }.

next_var(#env{ vars = Vars }) ->
    1 + lists:max([-1 | [J || {_, {var, J}} <- Vars]]).

bind_var(Name, Var, Env = #env{ vars = Vars }) ->
    Env#env{ vars = [{Name, Var} | Vars] }.

bind_local(Name, Env) ->
    I = next_var(Env),
    {I, bind_var(Name, {var, I}, Env)}.

notail(Env) -> Env#env{ tailpos = false }.

lookup_var(#env{vars = Vars}, X) ->
    case lists:keyfind(X, 1, Vars) of
        {_, Var} -> Var;
        false    -> code_error({unbound_variable, X, Vars})
    end.

%% -- The compiler --

serialize_contract_code(Env, C) ->
    Cache = case get(contract_code_cache) of
                undefined -> put(contract_code_cache, #{}), #{};
                Res       -> Res
            end,
    case maps:get(C, Cache, none) of
        none ->
            Options = Env#env.options,
            SavedFreshNames = Env#env.saved_fresh_names,
            FCode = maps:get(C, Env#env.child_contracts),
            FateCode = compile(Env#env.child_contracts, FCode, SavedFreshNames, Options),
            ByteCode = gmb_fate_code:serialize(FateCode, []),
            {ok, Version} = so_compiler:version(),
            OriginalSourceCode = proplists:get_value(original_src, Options, ""),
            Code = #{byte_code => ByteCode,
                     compiler_version => Version,
                     source_hash => crypto:hash(sha256, OriginalSourceCode ++ [0] ++ C),
                     type_info => [],
                     abi_version => gmb_fate_abi:abi_version(),
                     payable => maps:get(payable, FCode)
                   },
            Serialized = gmser_contract_code:serialize(Code),
            put(contract_code_cache, maps:put(C, Serialized, Cache)),
            Serialized;
        Serialized -> Serialized
    end.

lit_to_fate(Env, L) ->
    case L of
        {int, N}             -> gmb_fate_data:make_integer(N);
        {string, S}          -> gmb_fate_data:make_string(S);
        {bytes, B}           -> gmb_fate_data:make_bytes(B);
        {bool, B}            -> gmb_fate_data:make_boolean(B);
        {account_pubkey, K}  -> gmb_fate_data:make_address(K);
        {signature, S}       -> gmb_fate_data:make_bytes(S);
        {contract_pubkey, K} -> gmb_fate_data:make_contract(K);
        {contract_code, C}   -> gmb_fate_data:make_contract_bytearray(serialize_contract_code(Env, C));
        {typerep, T}         -> gmb_fate_data:make_typerep(type_to_scode(T))
     end.

term_to_fate(E) -> term_to_fate(#env{}, #{}, E).
term_to_fate(GlobEnv, E) -> term_to_fate(GlobEnv, #{}, E).

term_to_fate(GlobEnv, _Env, {lit, _, L}) ->
    lit_to_fate(GlobEnv, L);
%% negative literals are parsed as 0 - N
term_to_fate(_GlobEnv, _Env, {op, _, '-', [{lit, _, {int, 0}}, {lit, _, {int, N}}]}) ->
    gmb_fate_data:make_integer(-N);
term_to_fate(_GlobEnv, _Env, {nil, _}) ->
    gmb_fate_data:make_list([]);
term_to_fate(GlobEnv, Env, {op, _, '::', [Hd, Tl]}) ->
    %% The Tl will translate into a list, because FATE lists are just lists
    [term_to_fate(GlobEnv, Env, Hd) | term_to_fate(GlobEnv, Env, Tl)];
term_to_fate(GlobEnv, Env, {tuple, _, As}) ->
    gmb_fate_data:make_tuple(list_to_tuple([ term_to_fate(GlobEnv, Env, A) || A<-As]));
term_to_fate(GlobEnv, Env, {con, _, Ar, I, As}) ->
    FateAs = [ term_to_fate(GlobEnv, Env, A) || A <- As ],
    gmb_fate_data:make_variant(Ar, I, list_to_tuple(FateAs));
term_to_fate(_GlobEnv, _Env, {builtin, _, bits_all, []}) ->
    gmb_fate_data:make_bits(-1);
term_to_fate(_GlobEnv, _Env, {builtin, _, bits_none, []}) ->
    gmb_fate_data:make_bits(0);
term_to_fate(GlobEnv, _Env, {op, _, bits_set, [B, I]}) ->
    {bits, N} = term_to_fate(GlobEnv, B),
    J         = term_to_fate(GlobEnv, I),
    {bits, N bor (1 bsl J)};
term_to_fate(GlobEnv, _Env, {op, _, bits_clear, [B, I]}) ->
    {bits, N} = term_to_fate(GlobEnv, B),
    J         = term_to_fate(GlobEnv, I),
    {bits, N band bnot (1 bsl J)};
term_to_fate(GlobEnv, Env, {'let', _, X, E, Body}) ->
    Env1 = Env#{ X => term_to_fate(GlobEnv, Env, E) },
    term_to_fate(GlobEnv, Env1, Body);
term_to_fate(_GlobEnv, Env, {var, _, X}) ->
    case maps:get(X, Env, undefined) of
        undefined -> throw(not_a_fate_value);
        V         -> V
    end;
term_to_fate(_GlobEnv, _Env, {builtin, _, map_empty, []}) ->
    gmb_fate_data:make_map(#{});
term_to_fate(GlobEnv, Env, {op, _, map_set, [M, K, V]}) ->
    Map = term_to_fate(GlobEnv, Env, M),
    Map#{term_to_fate(GlobEnv, Env, K) => term_to_fate(GlobEnv, Env, V)};
term_to_fate(GlobEnv, Env, {builtin, _, bytes_to_any_size, [Bs]}) ->
    term_to_fate(GlobEnv, Env, Bs);
term_to_fate(_GlobEnv, _Env, _) ->
    throw(not_a_fate_value).

to_scode(Env, T) ->
    try term_to_fate(Env, T) of
        V ->
            FAnn = element(2, T),
            [dbg_loc(Env, FAnn), push(?i(V))]
    catch throw:not_a_fate_value ->
        to_scode1(Env, T)
    end.

to_scode1(Env, {lit, Ann, L}) ->
    [ dbg_loc(Env, Ann), push(?i(lit_to_fate(Env, L))) ];

to_scode1(Env, {nil, Ann}) ->
    [ dbg_loc(Env, Ann), gmb_fate_ops:nil(?a) ];

to_scode1(Env, {var, Ann, X}) ->
    [ dbg_loc(Env, Ann), push(lookup_var(Env, X)) ];

to_scode1(Env, {con, Ann, Ar, I, As}) ->
    N = length(As),
    [ dbg_loc(Env, Ann),
      [to_scode(notail(Env), A) || A <- As],
      gmb_fate_ops:variant(?a, ?i(Ar), ?i(I), ?i(N)) ];

to_scode1(Env, {tuple, Ann, As}) ->
    N = length(As),
    [ dbg_loc(Env, Ann),
      [ to_scode(notail(Env), A) || A <- As ],
      tuple(N) ];

to_scode1(Env, {proj, Ann, E, I}) ->
    [ dbg_loc(Env, Ann),
      to_scode(notail(Env), E),
      gmb_fate_ops:element_op(?a, ?i(I), ?a) ];

to_scode1(Env, {set_proj, Ann, R, I, E}) ->
    [ dbg_loc(Env, Ann),
      to_scode(notail(Env), E),
      to_scode(notail(Env), R),
      gmb_fate_ops:setelement(?a, ?i(I), ?a, ?a) ];

to_scode1(Env, {op, Ann, Op, Args}) ->
    [ dbg_loc(Env, Ann) | call_to_scode(Env, op_to_scode(Op), Args) ];

to_scode1(Env, {'let', Ann, X, {var, _, Y}, Body}) ->
    Env1 = bind_var(X, lookup_var(Env, Y), Env),
    [ dbg_loc(Env, Ann) | dbg_scoped_vars(Env1, [X], to_scode(Env1, Body)) ];
to_scode1(Env, {'let', Ann, X, Expr, Body}) ->
    {I, Env1} = bind_local(X, Env),
    SCode = [ to_scode(notail(Env), Expr),
              gmb_fate_ops:store({var, I}, {stack, 0}),
              to_scode(Env1, Body) ],
    [ dbg_loc(Env, Ann) | dbg_scoped_vars(Env1, [X], SCode) ];

to_scode1(Env = #env{ current_function = Fun, tailpos = true, debug_info = false }, {def, Ann, Fun, Args}) ->
    %% Tail-call to current function, f(e0..en). Compile to
    %%      [ let xi = ei ]
    %%      [ STORE argi xi ]
    %%      jump 0
    {Vars, Code, _Env} =
        lists:foldl(fun(Arg, {Is, Acc, Env1}) ->
                        {I, Env2} = bind_local("_", Env1),
                        ArgCode   = to_scode(notail(Env2), Arg),
                        Acc1 = [Acc, ArgCode,
                                gmb_fate_ops:store({var, I}, ?a)],
                        {[I | Is], Acc1, Env2}
                    end, {[], [], Env}, Args),
    [ dbg_loc(Env, Ann),
      Code,
      [ gmb_fate_ops:store({arg, I}, {var, J})
        || {I, J} <- lists:zip(lists:seq(0, length(Vars) - 1),
                               lists:reverse(Vars)) ],
      loop ];
to_scode1(Env, {def, Ann, Fun, Args}) ->
    FName = make_function_id(Fun),
    Lbl   = gmb_fate_data:make_string(FName),
    [ dbg_loc(Env, Ann) | call_to_scode(Env, local_call(Env, ?i(Lbl)), Args) ];
to_scode1(Env, {funcall, Ann, Fun, Args}) ->
    [ dbg_loc(Env, Ann) | call_to_scode(Env, [to_scode(Env, Fun), local_call(Env, ?a)], Args) ];

to_scode1(Env, {builtin, Ann, B, Args}) ->
    [ dbg_loc(Env, Ann) | builtin_to_scode(Env, B, Args) ];

to_scode1(Env, {remote, Ann, ArgsT, RetT, Ct, Fun, [Gas, Value, Protected | Args]}) ->
    Lbl = make_function_id(Fun),
    {ArgTypes, RetType0} = typesig_to_scode([{"_", T} || T <- ArgsT], RetT),
    ArgType = ?i(gmb_fate_data:make_typerep({tuple, ArgTypes})),
    RetType = ?i(gmb_fate_data:make_typerep(RetType0)),
    SCode = case Protected of
        {lit, _, {bool, false}} ->
            case Gas of
                {builtin, _, call_gas_left, _} ->
                    Call = gmb_fate_ops:call_r(?a, Lbl, ArgType, RetType, ?a),
                    call_to_scode(Env, Call, [Ct, Value | Args]);
                _ ->
                    Call = gmb_fate_ops:call_gr(?a, Lbl, ArgType, RetType, ?a, ?a),
                    call_to_scode(Env, Call, [Ct, Value, Gas | Args])
            end;
        {lit, _, {bool, true}} ->
            Call = gmb_fate_ops:call_pgr(?a, Lbl, ArgType, RetType, ?a, ?a, ?i(true)),
            call_to_scode(Env, Call, [Ct, Value, Gas | Args]);
        _ ->
            Call = gmb_fate_ops:call_pgr(?a, Lbl, ArgType, RetType, ?a, ?a, ?a),
            call_to_scode(Env, Call, [Ct, Value, Gas, Protected | Args])
    end,
    [ dbg_loc(Env, Ann) | SCode ];

to_scode1(Env, {get_state, Ann, Reg}) ->
    [ dbg_loc(Env, Ann), push(?s(Reg)) ];
to_scode1(Env, {set_state, Ann, Reg, Val}) ->
    [ dbg_loc(Env, Ann) | call_to_scode(Env, [{'STORE', ?s(Reg), ?a}, tuple(0)], [Val]) ];

to_scode1(Env, {closure, Ann, Fun, FVs}) ->
    [ to_scode(Env, {tuple, Ann, [{lit, Ann, {string, make_function_id(Fun)}}, FVs]}) ];

to_scode1(Env, {switch, Ann, Case}) ->
    [ dbg_loc(Env, Ann) | split_to_scode(Env, Case) ].

local_call( Env = #env{debug_info = false}, Fun) when Env#env.tailpos -> gmb_fate_ops:call_t(Fun);
local_call(_Env, Fun)                                                 -> gmb_fate_ops:call(Fun).

split_to_scode(Env, {nosplit, Renames, Expr}) ->
    [switch_body, dbg_scoped_vars(Env, Renames, to_scode(Env, Expr))];
split_to_scode(Env, {split, {tuple, _}, X, Alts}) ->
    {Def, Alts1} = catchall_to_scode(Env, X, Alts),
    Arg = lookup_var(Env, X),
    Alt = case [ {Xs, Split} || {'case', {tuple, Xs}, Split} <- Alts1 ] of
            []            -> missing;
            [{Xs, S} | _] ->
                {Code, Env1} = match_tuple(Env, Arg, Xs),
                [Code, split_to_scode(Env1, S)]
          end,
    case Def == missing andalso Alt /= missing of
       true  -> Alt;   % skip the switch if single tuple pattern
       false -> [{switch, Arg, tuple, [Alt], Def}]
    end;
split_to_scode(Env, {split, boolean, X, Alts}) ->
    {Def, Alts1} = catchall_to_scode(Env, X, Alts),
    GetAlt = fun(B) ->
                 case lists:keyfind({bool, B}, 2, Alts1) of
                     false          -> missing;
                     {'case', _, S} -> split_to_scode(Env, S)
                 end
             end,
    SAlts = [GetAlt(false), GetAlt(true)],
    Arg   = lookup_var(Env, X),
    [{switch, Arg, boolean, SAlts, Def}];
split_to_scode(Env, {split, {list, _}, X, Alts}) ->
    {Def, Alts1} = catchall_to_scode(Env, X, Alts),
    Arg = lookup_var(Env, X),
    GetAlt = fun(P) ->
                 case [C || C = {'case', Pat, _} <- Alts1, Pat == P orelse is_tuple(Pat) andalso element(1, Pat) == P] of
                     []      -> missing;
                     [{'case', nil, S} | _]           -> split_to_scode(Env, S);
                     [{'case', {'::', Y, Z}, S} | _] ->
                         {I, Env1} = bind_local(Y, Env),
                         {J, Env2} = bind_local(Z, Env1),
                         [gmb_fate_ops:hd({var, I}, Arg),
                          gmb_fate_ops:tl({var, J}, Arg),
                          split_to_scode(Env2, S)]
                 end
             end,
    SAlts = [GetAlt('::'), GetAlt(nil)],
    [gmb_fate_ops:is_nil(?a, Arg),
     {switch, ?a, boolean, SAlts, Def}];
split_to_scode(Env, {split, Type, X, Alts}) when Type == integer; Type == string ->
    {Def, Alts1} = catchall_to_scode(Env, X, Alts),
    literal_split_to_scode(Env, Type, lookup_var(Env, X), Alts1, Def);
split_to_scode(Env, {split, {variant, Cons}, X, Alts}) ->
    {Def, Alts1} = catchall_to_scode(Env, X, Alts),
    Arg = lookup_var(Env, X),
    GetAlt = fun(I) ->
                case [{Xs, S} || {'case', {con, _, J, Xs}, S} <- Alts1, I == J] of
                    [] -> missing;
                    [{Xs, S} | _] ->
                        {Code, Env1} = match_variant(Env, Arg, Xs),
                        [Code, split_to_scode(Env1, S)]
                end
             end,
    SType  = {variant, [length(Args) || Args <- Cons]},
    case {[GetAlt(I) || I <- lists:seq(0, length(Cons) - 1)], Def} of
        %% Skip the switch for single constructor datatypes (with no catchall)
        {[SAlt], missing} when SAlt /= missing -> SAlt;
        {SAlts, _} -> [{switch, Arg, SType, SAlts, Def}]
    end.

literal_split_to_scode(_Env, _Type, Arg, [], Def) ->
    {switch, Arg, boolean, [missing, missing], Def};
literal_split_to_scode(Env, Type, Arg, [{'case', Lit, Body} | Alts], Def) when Type == integer; Type == string ->
    True = split_to_scode(Env, Body),
    False =
        case Alts of
            [] -> missing;
            _  -> literal_split_to_scode(Env, Type, Arg, Alts, missing)
        end,
    SLit = case Lit of
               {int, N} -> N;
               {string, S} -> gmb_fate_data:make_string(S)
           end,
    [gmb_fate_ops:eq(?a, Arg, ?i(SLit)),
     {switch, ?a, boolean, [False, True], Def}].

catchall_to_scode(Env, X, Alts) -> catchall_to_scode(Env, X, Alts, []).

catchall_to_scode(Env, X, [{'case', {var, Y}, Split} | _], Acc) ->
    Env1 = bind_var(Y, lookup_var(Env, X), Env),
    {split_to_scode(Env1, Split), lists:reverse(Acc)};
catchall_to_scode(Env, X, [Alt | Alts], Acc) ->
    catchall_to_scode(Env, X, Alts, [Alt | Acc]);
catchall_to_scode(_, _, [], Acc) -> {missing, lists:reverse(Acc)}.

%% Tuple is in the accumulator. Arguments are the variable names.
match_tuple(Env, Arg, Xs) ->
    match_tuple(Env, 0, fun gmb_fate_ops:element_op/3, Arg, Xs).

match_variant(Env, Arg, Xs) ->
    Elem = fun(Dst, I, Val) -> gmb_fate_ops:variant_element(Dst, Val, I) end,
    match_tuple(Env, 0, Elem, Arg, Xs).

match_tuple(Env, I, Elem, Arg, ["_" | Xs]) ->
    match_tuple(Env, I + 1, Elem, Arg, Xs);
match_tuple(Env, I, Elem, Arg, [X | Xs]) ->
    {J,    Env1} = bind_local(X, Env),
    {Code, Env2} = match_tuple(Env1, I + 1, Elem, Arg, Xs),
    {[Elem({var, J}, ?i(I), Arg), Code], Env2};
match_tuple(Env, _, _, _, []) ->
    {[], Env}.

%% -- Builtins --

call_to_scode(Env, CallCode, Args) ->
    [[to_scode(notail(Env), A) || A <- lists:reverse(Args)],
     CallCode].

builtin_to_scode(Env, chain_event, Args) ->
    call_to_scode(Env, [erlang:apply(gmb_fate_ops, log, lists:duplicate(length(Args), ?a)),
                        tuple(0)], Args);
builtin_to_scode(_Env, map_empty, []) ->
    [gmb_fate_ops:map_empty(?a)];
builtin_to_scode(_Env, bits_none, []) ->
    [gmb_fate_ops:bits_none(?a)];
builtin_to_scode(_Env, bits_all, []) ->
    [gmb_fate_ops:bits_all(?a)];
builtin_to_scode(Env, bytes_to_int, [_] = Args) ->
    call_to_scode(Env, gmb_fate_ops:bytes_to_int(?a, ?a), Args);
builtin_to_scode(Env, bytes_to_str, [_] = Args) ->
    call_to_scode(Env, gmb_fate_ops:bytes_to_str(?a, ?a), Args);
builtin_to_scode(Env, bytes_concat, [_, _] = Args) ->
    call_to_scode(Env, gmb_fate_ops:bytes_concat(?a, ?a, ?a), Args);
builtin_to_scode(Env, bytes_split, [_, _] = Args) ->
    call_to_scode(Env, gmb_fate_ops:bytes_split(?a, ?a, ?a), Args);
builtin_to_scode(Env, bytes_split_any, [_, _] = Args) ->
    call_to_scode(Env, gmb_fate_ops:bytes_split_any(?a, ?a, ?a), Args);
builtin_to_scode(Env, bytes_to_fixed_size, [_, _] = Args) ->
    call_to_scode(Env, gmb_fate_ops:bytes_to_fixed_size(?a, ?a, ?a), Args);
builtin_to_scode(Env, bytes_to_any_size, [A]) ->
    [to_scode(Env, A)]; %% no_op!
builtin_to_scode(Env, bytes_size, [_] = Args) ->
    call_to_scode(Env, gmb_fate_ops:bytes_size(?a, ?a), Args);
builtin_to_scode(Env, abort, [_] = Args) ->
    call_to_scode(Env, gmb_fate_ops:abort(?a), Args);
builtin_to_scode(Env, exit, [_] = Args) ->
    call_to_scode(Env, gmb_fate_ops:exit(?a), Args);
builtin_to_scode(Env, chain_spend, [_, _] = Args) ->
    call_to_scode(Env, [gmb_fate_ops:spend(?a, ?a),
                        tuple(0)], Args);
builtin_to_scode(Env, chain_balance, [_] = Args) ->
    call_to_scode(Env, gmb_fate_ops:balance_other(?a, ?a), Args);
builtin_to_scode(Env, chain_block_hash, [_] = Args) ->
    call_to_scode(Env, gmb_fate_ops:blockhash(?a, ?a), Args);
builtin_to_scode(_Env, chain_coinbase, []) ->
    [gmb_fate_ops:beneficiary(?a)];
builtin_to_scode(_Env, chain_timestamp, []) ->
    [gmb_fate_ops:timestamp(?a)];
builtin_to_scode(_Env, chain_block_height, []) ->
    [gmb_fate_ops:generation(?a)];
builtin_to_scode(_Env, chain_difficulty, []) ->
    [gmb_fate_ops:difficulty(?a)];
builtin_to_scode(_Env, chain_gas_limit, []) ->
    [gmb_fate_ops:gaslimit(?a)];
builtin_to_scode(_Env, chain_network_id, []) ->
    [gmb_fate_ops:network_id(?a)];
builtin_to_scode(_Env, contract_balance, []) ->
    [gmb_fate_ops:balance(?a)];
builtin_to_scode(_Env, contract_address, []) ->
    [gmb_fate_ops:address(?a)];
builtin_to_scode(_Env, contract_creator, []) ->
    [gmb_fate_ops:contract_creator(?a)];
builtin_to_scode(_Env, call_origin, []) ->
    [gmb_fate_ops:origin(?a)];
builtin_to_scode(_Env, call_caller, []) ->
    [gmb_fate_ops:caller(?a)];
builtin_to_scode(_Env, call_value, []) ->
    [gmb_fate_ops:call_value(?a)];
builtin_to_scode(_Env, call_gas_price, []) ->
    [gmb_fate_ops:gasprice(?a)];
builtin_to_scode(_Env, call_fee, []) ->
    [gmb_fate_ops:fee(?a)];
builtin_to_scode(_Env, call_gas_left, []) ->
    [gmb_fate_ops:gas(?a)];
builtin_to_scode(Env, address_is_contract, [_] = Args) ->
    call_to_scode(Env, gmb_fate_ops:is_contract(?a, ?a), Args);
builtin_to_scode(Env, address_is_payable, [_] = Args) ->
    call_to_scode(Env, gmb_fate_ops:is_payable(?a, ?a), Args);
builtin_to_scode(Env, aens_resolve, [_Name, _Key, _Type] = Args) ->
    call_to_scode(Env, gmb_fate_ops:aens_resolve(?a, ?a, ?a, ?a), Args);
builtin_to_scode(Env, aens_preclaim, [_Sign, _Account, _Hash] = Args) ->
    call_to_scode(Env, [gmb_fate_ops:aens_preclaim(?a, ?a, ?a),
                        tuple(0)], Args);
builtin_to_scode(Env, aens_claim, [_Sign, _Account, _NameString, _Salt, _NameFee] = Args) ->
    call_to_scode(Env, [gmb_fate_ops:aens_claim(?a, ?a, ?a, ?a, ?a),
                        tuple(0)], Args);
builtin_to_scode(Env, aens_transfer, [_Sign, _From, _To, _Name] = Args) ->
    call_to_scode(Env, [gmb_fate_ops:aens_transfer(?a, ?a, ?a, ?a),
                        tuple(0)], Args);
builtin_to_scode(Env, aens_revoke, [_Sign, _Account, _Name] = Args) ->
    call_to_scode(Env, [gmb_fate_ops:aens_revoke(?a, ?a, ?a),
                        tuple(0)], Args);
builtin_to_scode(Env, aens_update, [_Sign, _Account, _NameString, _TTL, _ClientTTL, _Pointers] = Args) ->
    call_to_scode(Env, [gmb_fate_ops:aens_update(?a, ?a, ?a, ?a, ?a, ?a),
                        tuple(0)], Args);
builtin_to_scode(Env, aens_lookup, [_Name] = Args) ->
    call_to_scode(Env, gmb_fate_ops:aens_lookup(?a, ?a), Args);
builtin_to_scode(_Env, auth_tx_hash, []) ->
    [gmb_fate_ops:auth_tx_hash(?a)];
builtin_to_scode(_Env, auth_tx, []) ->
    [gmb_fate_ops:auth_tx(?a)];
builtin_to_scode(Env, chain_bytecode_hash, [_Addr] = Args) ->
    call_to_scode(Env, gmb_fate_ops:bytecode_hash(?a, ?a), Args);
builtin_to_scode(Env, chain_clone,
                 [InitArgsT, GasCap, Value, Prot, Contract | InitArgs]) ->
    case GasCap of
        {builtin, _, call_gas_left, _} ->
            call_to_scode(Env, gmb_fate_ops:clone(?a, ?a, ?a, ?a),
                          [Contract, InitArgsT, Value, Prot | InitArgs]
                         );
        _ ->
            call_to_scode(Env, gmb_fate_ops:clone_g(?a, ?a, ?a, ?a, ?a),
                          [Contract, InitArgsT, Value, GasCap, Prot | InitArgs]
                         )
    end;

builtin_to_scode(Env, chain_create,
                 [ Code, InitArgsT, Value | InitArgs]) ->
    call_to_scode(Env, gmb_fate_ops:create(?a, ?a, ?a),
                  [Code, InitArgsT, Value | InitArgs]
                 ).

%% -- Operators --

op_to_scode('+')               -> gmb_fate_ops:add(?a, ?a, ?a);
op_to_scode('-')               -> gmb_fate_ops:sub(?a, ?a, ?a);
op_to_scode('*')               -> gmb_fate_ops:mul(?a, ?a, ?a);
op_to_scode('/')               -> gmb_fate_ops:divide(?a, ?a, ?a);
op_to_scode(mod)               -> gmb_fate_ops:modulo(?a, ?a, ?a);
op_to_scode('^')               -> gmb_fate_ops:pow(?a, ?a, ?a);
op_to_scode('++')              -> gmb_fate_ops:append(?a, ?a, ?a);
op_to_scode('::')              -> gmb_fate_ops:cons(?a, ?a, ?a);
op_to_scode('<')               -> gmb_fate_ops:lt(?a, ?a, ?a);
op_to_scode('>')               -> gmb_fate_ops:gt(?a, ?a, ?a);
op_to_scode('=<')              -> gmb_fate_ops:elt(?a, ?a, ?a);
op_to_scode('>=')              -> gmb_fate_ops:egt(?a, ?a, ?a);
op_to_scode('==')              -> gmb_fate_ops:eq(?a, ?a, ?a);
op_to_scode('!=')              -> gmb_fate_ops:neq(?a, ?a, ?a);
op_to_scode('!')               -> gmb_fate_ops:not_op(?a, ?a);
op_to_scode('bnot')            -> gmb_fate_ops:bin_not(?a, ?a);
op_to_scode('band')            -> gmb_fate_ops:bin_and(?a, ?a, ?a);
op_to_scode('bor')             -> gmb_fate_ops:bin_or(?a, ?a, ?a);
op_to_scode('bxor')            -> gmb_fate_ops:bin_xor(?a, ?a, ?a);
op_to_scode('<<')              -> gmb_fate_ops:bin_sl(?a, ?a, ?a);
op_to_scode('>>')              -> gmb_fate_ops:bin_sr(?a, ?a, ?a);
op_to_scode(map_get)           -> gmb_fate_ops:map_lookup(?a, ?a, ?a);
op_to_scode(map_get_d)         -> gmb_fate_ops:map_lookup(?a, ?a, ?a, ?a);
op_to_scode(map_set)           -> gmb_fate_ops:map_update(?a, ?a, ?a, ?a);
op_to_scode(map_from_list)     -> gmb_fate_ops:map_from_list(?a, ?a);
op_to_scode(map_to_list)       -> gmb_fate_ops:map_to_list(?a, ?a);
op_to_scode(map_delete)        -> gmb_fate_ops:map_delete(?a, ?a, ?a);
op_to_scode(map_member)        -> gmb_fate_ops:map_member(?a, ?a, ?a);
op_to_scode(map_size)          -> gmb_fate_ops:map_size_(?a, ?a);
op_to_scode(stringinternal_length)    -> gmb_fate_ops:str_length(?a, ?a);
op_to_scode(stringinternal_concat)    -> gmb_fate_ops:str_join(?a, ?a, ?a);
op_to_scode(stringinternal_to_bytes)  -> gmb_fate_ops:str_to_bytes(?a, ?a);
op_to_scode(stringinternal_to_list)   -> gmb_fate_ops:str_to_list(?a, ?a);
op_to_scode(stringinternal_from_list) -> gmb_fate_ops:str_from_list(?a, ?a);
op_to_scode(stringinternal_to_lower)  -> gmb_fate_ops:str_to_lower(?a, ?a);
op_to_scode(stringinternal_to_upper)  -> gmb_fate_ops:str_to_upper(?a, ?a);
op_to_scode(char_to_int)       -> gmb_fate_ops:char_to_int(?a, ?a);
op_to_scode(char_from_int)     -> gmb_fate_ops:char_from_int(?a, ?a);
op_to_scode(bits_set)          -> gmb_fate_ops:bits_set(?a, ?a, ?a);
op_to_scode(bits_clear)        -> gmb_fate_ops:bits_clear(?a, ?a, ?a);
op_to_scode(bits_test)         -> gmb_fate_ops:bits_test(?a, ?a, ?a);
op_to_scode(bits_sum)          -> gmb_fate_ops:bits_sum(?a, ?a);
op_to_scode(bits_intersection) -> gmb_fate_ops:bits_and(?a, ?a, ?a);
op_to_scode(bits_union)        -> gmb_fate_ops:bits_or(?a, ?a, ?a);
op_to_scode(bits_difference)   -> gmb_fate_ops:bits_diff(?a, ?a, ?a);
op_to_scode(address_to_str)    -> gmb_fate_ops:addr_to_str(?a, ?a);
op_to_scode(address_to_bytes)  -> gmb_fate_ops:addr_to_bytes(?a, ?a);
op_to_scode(int_to_str)        -> gmb_fate_ops:int_to_str(?a, ?a);
op_to_scode(int_to_bytes)      -> gmb_fate_ops:int_to_bytes(?a, ?a, ?a);
op_to_scode(int_mulmod)        -> gmb_fate_ops:mulmod(?a, ?a, ?a, ?a);
op_to_scode(contract_to_address)         -> gmb_fate_ops:contract_to_address(?a, ?a);
op_to_scode(address_to_contract)         -> gmb_fate_ops:address_to_contract(?a, ?a);
op_to_scode(crypto_verify_sig)           -> gmb_fate_ops:verify_sig(?a, ?a, ?a, ?a);
op_to_scode(crypto_verify_sig_secp256k1) -> gmb_fate_ops:verify_sig_secp256k1(?a, ?a, ?a, ?a);
op_to_scode(crypto_ecverify_secp256k1)   -> gmb_fate_ops:ecverify_secp256k1(?a, ?a, ?a, ?a);
op_to_scode(crypto_ecrecover_secp256k1)  -> gmb_fate_ops:ecrecover_secp256k1(?a, ?a, ?a);
op_to_scode(crypto_sha3)                 -> gmb_fate_ops:sha3(?a, ?a);
op_to_scode(crypto_sha256)               -> gmb_fate_ops:sha256(?a, ?a);
op_to_scode(crypto_blake2b)              -> gmb_fate_ops:blake2b(?a, ?a);
op_to_scode(crypto_poseidon)             -> gmb_fate_ops:poseidon(?a, ?a, ?a);
op_to_scode(stringinternal_sha3)         -> gmb_fate_ops:sha3(?a, ?a);
op_to_scode(stringinternal_sha256)       -> gmb_fate_ops:sha256(?a, ?a);
op_to_scode(stringinternal_blake2b)      -> gmb_fate_ops:blake2b(?a, ?a);
op_to_scode(mcl_bls12_381_g1_neg)        -> gmb_fate_ops:bls12_381_g1_neg(?a, ?a);
op_to_scode(mcl_bls12_381_g1_norm)       -> gmb_fate_ops:bls12_381_g1_norm(?a, ?a);
op_to_scode(mcl_bls12_381_g1_valid)      -> gmb_fate_ops:bls12_381_g1_valid(?a, ?a);
op_to_scode(mcl_bls12_381_g1_is_zero)    -> gmb_fate_ops:bls12_381_g1_is_zero(?a, ?a);
op_to_scode(mcl_bls12_381_g1_add)        -> gmb_fate_ops:bls12_381_g1_add(?a, ?a, ?a);
op_to_scode(mcl_bls12_381_g1_mul)        -> gmb_fate_ops:bls12_381_g1_mul(?a, ?a, ?a);
op_to_scode(mcl_bls12_381_g2_neg)        -> gmb_fate_ops:bls12_381_g2_neg(?a, ?a);
op_to_scode(mcl_bls12_381_g2_norm)       -> gmb_fate_ops:bls12_381_g2_norm(?a, ?a);
op_to_scode(mcl_bls12_381_g2_valid)      -> gmb_fate_ops:bls12_381_g2_valid(?a, ?a);
op_to_scode(mcl_bls12_381_g2_is_zero)    -> gmb_fate_ops:bls12_381_g2_is_zero(?a, ?a);
op_to_scode(mcl_bls12_381_g2_add)        -> gmb_fate_ops:bls12_381_g2_add(?a, ?a, ?a);
op_to_scode(mcl_bls12_381_g2_mul)        -> gmb_fate_ops:bls12_381_g2_mul(?a, ?a, ?a);
op_to_scode(mcl_bls12_381_gt_inv)        -> gmb_fate_ops:bls12_381_gt_inv(?a, ?a);
op_to_scode(mcl_bls12_381_gt_add)        -> gmb_fate_ops:bls12_381_gt_add(?a, ?a, ?a);
op_to_scode(mcl_bls12_381_gt_mul)        -> gmb_fate_ops:bls12_381_gt_mul(?a, ?a, ?a);
op_to_scode(mcl_bls12_381_gt_pow)        -> gmb_fate_ops:bls12_381_gt_pow(?a, ?a, ?a);
op_to_scode(mcl_bls12_381_gt_is_one)     -> gmb_fate_ops:bls12_381_gt_is_one(?a, ?a);
op_to_scode(mcl_bls12_381_pairing)       -> gmb_fate_ops:bls12_381_pairing(?a, ?a, ?a);
op_to_scode(mcl_bls12_381_miller_loop)   -> gmb_fate_ops:bls12_381_miller_loop(?a, ?a, ?a);
op_to_scode(mcl_bls12_381_final_exp)     -> gmb_fate_ops:bls12_381_final_exp(?a, ?a);
op_to_scode(mcl_bls12_381_int_to_fr)     -> gmb_fate_ops:bls12_381_int_to_fr(?a, ?a);
op_to_scode(mcl_bls12_381_int_to_fp)     -> gmb_fate_ops:bls12_381_int_to_fp(?a, ?a);
op_to_scode(mcl_bls12_381_fr_to_int)     -> gmb_fate_ops:bls12_381_fr_to_int(?a, ?a);
op_to_scode(mcl_bls12_381_fp_to_int)     -> gmb_fate_ops:bls12_381_fp_to_int(?a, ?a).

%% PUSH and STORE ?a are the same, so we use STORE to make optimizations
%% easier, and specialize to PUSH (which is cheaper) at the end.
push(A) -> {'STORE', ?a, A}.

tuple(0) -> push(?i({tuple, {}}));
tuple(N) -> gmb_fate_ops:tuple(?a, N).

%% -- Debug info functions --

dbg_contract(#env{debug_info = false}) ->
    [];
dbg_contract(#env{contract = Contract}) ->
    [{'DBG_CONTRACT', {immediate, Contract}}].

dbg_loc(#env{debug_info = false}, _) ->
    [];
dbg_loc(_Env, Ann) ->
    File = case proplists:get_value(file, Ann, no_file) of
                no_file -> "";
                F       -> F
            end,
    Line = proplists:get_value(line, Ann, undefined),
    case Line of
        undefined -> [];
        _         -> [{'DBG_LOC', {immediate, File}, {immediate, Line}}]
    end.

dbg_scoped_vars(#env{debug_info = false}, _, SCode) ->
    SCode;
dbg_scoped_vars(_Env, [], SCode) ->
    SCode;
dbg_scoped_vars(Env, [{SavedVarName, Var} | Rest], SCode) ->
    dbg_scoped_vars(Env, Rest, dbg_scoped_var(Env, SavedVarName, Var, SCode));
dbg_scoped_vars(Env = #env{saved_fresh_names = SavedFreshNames}, [Var | Rest], SCode) ->
    SavedVarName = maps:get(Var, SavedFreshNames, Var),
    dbg_scoped_vars(Env, Rest, dbg_scoped_var(Env, SavedVarName, Var, SCode)).

dbg_scoped_var(Env, SavedVarName, Var, SCode) ->
    case SavedVarName == "_" orelse is_fresh_name(SavedVarName) of
        true ->
            SCode;
        false ->
            Register = lookup_var(Env, Var),
            Def      = [{'DBG_DEF',   {immediate, SavedVarName}, Register}],
            Undef    = [{'DBG_UNDEF', {immediate, SavedVarName}, Register}],
            Def ++ dbg_undef(Undef, SCode)
    end.

is_fresh_name([$% | _]) ->
    true;
is_fresh_name(_) ->
    false.

dbg_undef(_Undef, missing) ->
    missing;
dbg_undef(Undef, loop) ->
    [Undef, loop];
dbg_undef(Undef, switch_body) ->
    [switch_body, Undef];
dbg_undef(Undef, {switch, Arg, Type, Alts, Catch}) ->
    NewAlts   = [ dbg_undef(Undef, Alt) || Alt <- Alts ],
    NewCatch  = dbg_undef(Undef, Catch),
    NewSwitch = {switch, Arg, Type, NewAlts, NewCatch},
    NewSwitch;
dbg_undef(Undef, SCode) when is_list(SCode) ->
    lists:droplast(SCode) ++ [dbg_undef(Undef, lists:last(SCode))];
dbg_undef(Undef, SCode) when is_tuple(SCode); is_atom(SCode) ->
    [Mnemonic | _] =
        case is_tuple(SCode) of
            true  -> tuple_to_list(SCode);
            false -> [SCode]
        end,
    Op = gmb_fate_opcodes:m_to_op(Mnemonic),
    case gmb_fate_opcodes:end_bb(Op) of
        true  -> [Undef, SCode];
        false -> [SCode, Undef]
    end.

%% -- Phase II ---------------------------------------------------------------
%%  Optimize

optimize_scode(Funs, Options) ->
    maps:map(fun(Name, Def) -> optimize_fun(Funs, Name, Def, Options) end,
            Funs).

flatten(missing) -> missing;
flatten(Code)    -> lists:map(fun flatten_s/1, lists:flatten(Code)).

flatten_s({switch, Arg, Type, Alts, Catch}) ->
    {switch, Arg, Type, [flatten(Alt) || Alt <- Alts], flatten(Catch)};
flatten_s(I) -> I.

-define(MAX_SIMPL_ITERATIONS, 10).

optimize_fun(_Funs, Name, {Attrs, Sig, Code}, Options) ->
    Code0 = flatten(Code),
    ?debug(opt, Options, "Optimizing ~s\n", [Name]),
    Code1 = simpl_loop(0, Code0, Options),
    Code2 = desugar(Code1),
    {Attrs, Sig, Code2}.

simpl_loop(N, Code, Options) when N >= ?MAX_SIMPL_ITERATIONS ->
    ?debug(opt, Options, "  No simpl_loop fixed_point after ~p iterations.\n\n", [N]),
    Code;
simpl_loop(N, Code, Options) ->
    ACode = annotate_code(Code),
    [ ?debug(opt, Options, "  annotated:\n~s\n", [pp_ann("    ", ACode)]) || N == 0 ],
    Code1 = simplify(ACode, Options),
    [ ?debug(opt, Options, "  optimized:\n~s\n", [pp_ann("    ", Code1)]) || Code1 /= ACode ],
    Code2 = unannotate(Code1),
    case Code == Code2 of
        true  ->
            ?debug(opt, Options, "  Reached simpl_loop fixed point after ~p iteration~s.\n\n",
                                [N, if N /= 1 -> "s"; true -> "" end]),
            Code2;
        false -> simpl_loop(N + 1, Code2, Options)
    end.

pp_ann(Ind, [{switch, Arg, Type, Alts, Def} | Code]) ->
    Tags =
        case Type of
            boolean       -> ["FALSE", "TRUE"];
            tuple         -> ["(_)"];
            {variant, Ar} -> ["C" ++ integer_to_list(I) || I <- lists:seq(0, length(Ar) - 1)]
        end,
    Ind1 = "  " ++ Ind,
    Ind2 = "  " ++ Ind1,
    [Ind, "SWITCH ", pp_arg(Arg), "\n",
     [[Ind1, Tag, " =>\n", pp_ann(Ind2, Alt)]
      || {Tag, Alt} <- lists:zip(Tags, Alts), Alt /= missing],
     [[Ind1, "_ =>\n", pp_ann(Ind2, Def)] || Def /= missing],
     pp_ann(Ind, Code)];
pp_ann(Ind, [{i, #{ live_in := In, live_out := Out }, I} | Code]) ->
    Fmt = fun([]) -> "()";
             (Xs) -> string:join([lists:flatten(pp_arg(X)) || X <- Xs], " ")
          end,
    Op  = [Ind, pp_op(desugar_args(I))],
    Ann = [["   % ", Fmt(In), " -> ", Fmt(Out)] || In ++ Out /= []],
    [io_lib:format("~-40s~s\n", [Op, Ann]),
     pp_ann(Ind, Code)];
pp_ann(_, []) -> [].

pp_op(switch_body) -> "SWITCH-BODY";
pp_op(loop) -> "LOOP";
pp_op(I)    ->
    gmb_fate_pp:format_op(I, #{}).

pp_arg(?i(I))    -> io_lib:format("~w", [I]);
pp_arg({arg, N}) -> io_lib:format("arg~p", [N]);
pp_arg(?s(N))    -> io_lib:format("store~p", [N]);
pp_arg({var, N}) -> io_lib:format("var~p", [N]);
pp_arg(?a)       -> "a".

%% -- Analysis --

annotate_code(Code) ->
    annotate_code(5, [], Code).

annotate_code(Fuel, LiveTop, Code) ->
    {Code1, LiveIn} = ann_live(LiveTop, Code, []),
    case LiveIn == LiveTop of
        true  -> Code1;
        false when Fuel =< 0 ->
            code_error(liveness_analysis_out_of_fuel);
        false -> annotate_code(Fuel - 1, LiveIn, Code)
    end.

ann_live(_LiveTop, missing, _LiveOut) -> {missing, []};
ann_live(_LiveTop, [], LiveOut)       -> {[], LiveOut};
ann_live(LiveTop, [I | Is], LiveOut) ->
    {Is1, LiveMid} = ann_live(LiveTop, Is, LiveOut),
    {I1, LiveIn}   = ann_live1(LiveTop, I, LiveMid),
    {[I1 | Is1], LiveIn}.

ann_live1(_LiveTop, switch_body, LiveOut) ->
    Ann = #{ live_in => LiveOut, live_out => LiveOut },
    {{i, Ann, switch_body}, LiveOut};
ann_live1(LiveTop, loop, _LiveOut) ->
    Ann = #{ live_in => LiveTop, live_out => [] },
    {{i, Ann, loop}, LiveTop};
ann_live1(LiveTop, {switch, Arg, Type, Alts, Def}, LiveOut) ->
    Read              = [Arg || is_reg(Arg)],
    {Alts1, LiveAlts} = lists:unzip([ ann_live(LiveTop, Alt, LiveOut) || Alt <- Alts ]),
    {Def1,  LiveDef}  = ann_live(LiveTop, Def, LiveOut),
    LiveIn = ordsets:union([Read, LiveDef | LiveAlts]),
    {{switch, Arg, Type, Alts1, Def1}, LiveIn};
ann_live1(_LiveTop, I, LiveOut) ->
    #{ read := Reads0, write := W } = attributes(I),
    Reads   = lists:filter(fun is_reg/1, Reads0),
    %% If we write it here it's not live in (unless we also read it)
    LiveIn = ordsets:union(LiveOut -- [W], Reads),
    Ann = #{ live_in => LiveIn, live_out => LiveOut },
    {{i, Ann, I}, LiveIn}.

is_reg(?a)             -> false;
is_reg(none)           -> false;
is_reg(pc)             -> false;
is_reg({immediate, _}) -> false;
is_reg({arg, _})       -> true;
is_reg({store, _})     -> true;
is_reg({var, _})       -> true.

%% Instruction attributes: reads, writes and purity (pure means no writing to the chain).
attributes(I) ->
    Set  = fun(L) when is_list(L) -> ordsets:from_list(L);
              (X)                 -> ordsets:from_list([X]) end,
    Attr = fun(W, R, P) -> #{read => Set(R), write => W, pure => P}  end,
    Pure   = fun(W, R) -> Attr(W, R, true) end,
    Impure = fun(W, R) -> Attr(W, R, false) end,
    case I of
        loop                                  -> Impure(pc, []);
        switch_body                           -> Pure(none, []);
        'RETURN'                              -> Impure(pc, []);
        {'DBG_LOC', _, _}                     -> Impure(none, []);
        {'DBG_DEF', _, _}                     -> Impure(none, []);
        {'DBG_UNDEF', _, _}                   -> Impure(none, []);
        {'DBG_CONTRACT', _}                   -> Impure(none, []);
        {'RETURNR', A}                        -> Impure(pc, A);
        {'CALL', A}                           -> Impure(?a, [A]);
        {'CALL_R', A, _, B, C, D}             -> Impure(?a, [A, B, C, D]);
        {'CALL_GR', A, _, B, C, D, E}         -> Impure(?a, [A, B, C, D, E]);
        {'CALL_PGR', A, _, B, C, D, E, F}     -> Impure(?a, [A, B, C, D, E, F]);
        {'CALL_T', A}                         -> Impure(pc, [A]);
        {'CALL_VALUE', A}                     -> Pure(A, []);
        {'JUMP', _}                           -> Impure(pc, []);
        {'JUMPIF', A, _}                      -> Impure(pc, A);
        {'SWITCH_V2', A, _, _}                -> Impure(pc, A);
        {'SWITCH_V3', A, _, _, _}             -> Impure(pc, A);
        {'SWITCH_VN', A, _}                   -> Impure(pc, A);
        {'PUSH', A}                           -> Pure(?a, A);
        'DUPA'                                -> Pure(?a, ?a);
        {'DUP', A}                            -> Pure(?a, A);
        {'POP', A}                            -> Pure(A, ?a);
        {'STORE', A, B}                       -> Pure(A, B);
        'INCA'                                -> Pure(?a, ?a);
        {'INC', A}                            -> Pure(A, A);
        'DECA'                                -> Pure(?a, ?a);
        {'DEC', A}                            -> Pure(A, A);
        {'ADD', A, B, C}                      -> Pure(A, [B, C]);
        {'SUB', A, B, C}                      -> Pure(A, [B, C]);
        {'MUL', A, B, C}                      -> Pure(A, [B, C]);
        {'DIV', A, B, C}                      -> Pure(A, [B, C]);
        {'MOD', A, B, C}                      -> Pure(A, [B, C]);
        {'POW', A, B, C}                      -> Pure(A, [B, C]);
        {'MULMOD', A, B, C, D}                -> Pure(A, [B, C, D]);
        {'BAND', A, B, C}                     -> Pure(A, [B, C]);
        {'BOR', A, B, C}                      -> Pure(A, [B, C]);
        {'BXOR', A, B, C}                     -> Pure(A, [B, C]);
        {'BNOT', A, B}                        -> Pure(A, [B]);
        {'BSL', A, B, C}                      -> Pure(A, [B, C]);
        {'BSR', A, B, C}                      -> Pure(A, [B, C]);
        {'LT', A, B, C}                       -> Pure(A, [B, C]);
        {'GT', A, B, C}                       -> Pure(A, [B, C]);
        {'EQ', A, B, C}                       -> Pure(A, [B, C]);
        {'ELT', A, B, C}                      -> Pure(A, [B, C]);
        {'EGT', A, B, C}                      -> Pure(A, [B, C]);
        {'NEQ', A, B, C}                      -> Pure(A, [B, C]);
        {'AND', A, B, C}                      -> Pure(A, [B, C]);
        {'OR', A, B, C}                       -> Pure(A, [B, C]);
        {'NOT', A, B}                         -> Pure(A, B);
        {'TUPLE', A, N}                       -> Pure(A, [?a || N > 0]);
        {'ELEMENT', A, B, C}                  -> Pure(A, [B, C]);
        {'SETELEMENT', A, B, C, D}            -> Pure(A, [B, C, D]);
        {'MAP_EMPTY', A}                      -> Pure(A, []);
        {'MAP_LOOKUP', A, B, C}               -> Pure(A, [B, C]);
        {'MAP_LOOKUPD', A, B, C, D}           -> Pure(A, [B, C, D]);
        {'MAP_UPDATE', A, B, C, D}            -> Pure(A, [B, C, D]);
        {'MAP_DELETE', A, B, C}               -> Pure(A, [B, C]);
        {'MAP_MEMBER', A, B, C}               -> Pure(A, [B, C]);
        {'MAP_FROM_LIST', A, B}               -> Pure(A, B);
        {'MAP_TO_LIST', A, B}                 -> Pure(A, B);
        {'MAP_SIZE', A, B}                    -> Pure(A, B);
        {'NIL', A}                            -> Pure(A, []);
        {'IS_NIL', A, B}                      -> Pure(A, B);
        {'CONS', A, B, C}                     -> Pure(A, [B, C]);
        {'HD', A, B}                          -> Pure(A, B);
        {'TL', A, B}                          -> Pure(A, B);
        {'LENGTH', A, B}                      -> Pure(A, B);
        {'APPEND', A, B, C}                   -> Pure(A, [B, C]);
        {'STR_JOIN', A, B, C}                 -> Pure(A, [B, C]);
        {'INT_TO_STR', A, B}                  -> Pure(A, B);
        {'INT_TO_BYTES', A, B, C}             -> Pure(A, [B, C]);
        {'ADDR_TO_STR', A, B}                 -> Pure(A, B);
        {'STR_REVERSE', A, B}                 -> Pure(A, B);
        {'STR_LENGTH', A, B}                  -> Pure(A, B);
        {'STR_TO_BYTES', A, B}                -> Pure(A, B);
        {'INT_TO_ADDR', A, B}                 -> Pure(A, B);
        {'VARIANT', A, B, C, D}               -> Pure(A, [?a, B, C, D]);
        {'VARIANT_TEST', A, B, C}             -> Pure(A, [B, C]);
        {'VARIANT_ELEMENT', A, B, C}          -> Pure(A, [B, C]);
        'BITS_NONEA'                          -> Pure(?a, []);
        {'BITS_NONE', A}                      -> Pure(A, []);
        'BITS_ALLA'                           -> Pure(?a, []);
        {'BITS_ALL', A}                       -> Pure(A, []);
        {'BITS_ALL_N', A, B}                  -> Pure(A, B);
        {'BITS_SET', A, B, C}                 -> Pure(A, [B, C]);
        {'BITS_CLEAR', A, B, C}               -> Pure(A, [B, C]);
        {'BITS_TEST', A, B, C}                -> Pure(A, [B, C]);
        {'BITS_SUM', A, B}                    -> Pure(A, B);
        {'BITS_OR', A, B, C}                  -> Pure(A, [B, C]);
        {'BITS_AND', A, B, C}                 -> Pure(A, [B, C]);
        {'BITS_DIFF', A, B, C}                -> Pure(A, [B, C]);
        {'SHA3', A, B}                        -> Pure(A, [B]);
        {'SHA256', A, B}                      -> Pure(A, [B]);
        {'BLAKE2B', A, B}                     -> Pure(A, [B]);
        {'POSEIDON', A, B, C}                 -> Pure(A, [B, C]);
        {'VERIFY_SIG', A, B, C, D}            -> Pure(A, [B, C, D]);
        {'VERIFY_SIG_SECP256K1', A, B, C, D}  -> Pure(A, [B, C, D]);
        {'ECVERIFY_SECP256K1', A, B, C, D}    -> Pure(A, [B, C, D]);
        {'ECRECOVER_SECP256K1', A, B, C}      -> Pure(A, [B, C]);
        {'CONTRACT_TO_ADDRESS', A, B}         -> Pure(A, [B]);
        {'ADDRESS_TO_CONTRACT', A, B}         -> Pure(A, [B]);
        {'ADDRESS_TO_BYTES', A, B}            -> Pure(A, [B]);
        {'AUTH_TX_HASH', A}                   -> Pure(A, []);
        {'AUTH_TX', A}                        -> Pure(A, []);
        {'BYTES_TO_INT', A, B}                -> Pure(A, [B]);
        {'BYTES_TO_STR', A, B}                -> Pure(A, [B]);
        {'BYTES_CONCAT', A, B, C}             -> Pure(A, [B, C]);
        {'BYTES_SPLIT', A, B, C}              -> Pure(A, [B, C]);
        {'BYTES_SPLIT_ANY', A, B, C}          -> Pure(A, [B, C]);
        {'BYTES_SIZE', A, B}                  -> Pure(A, B);
        {'BYTES_TO_FIXED_SIZE', A, B, C}      -> Pure(A, [B, C]);
        {'ORACLE_CHECK', A, B, C, D}          -> Pure(A, [B, C, D]);
        {'ORACLE_CHECK_QUERY', A, B, C, D, E} -> Pure(A, [B, C, D, E]);
        {'IS_ORACLE', A, B}                   -> Pure(A, [B]);
        {'IS_CONTRACT', A, B}                 -> Pure(A, [B]);
        {'IS_PAYABLE', A, B}                  -> Pure(A, [B]);
        {'CREATOR', A}                        -> Pure(A, []);
        {'ADDRESS', A}                        -> Pure(A, []);
        {'BALANCE', A}                        -> Pure(A, []);
        {'BALANCE_OTHER', A, B}               -> Pure(A, [B]);
        {'ORIGIN', A}                         -> Pure(A, []);
        {'CALLER', A}                         -> Pure(A, []);
        {'GASPRICE', A}                       -> Pure(A, []);
        {'FEE', A}                            -> Pure(A, []);
        {'BLOCKHASH', A, B}                   -> Pure(A, [B]);
        {'BENEFICIARY', A}                    -> Pure(A, []);
        {'TIMESTAMP', A}                      -> Pure(A, []);
        {'GENERATION', A}                     -> Pure(A, []);
        {'MICROBLOCK', A}                     -> Pure(A, []);
        {'DIFFICULTY', A}                     -> Pure(A, []);
        {'GASLIMIT', A}                       -> Pure(A, []);
        {'NETWORK_ID', A}                     -> Pure(A, []);
        {'GAS', A}                            -> Pure(A, []);
        {'LOG0', A}                           -> Impure(none, [A]);
        {'LOG1', A, B}                        -> Impure(none, [A, B]);
        {'LOG2', A, B, C}                     -> Impure(none, [A, B, C]);
        {'LOG3', A, B, C, D}                  -> Impure(none, [A, B, C, D]);
        {'LOG4', A, B, C, D, E}               -> Impure(none, [A, B, C, D, E]);
        'DEACTIVATE'                          -> Impure(none, []);
        {'SPEND', A, B}                       -> Impure(none, [A, B]);
        {'ORACLE_REGISTER', A, B, C, D, E, F, G} -> Impure(A, [B, C, D, E, F, G]);
        {'ORACLE_QUERY', A, B, C, D, E, F, G, H} -> Impure(A, [B, C, D, E, F, G, H]);
        {'ORACLE_RESPOND', A, B, C, D, E, F}  -> Impure(none, [A, B, C, D, E, F]);
        {'ORACLE_EXTEND', A, B, C}            -> Impure(none, [A, B, C]);
        {'ORACLE_GET_ANSWER', A, B, C, D, E}  -> Pure(A, [B, C, D, E]);
        {'ORACLE_GET_QUESTION', A, B, C, D, E}-> Pure(A, [B, C, D, E]);
        {'ORACLE_QUERY_FEE', A, B}            -> Pure(A, [B]);
        {'ORACLE_EXPIRY', A, B}               -> Impure(A, [B]);
        {'AENS_RESOLVE', A, B, C, D}          -> Impure(A, [B, C, D]);
        {'AENS_PRECLAIM', A, B, C}            -> Impure(none, [A, B, C]);
        {'AENS_CLAIM', A, B, C, D, E}         -> Impure(none, [A, B, C, D, E]);
        {'AENS_UPDATE', A, B, C, D, E, F}     -> Impure(none, [A, B, C, D, E, F]);
        {'AENS_TRANSFER', A, B, C, D}         -> Impure(none, [A, B, C, D]);
        {'AENS_REVOKE', A, B, C}              -> Impure(none, [A, B, C]);
        {'AENS_LOOKUP', A, B}                 -> Impure(A, [B]);
        {'BLS12_381_G1_NEG', A, B}            -> Pure(A, [B]);
        {'BLS12_381_G1_NORM', A, B}           -> Pure(A, [B]);
        {'BLS12_381_G1_VALID', A, B}          -> Pure(A, [B]);
        {'BLS12_381_G1_IS_ZERO', A, B}        -> Pure(A, [B]);
        {'BLS12_381_G1_ADD', A, B, C}         -> Pure(A, [B, C]);
        {'BLS12_381_G1_MUL', A, B, C}         -> Pure(A, [B, C]);
        {'BLS12_381_G2_NEG', A, B}            -> Pure(A, [B]);
        {'BLS12_381_G2_NORM', A, B}           -> Pure(A, [B]);
        {'BLS12_381_G2_VALID', A, B}          -> Pure(A, [B]);
        {'BLS12_381_G2_IS_ZERO', A, B}        -> Pure(A, [B]);
        {'BLS12_381_G2_ADD', A, B, C}         -> Pure(A, [B, C]);
        {'BLS12_381_G2_MUL', A, B, C}         -> Pure(A, [B, C]);
        {'BLS12_381_GT_INV', A, B}            -> Pure(A, [B]);
        {'BLS12_381_GT_ADD', A, B, C}         -> Pure(A, [B, C]);
        {'BLS12_381_GT_MUL', A, B, C}         -> Pure(A, [B, C]);
        {'BLS12_381_GT_POW', A, B, C}         -> Pure(A, [B, C]);
        {'BLS12_381_GT_IS_ONE', A, B}         -> Pure(A, [B]);
        {'BLS12_381_PAIRING', A, B, C}        -> Pure(A, [B, C]);
        {'BLS12_381_MILLER_LOOP', A, B, C}    -> Pure(A, [B, C]);
        {'BLS12_381_FINAL_EXP', A, B}         -> Pure(A, [B]);
        {'BLS12_381_INT_TO_FR', A, B}         -> Pure(A, [B]);
        {'BLS12_381_INT_TO_FP', A, B}         -> Pure(A, [B]);
        {'BLS12_381_FR_TO_INT', A, B}         -> Pure(A, [B]);
        {'BLS12_381_FP_TO_INT', A, B}         -> Pure(A, [B]);
        {'STR_TO_LIST', A, B}                 -> Pure(A, [B]);
        {'STR_FROM_LIST', A, B}               -> Pure(A, [B]);
        {'STR_TO_UPPER', A, B}                -> Pure(A, [B]);
        {'STR_TO_LOWER', A, B}                -> Pure(A, [B]);
        {'CHAR_TO_INT', A, B}                 -> Pure(A, [B]);
        {'CHAR_FROM_INT', A, B}               -> Pure(A, [B]);
        {'CREATE', A, B, C}                   -> Impure(?a, [A, B, C]);
        {'CLONE', A, B, C, D}                 -> Impure(?a, [A, B, C, D]);
        {'CLONE_G', A, B, C, D, E}            -> Impure(?a, [A, B, C, D, E]);
        {'BYTECODE_HASH', A, B}               -> Impure(A, [B]);
        {'ABORT', A}                          -> Impure(pc, A);
        {'EXIT', A}                           -> Impure(pc, A);
        'NOP'                                 -> Pure(none, [])
    end.

var_writes({i, _, I}) -> var_writes(I);
var_writes(I) ->
    #{ write := W } = attributes(I),
    case W of
        {var, _}   -> [W];
        {arg, _}   -> [W];
        {store, _} -> [W];
        {stack, _} -> [];
        none       -> [];
        pc         -> []
    end.

-spec independent(sinstr_a(), sinstr_a()) -> boolean().
%% independent({switch, _, _, _, _}, _) -> false;       %% Commented due to Dialyzer whinging
independent(_, {switch, _, _, _, _}) -> false;
independent({i, _, I}, {i, _, J}) ->
    #{ write := WI, read := RI, pure := PureI } = attributes(I),
    #{ write := WJ, read := RJ, pure := PureJ } = attributes(J),

    StackI = lists:member(?a, [WI | RI]),
    StackJ = lists:member(?a, [WJ | RJ]),

    ReadStoreI = [] /= [ x || {store, _} <- RI ],
    ReadStoreJ = [] /= [ x || {store, _} <- RJ ],

    if  WI == pc; WJ == pc       -> false;  %% no jumps
        not (PureI or PureJ)     -> false;  %% at least one is pure
        StackI and StackJ        -> false;  %% cannot both use the stack
        WI == WJ                 -> false;  %% cannot write to the same register
        ReadStoreI and not PureJ -> false;  %% can't read store/state if other is impure
        ReadStoreJ and not PureI -> false;  %% can't read store/state if other is impure
        true                     ->
            %% and cannot write to each other's inputs
            not lists:member(WI, RJ) andalso
            not lists:member(WJ, RI)
    end.

merge_ann(#{ live_in := LiveIn }, #{ live_out := LiveOut }) ->
    #{ live_in => LiveIn, live_out => LiveOut }.

%% Swap two instructions. Precondition: the instructions are independent/2.
swap_instrs({i, #{ live_in := Live1 }, I}, {i, #{ live_in := Live2, live_out := Live3 }, J}) ->
    %% Since I and J are independent the J can't read or write anything in
    %% that I writes.
    WritesI = ordsets:subtract(Live2, Live1),
    %% Any final reads by J, that I does not read should be removed from Live2.
    #{ read := ReadsI } = attributes(I),
    ReadsJ  = ordsets:subtract(Live2, ordsets:union(Live3, ReadsI)),
    Live2_  = ordsets:subtract(ordsets:union([Live1, Live2, Live3]), ordsets:union(WritesI, ReadsJ)),
    {{i, #{ live_in => Live1,  live_out => Live2_ }, J},
     {i, #{ live_in => Live2_, live_out => Live3  }, I}}.

live_in({store, _}, _) -> true;
live_in(R, #{ live_in  := LiveIn  }) -> ordsets:is_element(R, LiveIn);
live_in(R, {i, Ann, _}) -> live_in(R, Ann);
live_in(R, [I = {i, _, _} | _]) -> live_in(R, I);
live_in(R, [{switch, A, _, Alts, Def} | _]) ->
    R == A orelse lists:any(fun(Code) -> live_in(R, Code) end, [Def | Alts]);
live_in(_, missing) -> false;
live_in(_, []) -> false.

live_out({store, _}, _) -> true;
live_out(R, #{ live_out := LiveOut }) -> ordsets:is_element(R, LiveOut).

%% -- Optimizations --

simplify([], _) -> [];
simplify(missing, _) -> missing;
simplify([I | Code], Options) ->
    simpl_top(simpl_s(I, Options), simplify(Code, Options), Options).

simpl_s({switch, Arg, Type, Alts, Def}, Options) ->
    {switch, Arg, Type, [simplify(A, Options) || A <- Alts], simplify(Def, Options)};
simpl_s(I, _) -> I.

%% Safe-guard against loops in the rewriting. Shouldn't happen so throw an
%% error if we run out.
-define(SIMPL_FUEL, 5000).

simpl_top(I, Code, Options) ->
    simpl_top(?SIMPL_FUEL, I, Code, Options).

simpl_top(0, I, Code, _Options) ->
    code_error({optimizer_out_of_fuel, I, Code});
simpl_top(Fuel, I, Code, Options) ->
    Rules = [R || R = {Rule, _} <- rules(), proplists:get_value(Rule, Options, true)],
    apply_rules(Fuel, Rules, I, Code, Options).

apply_rules(Fuel, Rules, I, Code, Options) ->
    Cons = fun(X, Xs) -> simpl_top(Fuel - 1, X, Xs, Options) end,
    case apply_rules_once(Rules, I, Code) of
        false -> [I | Code];
        {RName, New, Rest} ->
            case is_debug(opt_rules, Options) of
                true ->
                    {OldCode, NewCode} = drop_common_suffix([I | Code], New ++ Rest),
                    ?debug(opt_rules, Options, "  Applied ~p:\n~s  ==>\n~s\n", [RName, pp_ann("    ", OldCode), pp_ann("    ", NewCode)]);
                false -> ok
            end,
            lists:foldr(Cons, Rest, New)
    end.

apply_rules_once([], _, _) ->
    false;
apply_rules_once([{RName, Rule} | Rules], I, Code) ->
    case Rule(I, Code) of
        false       -> apply_rules_once(Rules, I, Code);
        {New, Rest} -> {RName, New, Rest}
    end.

-define(RULE(Name), {Name, fun Name/2}).

merge_rules() ->
    [?RULE(optimize_push_consume),
     ?RULE(optimize_one_shot_var),
     ?RULE(optimize_write_to_dead_var),
     ?RULE(optimize_inline_switch_target)
    ].

rules() ->
    merge_rules() ++
    [?RULE(optimize_swap_push),
     ?RULE(optimize_swap_pop),
     ?RULE(optimize_swap_write),
     ?RULE(optimize_constant_propagation),
     ?RULE(optimize_prune_impossible_branches),
     ?RULE(optimize_single_successful_branch),
     ?RULE(optimize_inline_store),
     ?RULE(optimize_float_switch_body)
    ].

%% Removing pushes that are immediately consumed.
optimize_push_consume({i, Ann1, {'STORE', ?a, A}}, Code) ->
    inline_push(Ann1, A, 0, Code, []);
%% Writing directly to memory instead of going through the accumulator.
optimize_push_consume({i, Ann1, I}, [{i, Ann2, {'STORE', R, ?a}} | Code]) ->
    IsPush =
        case op_view(I) of
            {_, ?a, _} -> true;
            _          -> false
        end orelse
        case I of
            {'VARIANT', ?a, _, _, _} -> true;
            _                        -> false
        end,
    if IsPush -> {[{i, merge_ann(Ann1, Ann2), setelement(2, I, R)}], Code};
       true   -> false end;
optimize_push_consume(_, _) -> false.

inline_push(Ann, Arg, Stack, [{i, _, switch_body} = AI | Code], Acc) ->
    {AI1, {i, Ann1, _}} = swap_instrs({i, Ann, {'STORE', ?a, Arg}}, AI),
    inline_push(Ann1, Arg, Stack, Code, [AI1 | Acc]);
inline_push(Ann1, Arg, Stack, [{i, Ann2, I} = AI | Code], Acc) ->
    case op_view(I) of
        {Op, R, As} ->
            Consumes = length([ ?a || ?a <- As ]),
            Produces = length([ ?a || ?a == R  ]),
            case Consumes > Stack of
                true ->
                    {As0, As1} = split_stack_arg(Stack, As),
                    Acc1 = [{i, merge_ann(Ann1, Ann2), from_op_view(Op, R, As0 ++ [Arg] ++ As1)} | Acc],
                    {lists:reverse(Acc1), Code};
                false when Arg /= R ->
                    {AI1, {i, Ann1b, _}} = swap_instrs({i, Ann1, {'STORE', ?a, Arg}}, AI),
                    inline_push(Ann1b, Arg, Stack + Produces - Consumes, Code, [AI1 | Acc]);
                false -> false
            end;
        _ -> false
    end;
inline_push(_, _, _, _, _) -> false.

split_stack_arg(N, As) -> split_stack_arg(N, As, []).
split_stack_arg(0, [?a | As], Acc) ->
    {lists:reverse(Acc), As};
split_stack_arg(N, [A | As], Acc) ->
    N1 = if A == ?a -> N - 1;
            true    -> N end,
    split_stack_arg(N1, As, [A | Acc]).

%% Move PUSHes past non-stack instructions.
optimize_swap_push(Push = {i, _, PushI}, [I | Code]) ->
    case op_view(PushI) of
        {_, ?a, _} ->
            case independent(Push, I) of
                true ->
                    {I1, Push1} = swap_instrs(Push, I),
                    {[I1, Push1], Code};
                false -> false
            end;
        _ -> false
    end;
optimize_swap_push(_, _) -> false.

%% Move non-stack instruction past POPs.
optimize_swap_pop(IA = {i, _, I}, [JA = {i, _, J} | Code]) ->
    case independent(IA, JA) of
        true ->
            case {op_view(I), op_view(J)} of
                {false, _} -> false;
                {_, false} -> false;
                {{_, IR, IAs}, {_, RJ, JAs}} ->
                    NonStackI = not lists:member(?a, [IR | IAs]),
                    %% RJ /= ?a to not conflict with optimize_swap_push
                    PopJ      = RJ /= ?a andalso lists:member(?a, JAs),
                    case NonStackI andalso PopJ of
                        false -> false;
                        true  ->
                            {JA1, IA1} = swap_instrs(IA, JA),
                            {[JA1, IA1], Code}
                    end
            end;
        false -> false
    end;
optimize_swap_pop(_, _) -> false.

%% Match up writes to variables with instructions further down.
optimize_swap_write(I = {i, _, _}, [J | Code]) ->
    case {var_writes(I), independent(I, J)} of
        {[_], true} ->
            {J1, I1} = swap_instrs(I, J),
            optimize_swap_write([J1], I1, Code);
        _ -> false
    end;
optimize_swap_write(_, _) -> false.

optimize_swap_write(Pre, I, [{i, _, switch_body} = J | Code]) ->
    {J1, I1} = swap_instrs(I, J),
    optimize_swap_write([J1 | Pre], I1, Code);
optimize_swap_write(Pre, I, Code0 = [J | Code]) ->
    case apply_rules_once(merge_rules(), I, Code0) of
        {_Rule, New, Rest} ->
            {lists:reverse(Pre) ++ New, Rest};
        false ->
            case independent(I, J) of
                false -> false;
                true  ->
                    {J1, I1} = swap_instrs(I, J),
                    optimize_swap_write([J1 | Pre], I1, Code)
            end
    end;
optimize_swap_write(_, _, _) -> false.

%% Precompute instructions with known values
optimize_constant_propagation(Cons = {i, Ann1, {'CONS', R, X, Xs}}, [{i, Ann, {'IS_NIL', S, R}} | Code]) ->
    Store = {i, Ann, {'STORE', S, ?i(false)}},
    Cons1 = case R of
                ?a -> {i, Ann1, {'CONS', ?void, X, Xs}};
                _  -> Cons
            end,
    {[Cons1, Store], Code};
optimize_constant_propagation(Nil = {i, Ann1, {'NIL', R}}, [{i, Ann, {'IS_NIL', S, R}} | Code]) ->
    Store = {i, Ann, {'STORE', S, ?i(true)}},
    Nil1 = case R of
               ?a -> {i, Ann1, {'NIL', ?void}};
               _  -> Nil
           end,
    {[Nil1, Store], Code};
optimize_constant_propagation({i, Ann, I}, Code) ->
    case op_view(I) of
        false -> false;
        {Op, R, As} ->
            Vs = [V || ?i(V) <- As],
            case length(Vs) == length(As) of
                false -> false;
                true  ->
                    case eval_op(Op, Vs) of
                        no_eval -> false;
                        V       -> {[{i, Ann, {'STORE', R, ?i(V)}}], Code}
                    end
            end
    end;
optimize_constant_propagation(_, _) -> false.

eval_op('ADD', [X, Y]) when is_integer(X), is_integer(Y) -> X + Y;
eval_op('SUB', [X, Y]) when is_integer(X), is_integer(Y) -> X - Y;
eval_op('MUL', [X, Y]) when is_integer(X), is_integer(Y) -> X * Y;
eval_op('DIV', [X, Y]) when is_integer(X), is_integer(Y), Y /= 0 -> X div Y;
eval_op('MOD', [X, Y]) when is_integer(X), is_integer(Y), Y /= 0 -> X rem Y;
eval_op('POW', [_, _])  -> no_eval;
eval_op('LT', [X, Y])   -> X < Y;
eval_op('GT', [X, Y])   -> X > Y;
eval_op('EQ', [X, Y])   -> X =:= Y;
eval_op('ELT', [X, Y])  -> X =< Y;
eval_op('EGT', [X, Y])  -> X >= Y;
eval_op('NEQ', [X, Y])  -> X =/= Y;
eval_op('NOT', [true])  -> false;
eval_op('NOT', [false]) -> true;
eval_op(_, _)           -> no_eval.   %% TODO: bits?

%% Prune impossible branches from switches
optimize_prune_impossible_branches({switch, ?i(V), Type, Alts, missing}, Code) ->
    case pick_branch(Type, V, Alts) of
        false -> false;
        Alt   -> {Alt, Code}
    end;
optimize_prune_impossible_branches({switch, ?i(V), boolean, [False, True] = Alts, Def}, Code) when V == true; V == false ->
    Alts1 = [if V -> missing; true -> False   end,
             if V -> True;    true -> missing end],
    case Alts == Alts1 of
        true  -> false;
        false ->
            case Alts1 of
                [missing, missing] -> {Def, Code};
                _                  -> {[{switch, ?i(V), boolean, Alts1, Def}], Code}
            end
    end;
optimize_prune_impossible_branches(Variant = {i, _, {'VARIANT', R, ?i(_), ?i(Tag), ?i(_)}},
                            [{switch, R, Type = {variant, _}, Alts, missing} | Code]) when is_integer(Tag) ->
    case {R, lists:nth(Tag + 1, Alts)} of
        {_, missing} ->
            Alts1 = [missing || _ <- Alts],
            case Alts == Alts1 of
                true -> false;
                false -> {[Variant, {switch, R, Type, Alts1, missing}], Code}
            end;
        {?a, Alt} -> {Alt, Code};
        {_,  Alt} ->
            case live_in(R, Alt) of
                true  -> {[Variant | Alt], Code};
                false -> {Alt, Code}
            end
    end;
optimize_prune_impossible_branches(_, _) -> false.

pick_branch(boolean, V, [False, True]) when V == true; V == false ->
    Alt = if V -> True; true -> False end,
    case Alt of
        missing -> false;
        _       -> Alt
    end;
pick_branch(_Type, _V, _Alts) ->
    false.

%% If there's a single branch that doesn't abort we can push the code for that
%% out of the switch.
optimize_single_successful_branch({switch, R, Type, Alts, Def}, Code) ->
    case push_code_out_of_switch([Def | Alts]) of
        {_, none} -> false;
        {_, many} -> false;
        {_, [{i, _, switch_body}]} -> false;
        {[Def1 | Alts1], PushedOut} ->
            {[{switch, R, Type, Alts1, Def1} | PushedOut], Code}
    end;
optimize_single_successful_branch(_, _) -> false.

push_code_out_of_switch([]) -> {[], none};
push_code_out_of_switch([Alt | Alts]) ->
    {Alt1, PushedAlt}   = push_code_out_of_alt(Alt),
    {Alts1, PushedAlts} = push_code_out_of_switch(Alts),
    Pushed =
        case {PushedAlt, PushedAlts} of
            {none, _} -> PushedAlts;
            {_, none} -> PushedAlt;
            _         -> many
        end,
    {[Alt1 | Alts1], Pushed}.

push_code_out_of_alt(missing) -> {missing, none};
push_code_out_of_alt([Body = {i, _, switch_body} | Code]) ->
    case does_abort(Code) of
        true  -> {[Body | Code], none};
        false -> {[Body], [Body | Code]}  %% Duplicate the switch_body, in case we apply this in the middle of a switch
    end;
push_code_out_of_alt([{switch, R, Type, Alts, Def}]) ->
    {[Def1 | Alts1], Pushed} = push_code_out_of_switch([Def | Alts]),
    {[{switch, R, Type, Alts1, Def1}], Pushed};
push_code_out_of_alt(Code) ->
    {Code, many}. %% Conservative

does_abort([I | Code]) ->
    does_abort(I) orelse does_abort(Code);
does_abort({i, _, {'ABORT', _}}) -> true;
does_abort({i, _, {'EXIT', _}}) -> true;
does_abort(missing) -> true;
does_abort({switch, _, _, Alts, Def}) ->
    lists:all(fun does_abort/1, [Def | Alts]);
does_abort(_) -> false.

%% STORE R A, SWITCH R --> SWITCH A
optimize_inline_switch_target({i, Ann, {'STORE', R, A}}, [{switch, R, Type, Alts, Def} | Code]) ->
    Ann1   =
        case is_reg(A) of
            true  -> Ann#{ live_out := ordsets:add_element(A, maps:get(live_out, Ann)) };
            false -> Ann
        end,
    Store  = {i, Ann1, {'STORE', R, A}},
    Switch = {switch, A, Type, Alts, Def},
    case R of
        A        -> false;
        ?a       -> {[Switch], Code};
        {var, _} ->
            case lists:any(fun(Alt) -> live_in(R, Alt) end, [Def | Alts]) of
                false             -> {[Switch], Code};
                true when A /= ?a -> {[Store, Switch], Code};
                true              -> false
            end;
        _        -> false %% impossible
    end;
optimize_inline_switch_target(_, _) -> false.

%% Float switch-body to closest switch
optimize_float_switch_body(I = {i, _, _}, [J = {i, _, switch_body} | Code]) ->
    {J1, I1} = swap_instrs(I, J),
    {[], [J1, I1 | Code]};
optimize_float_switch_body(_, _) -> false.

%% Inline stores
optimize_inline_store({i, _, {'STORE', R, R}}, Code) ->
    {[], Code};
optimize_inline_store(I = {i, _, {'STORE', R = {var, _}, A}}, Code) ->
    %% Not when A is var unless updating the annotations properly.
    Inline = case A of
                 {arg, _}   -> true;
                 ?i(_)      -> true;
                 {store, _} -> true;
                 _          -> false
             end,
    if Inline -> optimize_inline_store([I], false, R, A, Code);
       true   -> false end;
optimize_inline_store(_, _) -> false.

optimize_inline_store(Acc, Progress, R, A, [I = {i, _, switch_body} | Code]) ->
    optimize_inline_store([I | Acc], Progress, R, A, Code);
optimize_inline_store(Acc, Progress, R, A, [{i, Ann, I} | Code]) ->
    #{ write := W } = attributes(I),
    Inl = fun(X) when X == R -> A; (X) -> X end,
    case live_in(R, Ann) of
        false -> false;  %% No more reads of R
        true  ->
            {I1, Progress1} =
                case op_view(I) of
                    {Op, S, As} ->
                        case lists:member(R, As) of
                            true  -> {from_op_view(Op, S, lists:map(Inl, As)), true};
                            false -> {I, Progress}
                        end;
                    _ -> {I, Progress}
                end,
            Acc1 = [{i, Ann, I1} | Acc],
            %% Stop if write to R or A
            case lists:member(W, [R, A]) of
                true when Progress1 -> {lists:reverse(Acc1), Code};
                true                -> false;
                false               -> optimize_inline_store(Acc1, Progress1, R, A, Code)
            end
    end;
optimize_inline_store(Acc, true, _, _, Code) -> {lists:reverse(Acc), Code};
optimize_inline_store(_, false, _, _, _) -> false.

%% Shortcut write followed by final read
optimize_one_shot_var({i, Ann1, I}, [{i, Ann2, J} | Code]) ->
    case op_view(I) of
        {Op, R = {var, _}, As} ->
            Copy = case J of
                       {'STORE', S, R} -> {write_to, S};
                       _               -> false
                   end,
            case {live_out(R, Ann2), Copy} of
                {false, {write_to, X}} ->
                    {[{i, merge_ann(Ann1, Ann2), from_op_view(Op, X, As)}], Code};
                _ -> false
            end;
        _ -> false
    end;
optimize_one_shot_var(_, _) -> false.

%% Remove writes to dead variables
optimize_write_to_dead_var({i, _, {'STORE', ?void, ?a}}, _) -> false; %% Avoid looping
optimize_write_to_dead_var({i, Ann, I}, Code) ->
    #{ pure := Pure } = attributes(I),
    case op_view(I) of
        {_Op, R, As} when R /= ?a, Pure ->
            case live_out(R, Ann) of
                false ->
                    %% Subtle: we still have to pop the stack if any of the arguments
                    %% came from there.
                    {[{i, Ann, {'STORE', ?void, ?a}} || X <- As, X == ?a], Code};
                true -> false
            end;
        _ -> false
    end;
optimize_write_to_dead_var(_, _) -> false.

op_view({'ABORT', R}) -> {'ABORT', none, [R]};
op_view({'EXIT', R}) -> {'EXIT', none, [R]};
op_view(T) when is_tuple(T) ->
    [Op, R | As] = tuple_to_list(T),
    CheckReads = fun(Rs, X) -> case [] == Rs -- [dst, src] of true -> X; false -> false end end,
    case attributes(list_to_tuple([Op, dst | [src || _ <- As]])) of
        #{ write := dst, read := Rs  } -> CheckReads(Rs, {Op, R, As});
        #{ write := none, read := Rs } -> CheckReads(Rs, {Op, none, [R | As]});
        _                              -> false
    end;
op_view(_) -> false.

from_op_view(Op, none, As) -> list_to_tuple([Op | As]);
from_op_view(Op, R, As) -> list_to_tuple([Op, R | As]).

%% Desugar and specialize and remove annotations
-spec unannotate(scode_a())  -> scode();
                (sinstr_a()) -> sinstr();
                (missing)    -> missing.
unannotate({switch, Arg, Type, Alts, Def}) ->
    [{switch, Arg, Type, [unannotate(A) || A <- Alts], unannotate(Def)}];
unannotate(missing) -> missing;
unannotate(Code) when is_list(Code) ->
    lists:flatmap(fun unannotate/1, Code);
unannotate({i, _Ann, I}) -> [I].

%% Desugar and specialize
desugar({'ADD', ?a, ?i(1), ?a}) -> [gmb_fate_ops:inc()];
desugar({'ADD', A,  ?i(1), A})  -> [gmb_fate_ops:inc(desugar_arg(A))];
desugar({'ADD', ?a, ?a, ?i(1)}) -> [gmb_fate_ops:inc()];
desugar({'ADD', A,  A,  ?i(1)}) -> [gmb_fate_ops:inc(desugar_arg(A))];
desugar({'SUB', ?a, ?a, ?i(1)}) -> [gmb_fate_ops:dec()];
desugar({'SUB', A, A, ?i(1)})   -> [gmb_fate_ops:dec(desugar_arg(A))];
desugar({'STORE', ?a, A})       -> [gmb_fate_ops:push(desugar_arg(A))];
desugar({'STORE', R, ?a})       -> [gmb_fate_ops:pop(desugar_arg(R))];
desugar({switch, Arg, Type, Alts, Def}) ->
    [{switch, desugar_arg(Arg), Type, [desugar(A) || A <- Alts], desugar(Def)}];
desugar(missing) -> missing;
desugar(Code) when is_list(Code) ->
    lists:flatmap(fun desugar/1, Code);
desugar(I) -> [desugar_args(I)].

desugar_args(I) when is_tuple(I) ->
    [Op | Args] = tuple_to_list(I),
    list_to_tuple([Op | lists:map(fun desugar_arg/1, Args)]);
desugar_args(I) -> I.

desugar_arg(?s(N)) -> {var, -N};
desugar_arg(A) -> A.

%% -- Phase III --------------------------------------------------------------
%%  Constructing basic blocks

to_basic_blocks(Funs) ->
    to_basic_blocks(maps:to_list(Funs), gmb_fate_code:new()).

to_basic_blocks([{Name, {Attrs, Sig, Code}}|Left], Acc) ->
    BB = bb(Name, Code ++ [gmb_fate_ops:return()]),
    to_basic_blocks(Left, gmb_fate_code:insert_fun(Name, Attrs, Sig, BB, Acc));
to_basic_blocks([], Acc) ->
    Acc.

bb(_Name, Code) ->
    Blocks0 = blocks(Code),
    Blocks1 = optimize_blocks(Blocks0),
    Blocks  = lists:flatmap(fun split_calls/1, Blocks1),
    Labels  = maps:from_list([ {Ref, I} || {I, {Ref, _}} <- with_ixs(Blocks) ]),
    BBs     = [ set_labels(Labels, B) || B <- Blocks ],
    maps:from_list(dbg_loc_filter(BBs)).

%% Filter DBG_LOC instructions to keep one instruction per line
dbg_loc_filter(BBs) ->
    dbg_loc_filter(BBs, [], [], sets:new()).

dbg_loc_filter([], _, AllBlocks, _) ->
    lists:reverse(AllBlocks);
dbg_loc_filter([{I, []} | Rest], AllOps, AllBlocks, DbgLocs) ->
    dbg_loc_filter(Rest, [], [{I, lists:reverse(AllOps)} | AllBlocks], DbgLocs);
dbg_loc_filter([{I, [Op = {'DBG_LOC', _, _} | Ops]} | Rest], AllOps, AllBlocks, DbgLocs) ->
    case sets:is_element(Op, DbgLocs) of
        true  -> dbg_loc_filter([{I, Ops} | Rest], AllOps, AllBlocks, DbgLocs);
        false -> dbg_loc_filter([{I, Ops} | Rest], [Op | AllOps], AllBlocks, sets:add_element(Op, DbgLocs))
    end;
dbg_loc_filter([{I, [Op | Ops]} | Rest], AllOps, AllBlocks, DbgLocs) ->
    dbg_loc_filter([{I, Ops} | Rest], [Op | AllOps], AllBlocks, DbgLocs).

%% -- Break up scode into basic blocks --

-type bbref() :: reference().

%% Code to be turned into blocks.
-record(blk, { ref             :: bbref(),        %% block id
               code            :: scode(),
               catchall = none :: bbref() | none  %% closest catchall
             }).

-type bb()     :: {bbref(), bcode()}.
-type bcode()  :: [binstr()].
-type binstr() :: {jump, bbref()}
                | {jumpif, bbref()}
                | tuple().   %% FATE instruction

-spec blocks(scode()) -> [bb()].
blocks(Code) ->
    Top = make_ref(),
    blocks([#blk{ref = Top, code = Code}], []).

-spec blocks([#blk{}], [bb()]) -> [bb()].
blocks([], Acc) ->
    lists:reverse(Acc);
blocks([Blk | Blocks], Acc) ->
    block(Blk, [], Blocks, Acc).

-spec block(#blk{}, bcode(), [#blk{}], [bb()]) -> [bb()].
block(#blk{ref = Ref, code = []}, CodeAcc, Blocks, BlockAcc) ->
    blocks(Blocks, [{Ref, lists:reverse(CodeAcc)} | BlockAcc]);
block(Blk = #blk{code = [switch_body | Code]}, Acc, Blocks, BlockAcc) ->
    %% Reached the body of a switch. Clear catchall ref.
    block(Blk#blk{code = Code, catchall = none}, Acc, Blocks, BlockAcc);
block(Blk = #blk{code     = [{switch, Arg, Type, Alts, Default} | Code],
                 catchall = Catchall}, Acc, Blocks, BlockAcc) ->
    FreshBlk = fun(C, Ca) ->
                   R = make_ref(),
                   {R, [#blk{ref = R, code = C, catchall = Ca}]}
               end,
    {RestRef, RestBlk} = FreshBlk(Code, Catchall),
    {DefRef, DefBlk} =
        case Default of
            missing when Catchall == none ->
                FreshBlk([gmb_fate_ops:abort(?i(<<"Incomplete patterns">>))], none);
            missing -> {Catchall, []};
            _       -> FreshBlk(Default ++ [{jump, RestRef}], Catchall)
                       %% ^ fall-through to the outer catchall
        end,
    %% If we don't generate a switch, we need to pop the argument if on the stack.
    Pop = [{'POP', ?void} || Arg == ?a],
    {Blk1, Code1, AltBlks} =
        case Type of
            boolean ->
                [FalseCode, TrueCode] = Alts,
                {ThenRef, ThenBlk} =
                    case TrueCode of
                        missing -> {DefRef, []};
                        _       -> FreshBlk(TrueCode ++ [{jump, RestRef}], DefRef)
                    end,
                ElseCode =
                    case FalseCode of
                        missing -> [{jump, DefRef}];
                        _       -> FalseCode ++ [{jump, RestRef}]
                    end,
                case lists:usort(Alts) == [missing] of
                    true  -> {Blk#blk{code = Pop ++ [{jump, DefRef}]}, [], []};
                    false ->
                        case Arg of
                            ?i(false) -> {Blk#blk{code = ElseCode}, [], ThenBlk};
                            ?i(true)  -> {Blk#blk{code = []}, [{jump, ThenRef}], ThenBlk};
                            _         -> {Blk#blk{code = ElseCode}, [{jumpif, Arg, ThenRef}], ThenBlk}
                        end
                end;
            tuple ->
                [TCode] = Alts,
                case TCode of
                    missing -> {Blk#blk{code = Pop ++ [{jump, DefRef}]}, [], []};
                    _       -> {Blk#blk{code = Pop ++ TCode ++ [{jump, RestRef}]}, [], []}
                end;
            {variant, [_]} ->
                %% [SINGLE_CON_SWITCH] Single constructor switches don't need a
                %% switch instruction.
                [AltCode] = Alts,
                case AltCode of
                    missing -> {Blk#blk{code = Pop ++ [{jump, DefRef}]}, [], []};
                    _       -> {Blk#blk{code = Pop ++ AltCode ++ [{jump, RestRef}]}, [], []}
                end;
            {variant, _Ar} ->
                case lists:usort(Alts) == [missing] of
                    true  -> {Blk#blk{code = Pop ++ [{jump, DefRef}]}, [], []};
                    false ->
                        MkBlk = fun(missing) -> {DefRef, []};
                                   (ACode)   -> FreshBlk(ACode ++ [{jump, RestRef}], DefRef)
                                end,
                        {AltRefs, AltBs} = lists:unzip(lists:map(MkBlk, Alts)),
                        {Blk#blk{code = []}, [{switch, Arg, AltRefs}], lists:append(AltBs)}
                end
        end,
    Blk2 = Blk1#blk{catchall = DefRef}, %% Update catchall ref
    block(Blk2, Code1 ++ Acc, DefBlk ++ RestBlk ++ AltBlks ++ Blocks, BlockAcc);
block(Blk = #blk{code = [I | Code]}, Acc, Blocks, BlockAcc) ->
    block(Blk#blk{code = Code}, [I | Acc], Blocks, BlockAcc).

%% -- Reorder, inline, and remove dead blocks --

optimize_blocks(Blocks) ->
    %% We need to look at the last instruction a lot, so reverse all blocks.
    Rev       = fun(Bs) -> [ {Ref, lists:reverse(Code)} || {Ref, Code} <- Bs ] end,
    RBlocks   = Rev(Blocks),
    RBlockMap = maps:from_list(RBlocks),
    RBlocks1  = reorder_blocks(RBlocks, []),
    RBlocks2  = [ {Ref, inline_block(RBlockMap, Ref, Code)} || {Ref, Code} <- RBlocks1 ],
    RBlocks3  = shortcut_jump_chains(RBlocks2),
    RBlocks4  = remove_dead_blocks(RBlocks3),
    RBlocks5  = [ {Ref, tweak_returns(Code)} || {Ref, Code} <- RBlocks4 ],
    Rev(RBlocks5).

%% Choose the next block based on the final jump.
reorder_blocks([], Acc) ->
    lists:reverse(Acc);
reorder_blocks([{Ref, Code} | Blocks], Acc) ->
    reorder_blocks(Ref, Code, Blocks, Acc).

reorder_blocks(Ref, Code, Blocks, Acc) ->
    Acc1 = [{Ref, Code} | Acc],
    case Code of
        ['RETURN'|_]          -> reorder_blocks(Blocks, Acc1);
        [{'RETURNR', _}|_]    -> reorder_blocks(Blocks, Acc1);
        [{'CALL_T', _}|_]     -> reorder_blocks(Blocks, Acc1);
        [{'EXIT', _}|_]       -> reorder_blocks(Blocks, Acc1);
        [{'ABORT', _}|_]      -> reorder_blocks(Blocks, Acc1);
        [{switch, _, _}|_]    -> reorder_blocks(Blocks, Acc1);
        [{jump, L}|_]         ->
            NotL = fun({L1, _}) -> L1 /= L end,
            case lists:splitwith(NotL, Blocks) of
                {Blocks1, [{L, Code1} | Blocks2]} ->
                    reorder_blocks(L, Code1, Blocks1 ++ Blocks2, Acc1);
                {_, []} -> reorder_blocks(Blocks, Acc1)
            end
    end.

%% Inline short blocks (≤ 2 instructions)
inline_block(BlockMap, Ref, [{jump, L} | Code] = Code0) when L /= Ref ->
    case maps:get(L, BlockMap, nocode) of
        Dest when length(Dest) < 3 ->
            %% Remove Ref to avoid infinite loops
            inline_block(maps:remove(Ref, BlockMap), L, Dest) ++ Code;
        _ -> Code0
    end;
inline_block(_, _, Code) -> Code.

%% Shortcut jumps to blocks with a single jump
shortcut_jump_chains(RBlocks) ->
    Subst = lists:foldl(fun({L1, [{jump, L2}]}, Sub) ->
                            Sub#{ L1 => maps:get(L2, Sub, L2) };
                           (_, Sub) -> Sub end, #{}, RBlocks),
    [ {Ref, update_labels(Subst, Code)} || {Ref, Code} <- RBlocks ].

update_labels(Sub, Ref) when is_reference(Ref) ->
    maps:get(Ref, Sub, Ref);
update_labels(Sub, L) when is_list(L) ->
    lists:map(fun(X) -> update_labels(Sub, X) end, L);
update_labels(Sub, T) when is_tuple(T) ->
    list_to_tuple(update_labels(Sub, tuple_to_list(T)));
update_labels(_, X) -> X.

%% Remove unused blocks
remove_dead_blocks(Blocks = [{Top, _} | _]) ->
    BlockMap   = maps:from_list(Blocks),
    LiveBlocks = chase_labels([Top], BlockMap, #{}),
    [ B || B = {L, _} <- Blocks, maps:is_key(L, LiveBlocks) ].

chase_labels([], _, Live) -> Live;
chase_labels([L | Ls], Map, Live) ->
    Code = maps:get(L, Map),
    Jump = fun({jump, A})       -> [A || not maps:is_key(A, Live)];
              ({jumpif, _, A})  -> [A || not maps:is_key(A, Live)];
              ({switch, _, As}) -> [A || A <- As, not maps:is_key(A, Live)];
              (_)               -> [] end,
    New  = lists:flatmap(Jump, Code),
    chase_labels(New ++ Ls, Map, Live#{ L => true }).

%% Replace PUSH, RETURN by RETURNR, drop returns after tail calls.
tweak_returns(['RETURN', {'PUSH', A} | Code])          -> [{'RETURNR', A} | Code];
tweak_returns(['RETURN' | Code = [{'CALL_T', _} | _]]) -> Code;
tweak_returns(['RETURN' | Code = [{'ABORT', _} | _]])  -> Code;
tweak_returns(['RETURN' | Code = [{'EXIT', _} | _]])   -> Code;
tweak_returns(['RETURN' | Code = [loop | _]])          -> Code;
tweak_returns(Code) -> Code.

%% -- Split basic blocks at CALL instructions --
%%  Calls can only return to a new basic block. Also splits at JUMPIF instructions.

split_calls({Ref, Code}) ->
    split_calls(Ref, Code, [], []).

split_calls(Ref, [], Acc, Blocks) ->
    lists:reverse([{Ref, lists:reverse(Acc)} | Blocks]);
split_calls(Ref, [I | Code], Acc, Blocks) when element(1, I) == 'CALL';
                                               element(1, I) == 'CALL_R';
                                               element(1, I) == 'CALL_GR';
                                               element(1, I) == 'CALL_PGR';
                                               element(1, I) == 'CREATE';
                                               element(1, I) == 'CLONE';
                                               element(1, I) == 'CLONE_G';
                                               element(1, I) == 'jumpif' ->
    split_calls(make_ref(), Code, [], [{Ref, lists:reverse([I | Acc])} | Blocks]);
split_calls(Ref, [{'ABORT', _} = I | _Code], Acc, Blocks) ->
    lists:reverse([{Ref, lists:reverse([I | Acc])} | Blocks]);
split_calls(Ref, [{'EXIT', _} = I | _Code], Acc, Blocks) ->
    lists:reverse([{Ref, lists:reverse([I | Acc])} | Blocks]);
split_calls(Ref, [I | Code], Acc, Blocks) ->
    split_calls(Ref, Code, [I | Acc], Blocks).

%% -- Translate label refs to indices --

set_labels(Labels, {Ref, Code}) when is_reference(Ref) ->
    {maps:get(Ref, Labels), [ set_labels(Labels, I) || I <- Code ]};
set_labels(_Labels, loop)              -> gmb_fate_ops:jump(0);
set_labels(Labels, {jump, Ref})        -> gmb_fate_ops:jump(maps:get(Ref, Labels));
set_labels(Labels, {jumpif, Arg, Ref}) -> gmb_fate_ops:jumpif(Arg, maps:get(Ref, Labels));
set_labels(Labels, {switch, Arg, Refs}) ->
    case [ maps:get(Ref, Labels) || Ref <- Refs ] of
        [R1, R2]     -> gmb_fate_ops:switch(Arg, R1, R2);
        [R1, R2, R3] -> gmb_fate_ops:switch(Arg, R1, R2, R3);
        Rs           -> gmb_fate_ops:switch(Arg, Rs)
    end;
set_labels(_, I) -> I.

%% -- Helpers ----------------------------------------------------------------

with_ixs(Xs) ->
    lists:zip(lists:seq(0, length(Xs) - 1), Xs).

drop_common_suffix(Xs, Ys) ->
    drop_common_suffix_r(lists:reverse(Xs), lists:reverse(Ys)).

drop_common_suffix_r([X | Xs], [X | Ys]) ->
    drop_common_suffix_r(Xs, Ys);
drop_common_suffix_r(Xs, Ys) ->
    {lists:reverse(Xs), lists:reverse(Ys)}.
