Source file bigstring.ml

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
open! Base
open! Import
module Int64 = Base.Int64
module Bigstring = Base_bigstring

module Kind = struct
  type ('a : any) t =
    | Char : char t
    | Int8 : int8 t
    | Int16 : int16 t
    | Int32 : int32 t
    | Int64 : int64 t
    | Float32 : float32 t
    | Float64 : float t
  [@@deriving sexp_of]

  let[@inline] width (type a : any) : a t -> int = function
    | Char -> 1
    | Int8 -> 1
    | Int16 -> 2
    | Int32 -> 4
    | Int64 -> 8
    | Float32 -> 4
    | Float64 -> 8
  ;;
end

type ('a : any) t =
  { kind : 'a Kind.t
  ; data : Bigstring.t
  }
[@@deriving sexp_of]

external length : Bigstring.t @ contended -> int @@ portable = "%caml_ba_dim_1"

let%template[@inline] with_kind_exn kind data =
  if Bigstring.length data % Kind.width kind <> 0
  then invalid_arg "Bigstring length not divisible by element width.";
  { kind; data }
[@@mode m = (uncontended, shared)]
;;

let[@inline] empty kind =
  { kind
  ; data = Portability_hacks.magic_uncontended__promise_deeply_immutable Bigstring.empty
  }
;;

let[@inline] create kind n = { kind; data = Bigstring.create (n * Kind.width kind) }
let[@inline] kind { kind; _ } = kind
let[@inline] data { data; _ } = data

let%template sub_shared (type a : any) ({ kind; data } : a t) ~pos ~len =
  let width = Kind.width kind in
  let len = width * len in
  let pos = width * pos in
  { data = Stdlib.Bigarray.Array1.sub (Obj.magic_uncontended data) pos len; kind }
[@@mode m = (uncontended, shared, contended)]
;;

let%template[@inline] copy { kind; data } =
  { kind; data = Bigstring.copy (Obj.magic_uncontended data) }
[@@mode m = (uncontended, shared)]
;;

let[@inline] of_string ?pos ?len s =
  { kind = Char; data = Bigstring.of_string ?pos ?len s }
;;

let[@inline] length { kind; data } = length data / Kind.width kind

(* Bigstrings are usually allocated by [malloc], which provides 16-byte alignment.
   However, custom bigstrings or slices of bigstrings may not be aligned, so we must use
   unaligned instructions. Fortunately, they have no penalty when the address is aligned,
   which is the common case for 16-byte operations. *)

let%template[@inline] get (type a : value_or_null) { kind : a Kind.t; data } pos : a =
  let pos = pos * Kind.width kind in
  match kind with
  | Char -> Bigstring.get data pos
  | Int8 -> (Bigstring.get_int8 [@inlined]) data ~pos |> Int8.of_int
  | Int16 -> (Bigstring.get_int16_le [@inlined]) data ~pos |> Int16.of_int
  | Int32 -> (Bigstring.get_int32_t_le [@inlined]) data ~pos
  | Int64 -> (Bigstring.get_int64_t_le [@inlined]) data ~pos
  | Float32 -> Float32.Bigstring.get data ~pos
  | Float64 -> Int64.float_of_bits ((Bigstring.get_int64_t_le [@inlined]) data ~pos)
;;

let[@inline] set (type a : value_or_null) { kind : a Kind.t; data } pos (a : a) =
  let pos = pos * Kind.width kind in
  match kind with
  | Char -> Bigstring.set data pos a
  | Int8 -> (Bigstring.set_int8_exn [@inlined]) data ~pos (Int8.to_int a)
  | Int16 -> (Bigstring.set_int16_le_exn [@inlined]) data ~pos (Int16.to_int a)
  | Int32 -> (Bigstring.set_int32_t_le [@inlined]) data ~pos a
  | Int64 -> (Bigstring.set_int64_t_le [@inlined]) data ~pos a
  | Float32 -> Float32.Bigstring.set data ~pos a
  | Float64 -> (Bigstring.set_int64_t_le [@inlined]) data ~pos (Int64.bits_of_float a)
;;

let[@inline] unsafe_get (type a : value_or_null) { kind : a Kind.t; data } pos : a =
  let pos = pos * Kind.width kind in
  match kind with
  | Char -> Bigstring.unsafe_get data pos
  | Int8 -> (Bigstring.unsafe_get_int8 [@inlined]) data ~pos |> Int8.of_int
  | Int16 -> (Bigstring.unsafe_get_int16_le [@inlined]) data ~pos |> Int16.of_int
  | Int32 -> (Bigstring.unsafe_get_int32_t_le [@inlined]) data ~pos
  | Int64 -> (Bigstring.unsafe_get_int64_t_le [@inlined]) data ~pos
  | Float32 -> Float32.Bigstring.unsafe_get data ~pos
  | Float64 ->
    Int64.float_of_bits ((Bigstring.unsafe_get_int64_t_le [@inlined]) data ~pos)
;;

let[@inline] unsafe_set (type a : value_or_null) { kind : a Kind.t; data } pos (a : a) =
  let pos = pos * Kind.width kind in
  match kind with
  | Char -> Bigstring.unsafe_set data pos a
  | Int8 -> (Bigstring.unsafe_set_int8 [@inlined]) data ~pos (Int8.to_int a)
  | Int16 -> (Bigstring.unsafe_set_int16_le [@inlined]) data ~pos (Int16.to_int a)
  | Int32 -> (Bigstring.unsafe_set_int32_t_le [@inlined]) data ~pos a
  | Int64 -> (Bigstring.unsafe_set_int64_t_le [@inlined]) data ~pos a
  | Float32 -> Float32.Bigstring.unsafe_set data ~pos a
  | Float64 ->
    (Bigstring.unsafe_set_int64_t_le [@inlined]) data ~pos (Int64.bits_of_float a)
;;

module Expert = struct
  let[@inline] unsafe_racy_get_contended t pos = unsafe_get (Obj.magic_uncontended t) pos

  let[@inline] unsafe_racy_set_contended t pos a =
    unsafe_set (Obj.magic_uncontended t) pos a
  ;;
end