10
1
mirror of https://gitlab.com/scemama/QCaml.git synced 2024-12-22 20:33:36 +01:00
QCaml/common/lib/bitstring.ml

217 lines
4.6 KiB
OCaml

(** Single-integer implementation *)
module One = struct
let of_int x = assert (x > 0); x
let numbits _ = 63
let zero = 0
let is_zero x = x = 0
let shift_left x i = x lsl i
let shift_right x i = x lsr i
let shift_left_one i = 1 lsl i
let testbit x i = ( (x lsr i) land 1 ) = 1
let logor a b = a lor b
let neg a = - a
let logxor a b = a lxor b
let logand a b = a land b
let lognot a = lnot a
let minus_one a = a - 1
let plus_one a = a + 1
let popcount = function
| 0 -> 0
| r -> Util.popcnt (Int64.of_int r)
let trailing_zeros r =
Util.trailz (Int64.of_int r)
let hamdist a b =
a lxor b
|> popcount
let pp ppf s =
Format.fprintf ppf "@[@[%a@]@]" (Util.pp_bitstring 64)
(Z.of_int s)
end
(** Zarith implementation *)
module Many = struct
let of_z x = x
let zero = Z.zero
let is_zero x = x = Z.zero
let shift_left = Z.shift_left
let shift_right = Z.shift_right
let shift_left_one i = Z.shift_left Z.one i
let testbit = Z.testbit
let logor = Z.logor
let logxor = Z.logxor
let logand = Z.logand
let lognot = Z.lognot
let neg = Z.neg
let minus_one = Z.pred
let plus_one = Z.succ
let trailing_zeros = Z.trailing_zeros
let hamdist = Z.hamdist
let numbits i = max (Z.numbits i) 64
let popcount z =
if z = Z.zero then 0 else Z.popcount z
let pp ppf s =
Format.fprintf ppf "@[@[%a@]@]" (Util.pp_bitstring (Z.numbits s)) s
end
type t =
| One of int
| Many of Z.t
let of_int x =
One (One.of_int x)
let of_z x =
if Z.numbits x < 64 then One (Z.to_int x) else Many (Many.of_z x)
let zero = function
| n when n < 64 -> One (One.zero)
| _ -> Many (Many.zero)
let numbits = function
| One x -> One.numbits x
| Many x -> Many.numbits x
let is_zero = function
| One x -> One.is_zero x
| Many x -> Many.is_zero x
let neg = function
| One x -> One (One.neg x)
| Many x -> Many (Many.neg x)
let shift_left x i = match x with
| One x -> One (One.shift_left x i)
| Many x -> Many (Many.shift_left x i)
let shift_right x i = match x with
| One x -> One (One.shift_right x i)
| Many x -> Many (Many.shift_right x i)
let shift_left_one = function
| n when n < 64 -> fun i -> One (One.shift_left_one i)
| _ -> fun i -> Many (Many.shift_left_one i)
let testbit = function
| One x -> One.testbit x
| Many x -> Many.testbit x
(** General implementation *)
let logor a b =
match a,b with
| One a, One b -> One (One.logor a b)
| Many a, Many b -> Many (Many.logor a b)
| _ -> invalid_arg "Bitstring.logor"
let logxor a b =
match a,b with
| One a, One b -> One (One.logxor a b)
| Many a, Many b -> Many (Many.logxor a b)
| _ -> invalid_arg "Bitstring.logxor"
let logand a b =
match a,b with
| One a, One b -> One (One.logand a b)
| Many a, Many b -> Many (Many.logand a b)
| _ -> invalid_arg "Bitstring.logand"
let lognot = function
| One x -> One (One.lognot x)
| Many x -> Many (Many.lognot x)
let minus_one = function
| One x -> One (One.minus_one x)
| Many x -> Many (Many.minus_one x)
let plus_one = function
| One x -> One (One.plus_one x)
| Many x -> Many (Many.plus_one x)
let trailing_zeros = function
| One x -> One.trailing_zeros x
| Many x -> Many.trailing_zeros x
let hamdist a b = match a, b with
| One a, One b -> One.hamdist a b
| Many a, Many b -> Many.hamdist a b
| _ -> invalid_arg "Bitstring.hamdist"
let popcount = function
| One x -> One.popcount x
| Many x -> Many.popcount x
(* [[file:~/QCaml/common/bitstring.org::*General implementation][General implementation:10]] *)
let rec to_list ?(accu=[]) = function
| t when (is_zero t) -> List.rev accu
| t -> let newlist =
(trailing_zeros t + 1)::accu
in
let b = logand t @@ minus_one t in
(to_list [@tailcall]) ~accu:newlist b
(* General implementation:10 ends here *)
(* Example:
* #+begin_example
* Bitstring.(to_list (of_int 45));;
* - : int list = [1; 3; 4; 6]
* #+end_example *)
let permutations m n =
let rec aux k u rest =
if k=1 then
List.rev (u :: rest)
else
let t = logor u @@ minus_one u in
let t' = plus_one t in
let not_t = lognot t in
let neg_not_t = neg not_t in
let t'' = shift_right (minus_one @@ logand not_t neg_not_t) (trailing_zeros u + 1) in
(*
let t'' = shift_right (minus_one (logand (lognot t) t')) (trailing_zeros u + 1) in
*)
(aux [@tailcall]) (k-1) (logor t' t'') (u :: rest)
in
aux (Util.binom n m) (minus_one (shift_left_one n m)) []
let pp ppf = function
| One x -> One.pp ppf x
| Many x -> Many.pp ppf x