opam install qcheck-alcotest

open QCheck2

(* 基本属性测试 *)
let prop_reverse_length =
  Test.make ~name:"reverse 保持长度"
    (list int)
    (fun l -> List.length (List.rev l) = List.length l)

let prop_reverse_reverse =
  Test.make ~name:"reverse(reverse(l)) = l"
    (list int)
    (fun l -> List.rev (List.rev l) = l)

let prop_sort_preserves_length =
  Test.make ~name:"排序保持长度"
    (list int)
    (fun l -> List.length (List.sort compare l) = List.length l)

let prop_sort_is_sorted =
  Test.make ~name:"排序结果是有序的"
    (list int)
    (fun l ->
      let sorted = List.sort compare l in
      let rec is_sorted = function
        | [] | [_] -> true
        | a :: (b :: _ as rest) -> a <= b && is_sorted rest
      in
      is_sorted sorted)

let prop_sort_idempotent =
  Test.make ~name:"排序是幂等的"
    (list int)
    (fun l ->
      let s1 = List.sort compare l in
      let s2 = List.sort compare s1 in
      s1 = s2)

(* 运行测试 *)
let () =
  let open Alcotest in
  run "属性测试" [
    "列表属性", [
      QCheck_alcotest.to_alcotest prop_reverse_length;
      QCheck_alcotest.to_alcotest prop_reverse_reverse;
      QCheck_alcotest.to_alcotest prop_sort_preserves_length;
      QCheck_alcotest.to_alcotest prop_sort_is_sorted;
      QCheck_alcotest.to_alcotest prop_sort_idempotent;
    ];
  ]

open QCheck2.Gen

(* 基本生成器 *)
let gen_int = int
let gen_nat = nat
let gen_bool = bool
let gen_float = float
let gen_char = char
let gen_string = string

(* 范围生成器 *)
let gen_small_int = int_range 0 100
let gen_percentage = float_range 0.0 100.0
let gen_alpha = char_range 'a' 'z'

(* 组合生成器 *)
let gen_pair = pair gen_small_int gen_string
let gen_triple = triple gen_int gen_bool gen_string
let gen_list = list_size (int_range 0 50) gen_small_int
let gen_nonempty_list = list_size (int_range 1 50) gen_small_int

(* 选项生成器 *)
let gen_option = option gen_small_int
let gen_result = result gen_small_int gen_string

(* 频率生成器 *)
let gen_tree =
  let rec gen_tree size =
    if size <= 0 then pure Leaf
    else frequency [
      1, pure Leaf;
      3, map3
        (fun left v right -> Node (left, v, right))
        (gen_tree (size / 2))
        gen_small_int
        (gen_tree (size / 2));
    ]
  in
  sized gen_tree

(* 自定义生成器 *)
type color = Red | Green | Blue | RGB of int * int * int

let gen_color =
  frequency [
    3, pure Red;
    3, pure Green;
    3, pure Blue;
    1, map3
      (fun r g b -> RGB (r, g, b))
      (int_range 0 255)
      (int_range 0 255)
      (int_range 0 255);
  ]

(* 日期生成器 *)
type date = { year: int; month: int; day: int }

let gen_date =
  let gen_year = int_range 1900 2100 in
  let gen_month = int_range 1 12 in
  let gen_day month =
    let max_day = match month with
      | 2 -> 28
      | 4 | 6 | 9 | 11 -> 30
      | _ -> 31
    in
    int_range 1 max_day
  in
  gen_year >>= fun year ->
  gen_month >>= fun month ->
  gen_day month >>= fun day ->
  pure { year; month; day }

(* 字符串生成器 *)
let gen_identifier =
  let gen_first = char_range 'a' 'z' in
  let gen_rest = frequency [
    5, char_range 'a' 'z';
    3, char_range '0' '9';
    1, pure '_';
  ] in
  string_size ~gen:gen_first (int_range 0 0) >>= fun first ->
  string_size ~gen:gen_rest (int_range 0 20) >>= fun rest ->
  pure (first ^ rest)

(* 从列表中随机选择 *)
let gen_from_list xs =
  map (List.nth xs) (int_range 0 (List.length xs - 1))

(* 收缩器：找到最小失败用例 *)
open QCheck2.Shrink

(* 列表收缩 *)
let shrink_list shrink_elem l =
  (* 尝试删除元素 *)
  let rec remove_at i = function
    | [] -> []
    | _ :: xs when i = 0 -> xs
    | x :: xs -> x :: remove_at (i - 1) xs
  in
  let n = List.length l in
  if n = 0 then Seq.empty
  else
    (* 尝试删除每个位置 *)
    Seq.append
      (Seq.init n (fun i -> remove_at i l))
      (* 收缩每个元素 *)
      (List.to_seq l |> Seq.mapi (fun i x ->
        shrink_elem x |> Seq.map (fun x' ->
          List.mapi (fun j y -> if j = i then x' else y) l))
      |> Seq.concat)

(* 整数收缩 *)
let shrink_int n =
  if n = 0 then Seq.empty
  else
    let rec aux n =
      if n = 0 then Seq.return 0
      else Seq.cons (n / 2) (fun () -> aux (n / 2))
    in
    aux n

(* 自定义收缩器 *)
type expr =
  | Num of int
  | Add of expr * expr
  | Mul of expr * expr

let rec shrink_expr = function
  | Num n ->
    shrink_int n |> Seq.map (fun n' -> Num n')
  | Add (l, r) ->
    Seq.append
      (Seq.of_list [l; r])  (* 简化为子表达式 *)
      (Seq.append
        (shrink_expr l |> Seq.map (fun l' -> Add (l', r)))
        (shrink_expr r |> Seq.map (fun r' -> Add (l, r'))))
  | Mul (l, r) ->
    Seq.append
      (Seq.of_list [l; r])
      (Seq.append
        (shrink_expr l |> Seq.map (fun l' -> Mul (l', r)))
        (shrink_expr r |> Seq.map (fun r' -> Mul (l, r'))))

(* 带收缩的属性测试 *)
let prop_eval_nonneg =
  Test.make
    ~name:"求值结果非负"
    ~shrink:shrink_expr
    (make_expr_gen)
    (fun expr ->
      let result = eval expr in
      result >= 0)

(* 通用属性模式 *)

(* 1. 圆性（Roundtrip） *)
let prop_json_roundtrip =
  Test.make
    ~name:"JSON 圆性"
    gen_user
    (fun user ->
      let json = user_to_yojson user in
      let s = Yojson.Safe.to_string json in
      let parsed = Yojson.Safe.from_string s in
      match user_of_yojson parsed with
      | Ok user' -> user = user'
      | Error _ -> false)

(* 2. 不变量 *)
let prop_bst_invariant =
  Test.make
    ~name:"BST 插入后仍有序"
    (list (pair gen_small_int gen_string))
    (fun kvs ->
      let tree = List.fold_left (fun t (k, v) -> insert k v t) empty kvs in
      is_bst tree)

(* 3. 代数性质 *)
let prop_add_commutative =
  Test.make
    ~name:"加法交换律"
    (pair int int)
    (fun (a, b) -> add a b = add b a)

let prop_add_associative =
  Test.make
    ~name:"加法结合律"
    (triple int int int)
    (fun (a, b, c) -> add (add a b) c = add a (add b c))

let prop_add_identity =
  Test.make
    ~name:"加法单位元"
    int
    (fun a -> add a 0 = a)

(* 4. 蕴含关系 *)
let prop_filter_subset =
  Test.make
    ~name:"filter 是子集"
    (pair (list int) (fun1 int bool))
    (fun (l, f) ->
      let f = Fn.apply f in
      let filtered = List.filter f l in
      List.for_all f filtered)

(* 5. 幂等性 *)
let prop_normalize_idempotent =
  Test.make
    ~name:"normalize 是幂等的"
    gen_string
    (fun s ->
      normalize (normalize s) = normalize s)

(* 状态机测试：测试有状态的系统 *)
open QCheck2

type state = {
  counter: int;
  items: string list;
}

type command =
  | Increment
  | Decrement
  | Add_item of string
  | Remove_item of string
  | Get_counter
  | Get_items

let gen_command =
  Gen.frequency [
    3, Gen.pure Increment;
    2, Gen.pure Decrement;
    3, Gen.map (fun s -> Add_item s) Gen.string;
    2, Gen.map (fun s -> Remove_item s) Gen.string;
    2, Gen.pure Get_counter;
    1, Gen.pure Get_items;
  ]

(* 模型：理想行为 *)
let model_apply state = function
  | Increment -> { state with counter = state.counter + 1 }
  | Decrement -> { state with counter = state.counter - 1 }
  | Add_item s -> { state with items = s :: state.items }
  | Remove_item s -> { state with items = List.filter ((<>) s) state.items }
  | Get_counter | Get_items -> state

(* 实际系统 *)
let system_apply impl = function
  | Increment -> impl.increment ()
  | Decrement -> impl.decrement ()
  | Add_item s -> impl.add_item s
  | Remove_item s -> impl.remove_item s
  | Get_counter -> impl.get_counter ()
  | Get_items -> impl.get_items ()

(* 状态机属性 *)
let prop_state_machine =
  Test.make
    ~name:"状态机一致性"
    (Gen.list_size (Gen.int_range 0 100) gen_command)
    (fun commands ->
      let model = ref { counter = 0; items = [] } in
      let impl = create_impl () in
      List.for_all (fun cmd ->
        let expected = model_apply !model cmd in
        model := expected;
        let actual = system_apply impl cmd in
        expected.counter = actual.counter &&
        List.sort compare expected.items = List.sort compare actual.items
      ) commands)

(* 带前置条件的状态机测试 *)
let prop_state_machine_precond =
  Test.make
    ~name:"状态机（带前置条件）"
    (Gen.list_size (Gen.int_range 0 100) gen_command)
    (fun commands ->
      let model = ref { counter = 0; items = [] } in
      let impl = create_impl () in
      List.for_all (fun cmd ->
        match cmd with
        | Remove_item _ when !model.items = [] ->
          true  (* 跳过空列表的删除 *)
        | _ ->
          let expected = model_apply !model cmd in
          model := expected;
          let actual = system_apply impl cmd in
          expected = actual
      ) commands)

(* 简单的模型检查器 *)
module ModelChecker = struct
  type state = {
    data: int list;
    lock: bool;
    owner: string option;
  }

  type action =
    | Acquire of string
    | Release of string
    | Push of int
    | Pop

  let next_states state = function
    | Acquire who when not state.lock ->
      [{ state with lock = true; owner = Some who }]
    | Release who when state.owner = Some who ->
      [{ state with lock = false; owner = None }]
    | _ -> []

  let is_safe state =
    (* 不变量：持有锁时才能修改数据 *)
    true

  let check_invariant initial_state actions =
    let rec explore states visited =
      match states with
      | [] -> true
      | state :: rest ->
        if Set.mem state visited then
          explore rest visited
        else if not (is_safe state) then
          false
        else begin
          let next = List.concat_map (next_states state) actions in
          explore (next @ rest) (Set.add state visited)
        end
    in
    explore [initial_state] Set.empty
end

(* 使用 *)
let () =
  let initial = { data = []; lock = false; owner = None } in
  let actions = [
    Acquire "thread1";
    Acquire "thread2";
    Release "thread1";
    Release "thread2";
    Push 1;
    Pop;
  ] in
  let safe = ModelChecker.check_invariant initial actions in
  Printf.printf "系统安全: %b\n" safe

(* 测试数据结构不变量 *)

(* 红黑树不变量 *)
type color = Red | Black
type tree = Empty | Node of color * tree * int * tree

let rec check_rb_tree = function
  | Empty -> (true, 0)  (* (valid, black_height) *)
  | Node (color, left, value, right) ->
    let (left_valid, left_bh) = check_rb_tree left in
    let (right_valid, right_bh) = check_rb_tree right in
    let valid =
      left_valid && right_valid &&
      left_bh = right_bh &&  (* 黑高相等 *)
      (match color with
       | Red ->
         (* 红节点的子节点必须是黑的 *)
         (match left with Node (Red, _, _, _) -> false | _ -> true) &&
         (match right with Node (Red, _, _, _) -> false | _ -> true)
       | Black -> true)
    in
    let bh = (if color = Black then 1 else 0) + left_bh in
    (valid, bh)

(* 不变量属性测试 *)
let prop_rb_invariant =
  Test.make
    ~name:"红黑树不变量"
    (list int)
    (fun xs ->
      let tree = List.fold_left (fun t x -> rb_insert x t) Empty xs in
      let (valid, _) = check_rb_tree tree in
      valid)

(* 哈希表不变量 *)
let prop_hashtbl_invariant =
  Test.make
    ~name:"哈希表操作后仍正确"
    (list (pair gen_small_int gen_string))
    (fun kvs ->
      let tbl = Hashtbl.create 16 in
      List.iter (fun (k, v) -> Hashtbl.replace tbl k v) kvs;
      (* 检查所有插入的键都存在 *)
      List.for_all (fun (k, v) ->
        Hashtbl.find_opt tbl k = Some v
      ) (List.sort_uniq compare kvs))

(* 栈不变量 *)
let prop_stack_invariant =
  Test.make
    ~name:"栈 LIFO 性质"
    (pair (list gen_small_int) (list gen_small_int))
    (fun (pushes, pops) ->
      let stack = Stack.create () in
      List.iter (Stack.push stack) pushes;
      let popped = List.init (min (List.length pushes) (List.length pops))
        (fun _ -> Stack.pop stack) in
      let expected = List.rev (List.init (List.length popped)
        (fun i -> List.nth pushes (List.length pushes - 1 - i))) in
      popped = expected)

(* 契约式设计 *)
module Contract = struct
  let require name pred x =
    if not (pred x) then
      failwith (Printf.sprintf "前置条件违反: %s" name)

  let ensure name pred result =
    if not (pred result) then
      failwith (Printf.sprintf "后置条件违反: %s" name)

  let invariant name pred state =
    if not (pred state) then
      failwith (Printf.sprintf "不变量违反: %s" name)
end

(* 带契约的栈 *)
module SafeStack = struct
  type 'a t = {
    mutable elements: 'a list;
    mutable size: int;
  }

  let create () = { elements = []; size = 0 }

  let push stack x =
    Contract.invariant "size >= 0" (fun s -> s.size >= 0) stack;
    stack.elements <- x :: stack.elements;
    stack.size <- stack.size + 1;
    Contract.invariant "size > 0" (fun s -> s.size > 0) stack

  let pop stack =
    Contract.require "非空" (fun s -> s.size > 0) stack;
    match stack.elements with
    | x :: rest ->
      stack.elements <- rest;
      stack.size <- stack.size - 1;
      Contract.invariant "size >= 0" (fun s -> s.size >= 0) stack;
      x
    | [] -> failwith "不可能"

  let size stack = stack.size
  let is_empty stack = stack.size = 0
end

(* 带契约的队列 *)
module SafeQueue = struct
  type 'a t = {
    mutable front: 'a list;
    mutable rear: 'a list;
    mutable size: int;
  }

  let create () = { front = []; rear = []; size = 0 }

  let enqueue q x =
    q.rear <- x :: q.rear;
    q.size <- q.size + 1

  let dequeue q =
    Contract.require "非空" (fun q -> q.size > 0) q;
    match q.front with
    | x :: rest ->
      q.front <- rest;
      q.size <- q.size - 1;
      x
    | [] ->
      q.front <- List.rev q.rear;
      q.rear <- [];
      match q.front with
      | x :: rest ->
        q.front <- rest;
        q.size <- q.size - 1;
        x
      | [] -> failwith "不可能"
end

(* 案例1：二分搜索 Bug *)
let binary_search_bad arr target =
  let lo = ref 0 in
  let hi = ref (Array.length arr - 1) in
  while !lo <= !hi do
    let mid = (!lo + !hi) / 2 in  (* 溢出风险！ *)
    if arr.(mid) = target then
      raise (Found mid)
    else if arr.(mid) < target then
      lo := mid + 1
    else
      hi := mid - 1
  done;
  -1

(* QCheck 发现 *)
let prop_binary_search =
  Test.make
    ~name:"二分搜索找到存在的元素"
    (pair (array_size (int_range 0 100) (int_range 0 100)) (int_range 0 100))
    (fun (arr, target) ->
      Array.sort compare arr;
      let result = binary_search arr target in
      result = -1 || arr.(result) = target)

(* 修复：避免溢出 *)
let binary_search_good arr target =
  let lo = ref 0 in
  let hi = ref (Array.length arr - 1) in
  while !lo <= !hi do
    let mid = !lo + (!hi - !lo) / 2 in  (* 安全的中点计算 *)
    if arr.(mid) = target then
      raise (Found mid)
    else if arr.(mid) < target then
      lo := mid + 1
    else
      hi := mid - 1
  done;
  -1

(* 案例2：列表去重 Bug *)
let dedup_bad xs =
  let seen = Hashtbl.create 16 in
  List.filter (fun x ->
    if Hashtbl.mem seen x then false
    else begin
      Hashtbl.add seen x ();
      true
    end
  ) xs

(* QCheck 发现：依赖 Hashtbl 的遍历顺序 *)
let prop_dedup_length =
  Test.make
    ~name:"dedup 不增加长度"
    (list int)
    (fun xs ->
      List.length (dedup xs) <= List.length xs)