content: {-# 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"
  