From mininix Require Import nix.wp nix.notations.
From stdpp Require Import options.
Local Open Scope Z_scope.

Definition test αs :=
  let: "x" := 1 in
  with: EAttr αs in
  with: EAttr {[ "y" := AttrN 2 ]} in
  "x" =: "y".

Example test_wp μ αs :
  no_recs αs →
  wp μ (test αs) (.= false).
Proof.
  intros Hαs. rewrite /test. apply LetAttr_no_recs_wp.
  { by apply no_recs_insert. }
  rewrite /= !map_fmap_singleton /= right_id_L lookup_singleton lookup_singleton_ne //=.
  apply LetAttr_no_recs_wp.
  { by apply no_recs_attr_subst. }
  rewrite /= !map_fmap_singleton /= right_id_L.
  rewrite (map_fmap_compose attr_expr) lookup_fmap union_kinded_abs.
  rewrite !lookup_fmap.
  apply LetAttr_no_recs_wp.
  { by apply no_recs_insert. }
  rewrite /= map_fmap_singleton lookup_singleton lookup_singleton_ne //=.
  rewrite union_kinded_with.
  apply BinOp_wp.
  apply Id_wp, Lit_wp; first done. eexists; split; [constructor|].
  apply Id_wp, Lit_wp; first done.
  eexists; split; [done|]. by apply Lit_wp.
Qed.

Definition neg := λ: "b", if: "b" then false else true.

Lemma neg_wp μ (Φ : expr → Prop) e :
  wp SHALLOW e (λ e', ∃ b : bool, e' = b ∧ Φ (negb b)) →
  wp μ (neg e) Φ.
Proof.
  intros Hwp. apply β_wp. rewrite /= lookup_singleton /=.
  apply If_wp, Id_wp. eapply wp_mono; [done|].
  intros ? (b & -> & ?). exists b; split; [done|].
  destruct b; by apply Lit_wp.
Qed.

(* rec { f = x: if x = 0 then true else !(f (x - 1)); }.f n *)
Definition even_rec_attr :=
  EAttr {[ "f" := AttrR (λ: "x", if: "x" =: 0 then true else neg ("f" ("x" -: 1))) ]} .: "f".

Lemma even_rec_attr_wp e n :
  0 ≤ n ≤ int_max →
  wp SHALLOW e (.= n) →
  wp SHALLOW (even_rec_attr e) (.= Z.even n).
Proof.
  intros Hn Hwp. apply App_wp.
  revert e Hwp. induction (Z.lt_wf 0 n) as [n _ IH]; intros e Hwp.
  apply BinOp_wp. apply Attr_wp_shallow.
  eexists; split; [by constructor|].
  apply Lit_wp; [done|]. eexists; split; [by eexists|].
  rewrite /=. apply Abs_wp, β_wp.
  rewrite /= !lookup_singleton /= !lookup_singleton_ne //=.
  rewrite !union_with_None_l !union_with_None_r.
  rewrite /indirects map_imap_insert map_imap_empty lookup_insert.
  rewrite -/even_rec_attr -/neg.
  apply If_wp, BinOp_wp, Id_wp.
  eapply wp_mono; [apply Hwp|]; intros ? ->.
  eexists; split; [by constructor|].
  apply Lit_wp; [done|]. eexists; split; [by eexists|]. simpl.
  destruct (n =? 0) eqn:Hn0; (apply Lit_wp; [done|]; eexists; split; [done|]; simpl).
  { apply Lit_wp; [done|]. by apply Z.eqb_eq in Hn0 as ->. }
  apply neg_wp, App_wp, Id_wp.
  eapply wp_mono; [apply (IH (n-1))|]; [lia..| |].
  2:{ intros e' He'. eapply wp_mono; [apply He'|].
      intros ? ->. eexists; split; [done|].
      by rewrite Z.negb_even Z.sub_1_r Z.odd_pred. }
  eapply BinOp_wp, Id_wp. eapply wp_mono; [apply Hwp|]. intros ? ->.
  eexists; split; [by constructor|]. apply Lit_wp; [done|].
  eexists; split; [eexists _, _; split_and!; [done| |done]|].
  - rewrite /= option_guard_True //. apply bool_decide_pack.
    rewrite /int_min Z.shiftl_mul_pow2 //. lia.
  - apply Lit_wp; [|done]. apply bool_decide_pack.
    rewrite /int_min Z.shiftl_mul_pow2 //. lia.
Qed.

Lemma even_rec_attr_wp' n :
  0 ≤ n ≤ int_max →
  wp SHALLOW (even_rec_attr n) (.= Z.even n).
Proof.
  intros ?. apply even_rec_attr_wp; [done|]. apply Lit_wp; [|done].
  rewrite /= /int_ok. apply bool_decide_pack.
  rewrite /int_min Z.shiftl_mul_pow2 //. lia.
Qed.

(* { "__functor " = r: x: if x == 0 then true else !(r (x - 1)); } n *)
Definition even_rec_functor :=
  EAttr {[ "__functor" :=
    AttrN (λ: "r" "x", if: "x" =: 0 then true else neg ("r" ("x" -: 1))) ]}.

Lemma even_rec_functor_wp e n :
  0 ≤ n ≤ int_max →
  wp SHALLOW e (.= n) →
  wp SHALLOW (even_rec_functor e) (.= Z.even n).
Proof.
  intros Hn Hwp. apply App_wp.
  revert e Hwp. induction (Z.lt_wf 0 n) as [n _ IH]; intros e Hwp.
  apply Attr_wp_shallow. rewrite map_fmap_singleton /=. eapply Functor_wp.
  { by apply no_recs_insert. }
  { done. }
  apply App_wp. apply β_wp.
  rewrite /= !lookup_singleton !lookup_singleton_ne //=. apply Abs_wp, β_wp.
  rewrite /= !lookup_singleton /= !lookup_singleton_ne //=.
  rewrite -/even_rec_functor -/neg.
  apply If_wp, BinOp_wp, Id_wp.
  eapply wp_mono; [apply Hwp|]; intros ? ->.
  eexists; split; [by constructor|].
  apply Lit_wp; [done|]. eexists; split; [by eexists|]. simpl.
  destruct (n =? 0) eqn:Hn0; (apply Lit_wp; [done|]; eexists; split; [done|]; simpl).
  { apply Lit_wp; [done|]. by apply Z.eqb_eq in Hn0 as ->. }
  apply neg_wp, App_wp, Id_wp.
  eapply wp_mono; [apply (IH (n-1))|]; [lia..| |].
  2:{ intros e' He'. eapply wp_mono; [apply He'|].
      intros ? ->. eexists; split; [done|].
      by rewrite Z.negb_even Z.sub_1_r Z.odd_pred. }
  eapply BinOp_wp, Id_wp. eapply wp_mono; [apply Hwp|]. intros ? ->.
  eexists; split; [by constructor|]. apply Lit_wp; [done|].
  eexists; split; [eexists _, _; split_and!; [done| |done]|].
  - rewrite /= option_guard_True //. apply bool_decide_pack.
    rewrite /int_min Z.shiftl_mul_pow2 //. lia.
  - apply Lit_wp; [|done]. apply bool_decide_pack.
    rewrite /int_min Z.shiftl_mul_pow2 //. lia.
Qed.

Lemma even_rec_functor_wp' n :
  0 ≤ n ≤ int_max →
  wp SHALLOW (even_rec_functor n) (.= Z.even n).
Proof.
  intros ?. apply even_rec_functor_wp; [done|]. apply Lit_wp; [|done].
  rewrite /= /int_ok. apply bool_decide_pack.
  rewrite /int_min Z.shiftl_mul_pow2 //. lia.
Qed.

(* ({ f ? (x: if x == 0 then true else !(f (x - 1))) }: f) {} n *)
Definition even_rec_default :=
  (λattr:
    {[ "f" := Some (λ: "x", if: "x" =: 0 then true else neg ("f" ("x" -: 1))) ]}, "f")
  (EAttr ∅).

Lemma even_rec_default_wp e n :
  0 ≤ n ≤ int_max →
  wp SHALLOW e (.= n) →
  wp SHALLOW (even_rec_default e) (.= Z.even n).
Proof.
  intros Hn Hwp. apply App_wp.
  eapply βMatch_wp; [done|repeat econstructor|]. simplify_map_eq.
  rewrite -/even_rec_attr. by apply Id_wp, App_wp, even_rec_attr_wp.
Qed.

Lemma even_rec_default_wp' n :
  0 ≤ n ≤ int_max →
  wp SHALLOW (even_rec_default n) (.= Z.even n).
Proof.
  intros ?. apply even_rec_default_wp; [done|]. apply Lit_wp; [|done].
  rewrite /= /int_ok. apply bool_decide_pack.
  rewrite /int_min Z.shiftl_mul_pow2 //. lia.
Qed.