Add support for functions in untagged variants.

This was done at speed: need to double check that there are no corner cases missing.

Fixes #6278

Author: cristianoc

jscomp/core/js_exp_make.ml

@@ -769,6 +769,7 @@ let tag_type = function
   | Undefined -> undefined
   | Untagged IntType -> str "number"
   | Untagged FloatType -> str "number"
+  | Untagged FunctionType -> str "function"
   | Untagged StringType -> str "string"
   | Untagged ArrayType -> str "Array" ~delim:DNoQuotes
   | Untagged ObjectType -> str "object"

jscomp/frontend/ast_core_type.ml

@@ -125,7 +125,7 @@ let get_uncurry_arity (ty : t) =
   | _ -> None
 
 let get_curry_arity (ty : t) =
-  if Ast_uncurried.typeIsUncurriedFun ty then
+  if Ast_uncurried.coreTypeIsUncurriedFun ty then
     let arity, _ = Ast_uncurried.typeExtractUncurriedFun ty in
     arity
   else get_uncurry_arity_aux ty 0

jscomp/frontend/ast_external_process.ml

@@ -68,7 +68,7 @@ let spec_of_ptyp (nolabel : bool) (ptyp : Parsetree.core_type) :
     | _ -> Bs_syntaxerr.err ptyp.ptyp_loc Invalid_bs_unwrap_type)
   | `Uncurry opt_arity -> (
     let real_arity =
-      if Ast_uncurried.typeIsUncurriedFun ptyp then
+      if Ast_uncurried.coreTypeIsUncurriedFun ptyp then
         let arity, _ = Ast_uncurried.typeExtractUncurriedFun ptyp in
         Some arity
       else Ast_core_type.get_uncurry_arity ptyp

jscomp/ml/ast_uncurried.ml

@@ -63,12 +63,19 @@ let exprExtractUncurriedFun (expr : Parsetree.expression) =
   | Pexp_construct ({ txt = Lident "Function$" }, Some e) -> e
   | _ -> assert false
 
-let typeIsUncurriedFun (typ : Parsetree.core_type) =
+let coreTypeIsUncurriedFun (typ : Parsetree.core_type) =
   match typ.ptyp_desc with
   | Ptyp_constr ({txt = Lident "function$"}, [{ptyp_desc = Ptyp_arrow _}; _]) ->
     true
   | _ -> false
 
+let typeIsUncurriedFun (typ : Types.type_expr) =
+  match typ.desc with
+  | Tconstr (Pident {name = "function$"}, [{desc = Tarrow _}; _], _) ->
+    true
+  | _ -> false
+
+
 let typeExtractUncurriedFun (typ : Parsetree.core_type) =
   match typ.ptyp_desc with
   | Ptyp_constr ({txt = Lident "function$"}, [tArg; tArity]) ->

jscomp/ml/ast_untagged_variants.ml

@@ -1,4 +1,4 @@
-type untaggedError = OnlyOneUnknown | AtMostOneObject | AtMostOneArray | AtMostOneString | AtMostOneNumber | DuplicateLiteral of string
+type untaggedError = OnlyOneUnknown | AtMostOneObject | AtMostOneArray | AtMostOneFunction | AtMostOneString | AtMostOneNumber | DuplicateLiteral of string
 type error =
   | InvalidVariantAsAnnotation
   | Duplicated_bs_as
@@ -22,14 +22,15 @@ let report_error ppf =
     | OnlyOneUnknown -> "An unknown case must be the only case with payloads."
     | AtMostOneObject -> "At most one case can be an object type."
     | AtMostOneArray -> "At most one case can be an array type."
+    | AtMostOneFunction -> "At most one case can be a function type."
     | AtMostOneString -> "At most one case can be a string type."
     | AtMostOneNumber -> "At most one case can be a number type (int or float)."
     | DuplicateLiteral s -> "Duplicate literal " ^ s ^ "."
     )
 
 (* Type of the runtime representation of an untagged block (case with payoad) *)
 type block_type =
-  | IntType | StringType | FloatType | ArrayType | ObjectType | UnknownType
+  | IntType | StringType | FloatType | ArrayType | FunctionType | ObjectType | UnknownType
 
 (*
   Type of the runtime representation of a tag.
@@ -116,6 +117,10 @@ let get_block_type ~env (cstr: Types.constructor_declaration) : block_type optio
       Some FloatType
   | true, Cstr_tuple [{desc = Tconstr (path, _, _)}] when Path.same path Predef.path_array ->
       Some ArrayType
+  | true, Cstr_tuple [{desc = Tconstr _}  as t] when Ast_uncurried.typeIsUncurriedFun t ->
+      Some FunctionType
+  | true, Cstr_tuple [{desc = Tarrow _} ] ->
+      Some FunctionType
   | true, Cstr_tuple [{desc = Tconstr (path, _, _)}] when Path. same path Predef.path_string ->
       Some StringType
   | true, Cstr_tuple [{desc = Tconstr _} as t] when type_is_builtin_object t ->
@@ -162,6 +167,7 @@ let checkInvariant ~isUntaggedDef ~(consts : (Location.t * tag) list) ~(blocks :
   let string_literals = ref StringSet.empty in
   let nonstring_literals = ref StringSet.empty in
   let arrayTypes = ref 0 in
+  let functionTypes = ref 0 in
   let objectTypes = ref 0 in
   let stringTypes = ref 0 in
   let numberTypes = ref 0 in
@@ -181,6 +187,8 @@ let checkInvariant ~isUntaggedDef ~(consts : (Location.t * tag) list) ~(blocks :
       then raise (Error (loc, InvalidUntaggedVariantDefinition AtMostOneObject));
     if !arrayTypes > 1
       then raise (Error (loc, InvalidUntaggedVariantDefinition AtMostOneArray));
+    if !functionTypes > 1
+      then raise (Error (loc, InvalidUntaggedVariantDefinition AtMostOneFunction));
     if !stringTypes > 1
       then raise (Error (loc, InvalidUntaggedVariantDefinition AtMostOneString));
     if !numberTypes > 1
@@ -214,6 +222,9 @@ let checkInvariant ~isUntaggedDef ~(consts : (Location.t * tag) list) ~(blocks :
     | Some ArrayType ->
       incr arrayTypes;
       invariant loc
+    | Some FunctionType ->
+      incr functionTypes;
+      invariant loc
     | Some (IntType | FloatType) ->
       incr numberTypes;
       invariant loc
@@ -266,6 +277,8 @@ module DynamicChecks = struct
   let nil = Null |> tag_type
   let undefined = Undefined |> tag_type
   let object_ = Untagged ObjectType |> tag_type
+
+  let function_ = Untagged FunctionType |> tag_type
   let string = Untagged StringType |> tag_type
   let number = Untagged IntType |> tag_type
 
@@ -298,6 +311,8 @@ module DynamicChecks = struct
       typeof e != number
     | ArrayType -> 
       not (is_array e)
+    | FunctionType ->
+      typeof e != function_
     | ObjectType when literals_overlaps_with_object () = false ->
       typeof e != object_
     | ObjectType (* overlap *) ->
@@ -341,9 +356,8 @@ module DynamicChecks = struct
   let add_runtime_type_check ~tag_type ~(block_cases: block_type list) x y =
     let has_array() = Ext_list.exists block_cases (fun t -> t = ArrayType) in
     match tag_type with
-    | Untagged IntType
-    | Untagged StringType
-    | Untagged FloatType -> typeof y == x
+    | Untagged (IntType | StringType | FloatType | FunctionType) ->
+      typeof y == x
     | Untagged ObjectType ->
       if has_array() then
         typeof y == x &&& not (is_array y)

jscomp/syntax/src/react_jsx_common.ml

@@ -45,7 +45,7 @@ let raiseErrorMultipleReactComponent ~loc =
 let optionalAttr = ({txt = "res.optional"; loc = Location.none}, PStr [])
 
 let extractUncurried typ =
-  if Ast_uncurried.typeIsUncurriedFun typ then
+  if Ast_uncurried.coreTypeIsUncurriedFun typ then
     let _arity, t = Ast_uncurried.typeExtractUncurriedFun typ in
     t
   else typ

jscomp/syntax/src/res_parens.ml

@@ -442,7 +442,7 @@ let includeModExpr modExpr =
 let arrowReturnTypExpr typExpr =
   match typExpr.Parsetree.ptyp_desc with
   | Parsetree.Ptyp_arrow _ -> true
-  | _ when Ast_uncurried.typeIsUncurriedFun typExpr -> true
+  | _ when Ast_uncurried.coreTypeIsUncurriedFun typExpr -> true
   | _ -> false
 
 let patternRecordRowRhs (pattern : Parsetree.pattern) =

jscomp/syntax/src/res_printer.ml

@@ -1591,7 +1591,7 @@ and printTypExpr ~(state : State.t) (typExpr : Parsetree.core_type) cmtTbl =
         let doc = printTypExpr ~state n cmtTbl in
         match n.ptyp_desc with
         | Ptyp_arrow _ | Ptyp_tuple _ | Ptyp_alias _ -> addParens doc
-        | _ when Ast_uncurried.typeIsUncurriedFun n -> addParens doc
+        | _ when Ast_uncurried.coreTypeIsUncurriedFun n -> addParens doc
         | _ -> doc
       in
       Doc.group
@@ -1652,7 +1652,7 @@ and printTypExpr ~(state : State.t) (typExpr : Parsetree.core_type) cmtTbl =
         let needsParens =
           match typ.ptyp_desc with
           | Ptyp_arrow _ -> true
-          | _ when Ast_uncurried.typeIsUncurriedFun typ -> true
+          | _ when Ast_uncurried.coreTypeIsUncurriedFun typ -> true
           | _ -> false
         in
         let doc = printTypExpr ~state typ cmtTbl in
@@ -1664,7 +1664,7 @@ and printTypExpr ~(state : State.t) (typExpr : Parsetree.core_type) cmtTbl =
     | Ptyp_object (fields, openFlag) ->
       printObject ~state ~inline:false fields openFlag cmtTbl
     | Ptyp_arrow _ -> printArrow ~uncurried:false typExpr
-    | Ptyp_constr _ when Ast_uncurried.typeIsUncurriedFun typExpr ->
+    | Ptyp_constr _ when Ast_uncurried.coreTypeIsUncurriedFun typExpr ->
       let arity, tArg = Ast_uncurried.typeExtractUncurriedFun typExpr in
       printArrow ~uncurried:true ~arity tArg
     | Ptyp_constr (longidentLoc, [{ptyp_desc = Ptyp_object (fields, openFlag)}])

jscomp/test/UntaggedVariants.js

@@ -293,3 +293,17 @@ module OptionUnboxingHeuristic = {
   type untaggedInlineMultinaryOption = A | B({x: option<int>, y?: string})
   let untaggedInlineMultinaryOption = (x: untaggedInlineMultinaryOption) => Some(x)
 }
+
+module TestFunctionCase = {
+  @unboxed
+  type t = Array(array<int>) | Record({x:int}) | Function((. int) => int)
+
+  let classify = v =>
+    switch v {
+    | Record({x}) => x
+    | Array(a) => a[0]
+    | Function(f) => f(. 3) 
+    }
+
+    let ff = Function((. x) => x+1)
+}