------------------------------------------------------------------------
-- The Agda standard library
--
-- Annotated reflected syntax.
--
-- NOTE: This file does not check under --cubical-compatible due to
--       restrictions in the termination checker. In particular
--       recursive functions over a universe of types is not supported
--       by --cubical-compatible.
------------------------------------------------------------------------

{-# OPTIONS --safe --with-K #-}

module Reflection.AnnotatedAST where

open import Level                        using (Level; 0ℓ; suc; _⊔_)
open import Effect.Applicative           using (RawApplicative)
open import Data.Bool.Base               using (Bool; false; true; if_then_else_)
open import Data.List.Base               using (List; []; _∷_)
open import Data.List.Relation.Unary.All using (All; _∷_; [])
open import Data.Product.Base            using (_×_; _,_; proj₁; proj₂)
open import Data.String.Base             using (String)

open import Reflection                   hiding (pure)
open import Reflection.AST.Universe

open Clause
open Pattern
open Sort

------------------------------------------------------------------------
-- Annotations and annotated types

-- An Annotation is a type indexed by a reflected term.
Annotation : ∀ ℓ → Set (suc ℓ)
Annotation ℓ = ∀ {u} → ⟦ u ⟧ → Set ℓ

-- An annotated type is a family over an Annotation and a reflected term.
Typeₐ : ∀ ℓ → Univ → Set (suc (suc ℓ))
Typeₐ ℓ u = Annotation ℓ → ⟦ u ⟧ → Set (suc ℓ)

private
  variable
    ℓ             : Level
    u             : Univ
    Ann Ann₁ Ann₂ : Annotation ℓ

-- ⟪_⟫ packs up an element of an annotated type with a top-level annotation.
infixr  ⟨_⟩_
data ⟪_⟫ {u} (Tyₐ : Typeₐ ℓ u) : Typeₐ ℓ u where
  ⟨_⟩_ : ∀ {t} → Ann t → Tyₐ Ann t → ⟪ Tyₐ ⟫ Ann t

ann : {Tyₐ : Typeₐ ℓ u} {t : ⟦ u ⟧} → ⟪ Tyₐ ⟫ Ann t → Ann t
ann (⟨ α ⟩ _) = α


------------------------------------------------------------------------
-- Annotated reflected syntax

-- Annotated lists are lists (All) of annotated values. No top-level
-- annotation added, since we don't want an annotation for every tail
-- of a list. Instead a top-level annotation is added on the outside.
-- See Argsₐ.
Listₐ : Typeₐ ℓ u → Typeₐ ℓ (⟨list⟩ u)
Listₐ Tyₐ Ann = All (Tyₐ Ann)

-- We define the rest of the annotated types in two variants:
-- The primed variant which has annotations on subterms, and the
-- non-primed variant which adds a top-level annotation to the primed
-- one.

data Absₐ′ (Tyₐ : Typeₐ ℓ u) : Typeₐ ℓ (⟨abs⟩ u) where
  abs : ∀ {t} x → Tyₐ Ann t → Absₐ′ Tyₐ Ann (abs x t)

Absₐ : Typeₐ ℓ u → Typeₐ ℓ (⟨abs⟩ u)
Absₐ Tyₐ = ⟪ Absₐ′ Tyₐ ⟫

data Argₐ′ (Tyₐ : Typeₐ ℓ u) : Typeₐ ℓ (⟨arg⟩ u) where
  arg : ∀ {t} i → Tyₐ Ann t → Argₐ′ Tyₐ Ann (arg i t)

Argₐ : Typeₐ ℓ u → Typeₐ ℓ (⟨arg⟩ u)
Argₐ Tyₐ = ⟪ Argₐ′ Tyₐ ⟫

data Namedₐ′ (Tyₐ : Typeₐ ℓ u) : Typeₐ ℓ (⟨named⟩ u) where
  _,_ : ∀ {t} x → Tyₐ Ann t → Namedₐ′ Tyₐ Ann (x , t)
infixr  _,_

Namedₐ : Typeₐ ℓ u → Typeₐ ℓ (⟨named⟩ u)
Namedₐ Tyₐ = ⟪ Namedₐ′ Tyₐ ⟫

-- Add a top-level annotation for Args.
Argsₐ : Typeₐ ℓ u → Typeₐ ℓ (⟨list⟩ (⟨arg⟩ u))
Argsₐ Tyₐ = ⟪ Listₐ (Argₐ Tyₐ) ⟫

mutual
  Termₐ : Typeₐ ℓ ⟨term⟩
  Termₐ = ⟪ Termₐ′ ⟫

  Patternₐ : Typeₐ ℓ ⟨pat⟩
  Patternₐ = ⟪ Patternₐ′ ⟫

  Sortₐ : Typeₐ ℓ ⟨sort⟩
  Sortₐ = ⟪ Sortₐ′ ⟫

  Clauseₐ : Typeₐ ℓ ⟨clause⟩
  Clauseₐ = ⟪ Clauseₐ′ ⟫

  Clausesₐ : Typeₐ ℓ (⟨list⟩ ⟨clause⟩)
  Clausesₐ = ⟪ Listₐ Clauseₐ ⟫

  Telₐ : Typeₐ ℓ ⟨tel⟩
  Telₐ = ⟪ Listₐ (Namedₐ (Argₐ Termₐ)) ⟫

  data Termₐ′ {ℓ} : Typeₐ ℓ ⟨term⟩ where
    var       : ∀ x {args} → Argsₐ Termₐ Ann args → Termₐ′ Ann (var x args)
    con       : ∀ c {args} → Argsₐ Termₐ Ann args → Termₐ′ Ann (con c args)
    def       : ∀ f {args} → Argsₐ Termₐ Ann args → Termₐ′ Ann (def f args)
    lam       : ∀ v {b}    → Absₐ  Termₐ Ann b    → Termₐ′ Ann (lam v b)
    pat-lam   : ∀ {cs args} → Clausesₐ Ann cs → Argsₐ Termₐ Ann args →
                  Termₐ′ Ann (pat-lam cs args)
    pi        : ∀ {a b} → Argₐ Termₐ Ann a → Absₐ Termₐ Ann b → Termₐ′ Ann (pi a b)
    agda-sort : ∀ {s} → Sortₐ Ann s → Termₐ′ Ann (agda-sort s)
    lit       : ∀ l → Termₐ′ Ann (lit l)
    meta      : ∀ x {args} → Argsₐ Termₐ Ann args → Termₐ′ Ann (meta x args)
    unknown   : Termₐ′ Ann unknown

  data Clauseₐ′ {ℓ} : Typeₐ ℓ ⟨clause⟩ where
    clause        : ∀ {tel ps t} → Telₐ Ann tel → Argsₐ Patternₐ Ann ps →
                      Termₐ Ann t → Clauseₐ′ Ann (clause tel ps t)
    absurd-clause : ∀ {tel ps} → Telₐ Ann tel → Argsₐ Patternₐ Ann ps →
                      Clauseₐ′ Ann (absurd-clause tel ps)

  data Sortₐ′ {ℓ} : Typeₐ ℓ ⟨sort⟩ where
    set     : ∀ {t} → Termₐ Ann t → Sortₐ′ Ann (set t)
    lit     : ∀ n → Sortₐ′ Ann (lit n)
    prop    : ∀ {t} → Termₐ Ann t → Sortₐ′ Ann (prop t)
    propLit : ∀ n → Sortₐ′ Ann (propLit n)
    inf     : ∀ n → Sortₐ′ Ann (inf n)
    unknown : Sortₐ′ Ann unknown

  data Patternₐ′ {ℓ} : Typeₐ ℓ ⟨pat⟩ where
    con    : ∀ c {ps} → Argsₐ Patternₐ Ann ps → Patternₐ′ Ann (con c ps)
    dot    : ∀ {t} → Termₐ Ann t → Patternₐ′ Ann (dot t)
    var    : ∀ x → Patternₐ′ Ann (var x)
    lit    : ∀ l → Patternₐ′ Ann (lit l)
    proj   : ∀ f → Patternₐ′ Ann (proj f)
    absurd : ∀ x → Patternₐ′ Ann (absurd x)


-- Mapping a code in the universe to its corresponding primed and
-- non-primed annotated type.

mutual
  Annotated′ : Typeₐ ℓ u
  Annotated′ {u = ⟨term⟩}    = Termₐ′
  Annotated′ {u = ⟨pat⟩}     = Patternₐ′
  Annotated′ {u = ⟨sort⟩}    = Sortₐ′
  Annotated′ {u = ⟨clause⟩}  = Clauseₐ′
  Annotated′ {u = ⟨list⟩ u}  = Listₐ Annotated
  Annotated′ {u = ⟨arg⟩ u}   = Argₐ′ Annotated
  Annotated′ {u = ⟨abs⟩ u}   = Absₐ′ Annotated
  Annotated′ {u = ⟨named⟩ u} = Namedₐ′ Annotated

  Annotated : Typeₐ ℓ u
  Annotated = ⟪ Annotated′ ⟫


------------------------------------------------------------------------
-- Annotating terms

-- An annotation function computes the top-level annotation given a
-- term annotated at all sub-terms.
AnnotationFun : Annotation ℓ → Set (suc ℓ)
AnnotationFun Ann = ∀ u {t : ⟦ u ⟧} → Annotated′ Ann t → Ann t


-- Given an annotation function we can do the bottom-up traversal of a
-- reflected term to compute an annotated version.
module _ (annFun : AnnotationFun Ann) where

  private
    annotated : {t : ⟦ u ⟧} → Annotated′ Ann t → Annotated Ann t
    annotated ps = ⟨ annFun _ ps ⟩ ps

  mutual
    annotate′ : (t : ⟦ u ⟧) → Annotated′ Ann t
    annotate′ {⟨term⟩}    (var x args)           = var x (annotate args)
    annotate′ {⟨term⟩}    (con c args)           = con c (annotate args)
    annotate′ {⟨term⟩}    (def f args)           = def f (annotate args)
    annotate′ {⟨term⟩}    (lam v t)              = lam v (annotate t)
    annotate′ {⟨term⟩}    (pat-lam cs args)      = pat-lam (annotate cs) (annotate args)
    annotate′ {⟨term⟩}    (pi a b)               = pi (annotate a) (annotate b)
    annotate′ {⟨term⟩}    (agda-sort s)          = agda-sort (annotate s)
    annotate′ {⟨term⟩}    (lit l)                = lit l
    annotate′ {⟨term⟩}    (meta x args)          = meta x (annotate args)
    annotate′ {⟨term⟩}    unknown                = unknown
    annotate′ {⟨pat⟩}     (con c ps)             = con c (annotate ps)
    annotate′ {⟨pat⟩}     (dot t)                = dot (annotate t)
    annotate′ {⟨pat⟩}     (var x)                = var x
    annotate′ {⟨pat⟩}     (lit l)                = lit l
    annotate′ {⟨pat⟩}     (proj f)               = proj f
    annotate′ {⟨pat⟩}     (absurd x)             = absurd x
    annotate′ {⟨sort⟩}    (set t)                = set (annotate t)
    annotate′ {⟨sort⟩}    (lit n)                = lit n
    annotate′ {⟨sort⟩}    (prop t)               = prop (annotate t)
    annotate′ {⟨sort⟩}    (propLit n)            = propLit n
    annotate′ {⟨sort⟩}    (inf n)                = inf n
    annotate′ {⟨sort⟩}    unknown                = unknown
    annotate′ {⟨clause⟩}  (clause tel ps t)      = clause (annotate tel) (annotate ps) (annotate t)
    annotate′ {⟨clause⟩}  (absurd-clause tel ps) = absurd-clause (annotate tel) (annotate ps)
    annotate′ {⟨abs⟩ u}   (abs x t)              = abs x (annotate t)
    annotate′ {⟨arg⟩ u}   (arg i t)              = arg i (annotate t)
    annotate′ {⟨list⟩ u}  []                     = []
    annotate′ {⟨list⟩ u}  (x ∷ xs)               = annotate x ∷ annotate′ xs
    annotate′ {⟨named⟩ u} (x , t)                = x , annotate t

    annotate : (t : ⟦ u ⟧) → Annotated Ann t
    annotate t = annotated (annotate′ t)


------------------------------------------------------------------------
-- Annotation function combinators

-- Mapping over annotations
mutual
  map′ : ∀ u → (∀ {u} {t : ⟦ u ⟧} → Ann₁ t → Ann₂ t) → {t : ⟦ u ⟧} → Annotated′ Ann₁ t → Annotated′ Ann₂ t
  map′ ⟨term⟩      f (var x args)         = var x (map f args)
  map′ ⟨term⟩      f (con c args)         = con c (map f args)
  map′ ⟨term⟩      f (def h args)         = def h (map f args)
  map′ ⟨term⟩      f (lam v b)            = lam v (map f b)
  map′ ⟨term⟩      f (pat-lam cs args)    = pat-lam (map f cs) (map f args)
  map′ ⟨term⟩      f (pi a b)             = pi (map f a) (map f b)
  map′ ⟨term⟩      f (agda-sort s)        = agda-sort (map f s)
  map′ ⟨term⟩      f (lit l)              = lit l
  map′ ⟨term⟩      f (meta x args)        = meta x (map f args)
  map′ ⟨term⟩      f unknown              = unknown
  map′ ⟨pat⟩       f (con c ps)           = con c (map f ps)
  map′ ⟨pat⟩       f (dot t)              = dot (map f t)
  map′ ⟨pat⟩       f (var x)              = var x
  map′ ⟨pat⟩       f (lit l)              = lit l
  map′ ⟨pat⟩       f (proj g)             = proj g
  map′ ⟨pat⟩       f (absurd x)           = absurd x
  map′ ⟨sort⟩      f (set t)              = set (map f t)
  map′ ⟨sort⟩      f (lit n)              = lit n
  map′ ⟨sort⟩      f (prop t)             = prop (map f t)
  map′ ⟨sort⟩      f (propLit n)          = propLit n
  map′ ⟨sort⟩      f (inf n)              = inf n
  map′ ⟨sort⟩      f unknown              = unknown
  map′ ⟨clause⟩    f (clause Γ ps args)   = clause (map f Γ) (map f ps) (map f args)
  map′ ⟨clause⟩    f (absurd-clause Γ ps) = absurd-clause (map f Γ) (map f ps)
  map′ (⟨list⟩ u)  f []                   = []
  map′ (⟨list⟩ u)  f (x ∷ xs)             = map f x ∷ map′ _ f xs
  map′ (⟨arg⟩ u)   f (arg i x)            = arg i (map f x)
  map′ (⟨abs⟩ u)   f (abs x t)            = abs x (map f t)
  map′ (⟨named⟩ u) f (x , t)              = x , map f t

  map : (∀ {u} {t : ⟦ u ⟧} → Ann₁ t → Ann₂ t) → ∀ {u} {t : ⟦ u ⟧} → Annotated Ann₁ t → Annotated Ann₂ t
  map f {u = u} (⟨ α ⟩ t) = ⟨ f α ⟩ map′ u f t


module _ {W : Set ℓ} (ε : W) (_∪_ : W → W → W) where

  -- This annotation function does nothing except combine ε's with _∪_.
  -- Lets you skip the boring cases when defining non-dependent
  -- annotation functions by adding a catch-all calling defaultAnn.
  defaultAnn : AnnotationFun (λ _ → W)
  defaultAnn ⟨term⟩      (var x args)         = ann args
  defaultAnn ⟨term⟩      (con c args)         = ann args
  defaultAnn ⟨term⟩      (def f args)         = ann args
  defaultAnn ⟨term⟩      (lam v b)            = ann b
  defaultAnn ⟨term⟩      (pat-lam cs args)    = ann cs ∪ ann args
  defaultAnn ⟨term⟩      (pi a b)             = ann a ∪ ann b
  defaultAnn ⟨term⟩      (agda-sort s)        = ann s
  defaultAnn ⟨term⟩      (lit l)              = ε
  defaultAnn ⟨term⟩      (meta x args)        = ann args
  defaultAnn ⟨term⟩      unknown              = ε
  defaultAnn ⟨pat⟩       (con c args)         = ann args
  defaultAnn ⟨pat⟩       (dot t)              = ann t
  defaultAnn ⟨pat⟩       (var x)              = ε
  defaultAnn ⟨pat⟩       (lit l)              = ε
  defaultAnn ⟨pat⟩       (proj f)             = ε
  defaultAnn ⟨pat⟩       (absurd x)           = ε
  defaultAnn ⟨sort⟩      (set t)              = ann t
  defaultAnn ⟨sort⟩      (lit n)              = ε
  defaultAnn ⟨sort⟩      (prop t)             = ann t
  defaultAnn ⟨sort⟩      (propLit n)          = ε
  defaultAnn ⟨sort⟩      (inf n)              = ε
  defaultAnn ⟨sort⟩      unknown              = ε
  defaultAnn ⟨clause⟩    (clause Γ ps t)      = ann Γ ∪ (ann ps ∪ ann t)
  defaultAnn ⟨clause⟩    (absurd-clause Γ ps) = ann Γ ∪ ann ps
  defaultAnn (⟨list⟩ u)  []                   = ε
  defaultAnn (⟨list⟩ u)  (x ∷ xs)             = ann x ∪ defaultAnn _ xs
  defaultAnn (⟨arg⟩ u)   (arg i x)            = ann x
  defaultAnn (⟨abs⟩ u)   (abs x t)            = ann t
  defaultAnn (⟨named⟩ u) (x , t)              = ann t


-- Cartisian product of two annotation functions.
infixr  _⊗_
_⊗_ : AnnotationFun Ann₁ → AnnotationFun Ann₂ → AnnotationFun (λ t → Ann₁ t × Ann₂ t)
(f ⊗ g) u t = f u (map′ u proj₁ t) , g u (map′ u proj₂ t)


------------------------------------------------------------------------
-- Annotation-driven traversal

-- Top-down applicative traversal of an annotated term. Applies an
-- action (without going into subterms) to terms whose annotation
-- satisfies a given criterion. Returns an unannotated term.

module Traverse {M : Set → Set} (appl : RawApplicative M) where

  open RawApplicative appl

  module _ (apply? : ∀ {u} {t : ⟦ u ⟧} → Ann t → Bool)
           (action : ∀ {u} {t : ⟦ u ⟧} → Annotated Ann t → M ⟦ u ⟧) where

    mutual
      traverse′ : {t : ⟦ u ⟧} → Annotated′ Ann t → M ⟦ u ⟧
      traverse′ {⟨term⟩}    (var x args)         = var x <$> traverse args
      traverse′ {⟨term⟩}    (con c args)         = con c <$> traverse args
      traverse′ {⟨term⟩}    (def f args)         = def f <$> traverse args
      traverse′ {⟨term⟩}    (lam v b)            = lam v <$> traverse b
      traverse′ {⟨term⟩}    (pat-lam cs args)    = pat-lam <$> traverse cs <*> traverse args
      traverse′ {⟨term⟩}    (pi a b)             = pi <$> traverse a <*> traverse b
      traverse′ {⟨term⟩}    (agda-sort s)        = agda-sort <$> traverse s
      traverse′ {⟨term⟩}    (lit l)              = pure (lit l)
      traverse′ {⟨term⟩}    (meta x args)        = meta x <$> traverse args
      traverse′ {⟨term⟩}    unknown              = pure unknown
      traverse′ {⟨pat⟩}     (con c args)         = con c <$> traverse args
      traverse′ {⟨pat⟩}     (dot t)              = dot <$> traverse t
      traverse′ {⟨pat⟩}     (var x)              = pure (var x)
      traverse′ {⟨pat⟩}     (lit l)              = pure (lit l)
      traverse′ {⟨pat⟩}     (proj f)             = pure (proj f)
      traverse′ {⟨pat⟩}     (absurd x)           = pure (absurd x)
      traverse′ {⟨sort⟩}    (set t)              = set <$> traverse t
      traverse′ {⟨sort⟩}    (lit n)              = pure (lit n)
      traverse′ {⟨sort⟩}    (prop t)             = prop <$> traverse t
      traverse′ {⟨sort⟩}    (propLit n)          = pure (propLit n)
      traverse′ {⟨sort⟩}    (inf n)              = pure (inf n)
      traverse′ {⟨sort⟩}    unknown              = pure unknown
      traverse′ {⟨clause⟩}  (clause Γ ps t)      = clause <$> traverse Γ <*> traverse ps <*> traverse t
      traverse′ {⟨clause⟩}  (absurd-clause Γ ps) = absurd-clause <$> traverse Γ <*> traverse ps
      traverse′ {⟨list⟩ u}  []                   = pure []
      traverse′ {⟨list⟩ u}  (x ∷ xs)             = _∷_ <$> traverse x <*> traverse′ xs
      traverse′ {⟨arg⟩ u}   (arg i x)            = arg i <$> traverse x
      traverse′ {⟨abs⟩ u}   (abs x t)            = abs x <$> traverse t
      traverse′ {⟨named⟩ u} (x , t)              = x ,_ <$> traverse t

      traverse : {t : ⟦ u ⟧} → Annotated Ann t → M ⟦ u ⟧
      traverse t@(⟨ α ⟩ tₐ) = if apply? α then action t else traverse′ tₐ