Ad-hoc support for simple map comprehensions

Summary:
Precise handling of all generators in all cases may get exponential quickly, but it's possible to handle the simple case of one map generator easily.

We may revisit and add support for more/generalize the approach later.

Reviewed By: ilya-klyuchnikov

Differential Revision: D66447457

fbshipit-source-id: 959e861378c7286b0d2a0e481458f03346dc1b13

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

@@ -390,11 +390,8 @@ final class Check(pipelineContext: PipelineContext) {
           val qEnv = elab.elabQualifiers(qualifiers, env)
           checkExpr(template, BinaryType, qEnv)
           env
-        case MComprehension(kTemplate, vTemplate, qualifiers) =>
-          val qEnv = elab.elabQualifiers(qualifiers, env)
-          val (kType, _) = elab.elabExpr(kTemplate, qEnv)
-          val (vType, _) = elab.elabExpr(vTemplate, qEnv)
-          val elabType = MapType(Map(), kType, vType)
+        case m: MComprehension =>
+          val (elabType, env1) = elab.elabExpr(m, env)
           if (!subtype.subType(elabType, resTy))
             diagnosticsInfo.add(ExpectedSubtype(expr.pos, expr, expected = resTy, got = elabType))
           env

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

@@ -8,6 +8,7 @@ package com.whatsapp.eqwalizer.tc
 
 import com.whatsapp.eqwalizer.ast.Exprs._
 import com.whatsapp.eqwalizer.ast.Guards.Guard
+import com.whatsapp.eqwalizer.ast.Pats.PatVar
 import com.whatsapp.eqwalizer.ast.Types._
 import com.whatsapp.eqwalizer.ast.{BinarySpecifiers, Filters, Pats, RemoteId, Vars}
 import com.whatsapp.eqwalizer.tc.TcDiagnostics._
@@ -483,6 +484,20 @@ final class Elab(pipelineContext: PipelineContext) {
         val qEnv = elabQualifiers(qualifiers, env)
         check.checkExpr(template, BinaryType, qEnv)
         (BinaryType, env)
+      case MComprehension(kTemplate, vTemplate, List(MGenerate(gk: PatVar, gv: PatVar, gExpr))) =>
+        val (gT, gEnv) = elabExprAndCheck(gExpr, env, MapType(Map(), AnyType, AnyType))
+        val mapT = narrow.asMapTypes(gT)
+        val kvTys = mapT.flatMap(narrow.getKVType)
+        var mapsAcc: Set[MapType] = Set()
+        for (TupleType(List(kTy, vTy)) <- kvTys) {
+          val (_, kPatEnv) = elabPat.elabPat(gk, kTy, gEnv)
+          val (_, vPatEnv) = elabPat.elabPat(gv, vTy, kPatEnv)
+          val (kType, _) = elabExpr(kTemplate, vPatEnv)
+          val (vType, _) = elabExpr(vTemplate, vPatEnv)
+          val mapTy = MapType(Map(), kType, vType)
+          mapsAcc = mapsAcc + mapTy
+        }
+        (narrow.joinAndMergeMaps(mapsAcc), env)
       case MComprehension(kTemplate, vTemplate, qualifiers) =>
         val qEnv = elabQualifiers(qualifiers, env)
         val (kType, _) = elabExpr(kTemplate, qEnv)