Source file shape.ml

(**************************************************************************)
(*                                                                        *)
(*  This file is part of CAISAR.                                          *)
(*                                                                        *)
(*  Copyright (C) 2024                                                    *)
(*    CEA (Commissariat à l'énergie atomique et aux énergies              *)
(*         alternatives)                                                  *)
(*                                                                        *)
(*  You can redistribute it and/or modify it under the terms of the GNU   *)
(*  Lesser General Public License as published by the Free Software       *)
(*  Foundation, version 2.1.                                              *)
(*                                                                        *)
(*  It is distributed in the hope that it will be useful,                 *)
(*  but WITHOUT ANY WARRANTY; without even the implied warranty of        *)
(*  MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the          *)
(*  GNU Lesser General Public License for more details.                   *)
(*                                                                        *)
(*  See the GNU Lesser General Public License version 2.1                 *)
(*  for more details (enclosed in the file licenses/LGPLv2.1).            *)
(*                                                                        *)
(**************************************************************************)

open Base

type t = int array [@@deriving ord, eq]

let to_array = Array.copy
let to_array_unsafe x = x
let of_array = Array.copy
let to_list = Array.to_list
let of_list = Array.of_list
let rank = Array.length
let size t = Array.fold t ~f:( * ) ~init:1
let pp fmt x = Fmt.pf fmt "[%a]" Fmt.(array ~sep:semi int) x
let show s = Fmt.str "%a" pp s
let get = Array.get

let set t k v =
  let t = Array.copy t in
  Array.set t k v;
  t

let row_major t a =
  assert (Array.length t = Array.length a);
  let r = ref 0 in
  for i = 0 to Array.length t - 1 do
    r := (!r * t.(i)) + a.(i)
  done;
  !r

let unrow_major t i =
  let r = ref i in
  let a = Array.create ~len:(Array.length t) 0 in
  for i = Array.length t - 1 downto 0 do
    a.(i) <- !r % t.(i);
    r := !r / t.(i)
  done;
  a