Scheme.hs

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

-- | Type schemes
module Hyper.Syntax.Scheme
    ( Scheme (..)
    , sForAlls
    , sTyp
    , W_Scheme (..)
    , QVars (..)
    , _QVars
    , HasScheme (..)
    , loadScheme
    , saveScheme
    , MonadInstantiate (..)
    , inferType
    , QVarInstances (..)
    , _QVarInstances
    , makeQVarInstances
    ) where

import qualified Control.Lens as Lens
import Control.Monad.Trans.Class (MonadTrans (..))
import Control.Monad.Trans.State (StateT (..))
import qualified Data.Map as Map
import Hyper
import Hyper.Class.Optic (HNodeLens (..))
import Hyper.Infer
import Hyper.Recurse
import Hyper.Unify
import Hyper.Unify.Generalize
import Hyper.Unify.New (newTerm)
import Hyper.Unify.QuantifiedVar (HasQuantifiedVar (..), MonadQuantify (..), OrdQVar)
import Hyper.Unify.Term (UTerm (..), uBody)
import Text.PrettyPrint ((<+>))
import qualified Text.PrettyPrint as Pretty
import Text.PrettyPrint.HughesPJClass (Pretty (..), maybeParens)

import Hyper.Internal.Prelude

-- | A type scheme representing a polymorphic type.
data Scheme varTypes typ h = Scheme
    { _sForAlls :: varTypes # QVars
    , _sTyp :: h :# typ
    }
    deriving (Generic)

newtype QVars typ
    = QVars
        (Map (QVar (GetHyperType typ)) (TypeConstraintsOf (GetHyperType typ)))
    deriving stock (Generic)

newtype QVarInstances h typ = QVarInstances (Map (QVar (GetHyperType typ)) (h typ))
    deriving stock (Generic)

makeLenses ''Scheme
makePrisms ''QVars
makePrisms ''QVarInstances
makeCommonInstances [''Scheme, ''QVars, ''QVarInstances]
makeHTraversableApplyAndBases ''Scheme

instance RNodes t => RNodes (Scheme v t)
instance (c (Scheme v t), Recursively c t) => Recursively c (Scheme v t)
instance (HTraversable (Scheme v t), RTraversable t) => RTraversable (Scheme v t)

instance
    ( Ord (QVar (GetHyperType typ))
    , Semigroup (TypeConstraintsOf (GetHyperType typ))
    ) =>
    Semigroup (QVars typ)
    where
    QVars m <> QVars n = QVars (Map.unionWith (<>) m n)

instance
    ( Ord (QVar (GetHyperType typ))
    , Semigroup (TypeConstraintsOf (GetHyperType typ))
    ) =>
    Monoid (QVars typ)
    where
    mempty = QVars mempty

instance
    (Pretty (varTypes # QVars), Pretty (h :# typ)) =>
    Pretty (Scheme varTypes typ h)
    where
    pPrintPrec lvl p (Scheme forAlls typ)
        | Pretty.isEmpty f = pPrintPrec lvl p typ
        | otherwise = f <+> pPrintPrec lvl 0 typ & maybeParens (p > 0)
        where
            f = pPrintPrec lvl 0 forAlls

instance
    (Pretty (TypeConstraintsOf typ), Pretty (QVar typ)) =>
    Pretty (QVars # typ)
    where
    pPrint (QVars qvars) =
        qvars ^@.. Lens.itraversed
            <&> (<> Pretty.text ".") . (Pretty.text "∀" <>) . printVar
            & Pretty.hsep
        where
            printVar (q, c)
                | cP == mempty = pPrint q
                | otherwise = pPrint q <> Pretty.text "(" <> cP <> Pretty.text ")"
                where
                    cP = pPrint c

type instance Lens.Index (QVars typ) = QVar (GetHyperType typ)
type instance Lens.IxValue (QVars typ) = TypeConstraintsOf (GetHyperType typ)

instance Ord (QVar (GetHyperType typ)) => Lens.Ixed (QVars typ)

instance Ord (QVar (GetHyperType typ)) => Lens.At (QVars typ) where
    at h = _QVars . Lens.at h

type instance InferOf (Scheme _ t) = HFlip GTerm t

class UnifyGen m t => MonadInstantiate m t where
    localInstantiations ::
        QVarInstances (UVarOf m) # t ->
        m a ->
        m a
    lookupQVar :: QVar t -> m (UVarOf m # t)

instance
    ( HasInferredValue typ
    , UnifyGen m typ
    , HTraversable varTypes
    , HNodesConstraint varTypes (MonadInstantiate m)
    , Infer m typ
    ) =>
    Infer m (Scheme varTypes typ)
    where
    {-# INLINE inferBody #-}
    inferBody (Scheme vars typ) =
        do
            foralls <- htraverse (Proxy @(MonadInstantiate m) #> makeQVarInstances) vars
            let withForalls =
                    hfoldMap
                        (Proxy @(MonadInstantiate m) #> (: []) . localInstantiations)
                        foralls
                        & foldl (.) id
            InferredChild typI typR <- inferChild typ & withForalls
            generalize (typR ^. inferredValue)
                <&> (Scheme vars typI,) . MkHFlip

inferType ::
    ( InferOf t ~ ANode t
    , HNodesConstraint t HasInferredValue
    , MonadInstantiate m t
    ) =>
    t # InferChild m h ->
    m (t # h, InferOf t # UVarOf m)
inferType x =
    case x ^? quantifiedVar of
        Just q -> lookupQVar q <&> (quantifiedVar # q,) . MkANode
        Nothing ->
            do
                xI <- htraverse (const inferChild) x
                hmap (Proxy @HasInferredValue #> (^. inType . inferredValue)) xI
                    & newTerm
                    <&> (hmap (const (^. inRep)) xI,) . MkANode

{-# INLINE makeQVarInstances #-}
makeQVarInstances ::
    Unify m typ =>
    QVars # typ ->
    m (QVarInstances (UVarOf m) # typ)
makeQVarInstances (QVars foralls) =
    traverse (newVar binding . USkolem) foralls <&> QVarInstances

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

class
    (UnifyGen m t, HNodeLens varTypes t, Ord (QVar t)) =>
    HasScheme varTypes m t
    where
    hasSchemeRecursive :: Proxy varTypes -> Proxy m -> RecMethod (HasScheme varTypes m) t
    {-# INLINE hasSchemeRecursive #-}
    default hasSchemeRecursive ::
        HNodesConstraint t (HasScheme varTypes m) =>
        Proxy varTypes ->
        Proxy m ->
        RecMethod (HasScheme varTypes m) t
    hasSchemeRecursive _ _ _ = Dict

instance Recursive (HasScheme varTypes m) where
    recurse = hasSchemeRecursive (Proxy @varTypes) (Proxy @m) . proxyArgument

-- | Load scheme into unification monad so that different instantiations share
-- the scheme's monomorphic parts -
-- their unification is O(1) as it is the same shared unification term.
{-# INLINE loadScheme #-}
loadScheme ::
    forall m varTypes typ.
    ( HTraversable varTypes
    , HNodesConstraint varTypes (UnifyGen m)
    , HasScheme varTypes m typ
    ) =>
    Pure # Scheme varTypes typ ->
    m (GTerm (UVarOf m) # typ)
loadScheme (Pure (Scheme vars typ)) =
    do
        foralls <- htraverse (Proxy @(UnifyGen m) #> makeQVarInstances) vars
        wrapM (Proxy @(HasScheme varTypes m) #>> loadBody foralls) typ

saveH ::
    forall typ varTypes m.
    HasScheme varTypes m typ =>
    GTerm (UVarOf m) # typ ->
    StateT (varTypes # QVars, [m ()]) m (Pure # typ)
saveH (GBody x) =
    htraverse (Proxy @(HasScheme varTypes m) #> saveH) x
        <&> (_Pure #)
        \\ hasSchemeRecursive (Proxy @varTypes) (Proxy @m) (Proxy @typ)
saveH (GMono x) =
    unwrapM (Proxy @(HasScheme varTypes m) #>> f) x & lift
    where
        f v =
            semiPruneLookup v
                <&> \case
                    (_, UTerm t) -> t ^. uBody
                    (_, UUnbound{}) -> error "saveScheme of non-toplevel scheme!"
                    _ -> error "unexpected state at saveScheme of monomorphic part"
saveH (GPoly x) =
    lookupVar binding x
        & lift
        >>= \case
            USkolem l ->
                do
                    r <-
                        scopeConstraints (Proxy @typ)
                            <&> (<> l)
                            >>= newQuantifiedVariable
                            & lift
                    Lens._1
                        . hNodeLens
                        %= (\v -> v & _QVars . Lens.at r ?~ generalizeConstraints l :: QVars # typ)
                    Lens._2 %= (bindVar binding x (USkolem l) :)
                    let result = _Pure . quantifiedVar # r
                    UResolved result & bindVar binding x & lift
                    pure result
            UResolved v -> pure v
            _ -> error "unexpected state at saveScheme's forall"

saveScheme ::
    ( HNodesConstraint varTypes OrdQVar
    , HPointed varTypes
    , HasScheme varTypes m typ
    ) =>
    GTerm (UVarOf m) # typ ->
    m (Pure # Scheme varTypes typ)
saveScheme x =
    do
        (t, (v, recover)) <-
            runStateT
                (saveH x)
                ( hpure (Proxy @OrdQVar #> QVars mempty)
                , []
                )
        _Pure # Scheme v t <$ sequence_ recover