Improve typing of maps:map when the input is a shape

Summary:
Given `F :: fun((AnyKey, AnyVal) -> Ret)`:
- Previously, `maps:map(F, #{Ka => Va, Kb => Vb, ...}) -> #{Ka | Kb | ... => Ret}`.
- After this diff, `maps:map(F, #{Ka => Va, Kb => Vb, ...}) -> #{Ka =>  Ret, Kb => Ret, ...}`.

Reviewed By: VLanvin

Differential Revision: D59867075

fbshipit-source-id: 5850c1515afc5ba27ec79d40239262e8021360b6

eqwalizer/src/main/scala/com/whatsapp/eqwalizer/tc/ElabApplyCustom.scala

@@ -363,7 +363,16 @@ class ElabApplyCustom(pipelineContext: PipelineContext) {
             val vTys = narrow.asFunType(funArgCoercedTy, 2).get.map(_.resTy)
             subtype.join(vTys)
         }
-        (DictMap(keyTy, resValTy), env1)
+        def mapValueType(argMapTy: Type, argValTy: Type): Type = argMapTy match {
+          case ShapeMap(props) =>
+            ShapeMap(props.map {
+              case ReqProp(key, ty) => ReqProp(key, argValTy)
+              case OptProp(key, ty) => OptProp(key, argValTy)
+            })
+          case UnionType(tys) => subtype.join(tys.map(mapValueType(_, argValTy)))
+          case _              => DictMap(keyTy, resValTy)
+        }
+        (mapValueType(mapType, resValTy), env1)
 
       case RemoteId("maps", "filtermap", 2) =>
         val List(funArg, map) = args

eqwalizer/test_projects/check/src/custom.erl

@@ -479,6 +479,22 @@ maps_map_2_10_neg() ->
     maps:map(F, non_kv),
     nok.
 
+-spec maps_map_2_11(
+    fun((dynamic(), dynamic()) -> ret),
+    #{ka => va, kb => vb}
+) ->
+    #{ka => ret, kb => ret}.
+maps_map_2_11(Fun, M) ->
+    maps:map(Fun, M).
+
+-spec maps_map_2_12(
+    fun((dynamic(), dynamic()) -> ret),
+    #{ka => va1, kb => vb1} | #{ka => va2, kb => vb2} | #{kb => vb3, kc => vc3}
+) ->
+    #{ka => ret, kb => ret} | #{kb => ret, kc => ret}.
+maps_map_2_12(Fun, M) ->
+    maps:map(Fun, M).
+
 -spec maps_fold_3_1()
         -> [number() | a | b].
 maps_fold_3_1() ->

eqwalizer/test_projects/check/src/custom.erl.check

@@ -572,8 +572,13 @@ maps_map_2_9_neg() ->                      | ERROR   |
     M = #{a => a, b => b},                 |         |
     F = fun erlang:'=:='/2,                |         |
     maps:map(F, M).                        |         | maps:map(F, M).
-                                           |         | Expression has type:   dict  map #D{'a' | 'b' => boolean()}
-                                           |         | Context expected type: shape map #S{a => 'a', b => 'b'}
+                                           |         | Expression has type:   #S{a := boolean(), b := boolean()}
+                                           |         | Context expected type: #S{a => 'a', b => 'b'}
+                                           |         | 
+                                           |         |   at shape key 'a':
+                                           |         |   #S{a := boolean(), b := boolean()} is not compatible with #S{a => 'a', b => 'b'}
+                                           |         |   because
+                                           |         |   boolean() is not compatible with 'a'
                                            |         |
 -spec maps_map_2_10_neg()                  |         |
         -> nok.                            |         |
@@ -584,6 +589,22 @@ maps_map_2_10_neg() ->                     | ERROR   |
                                            |         | Context expected type: #D{term() => term()}
     nok.                                   |         |
                                            |         |
+-spec maps_map_2_11(                       |         |
+    fun((dynamic(), dynamic()) -> ret),    |         |
+    #{ka => va, kb => vb}                  |         |
+) ->                                       |         |
+    #{ka => ret, kb => ret}.               |         |
+maps_map_2_11(Fun, M) ->                   | OK      |
+    maps:map(Fun, M).                      |         |
+                                           |         |
+-spec maps_map_2_12(                       |         |
+    fun((dynamic(), dynamic()) -> ret),    |         |
+    #{ka => va1, kb => vb1} | #{ka => va2,……         |
+) ->                                       |         |
+    #{ka => ret, kb => ret} | #{kb => ret,……         |
+maps_map_2_12(Fun, M) ->                   | OK      |
+    maps:map(Fun, M).                      |         |
+                                           |         |
 -spec maps_fold_3_1()                      |         |
         -> [number() | a | b].             |         |
 maps_fold_3_1() ->                         | OK      |