Nominal.hs

{-# LANGUAGE EmptyCase #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE TemplateHaskell #-}
{-# LANGUAGE UndecidableInstances #-}

-- | Nominal (named) types declaration, instantiation, construction, and access.
module Hyper.Syntax.Nominal
    ( NominalDecl (..)
    , nParams
    , nScheme
    , W_NominalDecl (..)
    , NominalInst (..)
    , nId
    , nArgs
    , ToNom (..)
    , tnId
    , tnVal
    , W_ToNom (..)
    , FromNom (..)
    , _FromNom
    , HasNominalInst (..)
    , NomVarTypes
    , MonadNominals (..)
    , LoadedNominalDecl
    , loadNominalDecl
    ) where

import Control.Applicative (Alternative (..))
import Control.Lens (Prism')
import qualified Control.Lens as Lens
import Control.Monad.Writer (WriterT (..), execWriterT)
import Generics.Constraints (Constraints)
import Hyper
import Hyper.Class.Context (HContext (..))
import Hyper.Class.Optic
import Hyper.Class.Traversable (ContainedH (..))
import Hyper.Class.ZipMatch (ZipMatch (..))
import Hyper.Infer
import Hyper.Recurse
import Hyper.Syntax.FuncType (FuncType (..))
import Hyper.Syntax.Map (TermMap (..), _TermMap)
import Hyper.Syntax.Scheme
import Hyper.Unify
import Hyper.Unify.Generalize (GTerm (..), instantiate, instantiateForAll, instantiateWith, _GMono)
import Hyper.Unify.New (newTerm)
import Hyper.Unify.QuantifiedVar (HasQuantifiedVar (..), OrdQVar)
import Hyper.Unify.Term (UTerm (..))
import qualified Text.PrettyPrint as P
import Text.PrettyPrint.HughesPJClass (Pretty (..), maybeParens)

import Hyper.Internal.Prelude

type family NomVarTypes (t :: HyperType) :: HyperType

-- | A declaration of a nominal type.
data NominalDecl typ h = NominalDecl
    { _nParams :: NomVarTypes typ # QVars
    , _nScheme :: Scheme (NomVarTypes typ) typ h
    }
    deriving (Generic)

-- | An instantiation of a nominal type
data NominalInst nomId varTypes h = NominalInst
    { _nId :: nomId
    , _nArgs :: varTypes # QVarInstances (GetHyperType h)
    }
    deriving (Generic)

-- | Nominal data constructor.
--
-- Wrap content with a data constructor
-- (analogues to a data constructor of a Haskell `newtype`'s).
--
-- Introduces the nominal's foralled type variables into the value's scope.
data ToNom nomId term h = ToNom
    { _tnId :: nomId
    , _tnVal :: h :# term
    }
    deriving (Generic)

-- | Access the data in a nominally typed value.
--
-- Analogues to a getter of a Haskell `newtype`.
newtype FromNom nomId (term :: HyperType) (h :: AHyperType) = FromNom nomId
    deriving newtype (Eq, Ord, Binary, NFData)
    deriving stock (Show, Generic)

-- | A nominal declaration loaded into scope in an inference monad.
data LoadedNominalDecl typ v = LoadedNominalDecl
    { lnParams :: NomVarTypes typ # QVarInstances (GetHyperType v)
    , lnForalls :: NomVarTypes typ # QVarInstances (GetHyperType v)
    , lnType :: GTerm (GetHyperType v) # typ
    }
    deriving (Generic)

makeLenses ''NominalDecl
makeLenses ''NominalInst
makeLenses ''ToNom
makePrisms ''FromNom
makeCommonInstances [''NominalDecl, ''NominalInst, ''ToNom, ''LoadedNominalDecl]
makeHTraversableAndBases ''NominalDecl
makeHTraversableApplyAndBases ''ToNom
makeHTraversableApplyAndBases ''FromNom
makeHMorph ''ToNom
makeZipMatch ''ToNom
makeZipMatch ''FromNom
makeHContext ''ToNom
makeHContext ''FromNom

instance HNodes v => HNodes (NominalInst n v) where
    type HNodesConstraint (NominalInst n v) c = HNodesConstraint v c
    type HWitnessType (NominalInst n v) = HWitnessType v
    {-# INLINE hLiftConstraint #-}
    hLiftConstraint (HWitness w) = hLiftConstraint @v (HWitness w)

instance HFunctor v => HFunctor (NominalInst n v) where
    {-# INLINE hmap #-}
    hmap f = nArgs %~ hmap (\(HWitness w) -> _QVarInstances . Lens.mapped %~ f (HWitness w))

instance HFoldable v => HFoldable (NominalInst n v) where
    {-# INLINE hfoldMap #-}
    hfoldMap f =
        hfoldMap (\(HWitness w) -> foldMap (f (HWitness w)) . (^. _QVarInstances)) . (^. nArgs)

instance HTraversable v => HTraversable (NominalInst n v) where
    {-# INLINE hsequence #-}
    hsequence (NominalInst n v) =
        htraverse (const (_QVarInstances (traverse runContainedH))) v
            <&> NominalInst n

instance
    ( Eq nomId
    , ZipMatch varTypes
    , HTraversable varTypes
    , HNodesConstraint varTypes ZipMatch
    , HNodesConstraint varTypes OrdQVar
    ) =>
    ZipMatch (NominalInst nomId varTypes)
    where
    {-# INLINE zipMatch #-}
    zipMatch (NominalInst xId x) (NominalInst yId y)
        | xId /= yId = Nothing
        | otherwise =
            zipMatch x y
                >>= htraverse
                    ( Proxy @ZipMatch #*#
                        Proxy @OrdQVar #>
                            \(QVarInstances c0 :*: QVarInstances c1) ->
                                zipMatch (TermMap c0) (TermMap c1)
                                    <&> QVarInstances . (^. _TermMap)
                    )
                <&> NominalInst xId

instance
    ( HFunctor varTypes
    , HContext varTypes
    , HNodesConstraint varTypes OrdQVar
    ) =>
    HContext (NominalInst nomId varTypes)
    where
    hcontext (NominalInst n args) =
        hcontext args
            & hmap
                ( Proxy @OrdQVar #>
                    \(HFunc c :*: x) ->
                        x
                            & _QVarInstances
                                . Lens.imapped
                                %@~ \k v ->
                                    HFunc
                                        ( \newV ->
                                            x
                                                & _QVarInstances . Lens.at k ?~ newV
                                                & c
                                                & getConst
                                                & NominalInst n
                                                & Const
                                        )
                                        :*: v
                )
            & NominalInst n

instance Constraints (ToNom nomId term h) Pretty => Pretty (ToNom nomId term h) where
    pPrintPrec lvl p (ToNom nomId term) =
        (pPrint nomId <> P.text "#") P.<+> pPrintPrec lvl 11 term
            & maybeParens (p > 10)

class (Pretty (QVar h), Pretty (outer :# h)) => PrettyConstraints outer h
instance (Pretty (QVar h), Pretty (outer :# h)) => PrettyConstraints outer h

instance
    ( Pretty nomId
    , HApply varTypes
    , HFoldable varTypes
    , HNodesConstraint varTypes (PrettyConstraints h)
    ) =>
    Pretty (NominalInst nomId varTypes h)
    where
    pPrint (NominalInst n vars) =
        pPrint n
            <> joinArgs
                (hfoldMap (Proxy @(PrettyConstraints h) #> mkArgs) vars)
        where
            joinArgs [] = mempty
            joinArgs xs = P.text "[" <> P.sep (P.punctuate (P.text ",") xs) <> P.text "]"
            mkArgs (QVarInstances m) =
                m ^@.. Lens.itraversed
                    <&> \(h, v) ->
                        (pPrint h <> P.text ":") P.<+> pPrint v

{-# ANN module "HLint: ignore Use camelCase" #-}
data W_LoadedNominalDecl t n where
    E_LoadedNominalDecl_Body :: HRecWitness t n -> W_LoadedNominalDecl t n
    E_LoadedNominalDecl_NomVarTypes :: HWitness (NomVarTypes t) n -> W_LoadedNominalDecl t n

instance (RNodes t, HNodes (NomVarTypes t)) => HNodes (LoadedNominalDecl t) where
    type
        HNodesConstraint (LoadedNominalDecl t) c =
            ( HNodesConstraint (NomVarTypes t) c
            , c t
            , Recursive c
            )
    type HWitnessType (LoadedNominalDecl t) = W_LoadedNominalDecl t
    {-# INLINE hLiftConstraint #-}
    hLiftConstraint (HWitness (E_LoadedNominalDecl_Body w)) = hLiftConstraint @(HFlip GTerm _) (HWitness w)
    hLiftConstraint (HWitness (E_LoadedNominalDecl_NomVarTypes w)) = hLiftConstraint w

instance
    (Recursively HFunctor typ, HFunctor (NomVarTypes typ)) =>
    HFunctor (LoadedNominalDecl typ)
    where
    {-# INLINE hmap #-}
    hmap f (LoadedNominalDecl mp mf t) =
        LoadedNominalDecl
            (onMap mp)
            (onMap mf)
            (t & hflipped %~ hmap (\(HWitness w) -> f (HWitness (E_LoadedNominalDecl_Body w))))
        where
            onMap = hmap (\w -> _QVarInstances . Lens.mapped %~ f (HWitness (E_LoadedNominalDecl_NomVarTypes w)))

instance
    (Recursively HFoldable typ, HFoldable (NomVarTypes typ)) =>
    HFoldable (LoadedNominalDecl typ)
    where
    {-# INLINE hfoldMap #-}
    hfoldMap f (LoadedNominalDecl mp mf t) =
        onMap mp
            <> onMap mf
            <> hfoldMap (\(HWitness w) -> f (HWitness (E_LoadedNominalDecl_Body w))) (_HFlip # t)
        where
            onMap =
                hfoldMap
                    ( \w ->
                        foldMap (f (HWitness (E_LoadedNominalDecl_NomVarTypes w)))
                            . (^. _QVarInstances)
                    )

instance
    (RTraversable typ, HTraversable (NomVarTypes typ)) =>
    HTraversable (LoadedNominalDecl typ)
    where
    {-# INLINE hsequence #-}
    hsequence (LoadedNominalDecl p f t) =
        LoadedNominalDecl
            <$> onMap p
            <*> onMap f
            <*> hflipped hsequence t
        where
            onMap = htraverse (const ((_QVarInstances . traverse) runContainedH))

{-# INLINE loadBody #-}
loadBody ::
    ( UnifyGen m typ
    , HNodeLens varTypes typ
    , Ord (QVar typ)
    ) =>
    varTypes # QVarInstances (UVarOf m) ->
    varTypes # QVarInstances (UVarOf m) ->
    typ # GTerm (UVarOf m) ->
    m (GTerm (UVarOf m) # typ)
loadBody params foralls x =
    case x ^? quantifiedVar >>= get of
        Just r -> GPoly r & pure
        Nothing ->
            case htraverse (const (^? _GMono)) x of
                Just xm -> newTerm xm <&> GMono
                Nothing -> GBody x & pure
    where
        get v =
            params ^? hNodeLens . _QVarInstances . Lens.ix v
                <|> foralls ^? hNodeLens . _QVarInstances . Lens.ix v

{-# INLINE loadNominalDecl #-}
loadNominalDecl ::
    forall m typ.
    ( HTraversable (NomVarTypes typ)
    , HNodesConstraint (NomVarTypes typ) (Unify m)
    , HasScheme (NomVarTypes typ) m typ
    ) =>
    Pure # NominalDecl typ ->
    m (LoadedNominalDecl typ # UVarOf m)
loadNominalDecl (Pure (NominalDecl params (Scheme foralls typ))) =
    do
        paramsL <- htraverse (Proxy @(Unify m) #> makeQVarInstances) params
        forallsL <- htraverse (Proxy @(Unify m) #> makeQVarInstances) foralls
        wrapM
            ( Proxy @(HasScheme (NomVarTypes typ) m) #>>
                loadBody paramsL forallsL
            )
            typ
            <&> LoadedNominalDecl paramsL forallsL

class MonadNominals nomId typ m where
    getNominalDecl :: nomId -> m (LoadedNominalDecl typ # UVarOf m)

class HasNominalInst nomId typ where
    nominalInst :: Prism' (typ # h) (NominalInst nomId (NomVarTypes typ) # h)

type instance InferOf (ToNom n e) = NominalInst n (NomVarTypes (TypeOf e))

instance
    ( MonadScopeLevel m
    , MonadNominals nomId (TypeOf expr) m
    , HTraversable (NomVarTypes (TypeOf expr))
    , HNodesConstraint (NomVarTypes (TypeOf expr)) (UnifyGen m)
    , UnifyGen m (TypeOf expr)
    , HasInferredType expr
    , Infer m expr
    ) =>
    Infer m (ToNom nomId expr)
    where
    {-# INLINE inferBody #-}
    inferBody (ToNom nomId val) =
        do
            LoadedNominalDecl params foralls gen <- getNominalDecl nomId

            -- Setup forall variables to instantiate to skolems.
            -- This means they aren't allow to be be unified,
            -- nor to propagate to outer scope.
            recoverForAlls <-
                htraverse_
                    ( Proxy @(UnifyGen m) #>
                        traverse_ (instantiateForAll USkolem) . (^. _QVarInstances)
                    )
                    foralls
                    & localLevel . execWriterT
            -- Setup params in outer scope
            (paramsT, recoverParams) <-
                htraverse
                    ( Proxy @(UnifyGen m) #>
                        (_QVarInstances . traverse) (instantiateForAll UUnbound)
                    )
                    params
                    & runWriterT

            typ <- instantiate gen & localLevel

            -- Restore loaded nominal to original reusable state
            sequence_ (recoverParams <> recoverForAlls)

            -- Term within is in inner level
            InferredChild valI valR <- inferChild val & localLevel

            -- Unify the inner term's type with the type inside the nominal
            _ <- unify typ (valR ^# inferredType (Proxy @expr))
            pure (ToNom nomId valI, NominalInst nomId paramsT)

type instance InferOf (FromNom _ e) = FuncType (TypeOf e)

instance
    ( Infer m expr
    , HasNominalInst nomId (TypeOf expr)
    , MonadNominals nomId (TypeOf expr) m
    , HTraversable (NomVarTypes (TypeOf expr))
    , HNodesConstraint (NomVarTypes (TypeOf expr)) (UnifyGen m)
    , UnifyGen m (TypeOf expr)
    ) =>
    Infer m (FromNom nomId expr)
    where
    {-# INLINE inferBody #-}
    inferBody (FromNom nomId) =
        do
            (LoadedNominalDecl params _ gen) <- getNominalDecl nomId
            let lookupParams = htraverse (Proxy @(UnifyGen m) #> (_QVarInstances . traverse) lookupParam) params
            (typ, paramsT) <- instantiateWith lookupParams UUnbound gen
            newTerm (nominalInst # NominalInst nomId paramsT)
                <&> (FromNom nomId,) . (`FuncType` typ)

lookupParam :: forall m t. UnifyGen m t => UVarOf m # t -> m (UVarOf m # t)
lookupParam v =
    lookupVar binding v
        >>= \case
            UInstantiated r -> pure r
            USkolem l ->
                -- This is a phantom-type, wasn't instantiated by `instantiate`.
                scopeConstraints (Proxy @t) <&> (<> l) >>= newVar binding . UUnbound
            _ -> error "unexpected state at nominal's parameter"