More countable stuff.

Also, use a different encoding of lists.

prelude/countable.v

@@ -9,6 +9,8 @@ Class Countable A `{∀ x y : A, Decision (x = y)} := {
   decode : positive → option A;
   decode_encode x : decode (encode x) = Some x
 }.
+Arguments encode : simpl never.
+Arguments decode : simpl never.
 
 Definition encode_nat `{Countable A} (x : A) : nat :=
   pred (Pos.to_nat (encode x)).
@@ -19,6 +21,8 @@ Proof.
   intros x y Hxy; apply (injective Some).
   by rewrite <-(decode_encode x), Hxy, decode_encode.
 Qed.
+Instance encode_nat_injective `{Countable A} : Injective (=) (=) encode_nat.
+Proof. unfold encode_nat; intros x y Hxy; apply (injective encode); lia. Qed.
 Lemma decode_encode_nat `{Countable A} x : decode_nat (encode_nat x) = Some x.
 Proof.
   pose proof (Pos2Nat.is_pos (encode x)).
@@ -26,6 +30,7 @@ Proof.
   by rewrite Pos2Nat.id, decode_encode.
 Qed.
 
+(** * Choice principles *)
 Section choice.
   Context `{Countable A} (P : A → Prop) `{∀ x, Decision (P x)}.
 
@@ -33,7 +38,6 @@ Section choice.
     | choose_step_None {p} : decode p = None → choose_step (Psucc p) p
     | choose_step_Some {p x} :
        decode p = Some x → ¬P x → choose_step (Psucc p) p.
-
   Lemma choose_step_acc : (∃ x, P x) → Acc choose_step 1%positive.
   Proof.
     intros [x Hx]. cut (∀ i p,
@@ -46,13 +50,11 @@ Section choice.
     constructor. intros j.
     inversion 1 as [? Hd|? y Hd]; subst; auto with lia.
   Qed.
-
   Fixpoint choose_go {i} (acc : Acc choose_step i) : A :=
     match Some_dec (decode i) with
     | inleft (x↾Hx) =>
       match decide (P x) with
-      | left _ => x
-      | right H => choose_go (Acc_inv acc (choose_step_Some Hx H))
+      | left _ => x | right H => choose_go (Acc_inv acc (choose_step_Some Hx H))
       end
     | inright H => choose_go (Acc_inv acc (choose_step_None H))
     end.
@@ -76,18 +78,18 @@ Proof.
   intros y. by rewrite (choose_correct (λ x, f x = y) (surjective f y)).
 Qed.
 
-(** ** Instances *)
+(** * Instances *)
+(** ** Option *)
 Program Instance option_countable `{Countable A} : Countable (option A) := {|
-  encode o :=
-    match o with None => 1 | Some x => Pos.succ (encode x) end;
-  decode p :=
-    if decide (p = 1) then Some None else Some <$> decode (Pos.pred p)
+  encode o := match o with None => 1 | Some x => Pos.succ (encode x) end;
+  decode p := if decide (p = 1) then Some None else Some <$> decode (Pos.pred p)
 |}.
 Next Obligation.
   intros ??? [x|]; simpl; repeat case_decide; auto with lia.
   by rewrite Pos.pred_succ, decode_encode.
 Qed.
 
+(** ** Sums *)
 Program Instance sum_countable `{Countable A} `{Countable B} :
   Countable (A + B)%type := {|
     encode xy :=
@@ -99,6 +101,7 @@ Program Instance sum_countable `{Countable A} `{Countable B} :
   |}.
 Next Obligation. by intros ?????? [x|y]; simpl; rewrite decode_encode. Qed.
 
+(** ** Products *)
 Fixpoint prod_encode_fst (p : positive) : positive :=
   match p with
   | 1 => 1
@@ -162,75 +165,82 @@ Proof.
 Qed.
 Program Instance prod_countable `{Countable A} `{Countable B} :
   Countable (A * B)%type := {|
-    encode xy := let (x,y) := xy in prod_encode (encode x) (encode y);
+    encode xy := prod_encode (encode (xy.1)) (encode (xy.2));
     decode p :=
      x ← prod_decode_fst p ≫= decode;
      y ← prod_decode_snd p ≫= decode; Some (x, y)
   |}.
 Next Obligation.
   intros ?????? [x y]; simpl.
-  rewrite prod_decode_encode_fst, prod_decode_encode_snd.
-  csimpl. by rewrite !decode_encode.
+  rewrite prod_decode_encode_fst, prod_decode_encode_snd; simpl.
+  by rewrite !decode_encode.
 Qed.
 
-Fixpoint list_encode_ (l : list positive) : positive :=
-  match l with [] => 1 | x :: l => prod_encode x (list_encode_ l) end.
-Definition list_encode (l : list positive) : positive :=
-  prod_encode (Pos.of_nat (S (length l))) (list_encode_ l).
-
-Fixpoint list_decode_ (n : nat) (p : positive) : option (list positive) :=
-  match n with
-  | O => guard (p = 1); Some []
-  | S n =>
-     x ← prod_decode_fst p; pl ← prod_decode_snd p;
-     l ← list_decode_ n pl; Some (x :: l)
+(** ** Lists *)
+(* Lists are encoded as 1 separated sequences of 0s corresponding to the unary
+representation of the elements. *)
+Fixpoint list_encode `{Countable A} (acc : positive) (l : list A) : positive :=
+  match l with
+  | [] => acc
+  | x :: l => list_encode (Nat.iter (encode_nat x) (~0) (acc~1)) l
   end.
-Definition list_decode (p : positive) : option (list positive) :=
-  pn ← prod_decode_fst p; pl ← prod_decode_snd p;
-  list_decode_ (pred (Pos.to_nat pn)) pl.
-
-Lemma list_decode_encode l : list_decode (list_encode l) = Some l.
+Fixpoint list_decode `{Countable A} (acc : list A)
+    (n : nat) (p : positive) : option (list A) :=
+  match p with
+  | 1 => Some acc
+  | p~0 => list_decode acc (S n) p
+  | p~1 => x ← decode_nat n; list_decode (x :: acc) O p
+  end.
+Lemma x0_iter_x1 n acc : Nat.iter n (~0) acc~1 = acc ++ Nat.iter n (~0) 3.
+Proof. by induction n; f_equal'. Qed.
+Lemma list_encode_app' `{Countable A} (l1 l2 : list A) acc :
+  list_encode acc (l1 ++ l2) = list_encode acc l1 ++ list_encode 1 l2.
 Proof.
-  cut (list_decode_ (length l) (list_encode_ l) = Some l).
-  { intros help. unfold list_decode, list_encode.
-    rewrite prod_decode_encode_fst, prod_decode_encode_snd; csimpl.
-    by rewrite Nat2Pos.id by done; simpl. }
-  induction l; simpl; auto.
-  by rewrite prod_decode_encode_fst, prod_decode_encode_snd;
-    simplify_option_equality.
+  revert acc; induction l1; simpl; auto.
+  induction l2 as [|x l IH]; intros acc; simpl; [by rewrite ?(left_id_L _ _)|].
+  by rewrite !(IH (Nat.iter _ _ _)), (associative_L _), x0_iter_x1.
 Qed.
-
-Program Instance list_countable `{Countable A} : Countable (list A) :=  {|
-  encode l := list_encode (encode <$> l);
-  decode p := list_decode p ≫= mapM decode
-|}.
+Program Instance list_countable `{Countable A} : Countable (list A) :=
+  {| encode := list_encode 1; decode := list_decode [] 0 |}.
 Next Obligation.
-  intros ??? l; simpl; rewrite list_decode_encode; simpl.
-  apply mapM_fmap_Some; auto using decode_encode.
+  intros A ??; simpl.
+  assert (∀ m acc n p, list_decode acc n (Nat.iter m (~0) p)
+    = list_decode acc (n + m) p) as decode_iter.
+  { induction m as [|m IH]; intros acc n p; simpl; [by rewrite Nat.add_0_r|].
+    by rewrite IH, Nat.add_succ_r. }
+  cut (∀ l acc, list_decode acc 0 (list_encode 1 l) = Some (l ++ acc))%list.
+  { by intros help l; rewrite help, (right_id_L _ _). }
+  induction l as [|x l IH] using @rev_ind; intros acc; [done|].
+  rewrite list_encode_app'; simpl; rewrite <-x0_iter_x1, decode_iter; simpl.
+  by rewrite decode_encode_nat; simpl; rewrite IH, <-(associative_L _).
 Qed.
-
-Program Instance pos_countable : Countable positive := {|
-  encode := id; decode := Some; decode_encode x := eq_refl
-|}.
+Lemma list_encode_app `{Countable A} (l1 l2 : list A) :
+  encode (l1 ++ l2)%list = encode l1 ++ encode l2.
+Proof. apply list_encode_app'. Qed.
+Lemma list_encode_cons `{Countable A} x (l : list A) :
+  encode (x :: l) = Nat.iter (encode_nat x) (~0) 3 ++ encode l.
+Proof. apply (list_encode_app' [_]). Qed.
+Lemma list_encode_suffix `{Countable A} (l k : list A) :
+  l `suffix_of` k → ∃ q, encode k = q ++ encode l.
+Proof. intros [l' ->]; exists (encode l'); apply list_encode_app. Qed.
+
+(** ** Numbers *)
+Instance pos_countable : Countable positive :=
+  {| encode := id; decode := Some; decode_encode x := eq_refl |}.
 Program Instance N_countable : Countable N := {|
   encode x := match x with N0 => 1 | Npos p => Pos.succ p end;
   decode p := if decide (p = 1) then Some 0%N else Some (Npos (Pos.pred p))
 |}.
 Next Obligation.
-  intros [|p]; simpl; repeat case_decide; auto with lia.
-  by rewrite Pos.pred_succ.
+  by intros [|p];simpl;[|rewrite decide_False,Pos.pred_succ by (by destruct p)].
 Qed.
 Program Instance Z_countable : Countable Z := {|
-  encode x :=
-    match x with Z0 => 1 | Zpos p => p~0 | Zneg p => p~1 end;
-  decode p := Some
-    match p with 1 => Z0 | p~0 => Zpos p | p~1 => Zneg p end
+  encode x := match x with Z0 => 1 | Zpos p => p~0 | Zneg p => p~1 end;
+  decode p := Some match p with 1 => Z0 | p~0 => Zpos p | p~1 => Zneg p end
 |}.
 Next Obligation. by intros [|p|p]. Qed.
-Program Instance nat_countable : Countable nat := {|
-  encode x := encode (N.of_nat x);
-  decode p := N.to_nat <$> decode p
-|}.
+Program Instance nat_countable : Countable nat :=
+  {| encode x := encode (N.of_nat x); decode p := N.to_nat <$> decode p |}.
 Next Obligation.
-  intros x; lazy beta; rewrite decode_encode; csimpl. by rewrite Nat2N.id.
+  by intros x; lazy beta; rewrite decode_encode; csimpl; rewrite Nat2N.id.
 Qed.