module T = (val Containers_testlib.make ~__FILE__ ())
include T;;

open Containers_bencode;;


eq ~printer:to_string_debug
  (map_of_list []) (Decode.of_string_exn "de");;
eq ~printer:to_string_debug
  (list [int 1; int 2; string "foo"]) (Decode.of_string_exn "li1ei2e3:fooe");;

module B = Containers_bencode

let rec size = function
  | Int _ | String _ -> 1
  | List l -> List.fold_left (fun n x -> n + size x) 0 l
  | Map m -> Str_map.fold(fun _ v n -> size v + n) m 0

let g_rand_b =
  Q.Gen.(
    sized_size (0--7) @@ fix @@ fun self n ->
      let str n = string_size ~gen:char (0 -- n) in
      let base = [
       int >|= B.int;
       str 100 >|= B.string;
     ] in
    match n with
   | 0 -> oneof base
   | n ->
     frequency @@
      List.map (fun x -> 2, x) base @
      [ 1, list_size (0 -- 10) (self (n-1)) >|= B.list;
        1, list_size (0 -- 10) (pair (str 10) (self (n-1)) ) >|= B.map_of_list;
      ]
  )

let rec shrink_b self = Q.(Iter.(
  match self with
  | Int i -> Shrink.int64 i >|= B.int64
  | String s -> Shrink.string s >|= B.string
  | List l -> Shrink.list ~shrink:shrink_b l >|= B.list
  | Map l ->
      let l = Str_map.fold (fun k v l -> (k,v) :: l) l [] in
      Shrink.list ~shrink:(fun (k,v) ->
        (Shrink.string k >|= fun k -> k,v) <+>
        (shrink_b v >|= fun v -> k,v))
        l
      >|= B.map_of_list
  ))

let rand_b = Q.make ~print:to_string_debug ~stats:["size", size]
~shrink:shrink_b g_rand_b ;;

q rand_b (fun b ->
    let s=Encode.to_string b in
    equal (Decode.of_string_exn s) b) ;;