(*---------------------------------------------------------------------------
  Copyright (c) 2026 The Raven authors. All rights reserved.
  SPDX-License-Identifier: ISC
  ---------------------------------------------------------------------------*)

open Windtrap
module L = Nx_array.Layout
module M = Nx_array.Move
open Nx_array_gen

let ( let* ) = Result.bind

(* An iarray equal to another only as the same value. *)
let same = Testable.make ~pp:pp_iints ~equal:( == )
let distinct ps = List.length (List.sort_uniq compare ps) = List.length ps
let check b reason = if b then Ok () else Error reason

(* [m]'s map from an index of its result to an index of its argument of shape
   [s], as Move's interface states it. *)
let source m s idx =
  let r = Array.length s in
  match m with
  | M.Reshape s' ->
      let s' = Iarray.to_array s' in
      let k = ref 0 in
      Array.iteri (fun i j -> k := (!k * s'.(i)) + j) idx;
      let src = Array.make r 0 in
      for i = r - 1 downto 0 do
        src.(i) <- !k mod s.(i);
        k := !k / s.(i)
      done;
      src
  | M.Broadcast s' ->
      let s' = Iarray.to_array s' in
      let off = Array.length s' - r in
      Array.init r (fun i -> if s.(i) = 1 then 0 else idx.(off + i))
  | M.Permute p ->
      let p = Iarray.to_array p in
      let src = Array.make r 0 in
      Array.iteri (fun i a -> src.(a) <- idx.(i)) p;
      src
  | M.Slice rs ->
      let rs = Iarray.to_array rs in
      Array.init r (fun i -> rs.(i).start + (idx.(i) * rs.(i).step))
  | M.Window ws ->
      let ws = Iarray.to_array ws in
      let src = Array.sub idx 0 r in
      Array.iteri
        (fun j (w : M.window) ->
          src.(w.axis) <- (idx.(w.axis) * w.step) + (idx.(r + j) * w.dilation))
        ws;
      src

(* Whether strides express the positions [t] gives the indices of [s]: the
   offset and strides are forced by the index 0 and the unit indices, so the one
   candidate is checked against every index. *)
let expressible s t =
  match indices s with
  | [] -> true
  | idxs ->
      let r = Array.length s in
      let zero = Array.make r 0 in
      let o = t zero in
      let st =
        Array.init r (fun i ->
            if s.(i) < 2 then 0
            else
              let e = Array.copy zero in
              e.(i) <- 1;
              t e - o)
      in
      List.for_all
        (fun idx ->
          let p = ref o in
          Array.iteri (fun i j -> p := !p + (j * st.(i))) idx;
          !p = t idx)
        idxs

(* Whether [l] is in canonical form: an extent-1 axis has stride 0, and a layout
   with no element has offset 0 and every stride 0. *)
let canonical l =
  let s = L.shape l and st = L.strides l in
  let unit_axes = Iarray.for_all2 (fun d t -> d <> 1 || t = 0) s st in
  if L.numel l = 0 then
    unit_axes && L.offset l = 0 && Iarray.for_all (( = ) 0) st
  else unit_axes

(* Construction, by the interface *)

let max_numel = L.max_numel

let add a b =
  match (a, b) with
  | Some a, Some b ->
      let c = a + b in
      if a >= 0 = (b >= 0) && c >= 0 <> (a >= 0) then None else Some c
  | _ -> None

(* Whether [(d - 1)·|t|], the reach of an axis, fits in an [int]. *)
let fits d t = d <= 1 || (t <> min_int && abs t <= max_int / (d - 1))

(* The offset and strides [v] makes of its arguments, or why it must raise:
   the end of the span, one past the greatest position, is at most
   [max_numel]. *)
let v_reference s strides offset =
  let r = Iarray.length s in
  let* () = check (r <= L.max_rank) "too many axes" in
  let* () = check (Iarray.for_all (fun d -> d >= 0) s) "a negative extent" in
  let* () = check (numel s <> None) "too many elements" in
  let* () = check (Iarray.length strides = r) "strides of another length" in
  let* () = check (Iarray.for_all2 fits s strides) "an axis' reach overflows" in
  if numel s = Some 0 then Ok (0, Iarray.init r (Fun.const 0))
  else
    let reach f =
      Iarray.fold_left add (Some offset)
        (Iarray.map2 (fun d t -> Some ((d - 1) * f t)) s strides)
    in
    let* () =
      match (reach (min 0), reach (max 0)) with
      | Some lo, Some hi when lo >= 0 && hi < max_numel -> Ok ()
      | Some lo, Some _ when lo >= 0 -> Error "a span past max_numel"
      | Some _, Some _ | None, _ -> Error "a negative position"
      | Some _, None -> Error "a span past max_numel"
    in
    Ok (offset, Iarray.map2 (fun d t -> if d = 1 then 0 else t) s strides)

(* Arguments of [v] about the bounds it checks. *)
let v_args =
  let open Gen in
  let stride =
    frequency
      [
        (6, int_range (-3) 3);
        ( 1,
          ints_of
            [
              max_int;
              max_int / 2;
              (max_int / 2) + 1;
              max_numel;
              max_numel - 1;
              -(max_numel - 1);
              (max_numel / 2) + 1;
              max_int / 3;
              -(max_int / 2);
              -max_int;
              min_int;
            ] );
      ]
  in
  let shapes =
    frequency
      [
        (8, shape);
        (1, array ~size:(int_range 1 3) (int_range (-1) 3));
        ( 1,
          array ~size:(int_range 1 3)
            (ints_of
               [ 0; 1; 2; 1 lsl 26; 1 lsl 27; max_numel; max_numel + 1 ]) );
        (* No element, the other extents multiplying past max_numel. *)
        ( 1,
          of_list ~pp:pp_ints
            [
              [| 0; max_numel + 1 |];
              [| 1 lsl 27; 0; 1 lsl 27 |];
              [| max_numel; 2; 0 |];
            ] );
      ]
  in
  (* Non-negative strides and the offset that puts the last position at
     [max_numel - 1], so the span ends at [max_numel]: the draws below seldom
     do. *)
  let at_end =
    let* s = shape in
    let+ strides = array ~size:(const (Array.length s)) (int_range 0 3) in
    let reach = ref 0 in
    Array.iteri (fun i d -> reach := !reach + (max 0 (d - 1) * strides.(i))) s;
    (Iarray.of_array s, Iarray.of_array strides, max_numel - 1 - !reach)
  in
  let drawn =
    let* s = shapes in
    let* n = frequency [ (8, const (Array.length s)); (1, int_range 0 5) ] in
    let+ strides = array ~size:(const n) stride
    and+ least =
      frequency
        [
          (6, int_range (-2) 12);
          ( 1,
            ints_of
              [
                max_int;
                max_int - 1;
                min_int;
                max_int / 2;
                max_numel;
                max_numel - 1;
                max_numel - 2;
              ] );
        ]
    and+ lifted = bool in
    (* Raise the offset by the reach of the negative strides, where it fits, so
       that a negative stride often meets an offset that keeps positions
       non-negative. *)
    let offset =
      if lifted && n = Array.length s then
        let up =
          Array.fold_left add (Some least)
            (Array.map2
               (fun d t ->
                 if d > 1 && t < 0 && t <> min_int then
                   if fits d t then Some ((d - 1) * -t) else None
                 else Some 0)
               s strides)
        in
        Option.value up ~default:least
      else least
    in
    (Iarray.of_array s, Iarray.of_array strides, offset)
  in
  with_pp
    (fun ppf (s, st, o) ->
      Format.fprintf ppf "~offset:%d ~strides:%a %a" o pp_iints st pp_iints s)
    (frequency [ (19, drawn); (1, at_end) ])

let reasons =
  [
    "a negative extent";
    "too many elements";
    "strides of another length";
    "an axis' reach overflows";
    "a negative position";
    "a span past max_numel";
  ]

let law_v (s, strides, offset) =
  let expected = v_reference s strides offset in
  List.iter (fun r -> cover r (expected = Error r)) reasons;
  cover "no element, other extents past max_numel"
    (expected = Error "too many elements" && Iarray.mem 0 s);
  cover "a span that ends at max_numel"
    (match expected with
    | Ok _ -> numel s <> Some 0 && snd (L.span (L.v ~offset ~strides s)) = max_numel
    | Error _ -> false);
  match expected with
  | Error _ -> raises_match Exn.invalid_arg (fun () -> L.v ~offset ~strides s)
  | Ok (o, st) ->
      let l = L.v ~offset ~strides s in
      cover "a negative stride"
        (Iarray.exists2 (fun d t -> d > 1 && t < 0) s strides);
      cover "no element" (numel s = Some 0);
      equal iints s (L.shape l);
      equal iints st (L.strides l);
      equal int o (L.offset l);
      equal int (Iarray.length s) (L.rank l);
      equal ~msg:"numel" (option int) (numel s) (Some (L.numel l));
      Iarray.iteri
        (fun i d ->
          equal ~msg:"dim" int d (L.dim l i);
          equal ~msg:"stride" int (Iarray.get st i) (L.stride l i))
        s

let law_contiguous s =
  let s = Iarray.of_array s in
  let l = L.contiguous s in
  let n = Option.get (numel s) in
  equal same s (L.shape l);
  equal (list int) (List.init n Fun.id) (positions l)

(* [fits] is the test its interface states, and {!L.contiguous} takes exactly
   the shapes it holds of. *)
let law_fits (s, _, _) =
  let natural = Iarray.for_all (fun d -> d >= 0) s in
  let expected = Iarray.length s <= L.max_rank && natural && numel s <> None in
  cover "a layout's shape" expected;
  cover "a negative extent" (not natural);
  cover "too many elements" (natural && numel s = None);
  let taken =
    match L.contiguous s with _ -> true | exception Invalid_argument _ -> false
  in
  equal (pair bool bool) (expected, expected) (L.fits s, taken)

(* [n] copies of [x]. *)
let make n x = Iarray.init n (Fun.const x)

let test_fits_rank () =
  equal (pair bool bool) (true, false)
    (L.fits (make L.max_rank 1), L.fits (make (L.max_rank + 1) 1))

let refusals =
  [
    ("contiguous of too many axes", fun () -> L.contiguous (make 33 1));
    ("contiguous of a negative extent", fun () -> L.contiguous [| 2; -1 |]);
    ("contiguous of too many elements", fun () -> L.contiguous [| max_int; 2 |]);
    ("v of too many axes", fun () -> L.v ~strides:(make 33 0) (make 33 1));
    ( "v of an axis whose reach overflows",
      fun () -> L.v ~strides:[| (max_int / 2) + 1 |] [| 3 |] );
    ( "v of too many elements",
      fun () -> L.v ~strides:[| 0; 0 |] [| max_int; 2 |] );
    ("contiguous past max_numel", fun () -> L.contiguous [| max_numel + 1 |]);
    ( "contiguous of no element whose other extents pass max_numel",
      fun () -> L.contiguous [| 0; 1 lsl 27; (1 lsl 26) + 1 |] );
    ( "contiguous of no element whose other extents overflow",
      fun () -> L.contiguous [| max_int; 2; 0 |] );
    ( "a broadcast v past max_numel",
      fun () -> L.v ~strides:[| 0; 0 |] [| (max_numel / 2) + 1; 2 |] );
    ( "v of no element whose other extents pass max_numel",
      fun () -> L.v ~strides:[| 0; 0; 0 |] [| max_numel; 2; 0 |] );
    ( "v whose span ends past max_numel",
      fun () -> L.v ~strides:[| max_numel |] [| 2 |] );
    ( "v whose offset is max_numel",
      fun () -> L.v ~offset:max_numel ~strides:[| 1 |] [| 1 |] );
    ( "v whose negative reach passes max_numel",
      fun () -> L.v ~offset:max_numel ~strides:[| -max_numel |] [| 2 |] );
  ]

let test_refuses (_, f) = raises_match Exn.invalid_arg f

let test_bounds () =
  equal int 32 L.max_rank;
  equal int 32 (L.rank (L.contiguous (make 32 1)));
  equal int 32 (L.rank (L.v ~strides:(make 32 0) (make 32 2)));
  equal int (1 lsl 53) max_numel;
  equal ~msg:"NX_MAX_NUMEL" int max_numel (Nx_array_support.max_numel ());
  equal int max_numel (L.numel (L.contiguous [| max_numel |]));
  equal (pair int int) (0, max_numel)
    (L.span (L.contiguous [| max_numel / 2; 2 |]));
  let reach = max_numel - 1 in
  let l = L.v ~strides:[| reach |] [| 2 |] in
  equal iints [| reach |] (L.strides l);
  equal (pair int int) (0, max_numel) (L.span l);
  let l = L.v ~offset:(max_numel - 1) ~strides:[| -reach |] [| 2 |] in
  equal (pair int int) (0, max_numel) (L.span l);
  let l = L.v ~offset:(max_numel - 1) ~strides:[| 0 |] [| 1 |] in
  equal (pair int int) (max_numel - 1, max_numel) (L.span l);
  equal int max_numel
    (L.numel (L.v ~strides:[| 0; 0 |] [| max_numel / 2; 2 |]))

let test_no_element () =
  let s : int iarray = [| max_numel / 2; 2; 0 |] in
  equal iints s (L.shape (L.contiguous s));
  equal iints s (L.shape (L.v ~strides:[| 0; 0; 0 |] s));
  equal int 0 (L.numel (L.contiguous s))

let test_broadcast_scalar () =
  let one = L.contiguous [||] in
  let l = Option.get (L.move (M.Broadcast [| 2 |]) one) in
  equal bool false (L.is_contiguous l);
  equal bool false (L.is_distinct l);
  equal bool true (L.is_contiguous (L.contiguous [| 0; 3 |]))

let test_axis_refuses () =
  let l = L.contiguous [| 2; 3 |] in
  List.iter
    (fun i ->
      raises_match ~msg:(string_of_int i) Exn.invalid_arg (fun () -> L.dim l i);
      raises_match ~msg:(string_of_int i) Exn.invalid_arg (fun () ->
          L.stride l i))
    [ -1; 2; min_int; max_int ]

(* The words [f] allocates over a thousand calls, beyond the loop's own. *)
let words f =
  let count g =
    let before = Gc.minor_words () in
    for _ = 1 to 1000 do
      ignore (Sys.opaque_identity (g ()))
    done;
    Gc.minor_words () -. before
  in
  count f -. count (fun () -> 0)

let test_no_allocation () =
  let l = L.contiguous [| 2; 3; 4; 5 |] in
  let l' = L.v ~offset:1 ~strides:[| 60; 20; 5; 1 |] [| 2; 3; 4; 5 |] in
  let zero = float 0.5 in
  equal ~msg:"rank" zero 0. (words (fun () -> L.rank l));
  equal ~msg:"dim" zero 0. (words (fun () -> L.dim l 2));
  equal ~msg:"stride" zero 0. (words (fun () -> L.stride l 2));
  equal ~msg:"offset" zero 0. (words (fun () -> L.offset l));
  equal ~msg:"numel" zero 0. (words (fun () -> L.numel l));
  equal ~msg:"shape" zero 0. (words (fun () -> L.shape l));
  equal ~msg:"strides" zero 0. (words (fun () -> L.strides l));
  equal ~msg:"is_contiguous" zero 0.
    (words (fun () -> Bool.to_int (L.is_contiguous l)));
  equal ~msg:"is_distinct" zero 0.
    (words (fun () -> Bool.to_int (L.is_distinct l)));
  equal ~msg:"equal" zero 0. (words (fun () -> Bool.to_int (L.equal l l')));
  equal ~msg:"hash" zero 0. (words (fun () -> L.hash l))

(* Movements *)

(* One view of each constructor, a movement Move.shape refuses, and a reshape
   no strides express: every outcome, on every run. *)
let move_examples =
  let l = L.contiguous [| 2; 3 |] in
  let r start count step = { M.start; count; step } in
  [
    (l, M.Reshape [| 3; 2 |]);
    (l, M.Broadcast [| 4; 2; 3 |]);
    (l, M.Permute [| 1; 0 |]);
    (l, M.Slice [| r 1 1 1; r 2 2 (-2) |]);
    (l, M.Window [| { M.axis = 1; size = 2; step = 1; dilation = 1 } |]);
    (l, M.Permute [| 0; 0 |]);
    (L.v ~strides:[| 1; 3 |] [| 3; 2 |], M.Reshape [| 6 |]);
  ]

let law_move (l, m) =
  let s = L.shape l in
  match M.shape m s with
  | exception Invalid_argument _ ->
      collect "refused";
      raises_match Exn.invalid_arg (fun () -> L.move m l)
  | s' -> (
      let s = Iarray.to_array s in
      let target idx = position l (source m s idx) in
      match L.move m l with
      | Some l' ->
          collect
            (match m with
            | M.Reshape _ -> "a reshape view"
            | M.Broadcast _ -> "broadcast"
            | M.Permute _ -> "permute"
            | M.Slice _ -> "slice"
            | M.Window _ -> "window");
          equal iints s' (L.shape l');
          List.iter
            (fun idx ->
              equal
                ~msg:(Format.asprintf "%a" pp_ints idx)
                int (target idx) (position l' idx))
            (indices (Iarray.to_array s'))
      | None ->
          collect "a reshape no strides express";
          (match m with
          | M.Reshape _ -> ()
          | _ -> failf "%a answered None" pp_move m);
          equal ~msg:"strides express the reshape" bool false
            (expressible (Iarray.to_array s') target))

(* Canonical form *)

let law_canonical_move (l, m) =
  match L.move m l with
  | exception Invalid_argument _ -> ()
  | None -> ()
  | Some l' ->
      cover "an extent-1 axis" (Iarray.mem 1 (L.shape l'));
      cover "no element" (L.numel l' = 0);
      if not (canonical l') then failf "%a is not canonical" L.pp l'

(* Two layouts of one shape over strides in [-1, 1], so that many map indices
   alike. *)
let similar =
  let open Gen in
  with_pp
    (fun ppf (a, b) -> Format.fprintf ppf "%a, %a" L.pp a L.pp b)
    (let* s = shape in
     let r = const (Array.length s) in
     let layout =
       let+ strides = array ~size:r (int_range (-1) 1)
       and+ least = int_range 2 3 in
       L.v ~offset:(lift s strides least) ~strides:(Iarray.of_array strides)
         (Iarray.of_array s)
     in
     pair layout layout)

let law_equal (a, b) =
  let same = positions a = positions b in
  cover "the same map" same;
  cover "different maps" (not same);
  equal bool same (L.equal a b)

let law_equivalence pair = Law.equivalence layout pair

let law_hash (a, b) =
  cover "equal" (L.equal a b);
  if L.equal a b then equal int (L.hash a) (L.hash b)

(* [v] on [l]'s own fields, with every extent-1 stride respelled. *)
let respell l =
  let s = L.shape l in
  let strides =
    Iarray.mapi (fun i t -> if Iarray.get s i = 1 then 7 else t) (L.strides l)
  in
  L.v ~offset:(L.offset l) ~strides s

let law_respelled l =
  equal layout l (respell l);
  equal int (L.hash l) (L.hash (respell l))

let test_equal_shapes () =
  equal bool false (L.equal (L.contiguous [| 1 |]) (L.contiguous [||]));
  equal bool false (L.equal (L.contiguous [| 0; 2 |]) (L.contiguous [| 2; 0 |]))

(* Flags and span *)

let law_contiguous_flag l =
  let ps = positions l in
  let c_order = List.mapi (fun k _ -> L.offset l + k) ps in
  cover "contiguous" (ps = c_order && ps <> []);
  cover "not contiguous" (ps <> c_order);
  equal bool (ps = c_order) (L.is_contiguous l)

(* The distinct test: axes of extent above 1 ordered by [|stride|], ties by
   axis, each stride past the reach of the axes before it. A layout with no
   element reaches no position twice, whatever its strides. *)
let distinct_rule l =
  let axes =
    List.filter (fun i -> L.dim l i > 1) (List.init (L.rank l) Fun.id)
  in
  let key i = (abs (L.stride l i), i) in
  let axes = List.sort (fun i j -> compare (key i) (key j)) axes in
  let rec go reach = function
    | [] -> true
    | i :: rest ->
        abs (L.stride l i) > reach
        && go (reach + ((L.dim l i - 1) * abs (L.stride l i))) rest
  in
  L.numel l = 0 || go 0 axes

let law_distinct_rule l =
  cover "distinct" (L.is_distinct l);
  cover "not distinct" (not (L.is_distinct l));
  cover "no element" (L.numel l = 0);
  equal bool (distinct_rule l) (L.is_distinct l)

let law_distinct_sound l =
  cover "distinct" (L.is_distinct l);
  cover "a repeated position" (not (distinct (positions l)));
  if L.is_distinct l then equal bool true (distinct (positions l))

let law_distinct_views l =
  cover "a strided view" ((not (L.is_contiguous l)) && L.numel l > 1);
  equal bool true (L.is_distinct l)

let law_span l =
  let lo, hi = L.span l in
  at_least ~msg:"lo" int ~than:0 lo;
  match positions l with
  | [] -> equal (pair int int) (0, 0) (lo, hi)
  | ps ->
      List.iter
        (fun p -> if p < lo || p >= hi then failf "%d outside [%d, %d)" p lo hi)
        ps

(* Every layout the constructors and movements make is well formed: strides
   of its shape's length, at most max_rank axes, a position no lower than
   0. The doors read layouts without checking this again. *)
let law_well_formed l =
  equal ~msg:"strides" int (L.rank l) (Iarray.length (L.strides l));
  at_most ~msg:"rank" int ~than:L.max_rank (L.rank l);
  at_least ~msg:"offset" int ~than:0 (L.offset l)

(* The descriptor nx_read fills for an array laid out by [l], over the bytes
   [l] reaches, holds what the accessors give: the rank, the flags, the offset,
   the extents and the strides. *)
let law_descriptor l =
  let flag b f = if b then f else 0 in
  let r = L.rank l in
  let flags =
    flag (L.is_contiguous l) 1
    lor flag (L.is_distinct l) 2
    lor flag (L.numel l = 0) 4
  in
  let bytes = Rig.Buffer.create Rig.host (max 1 (snd (L.span l))) in
  let a = Nx_array.v Nx_array.Dtype.Uint8 l bytes in
  equal ints
    (Array.concat
       [
         [| r; flags; L.offset l |];
         Iarray.to_array (L.shape l);
         Iarray.to_array (L.strides l);
       ])
    (Nx_array_support.descriptor a)

(* Coalescing *)

let same_shape =
  let open Gen in
  with_pp
    (fun ppf ls -> Format.pp_print_list L.pp ppf ls)
    (let* l = one_of [ strided; reached ~apart:false ] in
     let s = L.shape l in
     let other =
       one_of [ strided_of (Iarray.to_array s); constant (L.contiguous s) ]
     in
     let+ others = list ~size:(int_range 0 3) other in
     l :: others)

(* Whether axes [i] and [i + 1] of [l] lie as one run. *)
let one_run l i = L.stride l i = L.dim l (i + 1) * L.stride l (i + 1)

(* [coalesce ls] checked against what its interface states: the layouts it
   answers, of one shape. *)
let coalesced ls =
  let cs = L.coalesce (Array.of_list ls) in
  let s = L.shape cs.(0) in
  equal ~msg:"layouts" int (List.length ls) (Array.length cs);
  at_least ~msg:"rank" int ~than:1 (Iarray.length s);
  Array.iter (fun c -> equal ~msg:"one shape" iints s (L.shape c)) cs;
  if Iarray.length s > 1 && Iarray.mem 1 s then
    failf "the extent-1 axis of %a is kept" pp_iints s;
  for i = 0 to Iarray.length s - 2 do
    if Array.for_all (fun c -> one_run c i) cs then
      failf "axes %d and %d of %a lie as one run" i (i + 1) pp_iints s
  done;
  List.iteri (fun k l -> equal (list int) (positions l) (positions cs.(k))) ls;
  cs

let law_coalesce ls =
  let cs = coalesced ls in
  let rank = List.fold_left (fun r l -> max r (L.rank l)) 1 ls in
  cover "merges axes" (L.rank cs.(0) < rank)

(* Any number of layouts, one below the door's runs to many past them,
   merge where every one lays two axes out as one run: contiguous ones
   into one axis, and with a transposed one last, none. *)
let test_coalesce_counts () =
  let a = L.contiguous [| 2; 3 |] and t = L.v ~strides:[| 1; 2 |] [| 2; 3 |] in
  let layouts n last = List.init n (fun k -> if k = n - 1 then last else a) in
  equal ~msg:"max_operands" int 18 Nx_array.max_operands;
  List.iter
    (fun n ->
      let msg = Printf.sprintf "%d layouts" n in
      equal ~msg int 1 (L.rank (coalesced (layouts n a)).(0));
      if n > 1 then equal ~msg int 2 (L.rank (coalesced (layouts n t)).(0)))
    [ 1; 2; 17; 18; 19; 32; 100 ];
  raises_match Exn.invalid_arg (fun () -> L.coalesce [||]);
  raises_match Exn.invalid_arg (fun () ->
      L.coalesce [| a; L.contiguous [| 3; 2 |] |])

let test_pp () =
  let l = L.v ~offset:17 ~strides:[| 23; 19 |] [| 7; 5 |] in
  in_order ~subs:[ "7"; "5"; "23"; "19"; "17" ] (Format.asprintf "%a" L.pp l)

(* A refusal names the axis at fault. *)
let test_messages () =
  raises_match
    (Exn.invalid_arg ~substring:"Layout.v: axis 1 of [2; -1] is -1, below 0")
    (fun () -> L.v ~strides:[| 1; 1 |] [| 2; -1 |]);
  raises_match
    (Exn.invalid_arg
       ~substring:"Layout.v: axis 1 has extent 3 and stride 2305843009213693952")
    (fun () -> L.v ~strides:[| 1; (max_int / 2) + 1 |] [| 2; 3 |]);
  raises_match
    (Exn.invalid_arg ~substring:"Layout.contiguous: axis 0 of [-4] is -4")
    (fun () -> L.contiguous [| -4 |])

let tests =
  [
    group "construction"
      [
        test "refusals name the axis at fault" test_messages;
        test
          "max_rank is 32, max_numel 2^53, and layouts reach both bounds"
          test_bounds;
        prop ~count:1000
          "v keeps its arguments in canonical form, and raises iff they do not \
           fit"
          v_args law_v;
        prop
          "contiguous s places element k at position k, and its shape is s \
           itself"
          shape law_contiguous;
        prop ~count:1000 "fits is its test, and contiguous takes what it fits"
          v_args law_fits;
        test "fits takes max_rank axes and refuses one more" test_fits_rank;
        cases ~name:fst "refuses" refusals test_refuses;
        test "a shape with no element takes other extents up to max_numel"
          test_no_element;
        test "dim and stride refuse an axis out of range" test_axis_refuses;
        test
          "a scalar broadcast to two elements is neither contiguous nor \
           distinct"
          test_broadcast_scalar;
        test "queries, flags, equal and hash allocate nothing"
          test_no_allocation;
      ];
    group "movements"
      [
        prop ~count:500 ~examples:move_examples
          "a movement is its index map, None is a reshape no strides express, \
           and it raises as Move.shape does"
          layout_and_move law_move;
        prop "a movement's layout is canonical" layout_and_move
          law_canonical_move;
      ];
    group "equality"
      [
        prop "layouts of one shape are equal iff they map indices alike" similar
          law_equal;
        prop "equal is an equivalence" similar law_equivalence;
        prop "equal layouts hash alike" similar law_hash;
        prop
          "v of a layout's fields is the layout, whatever its extent-1 strides"
          any_layout law_respelled;
        test "layouts of other shapes differ" test_equal_shapes;
      ];
    group "flags"
      [
        prop "is_contiguous iff C order is offset + k" any_layout
          law_contiguous_flag;
        prop "is_distinct is the test its interface states" any_layout
          law_distinct_rule;
        prop "is_distinct is never true of a repeated position" any_layout
          law_distinct_sound;
        prop "is_distinct holds of views that keep positions apart"
          (Gen.with_pp L.pp (reached ~apart:true))
          law_distinct_views;
        prop "span holds every position, from 0 up" any_layout law_span;
        prop "every layout made is well formed" any_layout law_well_formed;
        prop "nx_read's descriptor holds the fields the accessors give"
          any_layout law_descriptor;
      ];
    group "coalesce"
      [
        prop
          "merges every run, drops extent-1 axes and keeps each layout's \
           positions in order"
          same_shape law_coalesce;
        test "takes any number of layouts of one shape" test_coalesce_counts;
      ];
    test "pp shows shape, strides and offset" test_pp;
  ]

let () = exit (run "nx_array.layout" tests)