-----------------------------------------------------------------------------
-- |
  Module       :   Disco .  . Unify
-- Copyright   :  disco team and contributors
-- Maintainer  :  
--
  SPDX - License - Identifier : BSD-3 - Clause
--
-- Unification.
--
-----------------------------------------------------------------------------

module Disco.Typecheck.Unify where

import           Unbound.Generics.LocallyNameless (Name, fv)

import           Control.Lens                     (anyOf)
import           Control.Monad.State
import qualified Data.Map                         as M
import           Data.Set                         (Set)
import qualified Data.Set                         as S

import           Disco.Subst
import           Disco.Types

-- XXX todo: might be better if unification took sorts into account
-- directly.  As it is, however, I think it works properly;
  e.g. suppose we have a with sort { sub } and we unify it with  .
-- unify will just return a substitution [a |-> Bool].  But then when
  we call extendSubst , and in particular substSortMap , the sort { sub }
  will be applied to  and decomposed which will throw an error .

-- | Given a list of equations between types, return a substitution
--   which makes all the equations satisfied (or fail if it is not
--   possible).
--
--   This is not the most efficient way to implement unification but
--   it is simple.
unify :: TyDefCtx -> [(Type, Type)] -> Maybe S
unify = unify' (==)

-- | Given a list of equations between types, return a substitution
--   which makes all the equations equal *up to* identifying all base
--   types.  So, for example, Int = Nat weakly unifies but Int = (Int
--   -> Int) does not.  This is used to check whether subtyping
--   constraints are structurally sound before doing constraint
--   simplification/solving, to ensure termination.
weakUnify :: TyDefCtx -> [(Type, Type)] -> Maybe S
weakUnify = unify' (\_ _ -> True)

-- | Given a list of equations between types, return a substitution
--   which makes all the equations satisfied (or fail if it is not
--   possible), up to the given comparison on base types.
unify' :: (BaseTy -> BaseTy -> Bool) -> TyDefCtx
       -> [(Type, Type)] -> Maybe S
unify' baseEq tyDefns eqs = evalStateT (go eqs) S.empty
  where
    go :: [(Type, Type)] -> StateT (Set (Type,Type)) Maybe S
    go [] = return idS
    go (e:es) = do
      u <- unifyOne e
      case u of
        Left sub    -> (@@ sub) <$> go (applySubst sub es)
        Right newEs -> go (newEs ++ es)

    unifyOne :: (Type, Type) -> StateT (Set (Type,Type)) Maybe (Either S [(Type, Type)])
    unifyOne pair = do
      seen <- get
      case pair `S.member` seen of
        True  -> return $ Left idS
        False -> unifyOne' pair

    unifyOne' :: (Type, Type) -> StateT (Set (Type,Type)) Maybe (Either S [(Type, Type)])

    unifyOne' (ty1, ty2)
      | ty1 == ty2 = return $ Left idS

    unifyOne' (TyVar x, ty2)
      | occurs x ty2 = mzero
      | otherwise    = return $ Left (x |-> ty2)
    unifyOne' (ty1, x@(TyVar _))
      = unifyOne (x, ty1)

    -- At this point we know ty2 isn't the same skolem nor a unification variable.
  Skolems do n't unify with anything .
    unifyOne' (TySkolem _, _) = mzero
    unifyOne' (_, TySkolem _) = mzero

  Unify two container types : unify the container descriptors as
    -- well as the type arguments
    unifyOne' p@(TyCon (CContainer ctr1) tys1, TyCon (CContainer ctr2) tys2) = do
      modify (S.insert p)
      return $ Right ((TyAtom ctr1, TyAtom ctr2) : zip tys1 tys2)

  If one of the types to be unified is a user - defined type ,
    -- unfold its definition before continuing the matching
    unifyOne' p@(TyCon (CUser t) tys1, ty2) = do
      modify (S.insert p)
      case M.lookup t tyDefns of
        Nothing                 -> mzero
        Just (TyDefBody _ body) -> return $ Right [(body tys1, ty2)]

    unifyOne' p@(ty1, TyCon (CUser t) tys2) = do
      modify (S.insert p)
      case M.lookup t tyDefns of
        Nothing                 -> mzero
        Just (TyDefBody _ body) -> return $ Right [(ty1, body tys2)]

    -- Unify any other pair of type constructor applications: the type
    -- constructors must match exactly
    unifyOne' p@(TyCon c1 tys1, TyCon c2 tys2)
      | c1 == c2  = do
          modify (S.insert p)
          return $ Right (zip tys1 tys2)
      | otherwise = mzero
    unifyOne' (TyAtom (ABase b1), TyAtom (ABase b2))
      | baseEq b1 b2 = return $ Left idS
      | otherwise    = mzero
    unifyOne' _ = mzero  -- Atom = Cons


equate :: TyDefCtx -> [Type] -> Maybe S
equate tyDefns tys = unify tyDefns eqns
  where
    eqns = zip tys (tail tys)

occurs :: Name Type -> Type -> Bool
occurs x = anyOf fv (==x)


unifyAtoms :: TyDefCtx -> [Atom] -> Maybe (Substitution Atom)
unifyAtoms tyDefns = fmap (fmap fromTyAtom) . equate tyDefns . map TyAtom
  where
    fromTyAtom (TyAtom a) = a
    fromTyAtom _          = error "fromTyAtom on non-TyAtom!"

unifyUAtoms :: TyDefCtx -> [UAtom] -> Maybe (Substitution UAtom)
unifyUAtoms tyDefns = fmap (fmap fromTyAtom) . equate tyDefns . map (TyAtom . uatomToAtom)
  where
    fromTyAtom (TyAtom (ABase b))    = UB b
    fromTyAtom (TyAtom (AVar (U v))) = UV v
    fromTyAtom _                     = error "fromTyAtom on wrong thing!"

---------------------------------------------------------------------------
 |
 Copyright   :  disco team and contributors
 Maintainer  :  


 Unification.

---------------------------------------------------------------------------
 XXX todo: might be better if unification took sorts into account
 directly.  As it is, however, I think it works properly;
 unify will just return a substitution [a |-> Bool].  But then when
 | Given a list of equations between types, return a substitution
   which makes all the equations satisfied (or fail if it is not
   possible).

   This is not the most efficient way to implement unification but
   it is simple.
 | Given a list of equations between types, return a substitution
   which makes all the equations equal *up to* identifying all base
   types.  So, for example, Int = Nat weakly unifies but Int = (Int
   -> Int) does not.  This is used to check whether subtyping
   constraints are structurally sound before doing constraint
   simplification/solving, to ensure termination.
 | Given a list of equations between types, return a substitution
   which makes all the equations satisfied (or fail if it is not
   possible), up to the given comparison on base types.
 At this point we know ty2 isn't the same skolem nor a unification variable.
 well as the type arguments
 unfold its definition before continuing the matching
 Unify any other pair of type constructor applications: the type
 constructors must match exactly
 Atom = Cons
  Module       :   Disco .  . Unify
  SPDX - License - Identifier : BSD-3 - Clause

module Disco.Typecheck.Unify where

import           Unbound.Generics.LocallyNameless (Name, fv)

import           Control.Lens                     (anyOf)
import           Control.Monad.State
import qualified Data.Map                         as M
import           Data.Set                         (Set)
import qualified Data.Set                         as S

import           Disco.Subst
import           Disco.Types

  e.g. suppose we have a with sort { sub } and we unify it with  .
  we call extendSubst , and in particular substSortMap , the sort { sub }
  will be applied to  and decomposed which will throw an error .

unify :: TyDefCtx -> [(Type, Type)] -> Maybe S
unify = unify' (==)

weakUnify :: TyDefCtx -> [(Type, Type)] -> Maybe S
weakUnify = unify' (\_ _ -> True)

unify' :: (BaseTy -> BaseTy -> Bool) -> TyDefCtx
       -> [(Type, Type)] -> Maybe S
unify' baseEq tyDefns eqs = evalStateT (go eqs) S.empty
  where
    go :: [(Type, Type)] -> StateT (Set (Type,Type)) Maybe S
    go [] = return idS
    go (e:es) = do
      u <- unifyOne e
      case u of
        Left sub    -> (@@ sub) <$> go (applySubst sub es)
        Right newEs -> go (newEs ++ es)

    unifyOne :: (Type, Type) -> StateT (Set (Type,Type)) Maybe (Either S [(Type, Type)])
    unifyOne pair = do
      seen <- get
      case pair `S.member` seen of
        True  -> return $ Left idS
        False -> unifyOne' pair

    unifyOne' :: (Type, Type) -> StateT (Set (Type,Type)) Maybe (Either S [(Type, Type)])

    unifyOne' (ty1, ty2)
      | ty1 == ty2 = return $ Left idS

    unifyOne' (TyVar x, ty2)
      | occurs x ty2 = mzero
      | otherwise    = return $ Left (x |-> ty2)
    unifyOne' (ty1, x@(TyVar _))
      = unifyOne (x, ty1)

  Skolems do n't unify with anything .
    unifyOne' (TySkolem _, _) = mzero
    unifyOne' (_, TySkolem _) = mzero

  Unify two container types : unify the container descriptors as
    unifyOne' p@(TyCon (CContainer ctr1) tys1, TyCon (CContainer ctr2) tys2) = do
      modify (S.insert p)
      return $ Right ((TyAtom ctr1, TyAtom ctr2) : zip tys1 tys2)

  If one of the types to be unified is a user - defined type ,
    unifyOne' p@(TyCon (CUser t) tys1, ty2) = do
      modify (S.insert p)
      case M.lookup t tyDefns of
        Nothing                 -> mzero
        Just (TyDefBody _ body) -> return $ Right [(body tys1, ty2)]

    unifyOne' p@(ty1, TyCon (CUser t) tys2) = do
      modify (S.insert p)
      case M.lookup t tyDefns of
        Nothing                 -> mzero
        Just (TyDefBody _ body) -> return $ Right [(ty1, body tys2)]

    unifyOne' p@(TyCon c1 tys1, TyCon c2 tys2)
      | c1 == c2  = do
          modify (S.insert p)
          return $ Right (zip tys1 tys2)
      | otherwise = mzero
    unifyOne' (TyAtom (ABase b1), TyAtom (ABase b2))
      | baseEq b1 b2 = return $ Left idS
      | otherwise    = mzero


equate :: TyDefCtx -> [Type] -> Maybe S
equate tyDefns tys = unify tyDefns eqns
  where
    eqns = zip tys (tail tys)

occurs :: Name Type -> Type -> Bool
occurs x = anyOf fv (==x)


unifyAtoms :: TyDefCtx -> [Atom] -> Maybe (Substitution Atom)
unifyAtoms tyDefns = fmap (fmap fromTyAtom) . equate tyDefns . map TyAtom
  where
    fromTyAtom (TyAtom a) = a
    fromTyAtom _          = error "fromTyAtom on non-TyAtom!"

unifyUAtoms :: TyDefCtx -> [UAtom] -> Maybe (Substitution UAtom)
unifyUAtoms tyDefns = fmap (fmap fromTyAtom) . equate tyDefns . map (TyAtom . uatomToAtom)
  where
    fromTyAtom (TyAtom (ABase b))    = UB b
    fromTyAtom (TyAtom (AVar (U v))) = UV v
    fromTyAtom _                     = error "fromTyAtom on wrong thing!"
