Bigarrays

vouillon · vouillon · commit 086b4567b36e · 2025-05-13T10:31:17.000+02:00
diff --git a/compiler/lib-wasm/gc_target.ml b/compiler/lib-wasm/gc_target.ml
@@ -422,6 +422,38 @@ module Type = struct
                     }
                   ])
           })
+
+  let int_array_type =
+    register_type "int_array" (fun () ->
+        return
+          { supertype = None
+          ; final = true
+          ; typ = W.Array { mut = true; typ = Value I32 }
+          })
+
+  let bigarray_type =
+    register_type "bigarray" (fun () ->
+        let* custom_operations = custom_operations_type in
+        let* int_array = int_array_type in
+        let* custom = custom_type in
+        return
+          { supertype = Some custom
+          ; final = true
+          ; typ =
+              W.Struct
+                [ { mut = false
+                  ; typ = Value (Ref { nullable = false; typ = Type custom_operations })
+                  }
+                ; { mut = true; typ = Value (Ref { nullable = false; typ = Extern }) }
+                ; { mut = false
+                  ; typ = Value (Ref { nullable = false; typ = Type int_array })
+                  }
+                ; { mut = false; typ = Packed I8 }
+                ; { mut = false; typ = Packed I8 }
+                ; { mut = false; typ = Packed I8 }
+                ; { mut = false; typ = Value I32 }
+                ]
+          })
 end
 
 module Value = struct
@@ -1367,6 +1399,56 @@ module Math = struct
   let exp2 x = power (return (W.Const (F64 2.))) x
 end
 
+module Bigarray = struct
+  let dim1 a =
+    let* ty = Type.bigarray_type in
+    Memory.wasm_array_get
+      ~ty:Type.int_array_type
+      (Memory.wasm_struct_get ty (Memory.wasm_cast ty a) 2)
+      (Arith.const 0l)
+
+  let get ~kind a i =
+    match kind with
+    | Typing.Bigarray.Int8_unsigned | Char ->
+        let* f =
+          register_import
+            ~import_module:"bindings"
+            ~name:"ta_get_ui8"
+            (Fun
+               { W.params = [ Ref { nullable = false; typ = Extern }; I32 ]
+               ; result = [ I32 ]
+               })
+        in
+        let* ty = Type.bigarray_type in
+        let* ta = Memory.wasm_struct_get ty (Memory.wasm_cast ty a) 1 in
+        let* i = Value.int_val i in
+        Value.val_int (return (W.Call (f, [ ta; i ])))
+    | _ -> assert false
+
+  let set ~kind a i v =
+    match kind with
+    | Typing.Bigarray.Int8_unsigned | Char ->
+        let* f =
+          register_import
+            ~import_module:"bindings"
+            ~name:"ta_set_ui8"
+            (Fun
+               { W.params =
+                   [ Ref { nullable = false; typ = Extern }
+                   ; I32
+                   ; Ref { nullable = false; typ = I31 }
+                   ]
+               ; result = []
+               })
+        in
+        let* ty = Type.bigarray_type in
+        let* ta = Memory.wasm_struct_get ty (Memory.wasm_cast ty a) 1 in
+        let* i = Value.int_val i in
+        let* v = cast I31 v in
+        instr (W.CallInstr (f, [ ta; i; v ]))
+    | _ -> assert false
+end
+
 module JavaScript = struct
   let anyref = W.Ref { nullable = true; typ = Any }
 
diff --git a/compiler/lib-wasm/generate.ml b/compiler/lib-wasm/generate.ml
@@ -705,7 +705,48 @@ module Generate (Target : Target_sig.S) = struct
                 let* x' = x' in
                 let* y' = y' in
                 return (W.Call (f, [ x'; y' ])))
-        | _ -> invalid_arity "caml_compare" l ~expected:2)
+        | _ -> invalid_arity "caml_compare" l ~expected:2);
+    register_prim "caml_ba_get_1" `Mutator (fun ctx context transl_prim_arg l ->
+        match l with
+        | [ x; y ] -> (
+            let x' = transl_prim_arg x in
+            let y' = transl_prim_arg y in
+            match get_type ctx x with
+            | Bigarray { kind = (Int8_unsigned | Char) as kind; layout = C } ->
+                seq
+                  (let* cond = Arith.uge (Value.int_val y') (Bigarray.dim1 x') in
+                   instr (W.Br_if (label_index context bound_error_pc, cond)))
+                  (Bigarray.get ~kind x' y')
+            | _ ->
+                let* f =
+                  register_import ~name:"caml_ba_get_1" (Fun (Type.primitive_type 2))
+                in
+                let* x' = x' in
+                let* y' = y' in
+                return (W.Call (f, [ x'; y' ])))
+        | _ -> invalid_arity "caml_ba_get_1" l ~expected:2);
+    register_prim "caml_ba_set_1" `Mutator (fun ctx context transl_prim_arg l ->
+        match l with
+        | [ x; y; z ] -> (
+            let x' = transl_prim_arg x in
+            let y' = transl_prim_arg y in
+            let z' = transl_prim_arg z in
+            match get_type ctx x with
+            | Bigarray { kind = (Int8_unsigned | Char) as kind; layout = C } ->
+                seq
+                  (let* cond = Arith.uge (Value.int_val y') (Bigarray.dim1 x') in
+                   let* () = instr (W.Br_if (label_index context bound_error_pc, cond)) in
+                   Bigarray.set ~kind x' y' z')
+                  Value.unit
+            | _ ->
+                let* f =
+                  register_import ~name:"caml_ba_set_1" (Fun (Type.primitive_type 3))
+                in
+                let* x' = x' in
+                let* y' = y' in
+                let* z' = z' in
+                return (W.Call (f, [ x'; y'; z' ])))
+        | _ -> invalid_arity "caml_ba_set_1" l ~expected:3)
 
   let rec translate_expr ctx context x e =
     match e with
@@ -942,7 +983,9 @@ module Generate (Target : Target_sig.S) = struct
                         | "caml_bytes_set"
                         | "caml_check_bound"
                         | "caml_check_bound_gen"
-                        | "caml_check_bound_float" )
+                        | "caml_check_bound_float"
+                        | "caml_ba_get_1"
+                        | "caml_ba_set_1" )
                     , _ ) ) -> fst n, true
             | Let
                 ( _
diff --git a/compiler/lib-wasm/target_sig.ml b/compiler/lib-wasm/target_sig.ml
@@ -252,6 +252,19 @@ module type S = sig
     val round : expression -> expression
   end
 
+  module Bigarray : sig
+    val dim1 : expression -> expression
+
+    val get : kind:Typing.Bigarray.kind -> expression -> expression -> expression
+
+    val set :
+         kind:Typing.Bigarray.kind
+      -> expression
+      -> expression
+      -> expression
+      -> unit Code_generation.t
+  end
+
   val internal_primitives :
     (string
     * Primitive.kind
diff --git a/compiler/lib-wasm/typing.ml b/compiler/lib-wasm/typing.ml
@@ -11,10 +11,89 @@ type number =
   | Nativeint
   | Float
 
+module Bigarray = struct
+  type kind =
+    | Float32
+    | Float64
+    | Int8_signed
+    | Int8_unsigned
+    | Int16_signed
+    | Int16_unsigned
+    | Int32
+    | Int64
+    | Int
+    | Nativeint
+    | Complex32
+    | Complex64
+    | Char
+    | Float16
+
+  type layout =
+    | C
+    | Fortran
+
+  type t =
+    { kind : kind
+    ; layout : layout
+    }
+
+  let make ~kind ~layout =
+    { kind =
+        (match kind with
+        | 0 -> Float32
+        | 1 -> Float64
+        | 2 -> Int8_signed
+        | 3 -> Int8_unsigned
+        | 4 -> Int16_signed
+        | 5 -> Int16_unsigned
+        | 6 -> Int32
+        | 7 -> Int64
+        | 8 -> Int
+        | 9 -> Nativeint
+        | 10 -> Complex32
+        | 11 -> Complex64
+        | 12 -> Char
+        | 13 -> Float16
+        | _ -> assert false)
+    ; layout =
+        (match layout with
+        | 0 -> C
+        | 1 -> Fortran
+        | _ -> assert false)
+    }
+
+  let print f { kind; layout } =
+    Format.fprintf
+      f
+      "bigarray{%s,%s}"
+      (match kind with
+      | Float32 -> "float32"
+      | Float64 -> "float64"
+      | Int8_signed -> "sint8"
+      | Int8_unsigned -> "uint8"
+      | Int16_signed -> "sint16"
+      | Int16_unsigned -> "uint16"
+      | Int32 -> "int32"
+      | Int64 -> "int64"
+      | Int -> "int"
+      | Nativeint -> "nativeint"
+      | Complex32 -> "complex32"
+      | Complex64 -> "complex64"
+      | Char -> "char"
+      | Float16 -> "float16")
+      (match layout with
+      | C -> "C"
+      | Fortran -> "Fortran")
+
+  let equal { kind; layout } { kind = kind'; layout = layout' } =
+    phys_equal kind kind' && phys_equal layout layout'
+end
+
 type typ =
   | Top
   | Number of number
   | Tuple of typ array
+  | Bigarray of Bigarray.t
   | Bot
 
 module Domain = struct
@@ -25,9 +104,17 @@ module Domain = struct
     | Bot, t | t, Bot -> t
     | Number n, Number n' -> if Poly.equal n n' then t else Top
     | Tuple t, Tuple t' ->
-        if Array.length t = Array.length t' then Tuple (Array.map2 ~f:join t t') else Top
+        let l = Array.length t in
+        let l' = Array.length t' in
+        Tuple
+          (if l = l'
+           then Array.map2 ~f:join t t'
+           else
+             Array.init (max l l') ~f:(fun i ->
+                 if i < l then if i < l' then join t.(i) t'.(i) else t.(i) else t'.(i)))
+    | Bigarray b, Bigarray b' when Bigarray.equal b b' -> t
     | Top, _ | _, Top -> Top
-    | Number _, Tuple _ | Tuple _, Number _ -> Top
+    | (Number _ | Tuple _ | Bigarray _), _ -> Top
 
   let join_set ?(others = false) f s =
     if others then Top else Var.Set.fold (fun x a -> join (f x) a) s Bot
@@ -38,7 +125,8 @@ module Domain = struct
     | Number t, Number t' -> Poly.equal t t'
     | Tuple t, Tuple t' ->
         Array.length t = Array.length t' && Array.for_all2 ~f:equal t t'
-    | (Top | Tuple _ | Number _ | Bot), _ -> false
+    | Bigarray b, Bigarray b' -> Bigarray.equal b b'
+    | (Top | Tuple _ | Number _ | Bigarray _ | Bot), _ -> false
 
   let bot = Bot
 
@@ -47,8 +135,15 @@ module Domain = struct
     | _, Top | Bot, _ -> true
     | Top, _ | _, Bot -> false
     | Number t, Number t' -> Poly.equal t t'
-    | Tuple t, Tuple t' -> Array.length t = Array.length t' && Array.for_all2 ~f:sub t t'
-    | Number _, _ | Tuple _, _ -> false
+    | Tuple t, Tuple t' ->
+        Array.length t <= Array.length t'
+        &&
+        let rec compare t t' i l =
+          i = l || (sub t.(i) t'.(i) && compare t t' (i + 1) l)
+        in
+        compare t t' 0 (Array.length t)
+    | Bigarray b, Bigarray b' -> Bigarray.equal b b'
+    | (Number _ | Tuple _ | Bigarray _), _ -> false
 
   let rec print f t =
     match t with
@@ -59,12 +154,30 @@ module Domain = struct
     | Number Int64 -> Format.fprintf f "int64"
     | Number Nativeint -> Format.fprintf f "nativeint"
     | Number Float -> Format.fprintf f "float"
+    | Bigarray b -> Bigarray.print f b
     | Tuple t ->
         Format.fprintf
           f
           "(%a)"
           (Format.pp_print_array ~pp_sep:(fun f () -> Format.fprintf f ",") print)
           t
+
+  let depth_treshold = 4
+
+  let rec depth t =
+    match t with
+    | Top | Bot | Number _ | Bigarray _ -> 0
+    | Tuple l -> 1 + Array.fold_left ~f:(fun acc t' -> max (depth t') acc) l ~init:0
+
+  let rec truncate depth t =
+    match t with
+    | Top | Bot | Number _ | Bigarray _ -> t
+    | Tuple l ->
+        if depth = 0
+        then Top
+        else Tuple (Array.map ~f:(fun t' -> truncate (depth - 1) t') l)
+
+  let limit t = if depth t > depth_treshold then truncate depth_treshold t else t
 end
 
 let update_deps st { blocks; _ } =
@@ -268,7 +381,7 @@ let propagate st approx x : Domain.t =
                  match st.state.mutable_fields.(Var.idx x) with
                  | All_fields -> Top
                  | Some_fields s when IntSet.mem i s -> Top
-                 | Some_fields _ | No_field -> Var.Tbl.get approx y)
+                 | Some_fields _ | No_field -> Domain.limit (Var.Tbl.get approx y))
                lst)
       | Field (y, n, _) -> (
           match Var.Tbl.get approx y with
@@ -316,7 +429,32 @@ let propagate st approx x : Domain.t =
                   | Expr (Closure (params, _, _))
                     when List.length args = List.length params ->
                       Domain.join_set
-                        (fun y -> Var.Tbl.get approx y)
+                        (fun y ->
+                          match st.state.defs.(Var.idx y) with
+                          | Expr
+                              (Prim (Extern "caml_ba_create", [ Pv kind; Pv layout; _ ]))
+                            -> (
+                              let m =
+                                List.fold_left2
+                                  ~f:(fun m p a -> Var.Map.add p a m)
+                                  ~init:Var.Map.empty
+                                  params
+                                  args
+                              in
+                              try
+                                match
+                                  ( st.state.defs.(Var.idx (Var.Map.find kind m))
+                                  , st.state.defs.(Var.idx (Var.Map.find layout m)) )
+                                with
+                                | Expr (Constant (Int kind)), Expr (Constant (Int layout))
+                                  ->
+                                    Bigarray
+                                      (Bigarray.make
+                                         ~kind:(Targetint.to_int_exn kind)
+                                         ~layout:(Targetint.to_int_exn layout))
+                                | _ -> raise Not_found
+                              with Not_found -> Var.Tbl.get approx y)
+                          | _ -> Var.Tbl.get approx y)
                         (Var.Map.find g st.state.return_values)
                   | Expr (Closure (_, _, _)) ->
                       (* The function is partially applied or over applied *)
@@ -368,6 +506,10 @@ let print_opt typ f e =
               l)
           [ Number Int; Number Int32; Number Int64; Number Nativeint; Number Float ]
       then Format.fprintf f " OPT"
+  | Prim (Extern ("caml_ba_get_1" | "caml_ba_set_1"), Pv x :: _) -> (
+      match Var.Tbl.get typ x with
+      | Bigarray _ -> Format.fprintf f " OPT"
+      | _ -> ())
   | _ -> ()
 
 let f ~state ~info p =
diff --git a/compiler/lib-wasm/typing.mli b/compiler/lib-wasm/typing.mli
