Require Import Hask.Data.Tuple.
Require Import Hask.Control.Monad.
Require Import Hask.Data.Functor.Identity.
Require Import Hask.Control.Monad.Cont.
Require Import Coq.Unicode.Utf8.
Require Import FunctionalExtensionality.
Require Import Pact.Data.Monoid.
Require Import Pact.Data.Either.

Generalizable All Variables.
Set Primitive Projections.

(******************************************************************************
 * The RWSET Monad transformer
 *)

Ltac rwse :=
  let r := fresh "r" in extensionality r;
  let s := fresh "s" in extensionality s.

Section RWSET.

Context {r w s e : Type}.
Context {m : Type → Type}.

Definition RWSET `{Monad m} `{Monoid w} (a : Type) : Type :=
  r → s → m (e + (a * (s * w)))%type.

Context `{Monad m}.
Context `{Monoid w}.

Definition ask : RWSET r :=
  λ r s, pure (inr (r, (s, mempty))).
Definition asks {a : Type} (f : r → a) : RWSET a :=
  λ r s, pure (inr (f r, (s, mempty))).
Definition local (f : r → r) `(x : RWSET a) : RWSET a :=
  λ r s, x (f r) s.
Definition get : RWSET s :=
  λ _ s, pure (inr (s, (s, mempty))).
Definition gets {a : Type} (f : s → a) : RWSET a :=
  λ _ s, pure (inr (f s, (s, mempty))).
Definition put (x : s) : RWSET unit :=
  λ _ _, pure (inr (tt, (x, mempty))).
Definition modify (f : s → s) : RWSET unit :=
  λ _ s, pure (inr (tt, (f s, mempty))).
Definition tell (v : w) : RWSET unit :=
  λ _ s, pure (inr (tt, (s, v))).
Definition throw {a : Type} (err : e) : RWSET a :=
  λ _ _, pure (inl err).

#[export]
Instance RWSET_Functor : Functor RWSET := {
  fmap := λ a _ f (x : RWSET a), λ r s,
    fmap[m] (fmap (first f)) (x r s)
}.

Definition RWSET_ap `(f : RWSET (a → b)) (x : RWSET a) :
  RWSET b := λ r s,
    f' <- f r s ;
    match f' with
    | inl e => pure (inl e)
    | inr (f'', (s', w')) =>
        x' <- fmap f'' x r s' ;
        pure (match x' with
              | inl e => inl e
              | inr (x'', (s'', w'')) =>
                  inr (x'', (s'', w' ⨂ w''))
              end)
    end.

#[export]
Instance RWSET_Applicative : Applicative RWSET := {
  is_functor := RWSET_Functor;
  pure := λ _ x, λ _ s, pure (inr (x, (s, mempty)));
  ap   := λ _ _, RWSET_ap
}.

Definition RWSET_join `(x : RWSET (RWSET a)) :
  RWSET a := λ r s,
    x' <- x r s ;
    match x' with
    | inl e => pure (inl e)
    | inr (y, (s', w')) =>
        y' <- y r s' ;
        pure (match y' with
              | inl e => inl e
              | inr (z, (s'', w'')) =>
                  inr (z, (s'', w' ⨂ w''))
              end)
    end.

#[export]
Instance RWSET_Monad : Monad RWSET | 1 := {
  is_applicative := RWSET_Applicative;
  join := λ _, RWSET_join
}.

End RWSET.

Arguments RWSET r w s e m {_ _}.
Arguments RWSET_ap {r w s e m _ _ a b} f x _ _ /.
Arguments RWSET_join {r w s e m _ _ a} x _ _ /.

#[export] Hint Unfold RWSET_ap : core.
#[export] Hint Unfold RWSET_join : core.
#[export] Hint Unfold Either_map : core.
#[export] Hint Unfold Tuple.first : core.
#[export] Hint Unfold id : core.

Definition RWSE r w s e `{Monoid w} := RWSET r w s e Identity.

Definition RWSEP r w s e `{Monoid w} := RWSET r w s e (Cont Prop).

Module RWSETLaws.

Module Import EL := EitherLaws.
Include MonadLaws.

Lemma first_id a z : first (a:=a) (b:=a) (z:=z) id = id.
Proof.
  unfold first.
  extensionality x.
  now destruct x.
Qed.

Lemma first_comp `(f : b → c) `(g : a → b) z x :
  first f (first g x) = first (z:=z) (f \o g) x.
Proof.
  unfold first.
  now destruct x.
Qed.

Section RWSETLaws.

Context {r w s e : Type}.
Context `{MonoidLaws w}.
Context `{MonadLaws m}.

#[global]
Program Instance RWSET_FunctorLaws :
  FunctorLaws (RWSET r w s e m).
Next Obligation.
  extensionality x.
  rwse.
  rewrite first_id.
  rewrite (fmap_id (FunctorLaws:=Either_FunctorLaws (e:=e))).
  now rewrite fmap_id.
Qed.
Next Obligation.
  extensionality x.
  rwse; simpl.
  rewrite fmap_comp_x.
  pose proof (fmap_comp (FunctorLaws:=Either_FunctorLaws (e:=e))).
  simpl in H3.
  unfold comp in H3.
  rewrite H3.
  repeat f_equal.
  extensionality x0.
  now apply first_comp.
Qed.

#[global]
Program Instance RWSET_ApplicativeLaws :
  ApplicativeLaws (RWSET r w s e m).
Next Obligation.
  extensionality x.
  rwse.
  unfold RWSET_ap.
  unfold bind, comp; simpl.
  rewrite fmap_pure_x.
  rewrite join_pure_x.
  rewrite fmap_comp_x.
  rewrite first_id; simpl.
  rewrite (fmap_id (FunctorLaws:=Either_FunctorLaws (e:=e))).
  unfold id.
  rewrite <- fmap_comp_x.
  rewrite join_fmap_pure_x.
  rewrite <- fmap_id_x.
  f_equal.
  extensionality y.
  destruct y; auto.
  destruct p; auto.
  destruct p; auto.
  now rewrite mempty_left.
Qed.
Next Obligation.
Admitted.
Next Obligation.
Admitted.
Next Obligation.
Admitted.
Next Obligation.
Admitted.

#[global]
Program Instance RWSET_MonadLaws :
  MonadLaws (RWSET r w s e m) := {|
    has_applicative_laws := RWSET_ApplicativeLaws
|}.
Next Obligation.
  unfold comp.
  extensionality x.
Admitted.
Next Obligation.
Admitted.
Next Obligation.
Admitted.
Next Obligation.
Admitted.

Notation "a >>[ H ] b" :=
  (bind (H:=H) (fun _ => b) a) (at level 90, right associativity) : monad_scope.

Lemma put_getM `(x : s) :
  (put x >> get (r:=r) (e:=e)) =
  (put x >> pure x).
Proof.
  simpl.
  rwse.
  unfold put, RWSET_join, bind, comp.
  now rewrite !fmap_comp_x, !fmap_pure_x.
Qed.

End RWSETLaws.

End RWSETLaws.

Require Import Coq.Arith.Arith.

#[export] Program Instance Unit_Semigroup : Semigroup unit.
#[export] Program Instance Unit_Monoid : Monoid unit.
Next Obligation. exact tt. Defined.

Definition sample {m : Type → Type} `{Monad m} :
  RWSET nat unit nat nat m nat :=
  x <- get ;
  put (x + 1) ;;
  if x <? 20
  then throw 100
  else pure 50.


(*
 * This is not in the Coq library, as far as I can tell.
 *
 * Some of this (without the typeclassery) was taken from
 * https://pdp7.org/blog/2011/01/the-maybe-monad-in-coq/
 *)

Class Monad (M : Type -> Type): Type := {
   ret {a : Type}: a -> M a;
   bind {a b : Type}: M a -> (a -> M b) -> M b;

   left_id: forall (a b : Type), forall (x : a) (f : a -> M b),
      bind (ret x) f = f x;
   right_id {a: Type}: forall (m : M a),
      bind m (fun (x : a) => ret x) = m;
   assoc {a b c: Type}: forall (m : M a) (n: a -> M b) (o : b -> M c),
      bind (bind m n) o = bind m (fun x => bind (n x) o);
}.

Notation "M >>= N" := (bind M N)
	(at level 42, left associativity).

Notation "'do' X <- M ; N" := (bind M (fun X => N))
	(at level 200, X ident, M at level 100, N at level 200).

Notation "'do' M ; N" := (bind M (fun _ => N))
	(at level 200, M at level 100, N at level 200).

Instance option_is_monad: Monad option := {
   (* It isn't clear to me why the implicit args need to be explicit. *)
   ret a x := Some x;
   bind a b ma fmb := match ma with
   | Some x => fmb x
   | None => None
   end;
}.
Proof.
   - intros. auto.
   - intros. destruct m; auto.
   - intros. destruct m; auto.
Qed.


Require Import Fiat.QueryStructure.Automation.MasterPlan.
Require Import Bedrock.Memory.
Definition boolToWord (b : bool) : W :=
  if b then Word.natToWord _ 1 else Word.natToWord _ 0.

Instance : Query_eq (Word.word n) :=
  { A_eq_dec := @Word.weq n }.
Opaque Word.weq.
Opaque Word.natToWord.
(* Our bookstore has two relations (tables):
   - The [Books] relation contains the books in the
     inventory, represented as a tuple with
     [Author], [Title], and [ISBN] attributes.
     The [ISBN] attribute is a key for the relation,
     specified by the [where attributes .. depend on ..]
     constraint.
   - The [Orders] relation contains the orders that
     have been placed, represented as a tuple with the
     [ISBN] and [Date] attributes.

   The schema for the entire query structure specifies that
   the [ISBN] attribute of [Orders] is a foreign key into
   [Books], specified by the [attribute .. of .. references ..]
   constraint.
 *)

(* Let's define some synonyms for strings we'll need,
 * to save on type-checking time. *)
Definition sBOOKS := "Books".
Definition sAUTHOR := "Authors".
Definition sTITLE := "Title".
Definition sISBN := "ISBN".
Definition sORDERS := "Orders".
Definition sDATE := "Date".

(* Now here's the actual schema, in the usual sense. *)
Definition BookStoreSchema :=
  Query Structure Schema
    [ relation sBOOKS has
              schema <sAUTHOR :: W,
                      sTITLE :: W,
                      sISBN :: W>
                      where UniqueAttribute ``sISBN;
      relation sORDERS has
              schema <sISBN :: W,
                      sDATE :: W> ]
    enforcing [attribute sISBN for sORDERS references sBOOKS].

(* Aliases for the tuples contained in Books and Orders, respectively. *)
Definition Book := TupleDef BookStoreSchema sBOOKS.
Definition Order := TupleDef BookStoreSchema sORDERS.

(* Our bookstore has two mutators:
   - [PlaceOrder] : Place an order into the 'Orders' table
   - [AddBook] : Add a book to the inventory

   Our bookstore has two observers:
   - [GetTitles] : The titles of books written by a given author
   - [NumOrders] : The number of orders for a given author
 *)

(* Now we write what the methods should actually do. *)

Definition BookStoreSpec : ADT _ :=
  Def ADT {
    rep := QueryStructure BookStoreSchema,

    Def Constructor0 "Init" : rep := empty,,

    Def Method1 "PlaceOrder" ( r : rep) (o : Order) : rep * bool :=
        Insert o into r!sORDERS,

    Def Method1 "AddBook" (r : rep) (b : Book) : rep * bool :=
        Insert b into r!sBOOKS ,

    Def Method1 "GetTitles" (r : rep) (author : W) : rep * list W :=
        titles <- For (b in r ! sBOOKS)
               Where (author = b!sAUTHOR)
               Return (b!sTITLE);
    ret (r, titles),

    Def Method1 "NumOrders" (r : rep) (author : W ) : rep * W :=
      count <- Count (For (o in r!sORDERS) (b in r!sBOOKS)
                              Where (author = b!sAUTHOR)
                              Where (o!sISBN = b!sISBN)
                              Return ());
      ret (r, Word.natToWord 32 count)
}%methDefParsing.

Theorem SharpenedBookStore :
  MostlySharpened BookStoreSpec.
Proof.

  start sharpening ADT.
  simpl; pose_string_hyps; pose_heading_hyps.
  start_honing_QueryStructure'.
  hone method "AddBook".
  { setoid_rewrite UniqueAttribute_symmetry.
    setoid_rewrite (@refine_uniqueness_check_into_query' BookStoreSchema Fin.F1 _ _ _ _).
    setoid_rewrite refine_For_rev.
    setoid_rewrite refine_Count.
    simplify with monad laws; simpl in *; subst.
    setoid_rewrite refine_pick_eq'.
    setoid_rewrite refine_bind_unit.
    setoid_rewrite refine_If_Then_Else_Duplicate.
    finish honing. }
  GenerateIndexesForAll EqExpressionAttributeCounter
  ltac:(fun attrlist =>
          let attrlist' := eval compute in (PickIndexes _ (CountAttributes' attrlist)) in
              make_simple_indexes attrlist'
                                  ltac:(LastCombineCase6 BuildEarlyEqualityIndex)
                                         ltac:(LastCombineCase5 BuildLastEqualityIndex)).
  + plan EqIndexUse createEarlyEqualityTerm createLastEqualityTerm
         EqIndexUse_dep createEarlyEqualityTerm_dep createLastEqualityTerm_dep.
  + setoid_rewrite refine_For_rev; simplify with monad laws.
    plan EqIndexUse createEarlyEqualityTerm createLastEqualityTerm
         EqIndexUse_dep createEarlyEqualityTerm_dep createLastEqualityTerm_dep.
  + setoid_rewrite refine_For_rev; simplify with monad laws.
    plan EqIndexUse createEarlyEqualityTerm createLastEqualityTerm
         EqIndexUse_dep createEarlyEqualityTerm_dep createLastEqualityTerm_dep.
  + setoid_rewrite refine_For_rev; simplify with monad laws.
    plan EqIndexUse createEarlyEqualityTerm createLastEqualityTerm
         EqIndexUse_dep createEarlyEqualityTerm_dep createLastEqualityTerm_dep.
  + setoid_rewrite refine_For_rev; simplify with monad laws.
    plan EqIndexUse createEarlyEqualityTerm createLastEqualityTerm
         EqIndexUse_dep createEarlyEqualityTerm_dep createLastEqualityTerm_dep.
  + hone method "PlaceOrder".
    { simpl in *; subst.
      setoid_rewrite refine_Count; simplify with monad laws.
      setoid_rewrite app_nil_r;
        setoid_rewrite map_map; simpl.
      unfold ilist2_hd at 1; simpl.
      setoid_rewrite rev_length.
      setoid_rewrite map_length.
      setoid_rewrite refine_pick_eq'; simplify with monad laws.
      repeat setoid_rewrite refine_If_Then_Else_Bind.
      repeat setoid_rewrite refineEquiv_bind_bind.
      repeat setoid_rewrite refineEquiv_bind_unit; simpl.
      unfold H1.
      eapply refine_under_bind; intros; set_evars.
      rewrite (CallBagFind_fst H0); simpl.
      finish honing.
    }
    hone method "AddBook".
    { simpl in *; subst; simplify with monad laws.
      setoid_rewrite app_nil_r;
        setoid_rewrite map_map; simpl.
      unfold ilist2_hd at 1; simpl.
      repeat setoid_rewrite rev_length.
      setoid_rewrite map_length.
      setoid_rewrite refine_pick_eq'; simplify with monad laws.
      repeat setoid_rewrite refine_If_Then_Else_Bind.
      repeat setoid_rewrite refineEquiv_bind_bind.
      repeat setoid_rewrite refineEquiv_bind_unit; simpl.
      unfold H1; eapply refine_under_bind; intros; set_evars.
      rewrite (CallBagFind_fst H0); simpl.
      finish honing.
    }
    hone method "GetTitles".
    { simpl in *; subst; simplify with monad laws.
      setoid_rewrite refine_pick_eq'; simplify with monad laws.
      simpl.
      unfold H1; eapply refine_under_bind; intros; set_evars.
      rewrite app_nil_r, map_map; unfold ilist2_hd; simpl.
      rewrite (CallBagFind_fst H0); simpl.
      finish honing.
    }
    hone method "NumOrders".
    { simpl in *; subst.
      setoid_rewrite refine_Count; simplify with monad laws.
      setoid_rewrite refine_pick_eq'; simplify with monad laws.
      simpl.
      unfold H1; eapply refine_under_bind; intros; set_evars.
      setoid_rewrite app_nil_r.
      rewrite (CallBagEnumerate_fst H0); simpl.
      etransitivity.
      eapply refine_under_bind_both.
      refine (@Join_Comp_Lists_eq BookStoreSchema Index Fin.F1 _ _ _ _ _).
      intros; finish honing.
      simplify with monad laws.
      unfold H2; apply refine_under_bind.
      intros.
      apply Join_Comp_Lists_eq' in H3; rewrite H3.
      finish honing.
    }
    simpl; eapply reflexivityT.
  + unfold CallBagFind, CallBagInsert.
    pose_headings_all.
    match goal with
    | |- appcontext[ @BuildADT (IndexedQueryStructure ?Schema ?Indexes) ] =>
      FullySharpenQueryStructure Schema Indexes
    end.

    Time Defined.

Time Definition BookstoreImpl :=
  Eval simpl in (fst (projT1 SharpenedBookStore)).

Print BookstoreImpl.