{-# LANGUAGE DataKinds             #-}
{-# LANGUAGE DeriveGeneric         #-}
{-# LANGUAGE DeriveTraversable     #-}
{-# LANGUAGE FlexibleContexts      #-}
{-# LANGUAGE FlexibleInstances     #-}
{-# LANGUAGE LambdaCase            #-}
{-# LANGUAGE MultiParamTypeClasses #-}
{-# LANGUAGE PatternSynonyms       #-}
{-# LANGUAGE PolyKinds             #-}
{-# LANGUAGE StandaloneDeriving    #-}
{-# LANGUAGE TemplateHaskell       #-}
{-# LANGUAGE TypeFamilies          #-}
{-# LANGUAGE TypeOperators         #-}
{-# LANGUAGE UndecidableInstances  #-}

-- | Annotating every node of a free foil term.
--
-- A node can be annotated with anything: a source position, a type, a memoised
-- normal form. 'AnnSig' turns a signature into an annotated one, and
-- 'AnnAST' is the annotated syntax it generates.
--
-- The annotation is a /functor of the term/, and that is what lets it depend on
-- the node's scope. In @'AST' binder sig n@ the signature's term parameter /is/
-- the AST at the node's own scope, so an annotation built from it holds terms in
-- that scope — which is what a type annotation in a dependent language needs.
-- An annotation that ignores the term (a source position, say) is @'Const' a@.
--
-- Whether two annotations must agree for their nodes to match is a property of
-- @ann@, not of 'AnnSig'. Two examples, and they are the two you will want:
--
-- An annotation that ignores the term (a source position) and /is/ compared:
--
-- > instance Eq a => ZipMatchK (Const (a :: Type)) where
-- >   zipMatchWithK _ (Const l) (Const r)
-- >     | l == r    = Just (Const l)
-- >     | otherwise = Nothing
--
-- An annotation that holds terms (a type, say) and is /ignored/, so that two
-- terms differing only in their types are α-equivalent. Make the held term
-- optional, so that pairing can fail /without/ failing the match:
--
-- > newtype TypeOf term = TypeOf (Maybe term)
-- >
-- > instance ZipMatchK TypeOf where
-- >   zipMatchWithK (f :^: M0) (TypeOf l) (TypeOf r) =
-- >     Just (TypeOf (do { a <- l; b <- r; f a b }))
--
-- __The instance must return 'Just' unconditionally, and lazily.__ Matching is
-- annotation-blind, so it must succeed whatever the annotations are; the paired
-- annotation is a thunk the (annotation-skipping) 'Bifoldable' never forces. A
-- /strict/ shape — @TypeOf '<$>' f l r@, or anything that yields 'Nothing' when
-- @f@ fails — is a footgun twice over: it breaks blindness (two nodes with
-- different types would fail to match), and it /diverges/ for a finite or
-- lazily-bottomed annotation (a universe tower ending in 'error', say), because
-- forcing the annotation runs off the end. This is why the held term is a 'Maybe':
-- a plain @term@ field would have no value to pair when @f@ fails, forcing a
-- bottom into the result.
--
-- __Do not reach for the generic instance here.__ It compares the annotation's
-- /shape/, so a node carrying a memoised normal form would fail to match the same
-- node without one. (Nor is 'zipMatchViaChooseLeft' available: an annotation
-- holding terms must /construct/ a value at the result index, and cannot simply
-- pick a side.)
--
-- __Note the asymmetry__ in the instances below: 'Bifunctor' and 'Bitraversable'
-- traverse the annotation, but 'Bifoldable' does /not/. This is deliberate — see
-- 'AnnSig'.
module Control.Monad.Free.Foil.Annotated (
  AnnSig(..),
  AnnAST,
  AnnScopedAST,
  pattern AnnNode,
  annotationOf,
  freeVarsOfAnnotated,
) where

import           Data.Bifoldable
import           Data.Bifunctor
import           Data.Bitraversable
import           Data.Kind             (Type)
import           Data.Maybe            (mapMaybe)
import           Data.ZipMatchK.TH     (deriveZipMatchK2)
import           Generics.Kind         (GenericK (..), Field, Var0, Var1, (:$:),
                                        Atom ((:@:)), (:*:))
import qualified GHC.Generics          as GHC

import qualified Control.Monad.Foil    as Foil
import           Control.Monad.Free.Foil

-- | A signature, with every node annotated by an @ann@ built from the node's term.
--
-- The annotation is applied to the /term/ parameter, so in an
-- @'AST' binder ('AnnSig' ann sig) n@ it holds terms in the node's own scope.
--
-- 'Bifoldable' deliberately __skips the annotation__. That is what makes
-- α-equivalence annotation-blind: 'Control.Monad.Free.Foil.alphaEquiv' consumes a
-- zipped node with 'bifoldMap', so a type annotation is never compared, and two
-- terms that differ only in their annotations are α-equivalent. The price is that
-- 'Control.Monad.Free.Foil.freeVarsOf', which is also 'Bifoldable', does not see
-- variables occurring inside annotations. Use 'freeVarsOfAnnotated' when those
-- matter.
data AnnSig (ann :: Type -> Type) (sig :: Type -> Type -> Type) scope term
  = AnnSig (ann term) (sig scope term)
  deriving ((forall x.
 AnnSig ann sig scope term -> Rep (AnnSig ann sig scope term) x)
-> (forall x.
    Rep (AnnSig ann sig scope term) x -> AnnSig ann sig scope term)
-> Generic (AnnSig ann sig scope term)
forall x.
Rep (AnnSig ann sig scope term) x -> AnnSig ann sig scope term
forall x.
AnnSig ann sig scope term -> Rep (AnnSig ann sig scope term) x
forall a.
(forall x. a -> Rep a x) -> (forall x. Rep a x -> a) -> Generic a
forall (ann :: * -> *) (sig :: * -> * -> *) scope term x.
Rep (AnnSig ann sig scope term) x -> AnnSig ann sig scope term
forall (ann :: * -> *) (sig :: * -> * -> *) scope term x.
AnnSig ann sig scope term -> Rep (AnnSig ann sig scope term) x
$cfrom :: forall (ann :: * -> *) (sig :: * -> * -> *) scope term x.
AnnSig ann sig scope term -> Rep (AnnSig ann sig scope term) x
from :: forall x.
AnnSig ann sig scope term -> Rep (AnnSig ann sig scope term) x
$cto :: forall (ann :: * -> *) (sig :: * -> * -> *) scope term x.
Rep (AnnSig ann sig scope term) x -> AnnSig ann sig scope term
to :: forall x.
Rep (AnnSig ann sig scope term) x -> AnnSig ann sig scope term
GHC.Generic)

deriving instance (Functor ann, Functor (sig scope))
  => Functor (AnnSig ann sig scope)
deriving instance (Foldable ann, Foldable (sig scope))
  => Foldable (AnnSig ann sig scope)
deriving instance (Traversable ann, Traversable (sig scope))
  => Traversable (AnnSig ann sig scope)

instance (Functor ann, Bifunctor sig) => Bifunctor (AnnSig ann sig) where
  bimap :: forall a b c d.
(a -> b) -> (c -> d) -> AnnSig ann sig a c -> AnnSig ann sig b d
bimap a -> b
f c -> d
g (AnnSig ann c
ann sig a c
sig) = ann d -> sig b d -> AnnSig ann sig b d
forall (ann :: * -> *) (sig :: * -> * -> *) scope term.
ann term -> sig scope term -> AnnSig ann sig scope term
AnnSig ((c -> d) -> ann c -> ann d
forall a b. (a -> b) -> ann a -> ann b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap c -> d
g ann c
ann) ((a -> b) -> (c -> d) -> sig a c -> sig b d
forall a b c d. (a -> b) -> (c -> d) -> sig a c -> sig b d
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap a -> b
f c -> d
g sig a c
sig)

-- | __The annotation is not folded.__ See 'AnnSig'.
instance Bifoldable sig => Bifoldable (AnnSig ann sig) where
  bifoldMap :: forall m a b.
Monoid m =>
(a -> m) -> (b -> m) -> AnnSig ann sig a b -> m
bifoldMap a -> m
f b -> m
g (AnnSig ann b
_ann sig a b
sig) = (a -> m) -> (b -> m) -> sig a b -> m
forall m a b. Monoid m => (a -> m) -> (b -> m) -> sig a b -> m
forall (p :: * -> * -> *) m a b.
(Bifoldable p, Monoid m) =>
(a -> m) -> (b -> m) -> p a b -> m
bifoldMap a -> m
f b -> m
g sig a b
sig

instance (Traversable ann, Bitraversable sig) => Bitraversable (AnnSig ann sig) where
  bitraverse :: forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c)
-> (b -> f d) -> AnnSig ann sig a b -> f (AnnSig ann sig c d)
bitraverse a -> f c
f b -> f d
g (AnnSig ann b
ann sig a b
sig) = ann d -> sig c d -> AnnSig ann sig c d
forall (ann :: * -> *) (sig :: * -> * -> *) scope term.
ann term -> sig scope term -> AnnSig ann sig scope term
AnnSig (ann d -> sig c d -> AnnSig ann sig c d)
-> f (ann d) -> f (sig c d -> AnnSig ann sig c d)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
<$> (b -> f d) -> ann b -> f (ann d)
forall (t :: * -> *) (f :: * -> *) a b.
(Traversable t, Applicative f) =>
(a -> f b) -> t a -> f (t b)
forall (f :: * -> *) a b.
Applicative f =>
(a -> f b) -> ann a -> f (ann b)
traverse b -> f d
g ann b
ann f (sig c d -> AnnSig ann sig c d)
-> f (sig c d) -> f (AnnSig ann sig c d)
forall a b. f (a -> b) -> f a -> f b
forall (f :: * -> *) a b. Applicative f => f (a -> b) -> f a -> f b
<*> (a -> f c) -> (b -> f d) -> sig a b -> f (sig c d)
forall (f :: * -> *) a c b d.
Applicative f =>
(a -> f c) -> (b -> f d) -> sig a b -> f (sig c d)
forall (t :: * -> * -> *) (f :: * -> *) a c b d.
(Bitraversable t, Applicative f) =>
(a -> f c) -> (b -> f d) -> t a b -> f (t c d)
bitraverse a -> f c
f b -> f d
g sig a b
sig

instance GenericK (AnnSig ann sig) where
  type RepK (AnnSig ann sig) =
    Field (ann :$: Var1) :*: Field ((sig :$: Var0) :@: Var1)

-- | Matching for 'AnnSig' is /derived/, not left to the generic default.
--
-- The generic default would rebuild the "Generics.Kind" view of every node on
-- every comparison, and comparing terms is most of what a typechecker does, so
-- for an annotated signature — where it lands on the hottest path — that is a
-- measurable cost. 'Data.ZipMatchK.TH.deriveZipMatchK2' generates the
-- written-out instance instead: the annotation matched with the term zipper (an
-- annotation is a functor of the term), the inner signature with both. Write
-- @deriveZipMatchK2 ''YourAnnSig@ for a bespoke annotated signature.
deriveZipMatchK2 ''AnnSig

-- | An annotated scope-safe term.
type AnnAST binder ann sig = AST binder (AnnSig ann sig)

-- | An annotated scope-safe term under a binder.
type AnnScopedAST binder ann sig = ScopedAST binder (AnnSig ann sig)

-- | An annotated node.
pattern AnnNode
  :: ann (AnnAST binder ann sig n)
  -> sig (AnnScopedAST binder ann sig n) (AnnAST binder ann sig n)
  -> AnnAST binder ann sig n
pattern $mAnnNode :: forall {r} {ann :: * -> *} {binder :: S -> S -> *}
       {sig :: * -> * -> *} {n :: S}.
AnnAST binder ann sig n
-> (ann (AnnAST binder ann sig n)
    -> sig (AnnScopedAST binder ann sig n) (AnnAST binder ann sig n)
    -> r)
-> ((# #) -> r)
-> r
$bAnnNode :: forall (ann :: * -> *) (binder :: S -> S -> *) (sig :: * -> * -> *)
       (n :: S).
ann (AnnAST binder ann sig n)
-> sig (AnnScopedAST binder ann sig n) (AnnAST binder ann sig n)
-> AnnAST binder ann sig n
AnnNode ann sig = Node (AnnSig ann sig)

{-# COMPLETE Var, AnnNode #-}

-- | The annotation of a term, unless it is a variable (which is not a node, so it
-- carries none).
annotationOf :: AnnAST binder ann sig n -> Maybe (ann (AnnAST binder ann sig n))
annotationOf :: forall (binder :: S -> S -> *) (ann :: * -> *) (sig :: * -> * -> *)
       (n :: S).
AnnAST binder ann sig n -> Maybe (ann (AnnAST binder ann sig n))
annotationOf = \case
  Var Name n
_         -> Maybe (ann (AnnAST binder ann sig n))
forall a. Maybe a
Nothing
  AnnNode ann (AnnAST binder ann sig n)
ann sig (AnnScopedAST binder ann sig n) (AnnAST binder ann sig n)
_ -> ann (AnnAST binder ann sig n)
-> Maybe (ann (AnnAST binder ann sig n))
forall a. a -> Maybe a
Just ann (AnnAST binder ann sig n)
ann

-- | The free variables of an annotated term, __including__ those occurring inside
-- annotations.
--
-- 'Control.Monad.Free.Foil.freeVarsOf' misses the latter, since 'Bifoldable' skips
-- the annotation (see 'AnnSig').
freeVarsOfAnnotated
  :: (Foil.Distinct n, Foil.CoSinkable binder, Bifoldable sig, Foldable ann)
  => AnnAST binder ann sig n -> [Foil.Name n]
freeVarsOfAnnotated :: forall (n :: S) (binder :: S -> S -> *) (sig :: * -> * -> *)
       (ann :: * -> *).
(Distinct n, CoSinkable binder, Bifoldable sig, Foldable ann) =>
AnnAST binder ann sig n -> [Name n]
freeVarsOfAnnotated = \case
  Var Name n
name -> [Name n
name]
  AnnNode ann (AnnAST binder ann sig n)
ann sig (AnnScopedAST binder ann sig n) (AnnAST binder ann sig n)
sig ->
    (AnnAST binder ann sig n -> [Name n])
-> ann (AnnAST binder ann sig n) -> [Name n]
forall m a. Monoid m => (a -> m) -> ann a -> m
forall (t :: * -> *) m a.
(Foldable t, Monoid m) =>
(a -> m) -> t a -> m
foldMap AnnAST binder ann sig n -> [Name n]
forall (n :: S) (binder :: S -> S -> *) (sig :: * -> * -> *)
       (ann :: * -> *).
(Distinct n, CoSinkable binder, Bifoldable sig, Foldable ann) =>
AnnAST binder ann sig n -> [Name n]
freeVarsOfAnnotated ann (AnnAST binder ann sig n)
ann
      [Name n] -> [Name n] -> [Name n]
forall a. Semigroup a => a -> a -> a
<> (AnnScopedAST binder ann sig n -> [Name n])
-> (AnnAST binder ann sig n -> [Name n])
-> sig (AnnScopedAST binder ann sig n) (AnnAST binder ann sig n)
-> [Name n]
forall m a b. Monoid m => (a -> m) -> (b -> m) -> sig a b -> m
forall (p :: * -> * -> *) m a b.
(Bifoldable p, Monoid m) =>
(a -> m) -> (b -> m) -> p a b -> m
bifoldMap AnnScopedAST binder ann sig n -> [Name n]
forall (n :: S) (binder :: S -> S -> *) (sig :: * -> * -> *)
       (ann :: * -> *).
(Distinct n, CoSinkable binder, Bifoldable sig, Foldable ann) =>
AnnScopedAST binder ann sig n -> [Name n]
freeVarsOfAnnotatedScoped AnnAST binder ann sig n -> [Name n]
forall (n :: S) (binder :: S -> S -> *) (sig :: * -> * -> *)
       (ann :: * -> *).
(Distinct n, CoSinkable binder, Bifoldable sig, Foldable ann) =>
AnnAST binder ann sig n -> [Name n]
freeVarsOfAnnotated sig (AnnScopedAST binder ann sig n) (AnnAST binder ann sig n)
sig

-- | The free variables of an annotated scoped term, including those inside
-- annotations.
freeVarsOfAnnotatedScoped
  :: (Foil.Distinct n, Foil.CoSinkable binder, Bifoldable sig, Foldable ann)
  => AnnScopedAST binder ann sig n -> [Foil.Name n]
freeVarsOfAnnotatedScoped :: forall (n :: S) (binder :: S -> S -> *) (sig :: * -> * -> *)
       (ann :: * -> *).
(Distinct n, CoSinkable binder, Bifoldable sig, Foldable ann) =>
AnnScopedAST binder ann sig n -> [Name n]
freeVarsOfAnnotatedScoped (ScopedAST binder n l
binder AST binder (AnnSig ann sig) l
body) =
  case binder n l -> DistinctEvidence l
forall (n :: S) (pattern :: S -> S -> *) (l :: S).
(Distinct n, CoSinkable pattern) =>
pattern n l -> DistinctEvidence l
Foil.assertDistinct binder n l
binder of
    DistinctEvidence l
Foil.Distinct ->
      (Name l -> Maybe (Name n)) -> [Name l] -> [Name n]
forall a b. (a -> Maybe b) -> [a] -> [b]
mapMaybe (binder n l -> Name l -> Maybe (Name n)
forall (pattern :: S -> S -> *) (n :: S) (l :: S).
(Distinct n, CoSinkable pattern) =>
pattern n l -> Name l -> Maybe (Name n)
Foil.unsinkNamePattern binder n l
binder) (AST binder (AnnSig ann sig) l -> [Name l]
forall (n :: S) (binder :: S -> S -> *) (sig :: * -> * -> *)
       (ann :: * -> *).
(Distinct n, CoSinkable binder, Bifoldable sig, Foldable ann) =>
AnnAST binder ann sig n -> [Name n]
freeVarsOfAnnotated AST binder (AnnSig ann sig) l
body)