From mininix Require Export utils.
From stdpp Require Import options.

Variant res A :=
  | Res (x : option A)
  | NoFuel.
Arguments Res {_} _.
Arguments NoFuel {_}.

Instance res_fail : MFail res := λ {A} _, Res None.

Instance res_mret : MRet res := λ {A} x, Res (Some x).

Instance res_mbind : MBind res := λ {A B} f rx,
  match rx with
  | Res mx => default mfail (f <$> mx)
  | NoFuel => NoFuel
  end.

Instance res_fmap : FMap res := λ {A B} f rx,
  match rx with
  | Res mx => Res (f <$> mx)
  | NoFuel => NoFuel
  end.

Instance Res_inj A : Inj (=) (=) (@Res A).
Proof. by injection 1. Qed.

Ltac simplify_res :=
  repeat match goal with
  | H : Res _ = mfail |- _ => apply (inj Res) in H
  | H : mfail = Res _ |- _ => apply (inj Res) in H
  | H : Res _ = mret _ |- _ => apply (inj Res) in H
  | H : mret _ = Res _ |- _ => apply (inj Res) in H
  | _ => progress simplify_eq/=
  end.

Lemma mapM_Res_impl {A B} (f g : A → res B) (xs : list A) ys :
  mapM f xs = Res ys →
  (∀ x y, f x = Res y → g x = Res y) →
  mapM g xs = Res ys.
Proof.
  intros Hxs Hf. revert ys Hxs.
  induction xs as [|x xs IH]; intros ys ?; simplify_res; [done|].
  destruct (f x) as [my|] eqn:?; simplify_res. rewrite (Hf x my) //=.
  destruct my as [y|]; simplify_res; [|done].
  destruct (mapM f _) as [mys|]; simplify_res; [|done..].
  by rewrite (IH _ eq_refl).
Qed.

Lemma map_mapM_sorted_Res_impl `{FinMap K M}
    (R : relation K) `{!RelDecision R, !PartialOrder R, !Total R}
    {A B} (f g : A → res B) (m1 : M A) m2 :
  map_mapM_sorted R f m1 = Res m2 →
  (∀ x y, f x = Res y → g x = Res y) →
  map_mapM_sorted R g m1 = Res m2.
Proof.
  intros Hm Hf. revert m2 Hm.
  induction m1 as [|i x m1 ?? IH] using (map_sorted_ind R); intros m2.
  { by rewrite !map_mapM_sorted_empty. }
  rewrite !map_mapM_sorted_insert //. intros.
  destruct (f x) as [my|] eqn:?; simplify_res. rewrite (Hf x my) //=.
  destruct my as [y|]; simplify_res; [|done].
  destruct (map_mapM_sorted _ f _) as [mm2'|]; simplify_res; [|done..].
  by rewrite (IH _ eq_refl).
Qed.

Lemma mapM_res_app {A B} (f : A → res B) xs1 xs2 :
  mapM f (xs1 ++ xs2) = ys1 ← mapM f xs1; ys2 ← mapM f xs2; mret (ys1 ++ ys2).
Proof.
  induction xs1 as [|x1 xs1 IH]; simpl.
  { by destruct (mapM f xs2) as [[]|]. }
  destruct (f x1) as [[y1|]|]; simpl; [|done..].
  rewrite IH. by destruct (mapM f xs1) as [[]|],  (mapM f xs2) as [[]|].
Qed.