[flow][perf] Almost disjoint union optimization

Summary:
Why am I doing this? [Hermes ESTree types](https://github.com/facebook/hermes/blob/main/tools/hermes-parser/js/hermes-estree/src/types.js) has almost disjoint union and it's very hard to change (we can't really change ESTree), and it's very slow to check. With this change, the ones in hermes-parser repo is still not optimized yet, since they use interfaces, but at least we are much closer.

In this diff, we start track the currently failing cases of disjoint union optimizations due to non-unique keys, and we group them together. Therefore, given the following union:

```
type Node = {type: 'lit', subType: 1} | {type: 'lit', subType: 2} | {type: '+', e1: Node, e2: Node}  | {type: '-', e1: Node, e2: Node}
```

the grouping will be:

```
type: {
'lit' => {type: 'lit', subType: 1}, {type: 'lit', subType: 2}
'+' => {type: '+', e1: Node, e2: Node}
'-' => {type: '-', e1: Node, e2: Node}
}
```

When there are only one item in the bucket, the usual optimization will apply and we are in a happy path. When it's not, then we will emit `Unknown` as the member test result, but we only pay the extra cost proportional to the number of non-unique keys.

We don't see a significant change in perf number, probably because such code that will can benefit from this optimization is either too small, or already impossible to land. However, I think the infrastructure built in this diff can still be helpful for other purposes. e.g. consider

```
const node: Node = {
  type: '+',
  e1,
  // still typing e2
}
```

right now we generate a big error with all the branches, but perhaps we can use the information to only show that it's incompatible with plus node.

In addition to this optimization, I still keep this case as an error with `$Flow$EnforceOptimize`, but give it a different error code. The infra built in this diff also allows us to show all the non-unique keys at once, instead of the first one we find.

Changelog: [internal]

Reviewed By: panagosg7

Differential Revision: D54650850

fbshipit-source-id: 67e084a098e4c2675eeecb654568b4075345956d

src/common/errors/error_codes.ml

@@ -156,6 +156,7 @@ type error_code =
   | UnderconstrainedImplicitInstantiation
   | UninitializedInstanceProperty
   | UnionUnoptimizable
+  | UnionPartiallyOptimizableNonUniqueKeys
   | UnnecessaryInvariant
   | UnnecessaryOptionalChain
   | UnnecessaryDeclareTypeOnlyExport
@@ -370,6 +371,7 @@ let string_of_code : error_code -> string = function
   | UnderconstrainedImplicitInstantiation -> "underconstrained-implicit-instantiation"
   | UninitializedInstanceProperty -> "uninitialized-instance-property"
   | UnionUnoptimizable -> "union-unoptimizable"
+  | UnionPartiallyOptimizableNonUniqueKeys -> "union-partially-optimizable-non-unique-keys"
   | UnnecessaryInvariant -> "unnecessary-invariant"
   | UnnecessaryOptionalChain -> "unnecessary-optional-chain"
   | UnnecessaryDeclareTypeOnlyExport -> "unnecessary-declare-type-only-export"

src/typing/debug_js.ml

@@ -1950,6 +1950,8 @@ let dump_error_message =
         (string_of_aloc loc)
         (SSet.to_string available_platforms)
         (SSet.to_string required_platforms)
+    | EUnionPartialOptimizationNonUniqueKey { loc; _ } ->
+      spf "EUnionPartialOptimizationNonUniqueKey (%s)" (string_of_aloc loc)
     | EUnionOptimization { loc; _ } -> spf "EUnionOptimization (%s)" (string_of_aloc loc)
     | EUnionOptimizationOnNonUnion { loc; _ } ->
       spf "EUnionOptimizationOnNonUnion (%s)" (string_of_aloc loc)

src/typing/errors/error_message.ml

@@ -692,6 +692,10 @@ and 'loc t' =
       available_platforms: SSet.t;
       required_platforms: SSet.t;
     }
+  | EUnionPartialOptimizationNonUniqueKey of {
+      loc: 'loc;
+      non_unique_keys: 'loc virtual_reason Nel.t Type.UnionRep.UnionEnumMap.t NameUtils.Map.t;
+    }
   | EUnionOptimization of {
       loc: 'loc;
       kind: 'loc Type.UnionRep.optimized_error;
@@ -1490,16 +1494,20 @@ let rec map_loc_of_error_message (f : 'a -> 'b) : 'a t' -> 'b t' =
     EInvalidTypeCastSyntax { loc = f loc; enabled_casting_syntax }
   | EMissingPlatformSupport { loc; available_platforms; required_platforms } ->
     EMissingPlatformSupport { loc = f loc; available_platforms; required_platforms }
+  | EUnionPartialOptimizationNonUniqueKey { loc; non_unique_keys } ->
+    EUnionPartialOptimizationNonUniqueKey
+      {
+        loc = f loc;
+        non_unique_keys =
+          NameUtils.Map.map (Type.UnionRep.UnionEnumMap.map (Nel.map map_reason)) non_unique_keys;
+      }
   | EUnionOptimization { loc; kind } ->
     let kind =
       let open UnionRep in
       match kind with
       | ContainsUnresolved r -> ContainsUnresolved (map_reason r)
       | NoCandidateMembers -> NoCandidateMembers
       | NoCommonKeys -> NoCommonKeys
-      | NonUniqueKeys map ->
-        NonUniqueKeys
-          (NameUtils.Map.map (fun (enum, r1, r2) -> (enum, map_reason r1, map_reason r2)) map)
     in
     EUnionOptimization { loc = f loc; kind }
   | EUnionOptimizationOnNonUnion { loc; arg } ->
@@ -1771,6 +1779,7 @@ let util_use_op_of_msg nope util = function
   | EInvalidRendersTypeArgument _
   | EInvalidTypeCastSyntax _
   | EMissingPlatformSupport _
+  | EUnionPartialOptimizationNonUniqueKey _
   | EUnionOptimization _
   | EUnionOptimizationOnNonUnion _ ->
     nope
@@ -1939,6 +1948,7 @@ let loc_of_msg : 'loc t' -> 'loc option = function
   | ERefComponentProp { spread = loc; _ }
   | ETypeGuardIncompatibleWithFunctionKind { loc; _ }
   | EMissingPlatformSupport { loc; _ }
+  | EUnionPartialOptimizationNonUniqueKey { loc; _ }
   | EUnionOptimization { loc; _ }
   | EUnionOptimizationOnNonUnion { loc; _ } ->
     Some loc
@@ -5609,7 +5619,7 @@ let friendly_message_of_msg loc_of_aloc msg =
       @ [text " is missing."]
     in
     Normal { features }
-  | EUnionOptimization { loc = _; kind } ->
+  | EUnionPartialOptimizationNonUniqueKey { loc = _; non_unique_keys } ->
     let string_of_union_enum =
       let open Type.UnionEnum in
       function
@@ -5620,19 +5630,32 @@ let friendly_message_of_msg loc_of_aloc msg =
       | Void -> code "undefined"
       | Null -> code "null"
     in
-    let string_of_non_unique_key (name, (enum, r, r')) =
-      [
-        text "Key ";
-        code (display_string_of_name name);
-        text " has value ";
-        string_of_union_enum enum;
-        text " in both ";
-        ref r;
-        text " and ";
-        ref r';
-        text ".";
-      ]
+    let string_of_non_unique_key (name, map) =
+      let ref_texts (r0, rs) = ref r0 :: List.concat_map (fun r -> [text ", "; ref r]) rs in
+      map
+      |> Type.UnionRep.UnionEnumMap.bindings
+      |> List.concat_map (fun (enum, rs) ->
+             [
+               text "\n - Key ";
+               code (display_string_of_name name);
+               text " has value ";
+               string_of_union_enum enum;
+               text " in ";
+             ]
+             @ ref_texts rs
+             @ [text "."]
+         )
+    in
+    let keys =
+      non_unique_keys |> NameUtils.Map.bindings |> Base.List.concat_map ~f:string_of_non_unique_key
+    in
+    let features =
+      text
+        "Union could not be fully optimized internally. The following keys have non-unique values:"
+      :: keys
     in
+    Normal { features }
+  | EUnionOptimization { loc = _; kind } ->
     let kind =
       let open Type.UnionRep in
       match kind with
@@ -5651,21 +5674,6 @@ let friendly_message_of_msg loc_of_aloc msg =
           text "boolean literal, void or null types.";
         ]
       | NoCommonKeys -> [text "There are no common keys among the members of the union."]
-      | NonUniqueKeys map ->
-        let bindings = NameUtils.Map.bindings map in
-        begin
-          match bindings with
-          | [] -> [text ""]
-          | [x] -> string_of_non_unique_key x
-          | xs ->
-            let keys =
-              xs
-              |> Base.List.map ~f:string_of_non_unique_key
-              |> Base.List.map ~f:(fun x -> [text " - "] @ x @ [text "\n"])
-              |> Base.List.concat
-            in
-            text "The following keys have non-unique values:\n" :: keys
-        end
     in
     let features = text "Union could not be optimized internally. " :: kind in
     Normal { features }
@@ -6017,5 +6025,6 @@ let error_code_of_message err : error_code option =
   | ETSSyntax _ -> Some TSSyntax
   | EInvalidTypeCastSyntax _ -> Some InvalidTypeCastSyntax
   | EMissingPlatformSupport _ -> Some MissingPlatformSupport
+  | EUnionPartialOptimizationNonUniqueKey _ -> Some UnionPartiallyOptimizableNonUniqueKeys
   | EUnionOptimization _ -> Some UnionUnoptimizable
   | EUnionOptimizationOnNonUnion _ -> Some UnionUnoptimizable

src/typing/speculation_kit.ml

@@ -799,12 +799,34 @@ module Make (Flow : INPUT) : OUTPUT = struct
           ~find_resolved:(Context.find_resolved cx)
           ~find_props:(Context.find_props cx)
       in
-      Base.Result.iter_error specialization ~f:(fun kind ->
+      begin
+        match specialization with
+        | Error kind ->
           add_output
             cx
             ~trace
             (Error_message.EUnionOptimization { loc = loc_of_reason reason; kind })
-      );
+        | Ok
+            ( UnionRep.AlmostDisjointUnionWithPossiblyNonUniqueKeys map
+            | UnionRep.PartiallyOptimizedAlmostDisjointUnionWithPossiblyNonUniqueKeys map ) ->
+          let non_unique_keys =
+            map
+            |> NameUtils.Map.map (fun map ->
+                   map
+                   |> UnionRep.UnionEnumMap.filter (fun _ (_, ts) -> ts <> [])
+                   |> UnionRep.UnionEnumMap.map (Nel.map TypeUtil.reason_of_t)
+               )
+            |> NameUtils.Map.filter (fun _ map -> not (UnionRep.UnionEnumMap.is_empty map))
+          in
+          if not (NameUtils.Map.is_empty non_unique_keys) then
+            add_output
+              cx
+              ~trace
+              (Error_message.EUnionPartialOptimizationNonUniqueKey
+                 { loc = loc_of_reason reason; non_unique_keys }
+              )
+        | Ok _ -> ()
+      end;
       true
     | UnionCases (use_op, l, rep, _ts) ->
       if not (UnionRep.is_optimized_finally rep) then

src/typing/type.ml

@@ -2374,16 +2374,17 @@ and UnionRep : sig
   type finally_optimized_rep =
     | EnumUnion of UnionEnumSet.t
     | PartiallyOptimizedUnionEnum of UnionEnumSet.t
-    | DisjointUnion of TypeTerm.t UnionEnumMap.t NameUtils.Map.t
-    | PartiallyOptimizedDisjointUnion of TypeTerm.t UnionEnumMap.t NameUtils.Map.t
+    | AlmostDisjointUnionWithPossiblyNonUniqueKeys of
+        TypeTerm.t Nel.t UnionEnumMap.t NameUtils.Map.t
+    | PartiallyOptimizedAlmostDisjointUnionWithPossiblyNonUniqueKeys of
+        TypeTerm.t Nel.t UnionEnumMap.t NameUtils.Map.t
     | Empty
     | Singleton of TypeTerm.t
 
   type 'loc optimized_error =
     | ContainsUnresolved of 'loc virtual_reason
     | NoCandidateMembers (* E.g. `number | string` *)
     | NoCommonKeys
-    | NonUniqueKeys of (UnionEnum.t * 'loc virtual_reason * 'loc virtual_reason) NameUtils.Map.t
 
   val optimize :
     t ->
@@ -2477,16 +2478,17 @@ end = struct
   type finally_optimized_rep =
     | EnumUnion of UnionEnumSet.t
     | PartiallyOptimizedUnionEnum of UnionEnumSet.t
-    | DisjointUnion of TypeTerm.t UnionEnumMap.t NameUtils.Map.t
-    | PartiallyOptimizedDisjointUnion of TypeTerm.t UnionEnumMap.t NameUtils.Map.t
+    | AlmostDisjointUnionWithPossiblyNonUniqueKeys of
+        TypeTerm.t Nel.t UnionEnumMap.t NameUtils.Map.t
+    | PartiallyOptimizedAlmostDisjointUnionWithPossiblyNonUniqueKeys of
+        TypeTerm.t Nel.t UnionEnumMap.t NameUtils.Map.t
     | Empty
     | Singleton of TypeTerm.t
 
   type 'loc optimized_error =
     | ContainsUnresolved of 'loc virtual_reason
     | NoCandidateMembers
     | NoCommonKeys
-    | NonUniqueKeys of (UnionEnum.t * 'loc virtual_reason * 'loc virtual_reason) NameUtils.Map.t
 
   (** union rep is:
       - list of members in declaration order, with at least 2 elements
@@ -2652,33 +2654,37 @@ end = struct
         ([], false)
         ts
     in
-    let unique_values =
-      let rec unique_values ~reason_of_t ~reasonless_eq idx = function
-        | [] -> Ok idx
+    let almost_unique_values =
+      let rec almost_unique_values
+          ~reason_of_t ~reasonless_eq (hybrid_idx : TypeTerm.t Nel.t UnionEnumMap.t) = function
+        | [] -> hybrid_idx
         | (enum, t) :: values -> begin
-          match UnionEnumMap.find_opt enum idx with
-          | None -> unique_values ~reason_of_t ~reasonless_eq (UnionEnumMap.add enum t idx) values
-          | Some t' ->
-            if reasonless_eq t t' then
-              (* This corresponds to the case
-               * type T = { f: "a" };
-               * type Union = T | T;
-               *)
-              unique_values ~reason_of_t ~reasonless_eq idx values
-            else
-              Error (enum, reason_of_t t, reason_of_t t')
+          let hybrid_idx =
+            UnionEnumMap.adjust
+              enum
+              (function
+                | None -> Nel.one t
+                | Some l when Nel.exists (reasonless_eq t) l ->
+                  (* This corresponds to the case
+                   * type T = { f: "a" };
+                   * type Union = T | T;
+                   *)
+                  l
+                | Some l -> Nel.cons t l)
+              hybrid_idx
+          in
+          almost_unique_values ~reason_of_t ~reasonless_eq hybrid_idx values
         end
       in
-      unique_values UnionEnumMap.empty
+      almost_unique_values UnionEnumMap.empty
     in
-    let unique ~reason_of_t ~reasonless_eq idx =
+    let almost_unique ~reason_of_t ~reasonless_eq idx =
       NameUtils.Map.fold
-        (fun key values (acc, err_acc) ->
-          match unique_values ~reason_of_t ~reasonless_eq values with
-          | Error err -> (acc, NameUtils.Map.add key err err_acc)
-          | Ok idx -> (NameUtils.Map.add key idx acc, err_acc))
+        (fun key values acc ->
+          let hybrid_idx = almost_unique_values ~reason_of_t ~reasonless_eq values in
+          NameUtils.Map.add key hybrid_idx acc)
         idx
-        (NameUtils.Map.empty, NameUtils.Map.empty)
+        NameUtils.Map.empty
     in
     let intersect_props (base_props, candidates) =
       (* Compute the intersection of properties of objects that have singleton types *)
@@ -2709,13 +2715,11 @@ end = struct
               Error NoCommonKeys
             else
               (* Ensure that enums map to unique types *)
-              let (map, err_map) = unique ~reason_of_t ~reasonless_eq idx in
-              if NameUtils.Map.is_empty map then
-                Error (NonUniqueKeys err_map)
-              else if partial then
-                Ok (PartiallyOptimizedDisjointUnion map)
+              let hybrid_map = almost_unique ~reason_of_t ~reasonless_eq idx in
+              if partial then
+                Ok (PartiallyOptimizedAlmostDisjointUnionWithPossiblyNonUniqueKeys hybrid_map)
               else
-                Ok (DisjointUnion map)
+                Ok (AlmostDisjointUnionWithPossiblyNonUniqueKeys hybrid_map)
         end
 
   let optimize_ rep ~reason_of_t ~reasonless_eq ~flatten ~find_resolved ~find_props =
@@ -2779,8 +2783,8 @@ end = struct
           Yes
         else
           Conditional t
-      | Some (DisjointUnion _) -> No
-      | Some (PartiallyOptimizedDisjointUnion _) -> Unknown
+      | Some (AlmostDisjointUnionWithPossiblyNonUniqueKeys _) -> No
+      | Some (PartiallyOptimizedAlmostDisjointUnionWithPossiblyNonUniqueKeys _) -> Unknown
       | Some (EnumUnion tset) ->
         if UnionEnumSet.mem tcanon tset then
           Yes
@@ -2794,7 +2798,7 @@ end = struct
     end
     | None -> failwith "quick_mem_enum is defined only for canonizable type"
 
-  let lookup_disjoint_union find_resolved prop_map ~partial map =
+  let lookup_almost_disjoint_union find_resolved prop_map ~partial map =
     NameUtils.Map.fold
       (fun key idx acc ->
         if acc <> Unknown then
@@ -2805,7 +2809,8 @@ end = struct
             match canon_prop find_resolved p with
             | Some enum -> begin
               match UnionEnumMap.find_opt enum idx with
-              | Some t' -> Conditional t'
+              | Some (t, []) -> Conditional t
+              | Some (_, _ :: _) -> Unknown
               | None ->
                 if partial then
                   Unknown
@@ -2835,9 +2840,10 @@ end = struct
           Yes
         else
           Conditional t
-      | Some (DisjointUnion map) -> lookup_disjoint_union find_resolved prop_map ~partial:false map
-      | Some (PartiallyOptimizedDisjointUnion map) ->
-        lookup_disjoint_union find_resolved prop_map ~partial:true map
+      | Some (AlmostDisjointUnionWithPossiblyNonUniqueKeys map) ->
+        lookup_almost_disjoint_union find_resolved prop_map ~partial:false map
+      | Some (PartiallyOptimizedAlmostDisjointUnionWithPossiblyNonUniqueKeys map) ->
+        lookup_almost_disjoint_union find_resolved prop_map ~partial:true map
       | Some (EnumUnion _) -> No
       | Some (PartiallyOptimizedUnionEnum _) -> Unknown
     end
@@ -2852,8 +2858,10 @@ end = struct
     | Some (EnumUnion _) -> "Enum"
     | Some Empty -> "Empty"
     | Some (Singleton _) -> "Singleton"
-    | Some (PartiallyOptimizedDisjointUnion _) -> "Partially Optimized Disjoint Union"
-    | Some (DisjointUnion _) -> "Disjoint Union"
+    | Some (PartiallyOptimizedAlmostDisjointUnionWithPossiblyNonUniqueKeys _) ->
+      "Partially Optimized Almost Disjoint Union with possibly non-unique keys"
+    | Some (AlmostDisjointUnionWithPossiblyNonUniqueKeys _) ->
+      "Almost Disjoint Union with possibly non-unique keys"
     | Some (PartiallyOptimizedUnionEnum _) -> "Partially Optimized Enum"
     | None -> "No Specialization"
 

tests/almost_disjoint_union_opt/.flowconfig

@@ -0,0 +1,3 @@
+[options]
+all=true
+casting_syntax=both

tests/almost_disjoint_union_opt/almost_disjoint_union_opt.exp

@@ -0,0 +1 @@
+Found 0 errors