(** Key for hastables that contain tuples of integers encoded in a Zarith integer *)


type kind_array = 
| Kind_3
| Kind_6
| Kind_12
| Kind_9
| Kind_4
| Kind_2
| Kind_1


(** Build a Zkey from an array or 1, 2, 3, 4, 6, 9, or 12 integers *)
let of_int_array ~kind a = 
  let (<|) x a = 
    Z.logor (Z.shift_left x 64) a
  in
  let (<<) x a = 
    Int64.logor (Int64.shift_left x 10) (Int64.of_int a)
  in
  let (<+) x a = 
    Int64.logor (Int64.shift_left x 16) (Int64.of_int a)
  in
  match kind with
  | Kind_3 -> (Int64.of_int a.(0)) << a.(1) << a.(2) |> Z.of_int64
  | Kind_6 -> (Int64.of_int a.(0)) << a.(1) << a.(2) << a.(3) << a.(4) << a.(5)
        |> Z.of_int64
  | Kind_12 ->
    let a =
      (Int64.of_int a.(0)) << a.(1) << a.(2) << a.(3) << a.(4) << a.(5)
      |> Z.of_int64 
     and b = 
      (Int64.of_int a.(6)) << a.(7) << a.(8) << a.(9) << a.(10) << a.(11)
      |> Z.of_int64 
     in a <| b
  | Kind_9 ->
    let a =
      (Int64.of_int a.(0)) << a.(1) << a.(2) << a.(3) << a.(4) << a.(5)
      |> Z.of_int64 
     and b = 
      (Int64.of_int a.(6)) << a.(7) << a.(8) 
      |> Z.of_int64 
     in a <| b
  | Kind_4 -> (Int64.of_int a.(0)) <+ a.(1) <+ a.(2) <+ a.(3) |> Z.of_int64
  | Kind_2 -> (Int64.of_int a.(0)) <+ a.(1) |> Z.of_int64
  | Kind_1 -> Z.of_int a.(0) 


type kind = 
  | One    of (int)
  | Two    of (int*int)
  | Three  of (int*int*int)
  | Four   of ((int*int)*(int*int))
  | Six    of ((int*int*int)*(int*int*int))
  | Nine   of ((int*int*int)*(int*int*int)*(int*int*int))
  | Twelve of ((int*int*int)*(int*int*int)*(int*int*int)*(int*int*int))

let (<|) x a = 
  Z.logor (Z.shift_left x 64) a

let (<<) x a = 
  Int64.logor (Int64.shift_left x 10) (Int64.of_int a)

let (<+) x a = 
  Int64.logor (Int64.shift_left x 16) (Int64.of_int a)

let of_int_tuple a = 
  match a with
  | One   (a)       -> Z.of_int a
  | Two   (a,b)     -> (Int64.of_int a) <+ b |> Z.of_int64
  | Three (a,b,c)   -> (Int64.of_int a) << b << c |> Z.of_int64
  | Four  ((a,b),(c,d)) ->
    (Int64.of_int a) <+ b <+ c <+ d |> Z.of_int64
  | Six   ((a,b,c),(d,e,f)) ->
    (Int64.of_int a) << b << c << d << e << f |> Z.of_int64
  | Twelve ((a,b,c),(d,e,f),(g,h,i),(j,k,l)) ->
    let a =
      (Int64.of_int a) << b << c << d << e << f |> Z.of_int64
    and b = 
      (Int64.of_int g) << h << i << j << k << l |> Z.of_int64
    in a <| b
  | Nine ((a,b,c),(d,e,f),(g,h,i)) -> 
    let a =
      (Int64.of_int a) << b << c << d << e << f |> Z.of_int64
    and b = 
      (Int64.of_int g) << h << i |> Z.of_int64
    in a <| b


let mask10 = Int64.of_int 0x3ff
and mask16 = Int64.of_int 0xffff
(** Transform the Zkey into an int array *)
let to_int_array ~kind a = 
  let open Int64 in
  match kind with
  | Kind_3  -> let x = Z.to_int64 a in
               [| to_int ( logand mask10 (shift_right x 20)) ;
                  to_int ( logand mask10 (shift_right x 10)) ;
                  to_int ( logand mask10 x) |]
  | Kind_6  -> let x = Z.to_int64 a in
               [| to_int ( logand mask10 (shift_right x 50)) ;
                  to_int ( logand mask10 (shift_right x 40)) ;
                  to_int ( logand mask10 (shift_right x 30)) ;
                  to_int ( logand mask10 (shift_right x 20)) ;
                  to_int ( logand mask10 (shift_right x 10)) ;
                  to_int ( logand mask10 x) |]
  | Kind_12 -> let x = Z.to_int64 @@ Z.extract a  0 60
               and y = Z.to_int64 @@ Z.extract a 64 60
               in
               [| to_int ( logand mask10 (shift_right y 114)) ;
                  to_int ( logand mask10 (shift_right y 104)) ;
                  to_int ( logand mask10 (shift_right y 94)) ;
                  to_int ( logand mask10 (shift_right y 84)) ;
                  to_int ( logand mask10 (shift_right y 74)) ;
                  to_int ( logand mask10 (shift_right y 64)) ;
                  to_int ( logand mask10 (shift_right x 50)) ;
                  to_int ( logand mask10 (shift_right x 40)) ;
                  to_int ( logand mask10 (shift_right x 30)) ;
                  to_int ( logand mask10 (shift_right x 20)) ;
                  to_int ( logand mask10 (shift_right x 10)) ;
                  to_int ( logand mask10 x) |]
  | Kind_9  -> let x = Z.to_int64 @@ Z.extract a  0 60
               and y = Z.to_int64 @@ Z.extract a 64 60
               in
               [| to_int ( logand mask10 (shift_right y 84)) ;
                  to_int ( logand mask10 (shift_right y 74)) ;
                  to_int ( logand mask10 (shift_right y 64)) ;
                  to_int ( logand mask10 (shift_right x 50)) ;
                  to_int ( logand mask10 (shift_right x 40)) ;
                  to_int ( logand mask10 (shift_right x 30)) ;
                  to_int ( logand mask10 (shift_right x 20)) ;
                  to_int ( logand mask10 (shift_right x 10)) ;
                  to_int ( logand mask10 x) |]
  | Kind_4  -> let x = Z.to_int64 a in
               [| to_int ( logand mask16 (shift_right x 48)) ;
                  to_int ( logand mask16 (shift_right x 32)) ;
                  to_int ( logand mask16 (shift_right x 16)) ;
                  to_int ( logand mask16 x) |]
  | Kind_2  -> let x = Z.to_int64 a in
               [| to_int ( logand mask16 (shift_right x 16)) ;
                  to_int ( logand mask16 x) |]
  | Kind_1  -> [| Z.to_int a |]


let to_string ~kind a =
  "< " ^ ( Z.to_string a ) ^ " | " ^ (
  to_int_array kind a 
  |> Array.map string_of_int
  |> Array.to_list
  |> String.concat ", "
  ) ^ " >"


(** Transform the Zkey into an int array *)
let to_int_tuple ~kind a = 
  match kind with
  | Kind_3  -> Three  ( Z.to_int @@ Z.extract a  20 10 ,
                        Z.to_int @@ Z.extract a  10 10 ,
                        Z.to_int @@ Z.extract a   0 10 )

  | Kind_6  -> Six    ((Z.to_int @@ Z.extract a  50 10 ,
                        Z.to_int @@ Z.extract a  40 10 ,
                        Z.to_int @@ Z.extract a  30 10),
                       (Z.to_int @@ Z.extract a  20 10 ,
                        Z.to_int @@ Z.extract a  10 10 ,
                        Z.to_int @@ Z.extract a   0 10))

  | Kind_12 -> Twelve ((Z.to_int @@ Z.extract a 114 10 ,
                        Z.to_int @@ Z.extract a 104 10 ,
                        Z.to_int @@ Z.extract a  94 10),
                       (Z.to_int @@ Z.extract a  84 10 ,
                        Z.to_int @@ Z.extract a  74 10 ,
                        Z.to_int @@ Z.extract a  64 10),
                       (Z.to_int @@ Z.extract a  50 10 ,
                        Z.to_int @@ Z.extract a  40 10 ,
                        Z.to_int @@ Z.extract a  30 10),
                       (Z.to_int @@ Z.extract a  20 10 ,
                        Z.to_int @@ Z.extract a  10 10 ,
                        Z.to_int @@ Z.extract a   0 10))

  | Kind_9 ->  Nine   ((Z.to_int @@ Z.extract a  84 10 ,
                        Z.to_int @@ Z.extract a  74 10 ,
                        Z.to_int @@ Z.extract a  64 10),
                       (Z.to_int @@ Z.extract a  50 10 ,
                        Z.to_int @@ Z.extract a  40 10 ,
                        Z.to_int @@ Z.extract a  30 10),
                       (Z.to_int @@ Z.extract a  20 10 ,
                        Z.to_int @@ Z.extract a  10 10 ,
                        Z.to_int @@ Z.extract a   0 10))

  | Kind_4  -> Four   ((Z.to_int @@ Z.extract a  48 16 ,
                        Z.to_int @@ Z.extract a  32 16),
                       (Z.to_int @@ Z.extract a  16 16 ,
                        Z.to_int @@ Z.extract a   0 16))

  | Kind_2  -> Two    ( Z.to_int @@ Z.extract a  16 16,
                        Z.to_int @@ Z.extract a   0 16 )

  | Kind_1  -> One    ( Z.to_int a )


include Z

(*
let debug () = 
  let k2  = of_int_array Kind_2  [| 1 ; 2 |]
  and k3  = of_int_array Kind_3  [| 1 ; 2 ; 3 |]
  and k4  = of_int_array Kind_4  [| 1 ; 2 ; 3; 4 |]
  and k6  = of_int_array Kind_6  [| 1 ; 2 ; 3; 4 ; 5; 6|]
  and k12 = of_int_array Kind_12 [| 1 ; 2 ; 3; 4 ; 5; 6 ; 7 ; 8 ; 9 ; 10 ; 11; 12|]
  in
  print_endline @@ to_string Kind_2 k2 ;
  print_endline @@ to_string Kind_3 k3 ;
  print_endline @@ to_string Kind_4 k4 ;
  print_endline @@ to_string Kind_6 k6 ;
  print_endline @@ to_string Kind_12 k12
  

let () = debug ()
*)