[typer] make Map comprehension syntax more robust (closes #6620)

src/typing/typer.ml

@@ -3437,12 +3437,38 @@ and type_expr ctx (e,p) (with_type:with_type) =
 			| EFor(it,e2) -> (EFor (it, map_compr e2),p)
 			| EWhile(cond,e2,flag) -> (EWhile (cond,map_compr e2,flag),p)
 			| EIf (cond,e2,None) -> (EIf (cond,map_compr e2,None),p)
+			| EIf (cond,e2,Some e3) -> (EIf (cond,map_compr e2,Some (map_compr e3)),p)
+			| ESwitch(e1,cases,eo) ->
+				let cases = List.map (fun (el,eg,e,p) -> el,eg,Option.map map_compr e,p) cases in
+				let eo = Option.map (fun (eo,p) -> Option.map map_compr eo,p) eo in
+				ESwitch(e1,cases,eo),p
+			| ETry(e1,catches) ->
+				let catches = List.map (fun (n,ct,e,p) -> n,ct,map_compr e,p) catches in
+				ETry(map_compr e1,catches),p
+			| ETernary (cond,e2,e3) -> (ETernary (cond,map_compr e2,map_compr e3),p)
 			| EBlock [e] -> (EBlock [map_compr e],p)
+			| EBlock el -> begin match List.rev el with
+				| e :: el -> (EBlock ((List.rev el) @ [map_compr e]),p)
+				| [] -> e,p
+				end
 			| EParenthesis e2 -> (EParenthesis (map_compr e2),p)
 			| EBinop(OpArrow,a,b) ->
 				et := (ENew(({tpackage=["haxe";"ds"];tname="Map";tparams=[];tsub=None},null_pos),[]),p);
 				(ECall ((EField ((EConst (Ident v.v_name),p),"set"),p),[a;b]),p)
-			| _ ->
+			| EContinue | EBreak | EReturn _ | EThrow _ ->
+				(e,p)
+			| ECast(e1,tho) ->
+				ECast(map_compr e1,tho),p
+			| EUntyped e1 ->
+				EUntyped(map_compr e1),p
+			| EDisplay(e1,b) ->
+				EDisplay(map_compr e1,b),p
+			| EMeta(m,e1) ->
+				EMeta(m,map_compr e1),p
+			| ECheckType(e1,ct) ->
+				ECheckType(map_compr e1,ct),p
+			| EConst _ | EArray _ | EBinop _ | EField _ | EObjectDecl _ | EArrayDecl _
+			| ECall _ | ENew _ | EUnop _ | EVars _ | EFunction _ | EDisplayNew _ ->
 				et := (EArrayDecl [],p);
 				(ECall ((EField ((EConst (Ident v.v_name),p),"push"),p),[(e,p)]),p)
 		in

tests/unit/src/unit/TestMain.hx

@@ -93,6 +93,7 @@ class TestMain {
 			#if ((dce == "full") && !interp && !as3)
 			new TestDCE(),
 			#end
+			new TestMapComprehension(),
 			// #if ( (java || neko) && !macro && !interp)
 			// new TestThreads(),
 			// #end

tests/unit/src/unit/TestMapComprehension.hx

@@ -0,0 +1,37 @@
+package unit;
+
+class TestMapComprehension extends unit.Test {
+	public function testBasic() {
+		mapEq([for (i in 0...2) i => i], [0 => 0, 1 => 1]);
+		mapEq([for (i in 0...2) (i => i)], [0 => 0, 1 => 1]);
+		mapEq([for (i in 0...2) cast i => i], [0 => 0, 1 => 1]);
+		mapEq([for (i in 0...2) untyped i => i], [0 => 0, 1 => 1]);
+		mapEq([for (i in 0...2) (i => i : Void)], [0 => 0, 1 => 1]);
+		mapEq([for (i in 0...2) if (i == 1) i => i], [1 => 1]);
+		mapEq([for (i in 0...2) if (i == 0) i => i else i => i * 2], [0 => 0, 1 => 2]);
+		mapEq([for (i in 0...2) (i == 0) ? i => i : i => i * 2], [0 => 0, 1 => 2]);
+		mapEq([for (i in 0...2) switch i { case 0: i => i; case _: i => i * 2; }], [0 => 0, 1 => 2]);
+		mapEq([for (i in 0...2) try i => i catch(e:Dynamic) i => i * 2], [0 => 0, 1 => 1]);
+		mapEq([for (i in 0...2) try { throw null; i => i; } catch(e:Dynamic) i => i * 2], [0 => 0, 1 => 2]);
+		mapEq([for (i in 0...2) try { throw null; } catch(e:Dynamic) i => i * 2], [0 => 0, 1 => 2]);
+
+		mapEq([for (i in 0...2) { continue; i => i; }], new Map<Int, Int>());
+	}
+
+	public function testMark() {
+		var specialValue = null;
+		function doSomething(i:Int) specialValue = ("special" + i);
+		var m = [for(i in 0...3) if (i != 1) i => 'number $i' else { doSomething(i); continue; }];
+		mapEq(m, [0 => "number 0", 2 => "number 2"]);
+		eq(specialValue, "special1");
+	}
+
+	function mapEq<K, V>(m1:Map<K, V>, m2:Map<K, V>, ?p:haxe.PosInfos) {
+		for (k1 in m1.keys()) {
+			eq(m1[k1], m2[k1], p);
+		}
+		for (k2 in m2.keys()) {
+			eq(m1[k2], m2[k2], p);
+		}
+	}
+}