-----------------------------------------------------------------------------
-- |
-- Module      :  Disco.Typecheck.Unify
-- Copyright   :  disco team and contributors
-- Maintainer  :  byorgey@gmail.com
--
-- 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 Bool.
-- 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 Bool 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 don'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!"