content: {-# LANGUAGE FlexibleContexts #-}
  
  -- | A class for unification
  module Hyper.Class.Unify
      ( Unify (..)
      , UVarOf
      , UnifyGen (..)
      , BindingDict (..)
      , applyBindings
      , semiPruneLookup
      , occursError
      ) where
  
  import Control.Monad (unless)
  import Control.Monad.Error.Class (MonadError (..))
  import Control.Monad.Trans.Class (MonadTrans (..))
  import Control.Monad.Trans.State (get, put, runStateT)
  import Data.Kind (Type)
  import Hyper.Class.Nodes (HNodes (..), (#>))
  import Hyper.Class.Optic (HSubset (..), HSubset')
  import Hyper.Class.Recursive
  import Hyper.Class.Traversable (htraverse)
  import Hyper.Class.ZipMatch (ZipMatch)
  import Hyper.Type (HyperType, type (#))
  import Hyper.Type.Pure (Pure, _Pure)
  import Hyper.Unify.Constraints
  import Hyper.Unify.Error (UnifyError (..))
  import Hyper.Unify.QuantifiedVar (HasQuantifiedVar (..), MonadQuantify (..))
  import Hyper.Unify.Term (UTerm (..), UTermBody (..), uBody)
  
  import Hyper.Internal.Prelude
  
  -- | Unification variable type for a unification monad
  type family UVarOf (m :: Type -> Type) :: HyperType
  
  -- | BindingDict implements unification variables for a type in a unification monad.
  --
  -- It is parameterized on:
  --
  -- * @v@: The unification variable 'HyperType'
  -- * @m@: The 'Monad' to bind in
  -- * @t@: The unified term's 'HyperType'
  --
  -- Has 2 implementations in hypertypes:
  --
  -- * 'Hyper.Unify.Binding.bindingDict' for pure state based unification
  -- * 'Hyper.Unify.Binding.ST.stBinding' for 'Control.Monad.ST.ST' based unification
  data BindingDict v m t = BindingDict
      { lookupVar :: !(v # t -> m (UTerm v # t))
      , newVar :: !(UTerm v # t -> m (v # t))
      , bindVar :: !(v # t -> UTerm v # t -> m ())
      }
  
  -- | @Unify m t@ enables 'Hyper.Unify.unify' to perform unification for @t@ in the 'Monad' @m@.
  --
  -- The 'unifyRecursive' method represents the constraint that @Unify m@ applies to all recursive child nodes.
  -- It replaces context for 'Unify' to avoid @UndecidableSuperClasses@.
  class
      ( Eq (UVarOf m # t)
      , RTraversable t
      , ZipMatch t
      , HasTypeConstraints t
      , HasQuantifiedVar t
      , Monad m
      , MonadQuantify (TypeConstraintsOf t) (QVar t) m
      ) =>
      Unify m t
      where
      -- | The implementation for unification variables binding and lookup
      binding :: BindingDict (UVarOf m) m t
  
      -- | Handles a unification error.
      --
      -- If 'unifyError' is called then unification has failed.
      -- A compiler implementation may present an error message based on the provided 'UnifyError' when this occurs.
      unifyError :: UnifyError t # UVarOf m -> m a
      default unifyError ::
          (MonadError (e # Pure) m, HSubset' e (UnifyError t)) =>
          UnifyError t # UVarOf m ->
          m a
      unifyError e =
          htraverse (Proxy @(Unify m) #> applyBindings) e
              >>= throwError
                  . (hSubset #)
              \\ unifyRecursive (Proxy @m) (Proxy @t)
  
      -- | What to do when top-levels of terms being unified do not match.
      --
      -- Usually this will cause a 'unifyError'.
      --
      -- Some AST terms could be equivalent despite not matching structurally,
      -- like record field extentions with the fields ordered differently.
      -- Those would override the default implementation to handle the unification of mismatching structures.
      structureMismatch ::
          (forall c. Unify m c => UVarOf m # c -> UVarOf m # c -> m (UVarOf m # c)) ->
          t # UVarOf m ->
          t # UVarOf m ->
          m ()
      structureMismatch _ x y = unifyError (Mismatch x y)
  
      -- TODO: Putting documentation here causes duplication in the haddock documentation
      unifyRecursive :: Proxy m -> RecMethod (Unify m) t
      {-# INLINE unifyRecursive #-}
      default unifyRecursive :: HNodesConstraint t (Unify m) => Proxy m -> RecMethod (Unify m) t
      unifyRecursive _ _ = Dict
  
  instance Recursive (Unify m) where
      {-# INLINE recurse #-}
      recurse = unifyRecursive (Proxy @m) . proxyArgument
  
  -- | A class for unification monads with scope levels
  class Unify m t => UnifyGen m t where
      -- | Get the current scope constraint
      scopeConstraints :: Proxy t -> m (TypeConstraintsOf t)
  
      unifyGenRecursive :: Proxy m -> RecMethod (UnifyGen m) t
      {-# INLINE unifyGenRecursive #-}
      default unifyGenRecursive ::
          HNodesConstraint t (UnifyGen m) => Proxy m -> RecMethod (UnifyGen m) t
      unifyGenRecursive _ _ = Dict
  
  instance Recursive (UnifyGen m) where
      {-# INLINE recurse #-}
      recurse = unifyGenRecursive (Proxy @m) . proxyArgument
  
  -- | Look up a variable, and return last variable pointing to result.
  -- Prunes all variables on way to point to the last variable
  -- (path-compression ala union-find).
  {-# INLINE semiPruneLookup #-}
  semiPruneLookup ::
      Unify m t =>
      UVarOf m # t ->
      m (UVarOf m # t, UTerm (UVarOf m) # t)
  semiPruneLookup v0 =
      lookupVar binding v0
          >>= \case
              UToVar v1 ->
                  do
                      (v, r) <- semiPruneLookup v1
                      bindVar binding v0 (UToVar v)
                      pure (v, r)
              t -> pure (v0, t)
  
  -- | Resolve a term from a unification variable.
  --
  -- Note that this must be done after
  -- all unifications involving the term and its children are done,
  -- as it replaces unification state with cached resolved terms.
  {-# INLINE applyBindings #-}
  applyBindings ::
      forall m t.
      Unify m t =>
      UVarOf m # t ->
      m (Pure # t)
  applyBindings v0 =
      do
          (v1, x) <- semiPruneLookup v0
          let result r = r <$ bindVar binding v1 (UResolved r)
          let quantify c =
                  newQuantifiedVariable c
                      <&> (_Pure . quantifiedVar #)
                      >>= result
          case x of
              UResolving t -> occursError v1 t
              UResolved t -> pure t
              UUnbound c -> quantify c
              USkolem c -> quantify c
              UTerm b ->
                  do
                      (r, anyChild) <-
                          htraverse
                              ( Proxy @(Unify m) #>
                                  \c ->
                                      do
                                          get >>= lift . (`unless` bindVar binding v1 (UResolving b))
                                          put True
                                          applyBindings c & lift
                              )
                              (b ^. uBody)
                              & (`runStateT` False)
                              \\ unifyRecursive (Proxy @m) (Proxy @t)
                      _Pure # r & if anyChild then result else pure
              UToVar{} -> error "lookup not expected to result in var"
              UConverted{} -> error "conversion state not expected in applyBindings"
              UInstantiated{} ->
                  -- This can happen in alphaEq,
                  -- where UInstantiated marks that var from one side matches var in the other.
                  quantify mempty
  
  -- | Format and throw an occurs check error
  occursError ::
      Unify m t =>
      UVarOf m # t ->
      UTermBody (UVarOf m) # t ->
      m a
  occursError v (UTermBody c b) =
      do
          q <- newQuantifiedVariable c
          bindVar binding v (UResolved (_Pure . quantifiedVar # q))
          unifyError (Occurs (quantifiedVar # q) b)
  