------------------------------------------------------------------------
-- Substitutions
------------------------------------------------------------------------

module SimplyTyped.Substitution where

open import SimplyTyped.TypeSystem
open import SimplyTyped.ContextExtension
import SimplyTyped.Environment as Env
open import Function hiding (_∋_)

------------------------------------------------------------------------
-- General substitution machinery

-- Uses idea from "Type-Preserving Renaming and Substitution" by Conor
-- McBride to enable defining both term and variable substitutions
-- without having to write several functions which traverse terms.

-- Definition of what a substitution is.

module SubstDef (_•_ : TmLike) -- The substitution replaces variables
                               -- with this kind of "term".
                where

  -- Substitutions of type Γ ⇒ Δ, when applied, take things with
  -- variables in Γ and give things with variables in Δ. (This concept
  -- is sometimes written with the arrow in the other direction.)

  -- Substitutions are defined in terms of environments.

  infix  _⇒_

  private
    module Dummy {Δ : Ctxt} where
      open Env (_•_ Δ) public
  open Dummy public using (∅; _▷_; lookup)

  _⇒_ : Ctxt → Ctxt → Set
  Γ ⇒ Δ = Dummy.Env {Δ} Γ

  -- The type of functions applying substitutions to things.

  Applier : TmLike → TmLike → Set
  Applier _•₁_ _•₂_ = ∀ {Γ Δ τ} → Γ •₁ τ → Γ ⇒ Δ → Δ •₂ τ

  -- You get lots of substitution machinery defined for you (see
  -- below) if you define a SubstKit, ...

  record SubstKit : Set where
    field
      vr     : ∀ {Γ τ}   → Γ ∋ τ → Γ • τ
      weaken : ∀ {Γ σ τ} → Γ • τ → (Γ ▻ σ) • τ
      tm     : ∀ {Γ τ}   → Γ • τ → Γ ⊢ τ

  -- ...and more if you define a SubstKit⁺.

  record SubstKit⁺ : Set where
    field
      kit  : SubstKit
      _/•_ : Applier _•_ _•_

open SubstDef using (SubstKit; SubstKit⁺) hiding (module SubstKit)

-- Various substitution-related functions.

module Subst {_•_ : TmLike}
             (kit : SubstKit _•_)
             where

  open SubstDef _•_ public
  open SubstKit kit

  infix   _↑
  infix   _[_]

  -- Weakening of substitutions.

  wk⇒ : ∀ {Γ Δ σ} → Γ ⇒ Δ → Γ ⇒ Δ ▻ σ
  wk⇒ ∅       = ∅
  wk⇒ (ρ ▷ t) = wk⇒ ρ ▷ weaken t

  -- Lifting of substitutions.

  _↑ : ∀ {Γ Δ σ} → Γ ⇒ Δ → Γ ▻ σ ⇒ Δ ▻ σ
  ρ ↑ = wk⇒ ρ ▷ vr vz

  -- Identity.

  idˢ : ∀ {Γ} → Γ ⇒ Γ
  idˢ {Γ = ε}     = ∅
  idˢ {Γ = Γ ▻ τ} = idˢ ↑

  -- Weakening.

  wk : ∀ {Γ σ} → Γ ⇒ Γ ▻ σ
  wk = wk⇒ idˢ

  -- Substitution of a single variable.

  sub : ∀ {Γ τ} → Γ • τ → Γ ▻ τ ⇒ Γ
  sub t = idˢ ▷ t

  -- Substitution application for terms.

  _[_] : Applier _⊢_ _⊢_
  var x   [ ρ ] = tm (lookup x ρ)
  ƛ t     [ ρ ] = ƛ (t [ ρ ↑ ])
  t₁ · t₂ [ ρ ] = (t₁ [ ρ ]) · (t₂ [ ρ ])

  -- Multiple weakenings.

  wk⇒⋆ : ∀ {Γ Δ} Δ⁺ → Γ ⇒ Δ → Γ ⇒ Δ ▻▻ Δ⁺
  wk⇒⋆ ε        ρ = ρ
  wk⇒⋆ (Δ⁺ ▻ σ) ρ = wk⇒ (wk⇒⋆ Δ⁺ ρ)

  wk⋆ : ∀ {Γ} Γ⁺ → Γ ⇒ Γ ▻▻ Γ⁺
  wk⋆ Γ⁺ = wk⇒⋆ Γ⁺ idˢ

-- Subst⁺ defines composition. This function uses a stronger
-- substitution kit, which is sometimes defined using the Subst module
-- (see the examples below); this is the reason for placing
-- composition in a separate module.

module Subst⁺ {_•_  : TmLike}
              (kit⁺ : SubstKit⁺ _•_)
              where

  open SubstDef.SubstKit⁺ kit⁺
  open Subst kit public

  infixl  _∘ˢ_

  _∘ˢ_ : ∀ {Γ Δ X} → Γ ⇒ Δ → Δ ⇒ X → Γ ⇒ X
  ∅        ∘ˢ ρ₂ = ∅
  (ρ₁ ▷ t) ∘ˢ ρ₂ = ρ₁ ∘ˢ ρ₂ ▷ (t /• ρ₂)

------------------------------------------------------------------------
-- Instantiations of the general structures above

-- Variables.

varKit : SubstKit _∋_
varKit = record
  { vr     = id
  ; weaken = vs
  ; tm     = var
  }

varKit⁺ : SubstKit⁺ _∋_
varKit⁺ = record { kit  = varKit
                 ; _/•_ = Subst.lookup varKit
                 }

module VarSubst where
  open Subst⁺ varKit⁺ public

-- Terms.

tmKit : SubstKit _⊢_
tmKit = record
  { vr     = var
  ; weaken = λ t → t [ wk ]
  ; tm     = id
  }
  where open VarSubst

tmKit⁺ : SubstKit⁺ _⊢_
tmKit⁺ = record { kit  = tmKit
                ; _/•_ = Subst._[_] tmKit
                }

module TmSubst where
  open Subst⁺ tmKit⁺ public