Introduce is_map_key/2 guard BIF

This complements the `map_get/2` guard BIF introduced in #1784.

Rationale.

`map_get/2` allows accessing map fields in guards, but it might be
problematic in more complex guard expressions, for example:

foo(X) when map_get(a, X) =:= 1 or is_list(X) -> ...

The `is_list/1` part of the guard could never succeed since the
`map_get/2` guard would fail the whole guard expression. In this
situation, this could be solved by using `;` instead of `or` to separate
the guards, but it is not possible in every case.

To solve this situation, this PR proposes a `is_map_key/2` guard that
allows to check if a map has key inside a guard before trying to access
that key. When combined with `is_map/1` this allows to construct a
purely boolean guard expression testing a value of a key in a map.

Implementation.

Given the use case motivating the introduction of this function, the PR
contains compiler optimisations that produce optimial code for the
following guard expression:

foo(X) when is_map(X) and is_map_key(a, X) and map_get(a, X) =:= 1 -> ok;
foo(_) -> error.

Given all three tests share the failure label, the `is_map_key/2` and
`is_map/2` tests are optimised away.

As with `map_get/2` the `is_map_key/2` BIF is allowed in match specs.

6 files changed, 54 insertions, 0 deletions

diff --git a/lib/compiler/src/beam_dead.erl b/lib/compiler/src/beam_dead.erl
index dbbaae05eb..762c7bdf9e 100644
--- a/lib/compiler/src/beam_dead.erl
+++ b/lib/compiler/src/beam_dead.erl
@@ -392,6 +392,26 @@ backward([{bif,'or',{f,To0},[Dst,{atom,false}],Dst}=I|Is], D,
 	_ ->
 	    backward(Is, D, [I|Acc])
     end;
+backward([{bif,map_get,{f,FF},[Key,Map],_}=I0,
+          {test,has_map_fields,{f,FT}=F,[Map|Keys0]}=I1|Is], D, Acc) when FF =/= 0 ->
+    case shortcut_label(FF, D) of
+        FT ->
+            case lists:delete(Key, Keys0) of
+                [] ->
+                    backward([I0|Is], D, Acc);
+                Keys ->
+                    Test = {test,has_map_fields,F,[Map|Keys]},
+                    backward([Test|Is], D, [I0|Acc])
+            end;
+        _ ->
+            backward([I1|Is], D, [I0|Acc])
+    end;
+backward([{bif,map_get,{f,FF},[_,Map],_}=I0,
+          {test,is_map,{f,FT},[Map]}=I1|Is], D, Acc) when FF =/= 0 ->
+    case shortcut_label(FF, D) of
+        FT -> backward([I0|Is], D, Acc);
+        _ -> backward([I1|Is], D, [I0|Acc])
+    end;
 backward([I|Is], D, Acc) ->
     backward(Is, D, [I|Acc]);
 backward([], _D, Acc) -> Acc.
diff --git a/lib/compiler/src/beam_peep.erl b/lib/compiler/src/beam_peep.erl
index eb3192fe8f..920fb00397 100644
--- a/lib/compiler/src/beam_peep.erl
+++ b/lib/compiler/src/beam_peep.erl
@@ -77,6 +77,12 @@ peep([{bif,tuple_size,_,[_]=Ops,Dst}=I|Is], SeenTests0, Acc) ->
     %% Kill all remembered tests that depend on the destination register.
     SeenTests = kill_seen(Dst, SeenTests1),
     peep(Is, SeenTests, [I|Acc]);
+peep([{bif,map_get,_,[Key,Map],Dst}=I|Is], SeenTests0, Acc) ->
+    %% Pretend that we have seen {test,has_map_fields,_,[Map,Key]}
+    SeenTests1 = gb_sets:add({has_map_fields,[Map,Key]}, SeenTests0),
+    %% Kill all remembered tests that depend on the destination register.
+    SeenTests = kill_seen(Dst, SeenTests1),
+    peep(Is, SeenTests, [I|Acc]);
 peep([{bif,_,_,_,Dst}=I|Is], SeenTests0, Acc) ->
     %% Kill all remembered tests that depend on the destination register.
     SeenTests = kill_seen(Dst, SeenTests0),
diff --git a/lib/compiler/src/beam_type.erl b/lib/compiler/src/beam_type.erl
index 28f36db399..12da8c9446 100644
--- a/lib/compiler/src/beam_type.erl
+++ b/lib/compiler/src/beam_type.erl
@@ -462,6 +462,9 @@ update({set,[D],[Index,Reg],{bif,element,_}}, Ts0) ->
               end,
     Ts = tdb_meet(Reg, {tuple,min_size,MinSize,[]}, Ts0),
     tdb_store(D, any, Ts);
+update({set,[D],[_Key,Map],{bif,map_get,_}}, Ts0) ->
+    Ts = tdb_meet(Map, map, Ts0),
+    tdb_store(D, any, Ts);
 update({set,[D],Args,{bif,N,_}}, Ts) ->
     Ar = length(Args),
     BoolOp = erl_internal:new_type_test(N, Ar) orelse
diff --git a/lib/compiler/src/erl_bifs.erl b/lib/compiler/src/erl_bifs.erl
index 70b36f029e..a7452aebc8 100644
--- a/lib/compiler/src/erl_bifs.erl
+++ b/lib/compiler/src/erl_bifs.erl
@@ -94,6 +94,7 @@ is_pure(erlang, is_function, 1) -> true;
 is_pure(erlang, is_integer, 1) -> true;
 is_pure(erlang, is_list, 1) -> true;
 is_pure(erlang, is_map, 1) -> true;
+is_pure(erlang, is_map_key, 2) -> true;
 is_pure(erlang, is_number, 1) -> true;
 is_pure(erlang, is_pid, 1) -> true;
 is_pure(erlang, is_port, 1) -> true;
diff --git a/lib/compiler/src/v3_codegen.erl b/lib/compiler/src/v3_codegen.erl
index 8e73b613a0..9652a8476d 100644
--- a/lib/compiler/src/v3_codegen.erl
+++ b/lib/compiler/src/v3_codegen.erl
@@ -589,6 +589,7 @@ is_gc_bif(element, 2) -> false;
 is_gc_bif(get, 1) -> false;
 is_gc_bif(tuple_size, 1) -> false;
 is_gc_bif(map_get, 2) -> false;
+is_gc_bif(is_map_key, 2) -> false;
 is_gc_bif(Bif, Arity) ->
     not (erl_internal:bool_op(Bif, Arity) orelse
 	 erl_internal:new_type_test(Bif, Arity) orelse
@@ -1620,6 +1621,11 @@ test_cg(is_boolean, [#k_atom{val=Val}], Fail, I, Vdb, Bef, St) ->
 	     false -> [{jump,{f,Fail}}]
 	 end,
     {Is,Aft,St};
+test_cg(is_map_key, As, Fail, I, Vdb, Bef, St) ->
+    [Key,Map] = cg_reg_args(As, Bef),
+    Aft = clear_dead(Bef, I, Vdb),
+    F = {f,Fail},
+    {[{test,is_map,F,[Map]},{test,has_map_fields,F,Map,{list,[Key]}}],Aft,St};
 test_cg(Test, As, Fail, I, Vdb, Bef, St) ->
     Args = cg_reg_args(As, Bef),
     Aft = clear_dead(Bef, I, Vdb),
diff --git a/lib/compiler/test/map_SUITE.erl b/lib/compiler/test/map_SUITE.erl
index e98c295da6..6badc7a8b8 100644
--- a/lib/compiler/test/map_SUITE.erl
+++ b/lib/compiler/test/map_SUITE.erl
@@ -1203,12 +1203,18 @@ t_guard_bifs(Config) when is_list(Config) ->
     true   = map_guard_empty_2(),
     true   = map_guard_head(#{a=>1}),
     false  = map_guard_head([]),
+    true   = map_get_head(#{a=>1}),
+    false  = map_get_head([]),
+    true   = map_is_key_head(#{a=>1}),
+    false  = map_is_key_head(#{}),
     true   = map_guard_body(#{a=>1}),
     false  = map_guard_body({}),
     true   = map_guard_pattern(#{a=>1, <<"hi">> => "hi" }),
     false  = map_guard_pattern("list"),
     true   = map_guard_tautology(),
     true   = map_guard_ill_map_size(),
+    true   = map_field_check_sequence(#{a=>1}),
+    false  = map_field_check_sequence(#{}),
     ok.
 
 map_guard_empty() when is_map(#{}); false -> true.
@@ -1218,6 +1224,12 @@ map_guard_empty_2() when true; #{} andalso false -> true.
 map_guard_head(M) when is_map(M) -> true;
 map_guard_head(_) -> false.
 
+map_get_head(M) when map_get(a, M) =:= 1 -> true;
+map_get_head(_) -> false.
+
+map_is_key_head(M) when is_map_key(a, M) -> true;
+map_is_key_head(M) -> false.
+
 map_guard_body(M) -> is_map(M).
 
 map_guard_pattern(#{}) -> true;
@@ -1227,6 +1239,12 @@ map_guard_tautology() when #{} =:= #{}; true -> true.
 
 map_guard_ill_map_size() when true; map_size(0) -> true.
 
+map_field_check_sequence(M)
+  when is_map(M) andalso is_map_key(a, M) andalso (map_get(a, M) == 1) ->
+    true;
+map_field_check_sequence(_) ->
+    false.
+
 t_guard_sequence(Config) when is_list(Config) ->
 	{1, "a"} = map_guard_sequence_1(#{seq=>1,val=>id("a")}),
 	{2, "b"} = map_guard_sequence_1(#{seq=>2,val=>id("b")}),