Better map update with union of keys

Summary: Unfold aliases and flatten unions using `narrow.asKeys` in `narrow.adjustMapType`. Also use `adjustMapType` in elaboration of `maps:find` for slightly better typing. Reviewed By: ilya-klyuchnikov Differential Revision: D65536147 fbshipit-source-id: 1b0f74d3ac8ac4291295c5181040e4189c9138a4
WhatsApp · Nov 6, 2024 · 4bef349 · 4bef349
1 parent 3bafe37
commit 4bef349
Show file tree

Hide file tree

Showing 2 changed files with 16 additions and 7 deletions.
diff --git a/eqwalizer/src/main/scala/com/whatsapp/eqwalizer/tc/ElabApplyCustom.scala b/eqwalizer/src/main/scala/com/whatsapp/eqwalizer/tc/ElabApplyCustom.scala
@@ -258,9 +258,9 @@ class ElabApplyCustom(pipelineContext: PipelineContext) {
         val List(keyTy, mapTy) = argTys
         val mapCoercedTy = coerce(map, mapTy, anyMapTy)
         val mapType = narrow.asMapType(mapCoercedTy)
-        val valTy = Key.fromType(keyTy) match {
-          case Some(key) =>
-            narrow.getValType(key, mapType)
+        val valTy = narrow.asKeys(keyTy) match {
+          case Some(keys) =>
+            subtype.join(keys.map(narrow.getValType(_, mapType)))
           case None =>
             narrow.getValType(mapType)
         }

diff --git a/eqwalizer/src/main/scala/com/whatsapp/eqwalizer/tc/Narrow.scala b/eqwalizer/src/main/scala/com/whatsapp/eqwalizer/tc/Narrow.scala
@@ -447,9 +447,17 @@ class Narrow(pipelineContext: PipelineContext) {
       case BoundedDynamicType(bound) =>
         BoundedDynamicType(adjustMapType(bound, keyT, valT))
       case mapType: MapType =>
-        Key.fromType(keyT) match {
-          case Some(key) =>
-            MapType(mapType.props.updated(key, Prop(req = true, valT)), mapType.kType, mapType.vType)
+        asKeys(keyT) match {
+          case Some(keys) if keys.size == 1 =>
+            MapType(mapType.props.updated(keys.head, Prop(req = true, valT)), mapType.kType, mapType.vType)
+          case Some(keys) =>
+            keys.foldLeft(mapType) { case (mapType, key) =>
+              val props = mapType.props.updatedWith(key) {
+                case Some(prop) => Some(Prop(prop.req, subtype.join(valT, prop.tp)))
+                case None       => Some(Prop(req = false, valT))
+              }
+              MapType(props, mapType.kType, mapType.vType)
+            }
           case None if subtype.isDynamicType(mapType.kType) && subtype.isDynamicType(mapType.vType) =>
             mapType
           case None =>
@@ -530,7 +538,8 @@ class Narrow(pipelineContext: PipelineContext) {
       case RemoteType(rid, args) =>
         val body = util.getTypeDeclBody(rid, args)
         asKeys(body)
-      case _ => Key.fromType(t).map(Set(_))
+      case NoneType => Some(Set())
+      case _        => Key.fromType(t).map(Set(_))
     }
 
   private def mergeMaps(s1: MapType, s2: MapType, inOrder: Boolean): MapType = {