content: {-# LANGUAGE FlexibleContexts #-}
  {-# LANGUAGE FlexibleInstances #-}
  {-# LANGUAGE TemplateHaskell #-}
  {-# LANGUAGE UndecidableInstances #-}
  
  -- | Generalization of type schemes
  module Hyper.Unify.Generalize
      ( generalize
      , instantiate
      , GTerm (..)
      , _GMono
      , _GPoly
      , _GBody
      , W_GTerm (..)
      , instantiateWith
      , instantiateForAll
        -- | Exports for @SPECIALIZE@ pragmas.
      , instantiateH
      ) where
  
  import Algebra.PartialOrd (PartialOrd (..))
  import qualified Control.Lens as Lens
  import Control.Monad.Trans.Class (MonadTrans (..))
  import Control.Monad.Trans.Writer (WriterT (..), tell)
  import Data.Monoid (All (..))
  import Hyper
  import Hyper.Class.Traversable
  import Hyper.Class.Unify
  import Hyper.Recurse
  import Hyper.Unify.Constraints
  import Hyper.Unify.New (newTerm)
  import Hyper.Unify.Term (UTerm (..), uBody)
  
  import Hyper.Internal.Prelude
  
  -- | An efficient representation of a type scheme arising from
  -- generalizing a unification term. Type subexpressions which are
  -- completely monomoprhic are tagged as such, to avoid redundant
  -- instantation and unification work
  data GTerm v ast
      = -- | Completely monomoprhic term
        GMono (v ast)
      | -- | Points to a quantified variable (instantiation will
        -- create fresh unification terms) (`Hyper.Unify.Term.USkolem`
        -- or `Hyper.Unify.Term.UResolved`)
        GPoly (v ast)
      | -- | Term with some polymorphic parts
        GBody (ast :# GTerm v)
      deriving (Generic)
  
  makePrisms ''GTerm
  makeCommonInstances [''GTerm]
  makeHTraversableAndBases ''GTerm
  
  instance RNodes a => HNodes (HFlip GTerm a) where
      type HNodesConstraint (HFlip GTerm a) c = (c a, Recursive c)
      type HWitnessType (HFlip GTerm a) = HRecWitness a
      {-# INLINE hLiftConstraint #-}
      hLiftConstraint (HWitness HRecSelf) = \_ x -> x
      hLiftConstraint (HWitness (HRecSub c n)) = hLiftConstraintH c n
  
  hLiftConstraintH ::
      forall a c b n r.
      (RNodes a, HNodesConstraint (HFlip GTerm a) c) =>
      HWitness a b ->
      HRecWitness b n ->
      Proxy c ->
      (c n => r) ->
      r
  hLiftConstraintH c n p f =
      (Proxy @RNodes #> (p #> (p #> f) (HWitness @(HFlip GTerm _) n)) c) c
          \\ recurse (Proxy @(c a))
          \\ recurse (Proxy @(RNodes a))
  
  instance Recursively HFunctor ast => HFunctor (HFlip GTerm ast) where
      {-# INLINE hmap #-}
      hmap f =
          _HFlip
              %~ \case
                  GMono x -> f (HWitness HRecSelf) x & GMono
                  GPoly x -> f (HWitness HRecSelf) x & GPoly
                  GBody x ->
                      hmap
                          ( Proxy @(Recursively HFunctor) #*#
                              \cw ->
                                  hflipped
                                      %~ hmap (f . (\(HWitness nw) -> HWitness (HRecSub cw nw)))
                          )
                          x
                          & GBody
                          \\ recursively (Proxy @(HFunctor ast))
  
  instance Recursively HFoldable ast => HFoldable (HFlip GTerm ast) where
      {-# INLINE hfoldMap #-}
      hfoldMap f =
          \case
              GMono x -> f (HWitness HRecSelf) x
              GPoly x -> f (HWitness HRecSelf) x
              GBody x ->
                  hfoldMap
                      ( Proxy @(Recursively HFoldable) #*#
                          \cw ->
                              hfoldMap (f . (\(HWitness nw) -> HWitness (HRecSub cw nw)))
                                  . (_HFlip #)
                      )
                      x
                      \\ recursively (Proxy @(HFoldable ast))
              . (^. _HFlip)
  
  instance RTraversable ast => HTraversable (HFlip GTerm ast) where
      {-# INLINE hsequence #-}
      hsequence =
          \case
              GMono x -> runContainedH x <&> GMono
              GPoly x -> runContainedH x <&> GPoly
              GBody x ->
                  -- HTraversable will be required when not implied by Recursively
                  htraverse (Proxy @RTraversable #> hflipped hsequence) x
                      \\ recurse (Proxy @(RTraversable ast))
                      <&> GBody
              & _HFlip
  
  -- | Generalize a unification term pointed by the given variable to a `GTerm`.
  -- Unification variables that are scoped within the term
  -- become universally quantified skolems.
  generalize ::
      forall m t.
      UnifyGen m t =>
      UVarOf m # t ->
      m (GTerm (UVarOf m) # t)
  generalize v0 =
      do
          (v1, u) <- semiPruneLookup v0
          c <- scopeConstraints (Proxy @t)
          case u of
              UUnbound l
                  | toScopeConstraints l `leq` c ->
                      GPoly v1
                          <$
                          -- We set the variable to a skolem,
                          -- so additional unifications after generalization
                          -- (for example hole resumptions where supported)
                          -- cannot unify it with anything.
                          bindVar binding v1 (USkolem (generalizeConstraints l))
              USkolem l | toScopeConstraints l `leq` c -> pure (GPoly v1)
              UTerm t ->
                  do
                      bindVar binding v1 (UResolving t)
                      b <- htraverse (Proxy @(UnifyGen m) #> generalize) (t ^. uBody)
                      bindVar binding v1 (UTerm t)
                      pure
                          ( if hfoldMap (Proxy @(UnifyGen m) #> All . Lens.has _GMono) b ^. Lens._Wrapped
                              then GMono v1
                              else GBody b
                          )
                      \\ unifyGenRecursive (Proxy @m) (Proxy @t)
              UResolving t -> GMono v1 <$ occursError v1 t
              _ -> pure (GMono v1)
  
  {-# INLINE instantiateForAll #-}
  instantiateForAll ::
      forall m t.
      UnifyGen m t =>
      (TypeConstraintsOf t -> UTerm (UVarOf m) # t) ->
      UVarOf m # t ->
      WriterT [m ()] m (UVarOf m # t)
  instantiateForAll cons x =
      lift (lookupVar binding x)
          >>= \case
              USkolem l ->
                  do
                      tell [bindVar binding x (USkolem l)]
                      r <- scopeConstraints (Proxy @t) >>= newVar binding . cons . (<> l) & lift
                      UInstantiated r & bindVar binding x & lift
                      pure r
              UInstantiated v -> pure v
              _ -> error "unexpected state at instantiate's forall"
  
  -- TODO: Better name?
  {-# INLINE instantiateH #-}
  instantiateH ::
      forall m t.
      UnifyGen m t =>
      (forall n. TypeConstraintsOf n -> UTerm (UVarOf m) # n) ->
      GTerm (UVarOf m) # t ->
      WriterT [m ()] m (UVarOf m # t)
  instantiateH _ (GMono x) = pure x
  instantiateH cons (GPoly x) = instantiateForAll cons x
  instantiateH cons (GBody x) =
      htraverse (Proxy @(UnifyGen m) #> instantiateH cons) x
          >>= lift
              . newTerm
          \\ unifyGenRecursive (Proxy @m) (Proxy @t)
  
  {-# INLINE instantiateWith #-}
  instantiateWith ::
      forall m t a.
      UnifyGen m t =>
      m a ->
      (forall n. TypeConstraintsOf n -> UTerm (UVarOf m) # n) ->
      GTerm (UVarOf m) # t ->
      m (UVarOf m # t, a)
  instantiateWith action cons g =
      do
          (r, recover) <- runWriterT (instantiateH cons g)
          action <* sequence_ recover <&> (r,)
  
  -- | Instantiate a generalized type with fresh unification variables
  -- for the quantified variables
  {-# INLINE instantiate #-}
  instantiate ::
      UnifyGen m t =>
      GTerm (UVarOf m) # t ->
      m (UVarOf m # t)
  instantiate g = instantiateWith (pure ()) UUnbound g <&> (^. Lens._1)
  