module STLC where

open import Prelude using (DecidableEq; _≟_)

open import Data.String using (String)
open import Data.List using (List; [_]; _++_; filter)
open import Data.Product.Base using (_×_; _,_; ∃; ∃-syntax; proj₂)
open import Function.Base using (case_of_)
open import Function.Bundles using (_⇔_; Equivalence)
open import Relation.Nullary using (Dec; yes; no; ¬?; contradiction)
open import Relation.Nullary.Negation.Core using (¬_)
open import Relation.Binary.Definitions using (DecidableEquality)
open import Relation.Binary.PropositionalEquality using (_≡_; _≢_; refl; trans; cong; ≢-sym)

open Equivalence using (from; to; to-cong; from-cong)

Name : Set
Name = String

infix   ``_
infixr  _⇒_

data Type : Set where
  ``_ : Name → Type
  _⇒_ : Type → Type → Type

⇒-cong₁ : ∀ {A B C} → A ≡ B → A ⇒ C ≡ B ⇒ C
⇒-cong₁ refl = refl

⇒-cong₂ : ∀ {A B C} → A ≡ B → C ⇒ A ≡ C ⇒ B
⇒-cong₂ refl = refl

⇒-inj₁ : ∀ {A B C D} → A ⇒ B ≡ C ⇒ D → A ≡ C
⇒-inj₁ refl = refl

⇒-inj₂ : ∀ {A B C D} → A ⇒ B ≡ C ⇒ D → B ≡ D
⇒-inj₂ refl = refl

``≢⇒ : ∀ {A B x} → `` x ≢ A ⇒ B
``≢⇒ ()

infixr  ƛ_∶_⇒_
infixl  _·_
infix   `_

data Term : Set where
  `_     : Name → Term
  ƛ_∶_⇒_ : Name → Type → Term → Term
  _·_    : Term → Term → Term

FV : Term → List Name
FV (` x) = [ x ]
FV (ƛ x ∶ A ⇒ M) = filter (λ y → ¬? (x ≟ y)) (FV M)
FV (M · N) = FV M ++ FV N

infixl  _,_∶_

data Context : Set where
  ∅     : Context
  _,_∶_ : Context → Name → Type → Context

infix  _∋_∶_

data _∋_∶_ : Context → Name → Type → Set where
  here  : ∀ {Γ x A} → Γ , x ∶ A ∋ x ∶ A
  there : ∀ {Γ x y A B}
    → x ≢ y
    -----------
    → Γ ∋ x ∶ A
    -------------------
    → Γ , y ∶ B ∋ x ∶ A

infix  _⊢_∶_

data _⊢_∶_ : Context → Term → Type → Set where
  ⊢`_ : ∀ {Γ x A}
    → Γ ∋ x ∶ A
    -------------
    → Γ ⊢ ` x ∶ A

  ⊢ƛ_ : ∀ {Γ x A B M}
    → Γ , x ∶ A ⊢ M ∶ B
    -----------------------------
    → Γ ⊢ (ƛ x ∶ A ⇒ M) ∶ (A ⇒ B)

  ⊢_·_ : ∀ {Γ A B M N}
    → Γ ⊢ M ∶ (A ⇒ B)
    -----------------
    → Γ ⊢ N ∶ A
    -----------------
    → Γ ⊢ (M · N) ∶ B

`-gen : ∀ {Γ x A} → (Γ ⊢ ` x ∶ A) ⇔ (Γ ∋ x ∶ A)
`-gen .to (⊢` x) = x
`-gen .from x = ⊢` x
`-gen .to-cong = cong (`-gen .to)
`-gen .from-cong = cong (`-gen .from)

·-gen : ∀ {Γ M N B} → (Γ ⊢ (M · N) ∶ B) ⇔ (∃[ A ] (Γ ⊢ M ∶ (A ⇒ B)) × (Γ ⊢ N ∶ A))
·-gen .to (⊢ M · N) = _ , M , N
·-gen .from (_ , M , N) = ⊢ M · N
·-gen .to-cong = cong (·-gen .to)
·-gen .from-cong = cong (·-gen .from)

ƛ-gen : ∀ {Γ x A C M} → (Γ ⊢ (ƛ x ∶ A ⇒ M) ∶ C) ⇔ (∃[ B ] (((Γ , x ∶ A) ⊢ M ∶ B) × (C ≡ (A ⇒ B))))
ƛ-gen .to (⊢ƛ M) = _ , M , refl
ƛ-gen .from (_ , M , refl) = ⊢ƛ M
ƛ-gen .to-cong = cong (ƛ-gen .to)
ƛ-gen .from-cong = cong (ƛ-gen .from)

Type-deq : DecidableEquality Type
Type-deq (`` x) (`` y) with x ≟ y
... | yes refl = yes refl
... | no n     = no λ{refl → n refl}
Type-deq (`` x) (A ⇒ B) = no λ()
Type-deq (A ⇒ B) (`` x) = no λ()
Type-deq (A ⇒ B) (A’ ⇒ B’) with Type-deq A A’ | Type-deq B B’
... | no na    | _        = no λ{refl → na refl}
... | yes _    | no nb    = no λ{refl → nb refl}
... | yes refl | yes refl = yes refl

instance
  DecidableEq-Type : DecidableEq Type
  DecidableEq-Type = record { _≟_ = Type-deq }

uniq-∋ : ∀ {Γ x A B} → Γ ∋ x ∶ A → Γ ∋ x ∶ B → A ≡ B
uniq-∋ here here               = refl
uniq-∋ here (there nx _)       = contradiction refl nx
uniq-∋ (there nx _) here       = contradiction refl nx
uniq-∋ (there _ x) (there _ y) = uniq-∋ x y

uniq : ∀ {Γ M A B} → Γ ⊢ M ∶ A → Γ ⊢ M ∶ B → A ≡ B
uniq (⊢` ∋x) (⊢` ∋y)       = uniq-∋ ∋x ∋y
uniq (⊢ƛ MA) (⊢ƛ MB)       = ⇒-cong₂ (uniq MA MB)
uniq (⊢ MA · _) (⊢ MB · _) = ⇒-inj₂ (uniq MA MB)

ext-∋ : ∀ {Γ x y B}
  → x ≢ y
  → ¬ (∃[ A ] Γ ∋ x ∶ A)
  → ¬ (∃[ A ] Γ , y ∶ B ∋ x ∶ A)
ext-∋ x≢y _ (_ , here)       = contradiction refl x≢y
ext-∋ _  ¬∃ (A , there _ ∋x) = ¬∃ (A , ∋x)

lookup : ∀ Γ x → Dec (∃[ A ] Γ ∋ x ∶ A)
lookup ∅ x = no λ()
lookup (Γ , y ∶ B) x with x ≟ y
... | yes refl = yes (B , here)
... | no nx with lookup Γ x
...   | yes (A , x) = yes (A , there nx x)
...   | no nex      = no (ext-∋ nx nex)

type-infer : ∀ Γ M → Dec (∃[ A ] Γ ⊢ M ∶ A)
type-infer Γ (` x) with lookup Γ x
... | no ¬∃        = no λ { (A , (⊢` ∋x)) → ¬∃ (A , ∋x) }
... | yes (A , ∋x) = yes (A , (⊢` ∋x))
type-infer Γ (ƛ x ∶ A ⇒ M) with type-infer (Γ , x ∶ A) M
... | no ¬MB       = no λ { (_ ⇒ B , (⊢ƛ MB)) → ¬MB (B , MB) }
... | yes (B , MB) = yes (A ⇒ B , (⊢ƛ MB))
type-infer Γ (M · N) with type-infer Γ M | type-infer Γ N
... | no _    | no ¬NA  = no λ { (_ , ⊢ _ · NA) → ¬NA (_ , NA) }
... | yes _   | no ¬NA  = no λ { (_ , ⊢ _ · NA) → ¬NA (_ , NA) }
... | no ¬MAB | yes _   = no λ { (_ , ⊢ MAB · _) → ¬MAB (_ , MAB) }
... | yes (`` _ ,  MA)  | yes _ = no λ { (_ , (⊢ MAB · _)) → contradiction (uniq MA MAB) ``≢⇒ }
... | yes (A ⇒ B , MAB) | yes (A' , NA') with A ≟ A'
...   | no A≢A' = no λ { (_ , (⊢ MA''B · NA'')) → A≢A' (trans (⇒-inj₁ (uniq MAB MA''B)) (uniq NA'' NA')) }
...   | yes refl = yes (B , ⊢ MAB · NA')

type-check : ∀ Γ M A → Dec (Γ ⊢ M ∶ A)
type-check Γ (` x) A with lookup Γ x
... | no ¬∃ = no λ xA → ¬∃ (A , `-gen .to xA)
... | yes (A’ , ∋x) with A’ ≟ A
...   | no A’≢A  = no λ xA → A’≢A (uniq-∋ ∋x (`-gen .to xA))
...   | yes refl = yes (⊢` ∋x)
type-check Γ (ƛ x ∶ A’ ⇒ M) N with N
... | (`` B) = no λ()
... | (A ⇒ B) with A’ ≟ A
...   | no  A’≢A = no λ{ (⊢ƛ _) → A’≢A refl }
...   | yes refl with type-check (Γ , x ∶ A) M B
...     | no ¬MB = no λ { (⊢ƛ MB) → ¬MB MB }
...     | yes MB = yes (⊢ƛ MB)
type-check Γ (M · N) B with type-infer Γ M
... | no ¬MAB         = no λ { (⊢ MAB · _) → ¬MAB (_ , MAB) }
... | yes (`` _ , MA) = no λ { (⊢ MAB · _) → contradiction (uniq MA MAB) ``≢⇒ }
... | yes (A ⇒ B' , MAB') with B' ≟ B
...   | no B'≢B = no λ { (⊢ MA'B · _) → B'≢B (⇒-inj₂ (uniq MAB' MA'B)) }
...   | yes refl with type-check Γ N A
...     | no ¬NA = no λ { (⊢ MA'B · NA') → case ⇒-inj₁ (uniq MAB' MA'B) of λ { refl → ¬NA NA' }}
...     | yes NA = yes (⊢ MAB' · NA)