Source file npy.ml

type _ dtype =
  | Int8 : int dtype
  | Uint8 : int dtype
  | Float32 : float dtype
  | Float64 : float dtype

type descr = D_int8 | D_uint8 | D_float32 | D_float64

type t = {
  shape : int array;
  fortran_order : bool;
  descr : descr;
  data : string;
}

let find_substring haystack needle =
  let nlen = String.length needle in
  let hlen = String.length haystack in
  let rec search i =
    if i + nlen > hlen then None
    else if String.sub haystack i nlen = needle then Some i
    else search (i + 1)
  in
  search 0

let extract_quoted_value header key =
  let pattern = "'" ^ key ^ "': " in
  match find_substring header pattern with
  | None -> Error (Printf.sprintf "missing key: %s" key)
  | Some i ->
    let start = i + String.length pattern in
    if start >= String.length header then Error (Printf.sprintf "truncated value for key: %s" key)
    else
      let c = header.[start] in
      if c = '\'' then
        (* quoted string value *)
        let value_start = start + 1 in
        (match find_substring (String.sub header value_start (String.length header - value_start)) "'" with
         | None -> Error (Printf.sprintf "unterminated string for key: %s" key)
         | Some len -> Ok (String.sub header value_start len))
      else
        (* unquoted value - read until comma or } *)
        let rec find_end j =
          if j >= String.length header then j
          else match header.[j] with
            | ',' | '}' | ')' -> j
            | _ -> find_end (j + 1)
        in
        let end_pos = find_end start in
        let value = String.trim (String.sub header start (end_pos - start)) in
        Ok value

let parse_descr s =
  match s with
  | "|i1" -> Ok D_int8
  | "|u1" -> Ok D_uint8
  | "<f4" -> Ok D_float32
  | "<f8" -> Ok D_float64
  | _ -> Error (Printf.sprintf "unsupported dtype: %s" s)

let parse_fortran_order s =
  match s with
  | "True" -> Ok true
  | "False" -> Ok false
  | _ -> Error (Printf.sprintf "invalid fortran_order: %s" s)

let parse_shape header =
  let pattern = "'shape': (" in
  match find_substring header pattern with
  | None -> Error "missing shape"
  | Some i ->
    let start = i + String.length pattern in
    (match find_substring (String.sub header start (String.length header - start)) ")" with
     | None -> Error "unterminated shape"
     | Some len ->
       let shape_str = String.sub header start len in
       let shape_str = String.trim shape_str in
       if shape_str = "" then Ok [||]
       else
         let parts = String.split_on_char ',' shape_str in
         let parts = List.filter (fun s -> String.trim s <> "") parts in
         let dims = List.map (fun s -> int_of_string (String.trim s)) parts in
         Ok (Array.of_list dims))

let of_string s =
  let len = String.length s in
  if len < 10 then Error "too short for .npy file"
  else if String.sub s 0 6 <> "\x93NUMPY" then Error "bad magic number"
  else
    let major = Char.code s.[6] in
    let _minor = Char.code s.[7] in
    let header_len, header_offset =
      if major = 1 then
        let hl = Char.code s.[8] lor (Char.code s.[9] lsl 8) in
        (hl, 10)
      else if major = 2 then
        if len < 12 then (0, 12)
        else
          let hl =
            Char.code s.[8]
            lor (Char.code s.[9] lsl 8)
            lor (Char.code s.[10] lsl 16)
            lor (Char.code s.[11] lsl 24)
          in
          (hl, 12)
      else (0, 10)
    in
    if header_offset + header_len > len then Error "truncated header"
    else
      let header = String.sub s header_offset header_len in
      match extract_quoted_value header "descr" with
      | Error e -> Error e
      | Ok descr_str ->
        match parse_descr descr_str with
        | Error e -> Error e
        | Ok descr ->
          match extract_quoted_value header "fortran_order" with
          | Error e -> Error e
          | Ok fo_str ->
            match parse_fortran_order fo_str with
            | Error e -> Error e
            | Ok fortran_order ->
              match parse_shape header with
              | Error e -> Error e
              | Ok shape ->
                let data_offset = header_offset + header_len in
                let data = String.sub s data_offset (len - data_offset) in
                Ok { shape; fortran_order; descr; data }

let shape t = t.shape

let fortran_order t = t.fortran_order

let data_int8 t =
  match t.descr with
  | D_int8 ->
    let n = String.length t.data in
    let ba = Bigarray.Array1.create Bigarray.int8_signed Bigarray.c_layout n in
    for i = 0 to n - 1 do
      let v = Char.code t.data.[i] in
      let v = if v >= 128 then v - 256 else v in
      Bigarray.Array1.set ba i v
    done;
    Some ba
  | _ -> None

let data_uint8 t =
  match t.descr with
  | D_uint8 ->
    let n = String.length t.data in
    let ba = Bigarray.Array1.create Bigarray.int8_unsigned Bigarray.c_layout n in
    for i = 0 to n - 1 do
      Bigarray.Array1.set ba i (Char.code t.data.[i])
    done;
    Some ba
  | _ -> None

let read_le_int32 s off =
  let b0 = Char.code s.[off] in
  let b1 = Char.code s.[off + 1] in
  let b2 = Char.code s.[off + 2] in
  let b3 = Char.code s.[off + 3] in
  Int32.logor
    (Int32.of_int b0)
    (Int32.logor
       (Int32.shift_left (Int32.of_int b1) 8)
       (Int32.logor
          (Int32.shift_left (Int32.of_int b2) 16)
          (Int32.shift_left (Int32.of_int b3) 24)))

let read_le_int64 s off =
  let b i = Int64.of_int (Char.code s.[off + i]) in
  let ( lor ) = Int64.logor in
  let ( lsl ) = Int64.shift_left in
  (b 0) lor ((b 1) lsl 8) lor ((b 2) lsl 16) lor ((b 3) lsl 24)
  lor ((b 4) lsl 32) lor ((b 5) lsl 40) lor ((b 6) lsl 48) lor ((b 7) lsl 56)

let data_float32 t =
  match t.descr with
  | D_float32 ->
    let n = String.length t.data / 4 in
    let ba = Bigarray.Array1.create Bigarray.float32 Bigarray.c_layout n in
    for i = 0 to n - 1 do
      let bits = read_le_int32 t.data (i * 4) in
      Bigarray.Array1.set ba i (Int32.float_of_bits bits)
    done;
    Some ba
  | _ -> None

let data_float64 t =
  match t.descr with
  | D_float64 ->
    let n = String.length t.data / 8 in
    let ba = Bigarray.Array1.create Bigarray.float64 Bigarray.c_layout n in
    for i = 0 to n - 1 do
      let bits = read_le_int64 t.data (i * 8) in
      Bigarray.Array1.set ba i (Int64.float_of_bits bits)
    done;
    Some ba
  | _ -> None