Require Import List. Import ListNotations.
Require Import FunctionalExtensionality.

Definition kat {a b c} (f : a -> option b) (g : b -> option c) : a -> option c :=
  fun x =>
    match f x with
    | None => None
    | Some y => g y
    end.

(*
Fixpoint omap {a b} (f : a -> option b) (xs : list a) : option (list b) :=
  match xs with
  | [] => Some []
  | x :: xs =>
    match f x, omap f xs with
    | Some y, Some ys => Some (y :: ys)
    | _, _ => None
    end
  end.
*)

Fixpoint filter_option {a b} (f : a -> option b) (xs : list a) : list b :=
  match xs with
  | [] => []
  | x :: xs =>
    match f x with
    | None => filter_option f xs
    | Some y => y :: filter_option f xs
    end
  end.

Definition omap {a b} (f : a -> option b) (xs : list a) : option (list b) :=
  Some (filter_option f xs).

Lemma omap_kat {a b c} (f : a -> option b) (g : b -> option c)
  : omap (kat f g) = kat (omap f) (omap g).
Proof.
  apply functional_extensionality.
  unfold omap, kat.
  intros xs. f_equal. induction xs; cbn.
  - reflexivity.
  - destruct f; cbn; auto.
    destruct g; cbn; auto.
    rewrite IHxs. reflexivity.
Qed.

Fixpoint tails {a} (xs : list a) : list (list a) :=
  xs :: match xs with
        | [] => []
        | _ :: xs => tails xs
        end.

Definition duplicate {a} (xs : list a) : option (list (list a)) :=
  Some (tails xs).

Definition extract {a} (xs : list a) : option a :=
  match xs with
  | [] => None
  | x :: _ => Some x
  end.

Lemma top_id {a} : kat (@duplicate a) extract = Some.
Proof.
  apply functional_extensionality; intros []; reflexivity.
Qed.

Lemma bot_id {a} : kat (@duplicate a) (omap extract) = Some.
Proof.
  apply functional_extensionality; intros xs.
  unfold kat, omap, duplicate. f_equal.
  induction xs.
  - reflexivity.
  - cbn. rewrite IHxs; auto.
Qed.

Lemma assoc {a} : kat (@duplicate a) duplicate = kat (@duplicate a) (omap duplicate).
Proof.
  apply functional_extensionality; intros xs.
  unfold kat, omap, duplicate. f_equal.
  induction xs; cbn.