{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE MonadComprehensions #-}

module Language.Haskell.Exference.Core.Internal.Unify
  ( unify
  , unifyOffset
  , unifyRight
  , unifyRightEqs
  , unifyRightOffset
  , TypeEq (..)
  )
where



import Language.Haskell.Exference.Core.Types
import Language.Haskell.Exference.Core.TypeUtils
import qualified Data.Map.Strict as M
import qualified Data.IntMap.Strict as IntMap
import Data.Maybe

-- import Debug.Hood.Observe
import Debug.Trace



data TypeEq = TypeEq !HsType
                     !HsType

data UniState1 = UniState1 [TypeEq] Substs
data UniState2 = UniState2 [TypeEq] Substs Substs

occursIn :: TVarId -> HsType -> Bool
occursIn i (TypeVar j)         = i==j
occursIn _ (TypeConstant _)    = False
occursIn _ (TypeCons _)        = False
occursIn i (TypeArrow t1 t2)   = occursIn i t1 || occursIn i t2
occursIn i (TypeApp t1 t2)     = occursIn i t1 || occursIn i t2
occursIn i (TypeForall js _ t) = (i `notElem` js) && occursIn i t

-- Maybe (Either (Either Subst Subst) [TypeEq])
data StepResult
  = StepFailed
  | StepClear
  | StepNewEqs ![TypeEq]
  | StepLeftSubst {-# UNPACK #-} !Subst
  | StepRightSubst {-# UNPACK #-} !Subst

-- unification of types.
-- returns two substitutions: one for variables in the first type,
-- one for variables in the second. In the symmetric case,
-- substituting the right-hand (second) type will be preferred.
-- examples:
-- unify v C -> ([v=>C], [])
-- unify C v -> ([], [v=>C])
-- unify v w -> ([], [w=>v])
{-# INLINE unify #-}
unify :: HsType -> HsType -> Maybe (Substs, Substs)
unify ut1 ut2 = unify' $ UniState2 [TypeEq ut1 ut2] IntMap.empty IntMap.empty
  where
    unify' :: UniState2 -> Maybe (Substs, Substs)
    unify' (UniState2 [] l r)      = Just (l, r)
    unify' (UniState2 (x:xr) l r) = case uniStep x of
      StepFailed                       -> Nothing
      StepClear                        -> unify' $ UniState2 xr l r
      StepNewEqs eqs                   -> unify' $ UniState2 (eqs++xr) l r
      StepLeftSubst  subst@(Subst i t) ->
        unify' $ let f = applySubst subst in UniState2
          [ TypeEq (f a) b | TypeEq a b <- xr ]
          (IntMap.insert i t $ IntMap.map f l)
          r
      StepRightSubst subst@(Subst i t) ->
        unify' $ let f = applySubst subst in UniState2
          [ TypeEq a (f b) | TypeEq a b <- xr ]
          l
          (IntMap.insert i t $ IntMap.map f r)
      {-
      \ms -> unify' $ case ms of
        Left (Left subst@(Subst i t)) -> let f = applySubst subst in UniState2
          [ TypeEq (f a) b | TypeEq a b <- xr ]
          (IntMap.insert i t $ IntMap.map f l)
          r
        Left (Right subst@(Subst i t)) -> let f = applySubst subst in UniState2
          [ TypeEq a (f b) | TypeEq a b <- xr ]
          l
          (IntMap.insert i t $ IntMap.map f r)
        Right eqs -> UniState2 (eqs++xr) l r
      ) -- up here is the same control flow again, with the deconsing and the growing of the tail... maybe make your own hof for now? oh wait you also modify the tail
      -}
    uniStep :: TypeEq -> StepResult
    -- uniStep (TypeEq (TypeVar i1) (TypeVar i2)) | i1==i2 = Just (Right [])
    uniStep (TypeEq t1 (TypeVar i2)) = if occursIn i2 t1
      then StepFailed
      else StepRightSubst $ Subst i2 t1
    uniStep (TypeEq (TypeVar i1) t2) = if occursIn i1 t2
      then StepFailed
      else StepLeftSubst $ Subst i1 t2
    uniStep (TypeEq (TypeConstant i1) (TypeConstant i2)) | i1 == i2 = StepClear
    uniStep (TypeEq (TypeCons s1) (TypeCons s2)) | s1 == s2 = StepClear
    uniStep (TypeEq (TypeArrow t1 t2) (TypeArrow t3 t4)) = StepNewEqs [TypeEq t1 t3, TypeEq t2 t4]
    uniStep (TypeEq (TypeApp t1 t2) (TypeApp t3 t4)) = StepNewEqs [TypeEq t1 t3, TypeEq t2 t4]
    -- TODO TypeForall unification
    -- THIS IS WRONG; WE IGNORE FORALLS FOR THE MOMENT
    -- uniStep (TypeEq (TypeForall _ _ t1) (t2)) = uniStep $ TypeEq t1 t2
    -- uniStep (TypeEq (t1) (TypeForall _ _ t2)) = uniStep $ TypeEq t1 t2
    uniStep (TypeEq TypeForall{} _) = error "this probably does not ever happen anymore, as foralls are treated previously."
    uniStep (TypeEq _ TypeForall{}) = error "this probably does not ever happen anymore, as foralls are treated previously."
    uniStep _ = StepFailed


{-# INLINE unifyOffset #-}                   -- left, rightOffset
unifyOffset :: HsType -> HsTypeOffset -> Maybe (Substs, Substs)
unifyOffset ut1 (HsTypeOffset ut2 offset) = unify' $ UniState2 [TypeEq ut1 ut2]
                                                               IntMap.empty
                                                               IntMap.empty
  where
    unify' :: UniState2 -> Maybe (Substs, Substs)
    unify' (UniState2 [] l r)      = Just (l, r)
    unify' (UniState2 (x:xr) l r) = uniStep x >>= (
      \ms -> unify' $ case ms of
        Left (Left subst@(Subst i t)) -> let f = applySubst subst in UniState2
          [ TypeEq (f a) b | TypeEq a b <- xr ]
          (IntMap.insert i t $ IntMap.map f l)
          r
        Left (Right (substInternal, Subst substExtI substExtT)) ->
          let f = applySubst substInternal
          in UniState2
            [ TypeEq a (f b) | TypeEq a b <- xr ]
            l
            (IntMap.insert substExtI substExtT $ IntMap.map f r)
        Right eqs -> UniState2 (eqs++xr) l r
      )
    uniStep :: TypeEq -> Maybe (Either (Either Subst (Subst, Subst)) [TypeEq])
    uniStep (TypeEq (TypeVar i1) (TypeVar i2)) | i1==offset+i2 = Just (Right [])
    uniStep (TypeEq (t1) (TypeVar i2)) = if occursIn (offset+i2) t1
      then Nothing
      else Just $ Left $ Right (Subst i2 t1, Subst (i2+offset) t1)
    uniStep (TypeEq (TypeVar i1) (t2)) = if occursIn (i1-offset) t2
      then Nothing
      else Just $ Left $ Left $ Subst i1 (incVarIds (+offset) t2)
    uniStep (TypeEq (TypeConstant i1) (TypeConstant i2)) | i1==i2 = Just (Right [])
    uniStep (TypeEq (TypeCons s1) (TypeCons s2)) | s1==s2 = Just (Right [])
    uniStep (TypeEq (TypeArrow t1 t2) (TypeArrow t3 t4)) = Just (Right [TypeEq t1 t3, TypeEq t2 t4])
    uniStep (TypeEq (TypeApp t1 t2) (TypeApp t3 t4)) = Just (Right [TypeEq t1 t3, TypeEq t2 t4])
    -- TODO TypeForall unification
    -- THIS IS WRONG; WE IGNORE FORALLS FOR THE MOMENT
    uniStep (TypeEq (TypeForall _ _ t1) (t2)) = uniStep $ TypeEq t1 t2
    uniStep (TypeEq (t1) (TypeForall _ _ t2)) = uniStep $ TypeEq t1 t2
    uniStep _ = Nothing

-- treats the variables in the first parameter as constants, and returns
-- the variable bindings for the second parameter that unify both types.
{-# INLINE unifyRight #-}
unifyRight :: HsType -> HsType -> Maybe Substs
unifyRight ut1 ut2 = unifyRightEqs [TypeEq ut1 ut2]

{-# INLINE unifyRightEqs #-}
unifyRightEqs :: [TypeEq] -> Maybe Substs
unifyRightEqs teqs = unify' $ UniState1 teqs IntMap.empty
  where
    unify' :: UniState1 -> Maybe Substs
    unify' (UniState1 [] x) = Just x
    unify' (UniState1 (x:xr) ss) = uniStepRight x >>= (
      \r -> unify' $ case r of
        Left subst@(Subst i t) -> let f = applySubst subst in UniState1
          [ TypeEq a (f b) | TypeEq a b <- xr]
          (IntMap.insert i t $ IntMap.map f ss)
        Right eqs -> UniState1 (eqs++xr) ss
      )


-- treats the variables in the first parameter as constants, and returns
-- the variable bindings for the second parameter that unify both types.
{-# INLINE unifyRightOffset #-}
unifyRightOffset :: HsType -> HsTypeOffset -> Maybe Substs
unifyRightOffset ut1 (HsTypeOffset ut2 offset) = unify' $ UniState1 [TypeEq ut1 ut2] IntMap.empty
  where
    unify' :: UniState1 -> Maybe Substs
    unify' (UniState1 [] x) = Just x
    unify' (UniState1 (x:xr) ss) = uniStep x >>= (
      \r -> unify' $ case r of
        Left (substInternal, Subst substExtI substExtT) -> UniState1
          [ TypeEq a (applySubst substInternal b) | TypeEq a b <- xr]
          (IntMap.insert substExtI substExtT ss)
        Right eqs -> UniState1 (eqs++xr) ss
      )
    uniStep :: TypeEq -> Maybe (Either (Subst, Subst) [TypeEq])
    uniStep (TypeEq (TypeVar i1) (TypeVar i2)) | i1==offset+i2 = Just (Right [])
    uniStep (TypeEq (t1) (TypeVar i2)) = if occursIn (i2+offset) t1
      then Nothing
      else Just $ Left (Subst i2 t1, Subst (i2+offset) t1)
    uniStep (TypeEq (TypeVar _) _) = Nothing
    uniStep (TypeEq (TypeConstant i1) (TypeConstant i2)) | i1==i2 = Just (Right [])
    uniStep (TypeEq (TypeCons s1) (TypeCons s2)) | s1==s2 = Just (Right [])
    uniStep (TypeEq (TypeArrow t1 t2) (TypeArrow t3 t4)) = Just (Right [TypeEq t1 t3, TypeEq t2 t4])
    uniStep (TypeEq (TypeApp t1 t2) (TypeApp t3 t4)) = Just (Right [TypeEq t1 t3, TypeEq t2 t4])
    -- TODO TypeForall unification
    -- THIS IS WRONG; WE IGNORE FORALLS FOR THE MOMENT
    uniStep (TypeEq (TypeForall _ _ t1) t2)   = uniStep $ TypeEq t1 t2
    uniStep (TypeEq (t1) (TypeForall _ _ t2)) = uniStep $ TypeEq t1 t2
    uniStep _ = Nothing


uniStepRight :: TypeEq -> Maybe (Either Subst [TypeEq])
uniStepRight (TypeEq (TypeVar i1) (TypeVar i2)) | i1==i2 = Just (Right [])
uniStepRight (TypeEq (t1) (TypeVar i2)) = if occursIn i2 t1
  then Nothing
  else Just $ Left $ Subst i2 t1
uniStepRight (TypeEq (TypeVar _) _) = Nothing
uniStepRight (TypeEq (TypeConstant i1) (TypeConstant i2)) | i1==i2 = Just (Right [])
uniStepRight (TypeEq (TypeCons s1) (TypeCons s2)) | s1==s2 = Just (Right [])
uniStepRight (TypeEq (TypeArrow t1 t2) (TypeArrow t3 t4)) = Just (Right [TypeEq t1 t3, TypeEq t2 t4])
uniStepRight (TypeEq (TypeApp t1 t2) (TypeApp t3 t4)) = Just (Right [TypeEq t1 t3, TypeEq t2 t4])
-- TODO TypeForall unification
-- THIS IS WRONG; WE IGNORE FORALLS FOR THE MOMENT
uniStepRight (TypeEq (TypeForall _ _ t1) t2)   = uniStepRight $ TypeEq t1 t2
uniStepRight (TypeEq (t1) (TypeForall _ _ t2)) = uniStepRight $ TypeEq t1 t2
uniStepRight _ = Nothing



--unifyDist :: HsType -> HsType -> Maybe Substs
--unifyDist t1 t2 = unify t1 $ distinctify t1 t2

-- tries to unify two types, under the assumption that one is supposed
-- to serve as the parameter (at some levels) to the other one (when at
-- least one is a function, that is.)
-- either is BROKEN, or does something other than i expected.
-- not used (or needed) anyway, so.. deletion candidate.
{-
inflateUnify :: HsType -> HsType -> [HsType]
inflateUnify t1 t2 =
  let d1 = arrowDepth t1
      d2 = arrowDepth t2
  in if d1 > d2
    then f t1 $ inflateTo d1 t2
    else f t2 $ inflateTo d2 t1
  where
    inflateTo :: Int -> HsType -> [HsType]
    inflateTo n t = map (createArrow t (map TypeVar [1000..998+n])) [0..n-1]
    createArrow t varIds i =
      let (l,r) = splitAt i varIds
          types = l++[t]++r
      in foldr1 TypeArrow types
    f :: HsType -> [HsType] -> [HsType]
    f ft1 ft2s = catMaybes [g ft1 (distinctify ft1 ft2) | ft2 <- ft2s]
    g :: HsType -> HsType -> Maybe HsType
    g gt1 gt2 = fmap (\subst -> reduceIds $ applySubsts subst gt1) (unify gt1 gt2)
-}