Add %ffi as variation of %raw with arity checking.

CHANGELOG.md

@@ -12,6 +12,10 @@
 
 # 11.0.0-beta.2 (Unreleased)
 
+#### :rocket: New Feature
+
+- Introduced a new `%ffi` extension that provides a more robust mechanism for JavaScript function interoperation by considering function arity in type constraints. This enhancement improves safety when dealing with JavaScript functions by enforcing type constraints based on the arity of the function. [PR #6251](https://github.com/rescript-lang/rescript-compiler/pull/6251)
+
 # 11.0.0-beta.1
 
 #### :rocket: Main New Feature

jscomp/frontend/ast_exp_extension.ml

@@ -26,6 +26,7 @@ open Ast_helper
 let handle_extension e (self : Bs_ast_mapper.mapper)
     (({txt; loc}, payload) : Parsetree.extension) =
   match txt with
+  | "ffi" -> Ast_exp_handle_external.handle_ffi ~loc ~payload
   | "bs.raw" | "raw" ->
     Ast_exp_handle_external.handle_raw ~kind:Raw_exp loc payload
   | "bs.re" | "re" ->

jscomp/frontend/ast_exp_handle_external.ml

@@ -72,7 +72,7 @@ let handle_debugger loc (payload : Ast_payload.t) =
     Location.raise_errorf ~loc "%%debugger extension doesn't accept arguments"
 
 let handle_raw ~kind loc payload =
-  let is_function = ref false in
+  let is_function = ref None in
   match Ast_payload.raw_as_string_exp_exn ~kind ~is_function payload with
   | None -> (
     match kind with
@@ -93,11 +93,57 @@ let handle_raw ~kind loc payload =
           ~pval_type:(Typ.arrow Nolabel (Typ.any ()) (Typ.any ()))
           [exp];
       pexp_attributes =
-        (if !is_function then
-         Ast_attributes.internal_expansive :: exp.pexp_attributes
-        else exp.pexp_attributes);
+        (match !is_function with
+        | None -> exp.pexp_attributes
+        | Some _ -> Ast_attributes.internal_expansive :: exp.pexp_attributes);
     }
 
+let handle_ffi ~loc ~payload =
+  let is_function = ref None in
+  let err () =
+    Location.raise_errorf ~loc
+      "%%ffi extension can only be applied to a string containing a JavaScript \
+       function such as \"(x) => ...\""
+  in
+  match
+    Ast_payload.raw_as_string_exp_exn ~kind:Raw_exp ~is_function payload
+  with
+  | None -> err ()
+  | Some exp ->
+    (* Wrap a type constraint based on arity.
+       E.g. for arity 2 constrain to type (_, _) => _ *)
+    let wrapTypeConstraint (e : Parsetree.expression) =
+      let loc = e.pexp_loc in
+      let any = Ast_helper.Typ.any ~loc:e.pexp_loc () in
+      let unit = Ast_literal.type_unit ~loc () in
+      let rec arrow ~arity =
+        if arity = 0 then Ast_helper.Typ.arrow ~loc Nolabel unit any
+        else if arity = 1 then Ast_helper.Typ.arrow ~loc Nolabel any any
+        else Ast_helper.Typ.arrow ~loc Nolabel any (arrow ~arity:(arity - 1))
+      in
+      match !is_function with
+      | Some arity ->
+        let type_ =
+          Ast_uncurried.uncurriedType ~loc
+            ~arity:(if arity = 0 then 1 else arity)
+            (arrow ~arity)
+        in
+        Ast_helper.Exp.constraint_ ~loc e type_
+      | _ -> err ()
+    in
+    wrapTypeConstraint
+      {
+        exp with
+        pexp_desc =
+          Ast_external_mk.local_external_apply loc ~pval_prim:["#raw_expr"]
+            ~pval_type:(Typ.arrow Nolabel (Typ.any ()) (Typ.any ()))
+            [exp];
+        pexp_attributes =
+          (match !is_function with
+          | None -> exp.pexp_attributes
+          | Some _ -> Ast_attributes.internal_expansive :: exp.pexp_attributes);
+      }
+
 let handle_raw_structure loc payload =
   match Ast_payload.raw_as_string_exp_exn ~kind:Raw_program payload with
   | Some exp ->

jscomp/frontend/ast_exp_handle_external.mli

@@ -26,6 +26,8 @@ val handle_external : Location.t -> string -> Parsetree.expression
 
 val handle_debugger : Location.t -> Ast_payload.t -> Parsetree.expression_desc
 
+val handle_ffi : loc:Location.t -> payload:Ast_payload.t -> Parsetree.expression
+
 val handle_raw :
   kind:Js_raw_info.raw_kind ->
   Location.t ->

jscomp/ml/ast_payload.ml

@@ -138,11 +138,11 @@ let raw_as_string_exp_exn ~(kind : Js_raw_info.raw_kind) ?is_function (x : t) :
            Location.raise_errorf ~loc
              "Syntax error: a valid JS regex literal expected");
         (match is_function with
-        | Some is_function -> (
-          match Classify_function.classify_exp prog with
-          | Js_function {arity = _; _} -> is_function := true
-          | _ -> ())
-        | None -> ());
+          | Some is_function -> (
+            match Classify_function.classify_exp prog with
+            | Js_function {arity; _} -> is_function := Some arity
+            | _ -> ())
+          | None -> ());
         errors
       | Raw_program -> snd (Parser_flow.parse_program false None str));
     Some {e with pexp_desc = Pexp_constant (Pconst_string (str, None))}

jscomp/ml/ast_payload.mli

@@ -47,7 +47,7 @@ val is_single_ident : t -> Longident.t option
 
 val raw_as_string_exp_exn :
   kind:Js_raw_info.raw_kind ->
-  ?is_function:bool ref ->
+  ?is_function:int option ref ->
   t ->
   Parsetree.expression option
 (** Convert %raw into expression *)

jscomp/test/FFI.res

@@ -0,0 +1,9 @@
+@@uncurried
+
+let canUseCanvas: unit => bool = %ffi(`
+  function canUseCanvas() {
+    return !!document.createElement('canvas').getContext;
+  }
+`)
+
+let add: (int, int) => int = %ffi(`(x,y)=>x+y`)