Distinguish float arrays

ocsigen · Aug 30, 2024 · c5594b1 · c5594b1
1 parent 732cb95
commit c5594b1
Show file tree

Hide file tree

Showing 6 changed files with 358 additions and 260 deletions.
diff --git a/compiler/lib/eval.ml b/compiler/lib/eval.ml
@@ -299,7 +299,17 @@ let eval_instr info ((x, loc) as i) =
           let c = Constant (Int c) in
           Flow.Info.update_def info x c;
           [ Let (x, c), loc ])
-  | Let (_, Prim (Extern ("caml_array_unsafe_get" | "caml_array_unsafe_set"), _)) ->
+  | Let
+      ( _
+      , Prim
+          ( ( Extern
+                ( "caml_array_unsafe_get"
+                | "caml_array_unsafe_set"
+                | "caml_floatarray_unsafe_get"
+                | "caml_floatarray_unsafe_set"
+                | "caml_array_unsafe_set_addr" )
+            | Array_get )
+          , _ ) ) ->
       (* Fresh parameters can be introduced for these primitives
            in Specialize_js, which would make the call to [the_const_of]
            below fail. *)
@@ -371,7 +381,7 @@ let the_cond_of info x =
     info
     (fun x ->
       match Flow.Info.def info x with
-      | Some (Constant (Int 0l | Int32 0l | NativeInt 0n)) -> Zero
+      | Some (Constant (Int 0l)) -> Zero
       | Some
           (Constant
             ( Int _

diff --git a/compiler/lib/generate.ml b/compiler/lib/generate.ml
@@ -2045,7 +2045,10 @@ let init () =
     ; "caml_array_unsafe_get_float", "caml_array_unsafe_get"
     ; "caml_floatarray_unsafe_get", "caml_array_unsafe_get"
     ; "caml_array_unsafe_set_float", "caml_array_unsafe_set"
+    ; "caml_array_unsafe_set_addr", "caml_array_unsafe_set"
     ; "caml_floatarray_unsafe_set", "caml_array_unsafe_set"
+    ; "caml_check_bound_gen", "caml_check_bound"
+    ; "caml_check_bound_float", "caml_check_bound"
     ; "caml_alloc_dummy_float", "caml_alloc_dummy"
     ; "caml_make_array", "%identity"
     ; "caml_ensure_stack_capacity", "%identity"

diff --git a/compiler/lib/global_flow.ml b/compiler/lib/global_flow.ml
@@ -153,7 +153,15 @@ let expr_deps blocks st x e =
   | Constant _ | Prim ((Vectlength | Not | IsInt | Eq | Neq | Lt | Le | Ult), _) | Block _
     -> ()
   | Special _ -> ()
-  | Prim ((Extern ("caml_check_bound" | "caml_array_unsafe_get") | Array_get), l) ->
+  | Prim
+      ( ( Extern
+            ( "caml_check_bound"
+            | "caml_check_bound_float"
+            | "caml_check_bound_gen"
+            | "caml_array_unsafe_get"
+            | "caml_floatarray_unsafe_get" )
+        | Array_get )
+      , l ) ->
       (* The analysis knowns about these primitives, and will compute
          an approximation of the value they return based on an
          approximation of their arguments *)
@@ -424,8 +432,12 @@ let propagate st ~update approx x =
                   | Phi _ | Expr _ -> assert false)
                 known
           | Top -> Top)
-      | Prim (Extern "caml_check_bound", [ Pv y; _ ]) -> Var.Tbl.get approx y
-      | Prim ((Array_get | Extern "caml_array_unsafe_get"), [ Pv y; _ ]) -> (
+      | Prim
+          ( Extern ("caml_check_bound" | "caml_check_bound_float" | "caml_check_bound_gen")
+          , [ Pv y; _ ] ) -> Var.Tbl.get approx y
+      | Prim
+          ( (Array_get | Extern ("caml_array_unsafe_get" | "caml_floatarray_unsafe_get"))
+          , [ Pv y; _ ] ) -> (
           if st.fast
           then Domain.others
           else

diff --git a/compiler/lib/parse_bytecode.ml b/compiler/lib/parse_bytecode.ml
@@ -515,7 +515,8 @@ end = struct
     | Float_array _ -> false
     | Int64 _ -> false
     | Tuple _ -> false
-    | Int _ | Int32 _ | NativeInt _ -> true
+    | Int _ -> true
+    | Int32 _ | NativeInt _ -> false
 end
 
 let const i = Constant (Int i)
@@ -1554,7 +1555,17 @@ and compile infos pc state instrs =
         let x, state = State.fresh_var state loc in
 
         if debug_parser () then Format.printf "%a = %a[%d]@." Var.print x Var.print y n;
-        compile infos (pc + 2) state ((Let (x, Field (y, n)), loc) :: instrs)
+        compile
+          infos
+          (pc + 2)
+          state
+          (( Let
+               ( x
+               , Prim
+                   ( Extern "caml_floatarray_unsafe_get"
+                   , [ Pv y; Pc (Int (Int32.of_int n)) ] ) )
+           , loc )
+          :: instrs)
     | SETFIELD0 ->
         let y, _ = State.accu state in
         let z, _ = State.peek 0 state in
@@ -1628,7 +1639,13 @@ and compile infos pc state instrs =
           infos
           (pc + 2)
           (State.pop 1 state)
-          ((Let (x, const 0l), loc) :: (Set_field (y, n, z), loc) :: instrs)
+          (( Let
+               ( x
+               , Prim
+                   ( Extern "caml_floatarray_unsafe_set"
+                   , [ Pv y; Pc (Int (Int32.of_int n)); Pv z ] ) )
+           , loc )
+          :: instrs)
     | VECTLENGTH ->
         let y, _ = State.accu state in
         let x, state = State.fresh_var state loc in

diff --git a/compiler/lib/specialize_js.ml b/compiler/lib/specialize_js.ml
@@ -195,48 +195,90 @@ let specialize_instrs info l =
            the array access. The bound checking function returns the array,
            which allows to produce more compact code. *)
         match i with
-        | Let (x, Prim (Extern "caml_array_get", [ y; z ]))
-        | Let (x, Prim (Extern "caml_array_get_float", [ y; z ]))
-        | Let (x, Prim (Extern "caml_array_get_addr", [ y; z ])) ->
+        | Let
+            ( x
+            , Prim
+                ( Extern
+                    (( "caml_array_get"
+                     | "caml_array_get_float"
+                     | "caml_floatarray_get"
+                     | "caml_array_get_addr" ) as prim)
+                , [ y; z ] ) ) ->
             let idx =
               match the_int info z with
               | Some idx -> `Cst idx
               | None -> `Var z
             in
+            let instr y =
+              let prim =
+                match prim with
+                | "caml_array_get" -> Extern "caml_array_unsafe_get"
+                | "caml_array_get_float" | "caml_floatarray_get" ->
+                    Extern "caml_floatarray_unsafe_get"
+                | "caml_array_get_addr" -> Array_get
+                | _ -> assert false
+              in
+              Let (x, Prim (prim, [ y; z ])), loc
+            in
             if List.mem (y, idx) ~set:checks
             then
-              let acc =
-                (Let (x, Prim (Extern "caml_array_unsafe_get", [ y; z ])), loc) :: acc
-              in
+              let acc = instr y :: acc in
               aux info checks r acc
             else
+              let check =
+                match prim with
+                | "caml_array_get" -> "caml_check_bound_gen"
+                | "caml_array_get_float" | "caml_floatarray_get" ->
+                    "caml_check_bound_float"
+                | "caml_array_get_addr" -> "caml_check_bound"
+                | _ -> assert false
+              in
               let y' = Code.Var.fresh () in
               let acc =
-                (Let (x, Prim (Extern "caml_array_unsafe_get", [ Pv y'; z ])), loc)
-                :: (Let (y', Prim (Extern "caml_check_bound", [ y; z ])), noloc)
-                :: acc
+                instr (Pv y') :: (Let (y', Prim (Extern check, [ y; z ])), noloc) :: acc
               in
               aux info ((y, idx) :: checks) r acc
-        | Let (x, Prim (Extern "caml_array_set", [ y; z; t ]))
-        | Let (x, Prim (Extern "caml_array_set_float", [ y; z; t ]))
-        | Let (x, Prim (Extern "caml_array_set_addr", [ y; z; t ])) ->
+        | Let
+            ( x
+            , Prim
+                ( Extern
+                    (( "caml_array_set"
+                     | "caml_array_set_float"
+                     | "caml_floatarray_set"
+                     | "caml_array_set_addr" ) as prim)
+                , [ y; z; t ] ) ) ->
             let idx =
               match the_int info z with
               | Some idx -> `Cst idx
               | None -> `Var z
             in
+            let instr y =
+              let prim =
+                match prim with
+                | "caml_array_set" -> "caml_array_unsafe_set"
+                | "caml_array_set_float" | "caml_floatarray_set" ->
+                    "caml_floatarray_unsafe_set"
+                | "caml_array_set_addr" -> "caml_array_unsafe_set_addr"
+                | _ -> assert false
+              in
+              Let (x, Prim (Extern prim, [ y; z; t ])), loc
+            in
             if List.mem (y, idx) ~set:checks
             then
-              let acc =
-                (Let (x, Prim (Extern "caml_array_unsafe_set", [ y; z; t ])), loc) :: acc
-              in
+              let acc = instr y :: acc in
               aux info checks r acc
             else
+              let check =
+                match prim with
+                | "caml_array_set" -> "caml_check_bound_gen"
+                | "caml_array_set_float" | "caml_floatarray_set" ->
+                    "caml_check_bound_float"
+                | "caml_array_set_addr" -> "caml_check_bound"
+                | _ -> assert false
+              in
               let y' = Code.Var.fresh () in
               let acc =
-                (Let (x, Prim (Extern "caml_array_unsafe_set", [ Pv y'; z; t ])), loc)
-                :: (Let (y', Prim (Extern "caml_check_bound", [ y; z ])), noloc)
-                :: acc
+                instr (Pv y') :: (Let (y', Prim (Extern check, [ y; z ])), noloc) :: acc
               in
               aux info ((y, idx) :: checks) r acc
         | _ ->
@@ -270,6 +312,7 @@ let f_once p =
                      ( "caml_array_set"
                      | "caml_array_unsafe_set"
                      | "caml_array_set_float"
+                     | "caml_floatarray_set"
                      | "caml_array_set_addr"
                      | "caml_array_unsafe_set_float"
                      | "caml_floatarray_unsafe_set" )