-- Andreas Abel, 2019-10-30,
-- Parallel hereditary substitutions.
-- A riddle posed to me by Conor McBride during TYPES 2018.

-- This solution uses Agda's sized types to type substitutions
-- whose non-variable entries are bounded in the order of their type.
-- This bound is the main component of the termination measure
-- of parallel herediary substitutions.

-- Agda standard library imports

open import Size
open import Data.Product using (∃; _,_)

-- Lists (for typing contexts)
open import Data.List using (List; []; _∷_)

-- List membership _∈_ (for variables)
open import Data.List.Membership.Propositional using (_∈_)
open import Data.List.Relation.Unary.Any using (here; there)
open import Relation.Binary.PropositionalEquality using (refl)

pattern here! = here refl

-- Sublists _⊆_ (for weakening)
import Data.List.Relation.Binary.Sublist.Propositional as Sublist
open Sublist using (_⊆_; _∷_; _∷ʳ_; ⊆-refl)

wk1 : ∀ {a} {A : Set a} {x : A} {xs} → xs ⊆ x ∷ xs
wk1 = _ ∷ʳ ⊆-refl

-- Types are indexed by an upper bound on their order.

data Ty : (i : Size) → Set where
  o   : ∀{i} → Ty i
  _⇒_ : ∀{i} (a : Ty i) (b : Ty (↑ i)) → Ty (↑ i)

-- Agda's subtyping ensures the size index is just an upper bound.

test-Ty↑ : ∀{i} → Ty i → Ty (↑ i)
test-Ty↑ a = a

-- We do not expose the order of types in the type of the context.

TY = ∃ Ty

-- Typing contexts

Cxt = List TY

-- Normal forms

mutual

  -- Introductions and neutrals

  data Nf (Γ : Cxt) : ∀ i → Ty i → Set where

    ƛ   : ∀{i} {a : Ty i} {b : Ty (↑ i)}
        → (t : Nf ((i , a) ∷ Γ) (↑ i) b)
        → Nf Γ (↑ i) (a ⇒ b)

    app : ∀{i} {a : Ty i} {j} {b : Ty j}
        → (x : (i , a) ∈ Γ)
        → (s : Spine Γ i a j b)
        → Nf Γ j b

  -- Eliminations

  data Spine (Γ : Cxt) : ∀ i (a : Ty i) k (c : Ty k) → Set where
    []  : ∀{i} {c : Ty i} → Spine Γ i c i c
    _∷_ : ∀{i} {a : Ty i} {b : Ty (↑ i)} {k} {c : Ty k}
        → (t : Nf Γ i a)
        → (s : Spine Γ (↑ i) b k c)
        → Spine Γ (↑ i) (a ⇒ b) k c

-- Spine concatenation

_++_ : ∀{Γ i a j b k c} → Spine Γ i a j b → Spine Γ j b k c → Spine Γ i a k c
[]      ++ s' = s'
(t ∷ s) ++ s' = t ∷ (s ++ s')

-- Substitution entries

data Entry (i : Size) (Γ : Cxt) : ∀ j → Ty j → Set where
  var : ∀{j} {a : Ty j} (x : (j , a) ∈ Γ) → Entry i Γ j a
  nf  : ∀    {a : Ty i} (t : Nf Γ i a)    → Entry i Γ i a

-- Weakening of syntax

mutual

  weakNf : ∀{Γ Δ} (τ : Γ ⊆ Δ) {i} {a : Ty i}
         → Nf Γ i a
         → Nf Δ i a
  weakNf τ (ƛ t) = ƛ (weakNf (refl ∷ τ) t)
  weakNf τ (app x s) = app (Sublist.lookup τ x) (weakSp τ s)

  weakSp : ∀{Γ Δ} (τ : Γ ⊆ Δ) {i} {a : Ty i} {k} {c : Ty k}
         → Spine Γ i a k c
         → Spine Δ i a k c
  weakSp τ []      = []
  weakSp τ (t ∷ s) = weakNf τ t ∷ weakSp τ s

weakEntry : ∀{Γ Δ} (τ : Γ ⊆ Δ) {i j a }
        → Entry i Γ j a
        → Entry i Δ j a
weakEntry τ (var x) = var (Sublist.lookup τ x)
weakEntry τ (nf t)  = nf  (weakNf τ t)

-- Substitutions bounded by an order

data Subst i : ∀ (Γ Δ : Cxt) → Set where
  []   : ∀{Γ} → Subst i Γ []
  id   : ∀{Γ : Cxt} → Subst i Γ Γ

  lift : ∀{Γ Δ}
       → {a : TY}
       → (σ : Subst i Γ Δ)
       → Subst i (a ∷ Γ) (a ∷ Δ)

  _∷_  : ∀{Γ Δ} {a : Ty i}
       → (u : Nf Γ i a)
       → (σ : Subst i Γ Δ)
       → Subst i Γ ((i , a) ∷ Δ)

-- Substitutions (Subst i Γ Δ) are morally lists of entries
-- (All (uncurry (Entry i Γ)) Δ).

lookup : ∀{i Γ Δ} (σ : Subst i Γ Δ) {j} {a : Ty j} (x : (j , a) ∈ Δ) → Entry i Γ j a
lookup id       x         = var x
lookup (lift σ) here!     = var here!
lookup (lift σ) (there x) = weakEntry wk1 (lookup σ x)
lookup (u ∷ σ)  here!     = nf u
lookup (u ∷ σ)  (there x) = lookup σ x

-- Parallel hereditary substitutions and accumulating application of Nfs

mutual

  -- Substitution into a normal form

  subst : ∀{i Γ j a Δ} (t : Nf Γ j a) (σ : Subst i Δ Γ) → Nf Δ j a
  subst (ƛ t)     σ = ƛ (subst t (lift σ))
  subst (app x s) σ = applyEntry (lookup σ x) (substSp s σ)

  -- Substitution into a spine (pointwise)

  substSp : ∀{i Γ Δ j a k c} (s : Spine Γ j a k c) (σ : Subst i Δ Γ) → Spine Δ j a k c
  substSp []      σ = []
  substSp (u ∷ s) σ = subst u σ ∷ substSp s σ

  -- Applying a substitution entry to a spine

  applyEntry : ∀ {i Γ j a k c} (r : Entry i Γ j a) (s : Spine Γ j a k c) → Nf Γ k c
  applyEntry (var x) s = app x s
  applyEntry (nf t)  s = apply t s

  -- Application of a normal form to a spine

  apply : ∀ {Γ i a j b} (t : Nf Γ i a) (s : Spine Γ i a j b) → Nf Γ j b
  apply (app x s) s'      = app x (s ++ s')
  apply (ƛ t)     []      = ƛ t
  apply (ƛ t)     (u ∷ s) = applyClos t (u ∷ id) s

  -- Accumulating application

  applyClos : ∀ {i} {Γ} {Δ} {b} {k} {c}
            → (t : Nf Δ (↑ i) b)
            → (σ : Subst i Γ Δ)
            → (s : Spine Γ (↑ i) b k c)
            → Nf Γ k c
  applyClos (app x s) σ s'      = applyEntry (lookup σ x) (substSp s σ ++ s')
  applyClos (ƛ t)     σ []      = ƛ (subst t (lift σ))
  applyClos (ƛ t)     σ (u ∷ s) = applyClos t (u ∷ σ) s