Require Import ZArith List Bool Zwf Lia ZifyNat.
(* Import ListNotations. *)

Open Scope Z_scope.

 Ltac Zify.zify_post_hook ::= Z.div_mod_to_equations.

Definition app_Z (l1 l2 : list Z) := l1 ++ l2.

Definition Zseq (n : Z) : list Z :=
  if (n <=? 0) then nil else
  snd (Z.iter n (fun '(x, l) => (x + 1, app_Z l (x :: nil))) (0, nil)).

Definition Zfactorial (x : Z) : Z :=
	if (x <=? 0) then 1 else
	fold_right (fun u v => (u + 1) * v) 1 (Zseq x).

Definition Z_to_digits (x : Z) : list Z :=
  if x <=? 0 then
	  nil
	else
    snd (Z.iter (1 + (Z.log2_up x))
		  (fun '(x, l) =>
        if (x =? 0) then (x, l) else (x / 10, l ++ (x mod 10) :: nil))
			 (x, nil)).

Fixpoint digits_to_Z (l : list Z) : Z :=
  match l with
	| nil => 0
	| a :: tl => a + 10 * digits_to_Z tl
	end.

Fixpoint Znth {A : Type}(l : list A) (n : Z) : option A :=
  match l with
  | nil => None
  | a :: tl =>
    if n =? 0 then Some a else Znth tl (n - 1)
  end.

Fixpoint Zlist_min (l : list Z) : option Z :=
match l with
| nil => None
| a :: tl =>
  match Zlist_min tl with
  | None => Some a
  | Some b => if a <=? b then Some a else Some b
  end
end.

Module Private.

Inductive Znat_sandbox : Z -> Prop :=
  Zn0 : Znat_sandbox 0
| ZnS : forall n, Znat_sandbox n -> Znat_sandbox (n + 1).

Lemma Znat_sandbox_ge0 x : Znat_sandbox x -> 0 <= x.
Proof.
induction 1 as [ | x _ Ih]; lia.
Qed.

Lemma ge0_Znat_sandbox x : 0 <= x -> Znat_sandbox x.
Proof.
induction x as [ x Ih] using (well_founded_induction (Zwf_well_founded 0)).
intros xge0.
destruct (x =? 0) eqn:cmp0.
  replace x with 0 by lia; constructor.
replace x with ((x - 1) + 1) by ring.
apply ZnS.
apply Ih; unfold Zwf; lia.
Qed.

End Private.

Ltac Znat_induction h :=
  match type of h with
  | (0 <= ?x) =>
     let ind_hyp_name := fresh "Ih" in
     generalize (Private.ge0_Znat_sandbox x h); clear h;
     induction 1 as [ | ? h ind_hyp_name];
     [ | apply Private.Znat_sandbox_ge0 in h]
  end.

Lemma Znat_ind :
  forall P : Z -> Prop,
  P 0 ->
  (forall n, 0 <= n -> P n -> P (n + 1)) ->
  forall n, 0 <= n -> P n.
Proof.
intros P P0 Ps n nge0.
Znat_induction nge0.
  easy.
auto.
Qed.

Lemma Z_iter_unroll_left {A : Type} n (f : A -> A) (x : A) :
  0 <= n - 1 -> Z.iter n f x = f (Z.iter (n - 1) f x).
Proof.
intros ngt0.
assert (main : Z.abs_nat n = S (Z.abs_nat (n - 1))) by lia.
rewrite !iter_nat_of_Z; try lia.
now rewrite main.
Qed.

Lemma Z_iter_unroll_right {A : Type} n (f : A -> A) (x : A) :
  0 <= n - 1 -> Z.iter n f x = Z.iter (n - 1) f (f x).
Proof.
intros ngt0.
assert (main : Z.abs_nat n = S (Z.abs_nat (n - 1))) by lia.
rewrite !iter_nat_of_Z; try lia.
rewrite main.
clear main ngt0.
revert x.
induction (Z.abs_nat (n - 1)).
  easy.
intros x; simpl.
now rewrite <- IHn0.
Qed.

Lemma Z_iter_add {A : Type} (start : A) f x y : 0 <= x -> 0 <= y ->
  Z.iter (x + y) f start = Z.iter x f (Z.iter y f start).
Proof.
intros xge0 yge0.
Znat_induction xge0.
  easy.
rewrite (Z_iter_unroll_left (n + 1 + y));[ | lia].
rewrite (Z_iter_unroll_left (n + 1));[ | lia].
replace (n + 1 + y - 1) with (n + y) by lia.
replace (n + 1 - 1) with n by lia.
rewrite Ih.
easy.
Qed.

Lemma Zseq0 : Zseq 0 = nil.
Proof. easy. Qed.

(* Discovering here that Znat_ind is not recognized as a good induction
  principle by the induction tactic.  *)
Lemma Z_iter_loop_index {A : Type} (f : Z -> A -> A) (x : Z) (a0 : A) n :
  0 <= n -> fst (Z.iter n (fun '(y, v) => (y + 1, f y v)) (x, a0)) = n + x.
Proof.
intros nge0.
Znat_induction nge0.
  easy.
rewrite Z_iter_unroll_left;[ | lia].
replace (n + 1 - 1) with n by ring.
destruct Z.iter as [y v] eqn:iterq.
unfold fst in Ih.
rewrite Ih.
simpl.
ring.
Qed.

Lemma Zseq_unroll n : 0 <= n - 1 -> Zseq n = Zseq (n - 1) ++ (n - 1):: nil.
Proof.
unfold Zseq.
intros ngt0.
rewrite Z_iter_unroll_left;[ | lia].
destruct (n <=? 0) eqn:cmpn.
  lia.
destruct (n - 1 <=? 0) eqn:cmpp.
  replace (n - 1) with 0 by lia.
  easy.
simpl.
destruct Z.iter as [a b] eqn:iterq.
assert (aq : a = n - 1).
  replace a with (fst (Z.iter (n - 1)
    (fun '(x, l) => (x + 1, app_Z l (x :: nil))) (0, nil))).
    rewrite Z_iter_loop_index.
      ring.
    lia.
  now rewrite iterq.
now unfold app_Z; rewrite <- aq.
Qed.

Lemma Zseq_split m n : 0 <= m <= n ->
  Zseq n = Zseq m ++ map (fun x=> x + m) (Zseq (n - m)).
Proof.
intros [mge0 ngem].
assert (nge0 : 0 <= n) by lia.
revert m mge0 ngem.
Znat_induction nge0; intros m mge0 mlen.
  replace m with 0 by lia.
  easy.
destruct (m =? n + 1) eqn:cmpmnp1.
  replace m with (n + 1) by lia.
  replace (n + 1 - (n + 1)) with 0 by ring.
  (* Here it is very unpleasant that nil needs to be qualified. *)
  replace (map (fun x => x + (n + 1) - (n + 1)) (Zseq 0)) with
    (@nil Z) by easy.
  now rewrite app_nil_r.
replace (n + 1 - m) with ((n - m) + 1) by ring.
rewrite (Zseq_unroll (n + 1));[ | lia].
rewrite (Zseq_unroll (n - m + 1));[ | lia].
replace (n + 1 - 1) with n by ring.
replace (n - m + 1 - 1) with (n - m) by ring.
rewrite map_last.
replace (n - m + m) with n by ring.
rewrite app_assoc.
rewrite <- Ih.
    easy.
  easy.
lia.
Qed.

Fixpoint natural_bounded_search {A B : Type} (n : nat)
  (f : B -> A + B)
  (start : B) : A + B :=
  match n with
  | O => inr start
  | S m =>
    match f start with
      inl v => inl v
    | inr t => natural_bounded_search m f t
    end
  end.

Fixpoint positive_bounded_search {A B : Type} (p : positive)
  (f : B -> A + B)
  (start : B) : A + B :=
  match p with
  | xH => f start
  | xO q =>
    match positive_bounded_search q f start with
    | inl v => inl v
    | inr tmp => positive_bounded_search q f tmp
    end
  | xI q =>
    match positive_bounded_search q f start with
    | inl v => inl v
    | inr t1 =>
      match positive_bounded_search q f t1 with
      | inl v => inl v
      | inr t2 => f t2
      end
    end
  end.

Lemma natural_bounded_search_success_step {A B : Type} (n : nat)
  (f : B -> A + B) (start : B) r :
  natural_bounded_search n f start = inl r ->
  natural_bounded_search (S n) f start = inl r.
Proof.
revert start; induction n as [ | n Ih].
  discriminate.
simpl; intros start; destruct (f start) as [v | t].
  easy.
intros nprop; apply Ih in nprop; exact nprop.
Qed.

Lemma natural_bounded_search_success_monotone{A B : Type} (n : nat)
  (f : B -> A + B) (start : B) r :
  natural_bounded_search n f start = inl r ->
  forall m, (n <= m)%nat ->
  natural_bounded_search m f start = natural_bounded_search n f start.
Proof.
intros nprop m nlem.
assert (mq : (m = n + (m - n))%nat) by lia.
rewrite mq; clear mq nlem.
remember (m - n)%nat as m' eqn:Heqm'.
clear m Heqm'.
revert n start nprop.
induction m' as [ | m' Ih].
  intros n start nprop; rewrite Nat.add_0_r; easy.
intros n start nprop.
rewrite nprop.
rewrite Nat.add_succ_r, natural_bounded_search_success_step with (r := r).
  easy.
rewrite <- nprop.
now apply Ih.
Qed.

Lemma natural_bounded_search_progress_break_step {A B : Type} (n : nat)
  (f : B -> A + B) start res :
  natural_bounded_search (S n) f start = inr res ->
  exists v, natural_bounded_search n f start = inr v /\
    f v = inr res.
Proof.
revert start; induction n as [ | n Ih]; intros start.
  simpl.
  destruct (f start) as [ v | v1] eqn:fq ;[discriminate | ].
  intros vals; injection vals as v1q.
  now exists start; split;[ | rewrite fq, v1q].
simpl; destruct (f start) as [ v | v1] eqn:fq; [discriminate | ].
intros nprop; apply Ih; exact nprop.
Qed.

Lemma natural_bounded_search_progress_break {A B : Type} (n : nat)
  (f : B -> A + B) start res :
  natural_bounded_search n f start = inr res ->
  forall m, (m < n)%nat ->
  exists v1, natural_bounded_search m f start = inr v1 /\
    natural_bounded_search (n - m) f v1 = inr res.
Proof.
revert start res; induction n as [ | n Ih].
  intros start res _ m mlt0; lia.
intros start res nprop.
assert (nprop2 := nprop).
apply natural_bounded_search_progress_break_step in nprop2.
destruct (f start) as [ v | v] eqn:fstartq.
  simpl in nprop.
  rewrite fstartq in nprop; discriminate.
intros [ | m] SmltSn.
  exists start.
  generalize nprop; simpl; rewrite fstartq.
  now auto.
assert (m < n)%nat by lia.
generalize nprop.
now simpl; rewrite fstartq; intros nprop3; apply Ih.
Qed.

Lemma natural_bounded_search_progress {A B : Type} (n m : nat)
  (f : B -> A + B) start res :
  natural_bounded_search n f start = inr res ->
  natural_bounded_search (n + m) f start =
  natural_bounded_search m f res.
Proof.
revert m start.
induction n as [ | n Ih].
  now simpl; intros m start [= startq]; rewrite startq.
intros m start; simpl.
destruct (f start) as [v | v] eqn:fstartq.
  discriminate.
now intros nprop; apply (Ih m) in nprop.
Qed.

Lemma natural_bounded_search_correct {A B : Type} (n : nat)
  (f : B -> B) (p : B -> bool) (g : B -> A) (h : B -> A + B)
  start res :
  (forall x, h x = if p x then inl (g x) else inr (f x)) ->
  natural_bounded_search n h start = inl res ->
  exists k, (k < n)%nat /\
    p (Nat.iter k f start) = true /\
    (forall k', (k' < k)%nat -> p (Nat.iter k' f start) = false) /\
    g (Nat.iter k f start) = res.
Proof.
intros heq; revert start; induction n as [ | n Ih].
  discriminate.
simpl; intros start.
destruct (h start) as [v | v] eqn:hstartq.
  intros [= vq].
  generalize hstartq; rewrite heq.
  destruct (p start) eqn:pstartq;[ | discriminate].
  rewrite vq; intros [=gq].
  exists 0%nat.
  split; [lia | ].
  split; [exact pstartq | ].
  split; [ | exact gq].
  intros k klt0; lia.
generalize hstartq.
rewrite heq.
destruct (p start) eqn:pstartq;[discriminate | ].
intros [= fstartq] nprop.
assert (nprop2 := nprop).
apply Ih in nprop2.
destruct nprop2 as [k [kltn [pktrue [pk'false gq]]]].
exists (S k).
split;[lia | ].
split.
  now rewrite Nat.iter_succ_r, fstartq, pktrue.
split.
  intros [ | k'] Sk'ltSk.
    exact pstartq.
  rewrite Nat.iter_succ_r, fstartq; apply pk'false; lia.
now rewrite Nat.iter_succ_r, fstartq.
Qed.

Lemma bounded_search_nat_pos {A B : Type} (n : positive)
  (f : B -> A + B) (start : B) :
  positive_bounded_search n f start =
  natural_bounded_search (Pos.to_nat n) f start.
Proof.
revert start.
induction n as [n Ih | n Ih | ].
- intros start; simpl (positive_bounded_search _  _ _).
  assert (vn1 : Pos.to_nat (xI n) = (2 * Pos.to_nat n + 1)%nat).
    rewrite Pos2Nat.inj_xI; lia.
  destruct (positive_bounded_search n f start) as [a | b] eqn:cal1.
    assert (caln1 := cal1).
    rewrite Ih in caln1.
    assert (cmpn2n1 : (Pos.to_nat n <= Pos.to_nat n~1)%nat) by lia.
    assert (tmp := natural_bounded_search_success_monotone (Pos.to_nat n) f
         start a caln1 (Pos.to_nat (xI n)) cmpn2n1).
    now rewrite tmp.
  assert (caln1 := cal1).
  rewrite Ih in caln1.
  assert (tmp:= natural_bounded_search_progress (Pos.to_nat n) (Pos.to_nat n)
              f start b caln1).
  destruct (positive_bounded_search n f b) as [a2 | b2] eqn:cal2;
    assert (caln2 := cal2); rewrite Ih in caln2; rewrite caln2 in tmp.
    assert (cmp2n2n1 :
        (Pos.to_nat n + Pos.to_nat n <= Pos.to_nat (xI n))%nat) by lia.
    assert (tmp2:= natural_bounded_search_success_monotone
        (Pos.to_nat n + Pos.to_nat n) f start a2 tmp _ cmp2n2n1).
    now rewrite tmp2.
  assert (tmp2 := natural_bounded_search_progress (Pos.to_nat n + Pos.to_nat n)
    1 f start b2 tmp).
    replace (Pos.to_nat (xI n)) with (Pos.to_nat n + Pos.to_nat n + 1)%nat.
      now rewrite tmp2; simpl; case (f b2).
    rewrite Pos2Nat.inj_xI; lia.
- intros start; simpl (positive_bounded_search _  _ _).
  assert (vn2 : Pos.to_nat (xO n) = (2 * Pos.to_nat n)%nat).
    rewrite Pos2Nat.inj_xO; lia.
  destruct (positive_bounded_search n f start) as [a | b] eqn:cal1.
    assert (caln1 := cal1).
    rewrite Ih in caln1.
    assert (cmpn2n: (Pos.to_nat n <= Pos.to_nat (xO n))%nat) by lia.
    assert (tmp := natural_bounded_search_success_monotone (Pos.to_nat n) f
         start a caln1 (Pos.to_nat (xO n)) cmpn2n).
    now rewrite tmp.
  assert (caln1 := cal1).
  rewrite Ih in caln1.
  assert (tmp:= natural_bounded_search_progress (Pos.to_nat n) (Pos.to_nat n)
              f start b caln1).
  now rewrite Ih, vn2; simpl; rewrite Nat.add_0_r.
- now intros n; simpl; destruct (f n).
Qed.

Definition Z_loop_until {A : Type} (x : Z) (f : A -> A) (test : A -> bool)
  (start : A) : A + A :=
  match x with
  | Zpos p => positive_bounded_search p
      (fun t => if test t then inl t else inr (f t)) start
  | _ => inr start
  end.

Lemma Z_loop_until_induction {A : Type} (x : Z) (f : A -> A) (test : A-> bool)
  (start : A) (P : A -> Prop) (v : A) :
  0 <= x ->
  P start ->
  (forall y, P y -> test y = false -> P (f y)) ->
  Z_loop_until x f test start = inr v -> P v.
Proof.
intros xge0 P0 inv; unfold Z_loop_until.
destruct x as [ | x | x].
  1,3: intros [= startv]; rewrite <- startv; easy.
rewrite bounded_search_nat_pos.
revert v; induction (Pos.to_nat x) as [ | n IHn]; clear x xge0; intros v.
  simpl; intros [= startv]; rewrite <- startv; easy.
intros atSn.
apply natural_bounded_search_progress_break_step in atSn.
destruct atSn as [v0 [atn lastcomp]].
destruct (test v0) eqn:testp;[discriminate | ].
injection lastcomp; intros lastcomp'; rewrite <- lastcomp'.
apply IHn in atn.
now apply inv.
Qed.

Lemma Z_iter_to_Z_loop {A : Type} (x : Z) (f : A -> A) (test : A -> bool)
  (start : A) :
  0 <= x ->
  (forall k, 0 <= k < x -> test (Z.iter k f start) = false) ->
  Z_loop_until x f test start = inr (Z.iter x f start).
Proof.
intros xge0.
unfold Z_loop_until.
destruct x as [ | x | x];[ easy | | easy].
rewrite iter_nat_of_Z, bounded_search_nat_pos;[ | lia].
fold (Nat.iter (Z.abs_nat (Z.pos x)) f start).
intros no_test.
assert 
 (no_test' : forall k, (0 <= k < Pos.to_nat x)%nat -> test (Nat.iter k f start) = false).
  intros k kbounds.
  assert (kboundsz : 0 <= Z.of_nat k < Z.pos x) by lia.
  generalize (no_test _ kboundsz).
  rewrite iter_nat_of_Z, Zabs2Nat.id; easy.
clear no_test; revert start no_test'.
rewrite Zabs2Nat.inj_pos.
clear xge0.
induction (Pos.to_nat x) as [ | n IHn].
  easy.
intros start no_test; rewrite Nat.iter_succ_r; simpl.
destruct (test start) eqn:testp.
  assert (abs := no_test 0%nat); simpl in abs.
  rewrite abs in testp;[discriminate | lia].
change (f (Nat.iter n f start)) with (Nat.iter (S n) f start).
apply IHn.
intros k kbounds.
rewrite <- Nat.iter_succ_r.
apply no_test; lia.
Qed.

Definition Z_bounded_search {A B : Type} (x : Z)
  (f : B -> A + B) (start : B) :=
  match x with
  | Z0 => inr start
  | Zpos p => positive_bounded_search p f start
  | Zneg p => inr start
  end.

Lemma Z_loop_until_correct {A : Type} (n : Z) (f : A -> A) (p : A -> bool)
  start res:
  Z_loop_until n f p start = inl res ->
  exists k, 0 <= k < n /\ p (Z.iter k f start) = true /\
    (forall k', 0 <= k' < k -> p (Z.iter k' f start) = false) /\
    res = (Z.iter k f start).
Proof.
destruct n as [ | n | n]; [discriminate | | discriminate].
unfold Z_loop_until.
rewrite bounded_search_nat_pos; intros execq.
destruct (natural_bounded_search_correct (Pos.to_nat n) f p (fun x => x)
    _
        start res (fun x => eq_refl) execq) as 
          [k [kbound [pktrue [pk'false gq]]]].
exists (Z.of_nat k).
split;[lia | ].
split.
  rewrite iter_nat_of_Z;[ | lia].
  rewrite Zabs2Nat.id; exact pktrue.
split.
  intros k' k'bounds; rewrite iter_nat_of_Z;[ | lia].
  refine (pk'false _ _); lia.
rewrite iter_nat_of_Z;[| lia].
now rewrite Zabs2Nat.id.
Qed.

Lemma Z_loop_until_not_found {A : Type} (n : Z) (f : A -> A) (p : A -> bool)
  start res:
  Z_loop_until n f p start = inr res ->
  res = Z.iter n f start /\
  forall k, 0 <= k < n -> p (Z.iter k f start) = false.
Proof.
destruct n as [ | n | n]; simpl.
1,3: intros execq; split.
1,3: now injection execq.
1,2: intros k kbounds; lia.
unfold Z_bounded_search; rewrite bounded_search_nat_pos.
intros execq.
enough (main_nat : res = Nat.iter (Pos.to_nat n) f start /\
  forall k : nat, (k < (Pos.to_nat n))%nat ->
  p (Nat.iter k f start) = false).
  split.
    change (Pos.iter f start n) with (Z.iter (Z.pos n) f start).
    rewrite iter_nat_of_Z;[ | lia].
    change (res = Nat.iter (Z.abs_nat (Z.pos n)) f start).
    destruct main_nat as [it _]; exact it.
  intros k kbounds; rewrite iter_nat_of_Z;[ | lia].
  change (p (Nat.iter (Z.abs_nat k) f start) = false).
  apply main_nat.
  lia.
revert start execq; induction (Pos.to_nat n) as [ | n' Ih].
  intros start execq; split.
    now injection execq.
  intros k kbound; lia.
intros start; rewrite Nat.iter_succ_r; simpl.
destruct (p start) eqn:pstartq; [discriminate | ].
intros propn; assert (Ih' := Ih (f start) propn).
split;[tauto | ].
intros [ | k'] k'bound.
  exact pstartq.
rewrite Nat.iter_succ_r; apply Ih'; lia.
Qed.

(* Example usage of Z_loop_until: computing log_10, given input x, compute
  the smallest number k such that x < 10 ^ (k + 1) *)

Definition log_10 (x : Z) :=
  if x <=? 0 then 0 else 
    match Z_loop_until x (fun '(t, y) => (t + 1, y / 10))
      (fun '(_, y) => y <=? 9) (0, x) with
    | inl i => fst i
    | inr _ => 0
    end.

Lemma iter_incr x y : 0 <= y -> Z.iter y (fun x => x + 1) x = x + y.
Proof.
intros yge0; Znat_induction yge0.
  simpl; intros; ring.
rewrite Z_iter_unroll_left;[ | lia].
replace (n + 1 - 1) with n by ring.
rewrite Ih.
ring.
Qed.

Lemma iter_incr_context {A : Type} s0 y (f : A -> A) (pi : A -> Z) :
  (forall s : A, pi (f s) = pi s + 1) ->
  0 <= y -> pi (Z.iter y f s0) = pi s0 + y.
Proof.
intros fq yge0.
Znat_induction yge0.
  simpl; ring.
rewrite Z_iter_unroll_left; [ | lia].
rewrite fq.
replace (n + 1 - 1) with n by ring.
rewrite Ih.
ring.
Qed.

Lemma log_10_correct (x : Z) : 
  1 <= x -> 10 ^ log_10 x <= x < 10 ^ (log_10 x + 1).
Proof.
assert (tenn0 : 10 <> 0) by lia.
intros xge1.
unfold log_10.
destruct (x <=? 0) eqn:cmp0.
  lia.
clear cmp0.
assert (fst_body : forall s,
  fst (let '(u, y) := s in (u + 1, y / 10)) = fst s + 1).
  intros; destruct s as [u y].
  easy.
set (main_prop0 := fun '(t, z) => 0 <= t /\ z = x / 10 ^ t).
assert (main0 : forall i, 0 <= i ->
  main_prop0 (Z.iter i (fun '(t, y) =>(t + 1, y / 10)) (0, x))).
  intros i ige0.
  Znat_induction ige0.
    simpl.
    lia.
  rewrite Z_iter_unroll_left;[ | lia].
  replace (n + 1 - 1) with n by lia.
  unfold main_prop0 in Ih |- *.
  destruct (Z.iter n _ _) as [u v] eqn:iterq.
  destruct Ih as [ucond vcond].
  split; [lia | ].
  rewrite vcond.
  rewrite Z.div_div; try lia.
  replace (10 ^ u * 10) with (10 ^ (u + 1));[easy | ].
  rewrite Z.pow_add_r; lia.
set (main_prop := fun '(t, z) => 0 <= t /\ z = x / 10 ^  t /\
            (t <> 0 -> 10 <= x / 10 ^ (t - 1))).
assert (main:
      forall i j, 0 <= i < j ->
      (forall k, 0 <= k < i -> (let '(_, r) := Z.iter k 
         (fun '(t, y) => (t + 1, y / 10)) (0, x) in
        r <=? 9) = false) ->
      main_prop (Z.iter i (fun '(t, y) => (t + 1, y / 10)) (0, x))).
  intros i j iltj nfoundj.
  assert (tmp := Z_iter_to_Z_loop _ _ (fun '(_, y) => y <=? 9) _ 
            (proj1 iltj) nfoundj).
  refine (Z_loop_until_induction _ _
    (fun '(_, z) => z <=? 9) (0, x)
    main_prop _ (proj1 iltj) _ _ tmp).
    split.
      lia.
    split.
      lia.
    lia.
  intros (t, v) [vge0 [vxdiv prevv]] pre_vge10.
  split.
    lia.
  split.
    rewrite vxdiv, Z.div_div;[ | lia | lia].
    replace (10 ^ t * 10) with (10 ^ (t + 1)).
      easy.
    rewrite Z.pow_add_r, Z.pow_1_r; easy.
  intros _.
  replace (t + 1 - 1) with t by ring.
  lia.
destruct Z_loop_until eqn:loopq.
  destruct (Z_loop_until_correct _ _ _ _ _ loopq) as
    [i [iltx [itrue [ileast rq]]]].
    assert (main' : main_prop 
              (Z.iter i (fun '(t, y) => (t + 1, y / 10)) (0, x))).
      apply (main i x); easy.
  destruct (i =? 0) eqn:cmpi0.
    assert (i0 : i = 0) by lia.
    rewrite i0 in itrue, rq.
    simpl in itrue, rq.
    rewrite rq; simpl.
    replace (Z.pow_pos 10 1) with 10 by easy.
    lia.
  rewrite <- rq in main'.
  unfold main_prop in main'.
  destruct p as [t z]; simpl.
  rewrite <- rq in itrue.
  destruct main' as [tge0 [zq prezP]].
  assert (xval : x = z * (10 ^ t) + x mod (10 ^ t)).
    rewrite zq, Z.mul_comm, <-Z.div_mod; lia.
  assert (ti : t = i).
    generalize (iter_incr_context (A := Z * Z) (0, x) i
              (fun '(t, y) => (t + 1, y / 10)) fst fst_body (proj1 iltx)).
    rewrite <- rq.
    easy.
  set (z' := x / 10 ^ (t - 1)).
  fold z' in prezP.
  assert (zz' : z = z' / 10).
    unfold z'; rewrite Z.div_div; try lia.
    replace 10 with (10 ^ 1) at 2 by ring.
    rewrite <- Z.pow_add_r; try lia.
    replace (t - 1 + 1) with t by ring.
    easy.
  assert (z'ge0 : 10 <= z') by lia.
  assert (zge1 : 1 <= z) by lia.
  assert (bigmod_bounds: 0 <= x mod 10 ^ t < 10 ^ t) by lia.
  assert (10 ^ t <= x) by nia.
  replace (10 ^ (t + 1)) with (10 ^ t * 10).
    nia.
  rewrite Z.pow_add_r; nia.
apply Z_loop_until_not_found in loopq.
destruct loopq as [pq not_found].
assert (xm1bounds : 0 <= x - 1 < x) by lia.
assert (not_found' := not_found _ xm1bounds).
assert (xbounds : 0 <= x < x + 1) by lia.
assert (main' := main0 _ (proj1 xm1bounds)).
destruct (Z.iter (x - 1) _ _) as [u v] eqn:iterq.
unfold main_prop0 in main'.
assert (uq : u = x - 1).
  generalize (iter_incr_context (A := Z * Z) (0, x) (x - 1)
            (fun '(t, y) => (t + 1, y / 10)) fst fst_body (proj1 xm1bounds)).
  rewrite iterq.
  simpl.
  easy.
assert (abs : 1 <= x / 10 ^ x).
  replace (10 ^ x) with (10 ^ (x - 1) * 10).
    rewrite <- Z.div_div; try lia.
    rewrite <-uq, <- (proj2 main').
    assert (tmp := Z.div_mod v 10).
    assert (tmp2 := Z.mod_pos_bound v 10).
    lia.
  replace 10 with (10 ^ 1) at 2 by ring.
  rewrite <- Z.pow_add_r; try lia.
  replace (x - 1 + 1) with x by ring.
  easy.
assert (tmp3 := Z.pow_gt_lin_r 10 x).
assert (xlttenpx : x < 10 ^ x) by lia.
assert (tmp2 := Z.div_mod x (10 ^ x)).
assert (tmp := Z.mod_pos_bound x (10 ^ x)).
assert (tenpxgt0 : 0 < 10 ^ x).
  apply Z.pow_pos_nonneg; lia.
case (Zle_not_lt _ _ abs).
apply (Zmult_lt_reg_r _ _ (10 ^ x)).
  lia.
lia.
Qed.

Lemma Z_bounded_search_correct_aux {A B : Type} (n : Z)
  (f : B -> B) (p : B -> bool) (g : B -> A) (h : B -> A + B)
  start res :
  (forall x, h x = if p x then inl (g x) else inr (f x)) ->
  Z_bounded_search n h start = inl res ->
  exists k, 0 <= k < n /\
    p (Z.iter k f start) = true /\
    (forall k', 0 <= k' < k -> p (Z.iter k' f start) = false) /\
    g (Z.iter k f start) = res.
Proof.
intros heq.
destruct n as [ | n | n]; [discriminate | | discriminate].
unfold Z_bounded_search.
rewrite bounded_search_nat_pos; intros execq.
destruct (natural_bounded_search_correct (Pos.to_nat n) f p g h
        start res heq execq) as [k [kbound [pktrue [pk'false gq]]]].
exists (Z.of_nat k).
split;[lia | ].
split.
  rewrite iter_nat_of_Z;[ | lia].
  rewrite Zabs2Nat.id; exact pktrue.
split.
  intros k' k'bounds.
  rewrite iter_nat_of_Z;[ | lia].
  fold (Nat.iter (Z.abs_nat k') f start).
  apply pk'false.
  lia.
rewrite iter_nat_of_Z;[ | lia].
rewrite Zabs2Nat.id; exact gq.
Qed.

Lemma Z_bounded_search_correct {A B : Type} (n : Z)
  (f : B -> B) (p : B -> bool) (g : B -> A) start res:
  Z_bounded_search n (fun x => if p x then inl (g x) else inr (f x)) start
    = inl res ->
  exists k, 0 <= k < n /\
    p (Z.iter k f start) = true /\
    (forall k', 0 <= k' < k -> p (Z.iter k' f start) = false) /\
    g (Z.iter k f start) = res.
Proof.
exact (Z_bounded_search_correct_aux n f p g _ start res (fun x => eq_refl)).
Qed.

Lemma Z_bounded_search_not_found_aux {A B : Type} (n : Z)
  (f : B -> B) (p : B -> bool) (g : B -> A) (h : B -> A + B)
  start res : 
  (forall x, h x = if p x then inl (g x) else inr (f x)) ->
  Z_bounded_search n h start = inr res ->
  res = Z.iter n f start /\
  forall k, 0 <= k < n -> p (Z.iter k f start) = false.
Proof.
intros heq.
destruct n as [ | n | n]; simpl.
1,3: intros execq; split.
1,3: now injection execq.
1,2: intros k kbounds; lia.
unfold Z_bounded_search; rewrite bounded_search_nat_pos.
intros execq.
enough (main_nat : res = Nat.iter (Pos.to_nat n) f start /\
  forall k : nat, (k < (Pos.to_nat n))%nat ->
  p (Nat.iter k f start) = false).
  split.
    change (Pos.iter f start n) with (Z.iter (Z.pos n) f start).
    rewrite iter_nat_of_Z;[ | lia].
    change (res = Nat.iter (Z.abs_nat (Z.pos n)) f start).
    destruct main_nat as [it _]; exact it.
  intros k kbounds; rewrite iter_nat_of_Z;[ | lia].
  change (p (Nat.iter (Z.abs_nat k) f start) = false).
  apply main_nat.
  lia.
revert start execq; induction (Pos.to_nat n) as [ | n' Ih].
  intros start execq; split.
    now injection execq.
  intros k kbound; lia.
intros start; rewrite Nat.iter_succ_r; simpl.
rewrite heq.
destruct (p start) eqn:pstartq; [discriminate | ].
intros propn; assert (Ih' := Ih (f start) propn).
split;[tauto | ].
intros [ | k'] k'bound.
  exact pstartq.
rewrite Nat.iter_succ_r; apply Ih'; lia.
Qed.

Lemma Z_bounded_search_not_found {A B : Type} (n : Z)
  (f : B -> B) (p : B -> bool) (g : B -> A) start res :
  Z_bounded_search n (fun x => if p x then inl (g x) else inr (f x)) start =
    inr res ->
  res = Z.iter n f start /\
  forall k, 0 <= k < n -> p (Z.iter k f start) = false.
Proof.
exact (Z_bounded_search_not_found_aux n f p g _ start res (fun x => eq_refl)).
Qed.

Definition newton_sequence_sqrt_pre_approx (denominator target : Z) :=
  Z_bounded_search denominator
    (fun t : Z =>
      if ((t) ^ 2 <=?  target) &&
        (target <? (t + 1) ^ 2) then
        inl t
      else
        inr ((t + target / t) / 2))
    denominator.

Definition newton_sqrt (target : Z) :=
  let l := Z.log2 target in
  newton_sequence_sqrt_pre_approx (2 ^ (l / 2)) target.

Definition newton_sqrt_plain (x : Z) :=
  match Z_bounded_search (Z.log2_up x + 1)
    (fun t => let next := (t + x / t) / 2 in
       if (t =? next) then
         inl t else inr next)
    (Z.iter (Z.log2_up (Z.log2_up x) / 2 + 1)
      (fun t => (t + x / t) / 2) (2 ^ (Z.log2_up x / 2))) with
  inl l => l
  | inr _ => Z.sqrt x
  end.
Time Compute (_ (newton_sqrt_plain (2 * 10 ^  1000))).
Time Compute (Z.log2 (Z.sqrt (10 ^ 1000))).


Lemma newton_sequence_sqrt_partial_correct v r :
  newton_sqrt v = inl r -> r ^ 2 <= v < (r + 1) ^ 2.
Proof.
intros execq.
destruct (Z_bounded_search_correct (2 ^ (Z.log2 v / 2))
         (fun t => (t + v / t) / 2)
         (fun t => (t ^ 2 <=? v) && (v <? (t + 1) ^ 2))
         (fun t => t) _ _ execq) as [k [kbounds [pktrue [pk'false rq]]]].
rewrite rq in pktrue; lia.
Qed.

Lemma Z_iter_invariant {A : Type} (P : A -> Prop) (a0 : A) (f : A -> A)
  (P0 : P a0) (Pf : forall x : A, P x -> P (f x)) : 
  forall z : Z, P (Z.iter z f a0).
Proof.
intros z.
destruct (z <? 0) eqn:cmpz0.
  destruct z as [ | | ]; easy.
assert (zge0 : 0 <= z) by lia.
clear cmpz0.
Znat_induction zge0.
  easy.
rewrite Z_iter_unroll_left;[ | lia].
apply Pf.
replace (n + 1 - 1) with n by ring.
exact Ih.
Qed.

Lemma Z_iter_counted_invariant {A : Type} (P : Z -> A -> Prop) (a0 : A) (f : A -> A)
  (P0 : P 0 a0) (Pf : forall (n : Z) (a : A), P n a -> P (n + 1) (f a)) :
  forall z,  0 <= z -> P z (Z.iter z f a0).
Proof.
intros z zge0; Znat_induction zge0.
  easy.
rewrite Z_iter_unroll_left;[ | lia].
apply Pf.
replace (n + 1 - 1) with n by ring.
exact Ih.
Qed.