finitepoly.v

Set Primitive Projections.

Require Import Coq.Unicode.Utf8.
Require Coq.Program.Basics.
Require Import Coq.Program.Equality.
Require Import Coq.Classes.RelationClasses.
Require Import Coq.Classes.SetoidClass.
Require Import Coq.Lists.List.
Require Coq.Vectors.Fin.
Require Import Coq.Vectors.VectorDef.
Require Coq.Arith.PeanoNat Nat.
Require Coq.FSets.FMapAVL.
Require Coq.Structures.OrderedTypeEx.

Import IfNotations.
Import List.ListNotations.
Import Vector.VectorNotations.

Import EqNotations.

Module Vector.
  Fixpoint Σ {n}: (Fin.t n → nat) → nat :=
    if n is S n'
    then
      λ f, f Fin.F1 + Σ (λ x, f (Fin.FS x))
    else
      λ _, 0.

  Fixpoint build {A} {n}: (Fin.t n → A) → Vector.t A n :=
    match n with
    | S n' => λ f, f Fin.F1 :: build (λ x, f (Fin.FS x))
    | _ => λ _, []
    end.

  Lemma build_nth {A B} (x: Vector.t A B): build (Vector.nth x) = x.
  Proof.
    induction x.
    all: cbn.
    1: reflexivity.
    replace (build _) with x.
    1: reflexivity.
  Qed.

  Lemma nth_build {A B} (x: Fin.t A → B) y: Vector.nth (build x) y = x y.
  Proof.
    induction y.
    all: cbn.
    1: reflexivity.
    rewrite IHy.
    reflexivity.
  Qed.

  Definition act {A B C} (f: Fin.t B → Fin.t A) (e: Vector.t C A): Vector.t C B :=
    build (λ x, e [@ f x]).

  Lemma act_id {A B}: ∀ x: Vector.t A B, act id x = x.
  Proof.
    intro x.
    unfold act, id.
    apply build_nth.
  Qed.

  Fixpoint Forall_Prop {n A} (P: A → Prop) (e: Vector.t A n): Prop :=
    if e is h :: t
    then
      P h ∧ Forall_Prop P t
    else
      True.

  Fixpoint Forall {n A} (P: A → Type) (e: Vector.t A n): Type :=
    if e is h :: t
    then
      prod (P h) (Forall P t)
    else
      unit.

  Fixpoint Forall_Prop_build {A} (P: A → Prop) {n} (f: Fin.t n → _) (q: Forall_Prop P (build f)) i {struct i}: P (f i).
  Proof.
    destruct i.
    - cbn in q.
      destruct q.
      auto.
    - cbn in q.
      destruct q.
      apply (Forall_Prop_build _ _ _ _ H0 i).
  Qed.

  Lemma Forall_Prop_nth {A} (P: A → Prop) {n} (e: Vector.t _ n) (q: Forall_Prop P e): ∀ i, P (e [@ i]).
  Proof.
    assert (p := Forall_Prop_build P (Vector.nth e)).
    rewrite build_nth in p.
    apply p.
    auto.
  Qed.

  Fixpoint Forall_build {A} (P: A → Type) {n} (f: Fin.t n → _) (q: Forall P (build f)) i {struct i}: P (f i).
  Proof.
    destruct i.
    - cbn in q.
      destruct q.
      auto.
    - cbn in q.
      destruct q.
      apply (Forall_build _ _ _ _ f0 i).
  Qed.

  Lemma Forall_nth {A} (P: A → Type) {n} (e: Vector.t _ n) (q: Forall P e): ∀ i, P (e [@ i]).
  Proof.
    assert (p := Forall_build P (Vector.nth e)).
    rewrite build_nth in p.
    apply p.
    auto.
  Qed.
End Vector.

Module Poly.
  Record t := finpoly { s: Set ; π: s → nat ; }.

  Module PolyNotations.
    Infix "▹" := finpoly (at level 30).
  End PolyNotations.
End Poly.
Import Poly.PolyNotations.

Module Mor.
  Import Poly.

  Record t (A B: Poly.t) :=
    mor {
        s: s A → s B ;
        π x: Fin.t (π B (s x)) → Fin.t (π A x) ;
      }.

  Arguments mor {A B}.
  Arguments s {A B}.
  Arguments π {A B}.

  Definition id A: t A A := mor (λ x, x) (λ _, id).
  Definition compose {A B C} (f: t B C) (g: t A B): t A C :=
    mor (λ x, s f (s g x)) (λ _, Basics.compose (π g _) (π f _)).

  Module MorNotations.
    Infix "∘" := compose (at level 40, left associativity).
  End MorNotations.
End Mor.

Module Import Container.
  Variant format := stream | term.

  Definition container (f: format): Set → nat → Set :=
    match f with
    (* this encoding is primarily useful for induction *)
    | stream => λ V n, Fin.t n → V
    | term => Vector.t
    end.

  Definition map {F}: ∀ {A B: Set}, (A → B) → ∀ {n}, container F A n → container F B n :=
    match F with
    | stream => λ _ _ f _ g, Basics.compose f g
    | term => @Vector.map
    end.

  Definition act {F}: ∀ {A B C}, (Fin.t B → Fin.t A) → container F C A → container F C B :=
    match F with
    | stream => λ _ _ _ f g, Basics.compose g f
    | term => @Vector.act
    end.

  Definition nil {F}: ∀ {A}, container F A 0 :=
    match F with
    | stream => λ A x, match x with end
    | term => @Vector.nil
    end.
End Container.

Module Ext.
  Unset Elimination Schemes.

  Record t {f: format} {P: Poly.t} {V: Set}: Set := con { tag: _ ; field: container f V (Poly.π P tag) ; }.
  Arguments t: clear implicits.

  Arguments con {f} P {V} tag &.
  Arguments tag {f P V}.
  Arguments field {f P V}.

  Module Stream.
    Definition t := t stream.

    Definition t_ind {P: Poly.t} {A: Set}
               (Q: t P A → Prop)
               (f: ∀ s x, Q (con P s x))
               x: Q x := f (tag x) (field x).

    Definition t_rect {P: Poly.t} {A: Set}
               (Q: t P A → Type)
               (f: ∀ s x, Q (con P s x))
               x: Q x := f (tag x) (field x).

    Definition t_rec := @t_rect.
    Arguments t_rec {P A}.
  End Stream.

  Module Term.
    Definition t := t term.
  End Term.

  Definition map {F P} {A B: Set} (f: A → B) (e: t F P A): t F P B :=
    con P (tag e) (Container.map f (field e)).

  Definition map_Poly {F P Q} {A: Set} (f: Mor.t P Q) (e: t F P A): t F Q A :=
    con Q (Mor.s f (tag e)) (Container.act (Mor.π f (tag e)) (field e)).

  Definition bimap {F P Q} {A B: Set} (f: Mor.t P Q) (g: A → B) (e: t F P A): t F Q B :=
    con Q (Mor.s f (tag e)) (Container.act (Mor.π f (tag e)) (Container.map g (field e))).

  Definition stream {P A} (e: Term.t P A): Stream.t P A :=
    con P (tag e) (Vector.nth (field e)).

  Definition term {P A} (e: Stream.t P A): Term.t P A :=
    con P (tag e) (Vector.build (field e)).

  Coercion term: Stream.t >-> Term.t.
End Ext.

Module W.
  Module Stream.
    Unset Elimination Schemes.

    Inductive t {P: Poly.t}: Set := roll (e: Ext.Stream.t P t).

    Arguments t: clear implicits.
    Arguments roll {P}.

    Definition unroll {P} (e: t P) :=
      let 'roll e' := e in
      e'.

    Coercion unroll: t >-> Ext.Stream.t.

    Fixpoint t_ind
             (P: Poly.t)
             (Q: t P → Prop)
             (f: ∀ s x, (∀ i, Q (x i)) → Q (roll (Ext.con P s x)))
             x: Q x.
    Proof.
      destruct x.
      apply (f (Ext.tag e) (Ext.field e)).
      intro i.
      apply (t_ind P Q f).
    Defined.

    Fixpoint t_rect
             (P: Poly.t)
             (Q: t P → Type)
             (f: ∀ s x, (∀ i, Q (x i)) → Q (roll (Ext.con P s x)))
             x: Q x.
    Proof.
      destruct x.
      apply (f (Ext.tag e) (Ext.field e)).
      intro i.
      apply (t_rect P Q f).
    Defined.

    Definition t_rec: ∀ (P: Poly.t) (Q: t P → Set), _ := t_rect.

    (* implement induction via eta contraction? *)
    Fixpoint fold {P A} (f: Ext.Stream.t P A → A) (e: t P): A :=
      f (Ext.map (fold f) (unroll e)).

    Fixpoint map {P Q} (f: Mor.t P Q) (e: t P): t Q :=
      roll (Ext.bimap f (map f) (unroll e)).
  End Stream.

  Module Term.
    Unset Elimination Schemes.

    Inductive t {P: Poly.t}: Set := roll (e: Ext.Term.t P t).
    Arguments t: clear implicits.
    Arguments roll {P}.

    Definition unroll {P} (e: t P) :=
      let 'roll e' := e in
      e'.

    Coercion unroll: t >-> Ext.Term.t.

    Fixpoint fold {P A} (f: Ext.Term.t P A → A) (e: t P): A :=
      f (Ext.map (fold f) (unroll e)).

    Fixpoint map {P Q} (f: Mor.t P Q) (e: t P): t Q :=
      roll (Ext.bimap f (map f) (unroll e)).
  End Term.

  Definition t f :=
    match f with
    | stream => Stream.t
    | term => Term.t
    end.

  Definition roll {f}: ∀ {P}, Ext.t f P (t f P) → t f P :=
    match f with
    | stream => @Stream.roll
    | term => @Term.roll
    end.

  Definition unroll {f}: ∀ {P}, t f P → Ext.t f P (t f P) :=
    match f with
    | stream => @Stream.unroll
    | term => @Term.unroll
    end.

  Coercion unroll: t >-> Ext.t.

  Definition fold {F}: ∀ {P A}, (Ext.t F P A → A) → t F P → A :=
    match F with
    | stream => @Stream.fold
    | term => @Term.fold
    end.

  Definition map {F}: ∀ {P Q}, Mor.t P Q → t F P → t F Q :=
    match F with
    | stream => @Stream.map
    | term => @Term.map
    end.

  Fixpoint stream {P} (e: Term.t P) {struct e}: Stream.t P :=
    Stream.roll (Ext.stream (Ext.map stream (Term.unroll e))).

  Fixpoint term {P} (e: Stream.t P): Term.t P :=
    Term.roll (Ext.term (Ext.map term (Stream.unroll e))).

  Coercion term: Stream.t >-> Term.t.
End W.

Module Combinator.
  Import Poly.

  Definition e := unit ▹ λ _, 0.

  Definition app (A B: Poly.t): Poly.t :=
    Ext.t stream A (Poly.s B) ▹ (λ xy, Vector.Σ (λ x, Poly.π _ (Ext.field xy x))).

  Definition const A := A ▹ λ _, 0.

  Definition empty := const Empty_set.
  Definition unit := const unit.

  Definition sum A B := (Poly.s A + Poly.s B) ▹ λ x,
      match x with
      | inl x => Poly.π A x
      | inr x => Poly.π B x
      end.

  Definition prod A B := (s A * s B) ▹ (λ xy, π A (fst xy) + π B (snd xy)).

  Definition Σ {A: Set} (B: A → _) := { x: A & s (B x) } ▹ (λ xy, π (B (projT1 xy)) (projT2 xy)).

  (* FIXME define indexed vectors for dependent products *)

  (* Definition exp (A: nat) (B: Poly.t): Poly.t := *)
  (*   Vector.t (s B) A ▹ Vector.Σ (π B). *)

  Module CombinatorNotations.
    Infix "·" := app (at level 30).
    Infix "+" := sum.
    Infix "*" := prod.
  End CombinatorNotations.
End Combinator.

Module Free.
  Import Combinator.
  Import CombinatorNotations.

  Definition t F P A := W.t F (const A + P).

  Definition map {F P} {A B: Set} (f: A → B): t F P A → t F P B :=
    W.map
      {|
        Mor.s (x: Poly.s (const A + P)) := match x return Poly.s (const B + P) with
                                           | inl x => inl (f x)
                                           | inr x => inr x
                                           end ;
        Mor.π x := match x with
                   | inl _ => id
                   | inr _ => id
                   end
      |}.

  Definition pure {F P} {A: Set} (a: A): t F P A := W.roll (Ext.con (const A + _) (inl a) Container.nil).

  Module Stream.
    Definition t := t stream.
  End Stream.

  Module Term.
    Definition t := t term.
  End Term.

  Definition join {F P} {A: Set}: t F P (t F P A) → t F P A.
  Proof.
    refine (W.fold (λ e, _)).
    assert (f := Ext.field e).
    destruct (Ext.tag e).
    - exact s.
    - exact (W.roll (Ext.con (const _ + P) (inr s) f)).
  Defined.

  Definition bind {F P} {A B: Set} (f: A → t F P B): t F P A → t F P B.
  Proof.
    refine (W.fold (λ e, _)).
    assert (g := Ext.field e).
    destruct (Ext.tag e).
    - exact (f s).
    - exact (W.roll (Ext.con (const _ + P) (inr s) g)).
  Defined.

  (* hack to work around type checking *)
  Definition dobind {F P} {A B: Set} (e: t F P A) (f: A → t F P B): t F P B := bind f e.

  Module FreeNotations.
    Declare Scope term_scope.
    Delimit Scope term_scope with t.

    Bind Scope term_scope with t.
    Bind Scope term_scope with Stream.t.
    Bind Scope term_scope with Term.t.

    Notation "'do' n ← e0 ; e1" :=
      (dobind e0 (λ n, e1))
        (n pattern, at level 200, left associativity): term_scope.
  End FreeNotations.
End Free.

Module M := FMapAVL.Make OrderedTypeEx.Nat_as_OT.

Definition map (P: Poly.t) := M.t (Free.t stream P nat).
Fixpoint Forall {n}: (Fin.t n → bool) → bool :=
  match n with
  | S n' => λ f, f Fin.F1 && Forall (λ x, f (Fin.FS x))
  | _ => λ _, true
  end%bool.

Fixpoint Exists {n}: (Fin.t n → bool) → bool :=
  match n with
  | S n' => λ f, f Fin.F1 || Forall (λ x, f (Fin.FS x))
  | _ => λ _, false
  end%bool.

Definition state X A := X → option (X * A).

Definition ret {X A} (a: A): state X A := λ x, Some (x, a).
Definition bind {X A B} (e: state X A) (f: A → state X B): state X B :=
  λ x,
    if e x is Some (x', a)
    then
      f a x'
    else
      None.

Definition get {X}: state X X := λ x, Some (x, x).
Definition put {X} (x:X): state X unit := λ _, Some (x, tt).

Definition error {X A}: state X A := λ _, None.

(* FIXME do well founded recursion *)
Unset Guard Checking.
Section state.
  Import Combinator.CombinatorNotations.

  (* Mostly ripped from https://eli.thegreenplace.net/2018/unification/
     any errors are mine
   *)
  Notation "'do' n ← e0 ; e1" :=
    (bind e0 (λ n, e1))
      (n pattern, at level 200, left associativity): list_scope.

  Context {P: Poly.t}.
  Variable eq: ∀ x y: Poly.s P, {x = y} + {x ≠ y}.

  Definition lookup K: state (map P) (option _) :=
    do m ← get ;
    ret (M.find K m).

  Definition set K V: state (map P) unit :=
    do m ← get ;
    put (M.add K V m).

  Fixpoint iter {X} {n}: (Fin.t n → state X unit) → state X unit :=
    match n with
    | S n' => λ f,
        do _ ← f Fin.F1 ;
        iter (λ x, f (Fin.FS x))
    | _ => λ _, ret tt
    end.

  Fixpoint occurs (x: nat) (e: Free.t stream P nat) (m: map P) {struct e} : bool :=
    match Ext.tag (W.unroll e) with
    | inl y =>
        if Nat.eqb x y
        then
          false
        else
          if M.find y m is Some e'
          then
            occurs x e' m
          else
            true
    | inr s =>
        Exists (λ h, occurs x (Ext.field (W.unroll e)  h) m)
    end.

  Fixpoint unify (e1 e2: Free.t stream P nat) {struct e1}: state (map P) unit :=
      match Ext.tag (W.unroll e1), Ext.tag (W.unroll e2) with
      | inl x, inl y =>
          if Nat.eqb x y
          then
            ret tt
          else
            unify_var x (Free.pure y)
      | inl x, _ =>
          unify_var x e2
      | _, inl x =>
          unify_var x e1
      | inr s, inr t =>
          if eq s t is left p
          then
            iter (λ x, unify (Ext.field (W.unroll e1) x) (Ext.field (W.unroll e1) (rew [λ x, Fin.t (Poly.π (_ + P) _)] p in x)))
          else
            error
      end
  with unify_var (x: nat) (e: Free.t stream P nat) {struct e}: state (map P) unit :=
     do m ← lookup x ;
     if m is Some e'
     then
       unify e' e
     else
       if Ext.tag (W.unroll e) is inl y
       then
         do m' ← lookup y ;
         if m' is Some e'
         then
           unify_var x e'
         else
           error
       else
         do m ← get ;
         if occurs x e m
         then
           error
         else
           set x e
  .
End state.

Set Guard Checking.

Inductive prop {P: Poly.t} {V:Set}: Set :=
| true
| false
| and (P Q: prop)
| or (P Q: prop)

| eq (e e': Free.t term P V)
.

Arguments prop: clear implicits.

Notation "⊤" := true.
Notation "⊥" := false.

Infix "∧" := and.
Infix "∨" := or.

Infix "=" := eq.
Notation "x ≠ y" := (¬ (x = y)).

Inductive formula {P: Poly.t} {V: Set}: nat → Set :=
| colim {n} (P: Free.t term P V → formula n): formula (S n)
| prove (P: prop P V): formula 0.

Arguments formula: clear implicits.

Coercion prove: prop >-> formula.

Notation "'∃' x .. y , P" := (colim (fun x => .. (colim (fun y => P)) ..)).

Notation "'do' n ← e0 ; e1" :=
  (List.flat_map (λ n, e1) e0)
    (n pattern, at level 200, left associativity): list_scope.

Fixpoint solve {P: Poly.t} eq (e: prop P nat) (m: map P): list (map P) :=
  match e with
  | ⊤ => [m]
  | ⊥ => []
  | P ∧ Q =>
      do P' ← solve eq P m ;
      solve eq Q P'
  | P ∨ Q => solve eq P m ++ solve eq Q m
  | e = e' =>
      if unify eq (W.stream e) (W.stream e') m is Some (m', tt)
      then [m']
      else []
  end%list.

Fixpoint instance {P: Poly.t} {n} (e: formula P nat n): prop P nat :=
  match e with
  | prove P => P
  | @colim _ _ n' P => instance (P (Free.pure n'))
  end.

Definition eval {P: Poly.t} eq {n} (e: formula P nat n) :=
  List.map (λ x, M.elements (M.map W.term x)) (solve eq (instance e) (M.empty _)).

Variant PEANO := ZERO | SUCC | ADD.
Definition peano := PEANO ▹ λ x, match x with
                                 | ZERO => 0
                                 | SUCC => 1
                                 | ADD => 2
                                 end.

Module Import PeanoNotations.
  Import Combinator.
  Import CombinatorNotations.

  Notation "'O'" := (W.Term.roll (Ext.con (const _ + peano) (inr ZERO) [])).
  Notation "'S' e" := (W.Term.roll (Ext.con (const _ + peano) (inr SUCC) [e])) (at level 1).
  Notation "e + e'" := (W.Term.roll (Ext.con (const _ + peano) (inr ADD) [e; e'])).
End PeanoNotations.

Import Free.FreeNotations.

Fixpoint ofnat {A} (n: nat): Free.t term peano A :=
  if n is Datatypes.S n'
  then
    S (ofnat n')
  else
    O.

Definition PEANO_dec (x y: Poly.s peano): {x = y} + {x ≠ y}.
Proof.
  decide equality.
Defined.

(* still very buggy :| *)
Example foo {V}: formula _ V _ := ∃ x y z, x = O + O.

Compute eval PEANO_dec foo.