module GHC.TypeLits.Normalise.Unify
(
CoreSOP
, normaliseNat
, reifySOP
, SubstItem (..)
, TySubst
, CoreSubst
, substsSOP
, substsSubst
, UnifyResult (..)
, unifyNats
, unifiers
, fvSOP
)
where
import Data.Function (on)
import Data.List ((\\), intersect)
import Outputable (Outputable (..), (<+>), ($$), text)
import TcPluginM (TcPluginM, tcPluginTrace)
import TcRnMonad (Ct, ctEvidence, isGiven)
import TcTypeNats (typeNatAddTyCon, typeNatExpTyCon, typeNatMulTyCon,
typeNatSubTyCon)
import Type (TyVar, mkNumLitTy, mkTyConApp, mkTyVarTy, tcView)
import TypeRep (Type (..), TyLit (..))
import UniqSet (UniqSet, unionManyUniqSets, emptyUniqSet, unionUniqSets,
unitUniqSet)
import GHC.Type.Instances ()
import GHC.TypeLits.Normalise.SOP
import GHC.TypeLits (Nat)
type CoreSOP = SOP TyVar Type
type CoreProduct = Product TyVar Type
type CoreSymbol = Symbol TyVar Type
normaliseNat :: Type -> CoreSOP
normaliseNat ty | Just ty1 <- tcView ty = normaliseNat ty1
normaliseNat (TyVarTy v) = S [P [V v]]
normaliseNat (LitTy (NumTyLit i)) = S [P [I i]]
normaliseNat (TyConApp tc [x,y])
| tc == typeNatAddTyCon = mergeSOPAdd (normaliseNat x) (normaliseNat y)
| tc == typeNatSubTyCon = mergeSOPAdd (normaliseNat x)
(mergeSOPMul (S [P [I (1)]])
(normaliseNat y))
| tc == typeNatMulTyCon = mergeSOPMul (normaliseNat x) (normaliseNat y)
| tc == typeNatExpTyCon = normaliseExp (normaliseNat x) (normaliseNat y)
normaliseNat t = S [P [C t]]
reifySOP :: CoreSOP -> Type
reifySOP = combineP . map negateP . unS
where
negateP :: CoreProduct -> Either CoreProduct CoreProduct
negateP (P ((I i):ps)) | i < 0 = Left (P ps)
negateP ps = Right ps
combineP :: [Either CoreProduct CoreProduct] -> Type
combineP [] = mkNumLitTy 0
combineP [p] = either (\p' -> mkTyConApp typeNatSubTyCon
[mkNumLitTy 0, reifyProduct p'])
reifyProduct p
combineP (p:ps) = let es = combineP ps
in either (\x -> mkTyConApp typeNatSubTyCon
[es, reifyProduct x])
(\x -> mkTyConApp typeNatAddTyCon
[reifyProduct x, es])
p
reifyProduct :: CoreProduct -> Type
reifyProduct = foldr1 (\t1 t2 -> mkTyConApp typeNatMulTyCon [t1,t2])
. map reifySymbol . unP
reifySymbol :: CoreSymbol -> Type
reifySymbol (I i) = mkNumLitTy i
reifySymbol (C c) = c
reifySymbol (V v) = mkTyVarTy v
reifySymbol (E s p) = mkTyConApp typeNatExpTyCon [reifySOP s,reifyProduct p]
type CoreSubst = TySubst TyVar Type Ct
type TySubst v c n = [SubstItem v c n]
data SubstItem v c n = SubstItem { siVar :: v
, siSOP :: SOP v c
, siNote :: n
}
instance (Outputable v, Outputable c) => Outputable (SubstItem v c n) where
ppr si = ppr (siVar si) <+> text " := " <+> ppr (siSOP si)
substsSOP :: (Ord v, Ord c) => TySubst v c n -> SOP v c -> SOP v c
substsSOP [] u = u
substsSOP (si:s) u = substsSOP s (substSOP (siVar si) (siSOP si) u)
substSOP :: (Ord v, Ord c) => v -> SOP v c -> SOP v c -> SOP v c
substSOP tv e = foldr1 mergeSOPAdd . map (substProduct tv e) . unS
substProduct :: (Ord v, Ord c) => v -> SOP v c -> Product v c -> SOP v c
substProduct tv e = foldr1 mergeSOPMul . map (substSymbol tv e) . unP
substSymbol :: (Ord v, Ord c) => v -> SOP v c -> Symbol v c -> SOP v c
substSymbol _ _ s@(I _) = S [P [s]]
substSymbol _ _ s@(C _) = S [P [s]]
substSymbol tv e (V tv')
| tv == tv' = e
| otherwise = S [P [V tv']]
substSymbol tv e (E s p) = normaliseExp (substSOP tv e s) (substProduct tv e p)
substsSubst :: (Ord v, Ord c) => TySubst v c n -> TySubst v c n -> TySubst v c n
substsSubst s = map (\si -> si {siSOP = substsSOP s (siSOP si)})
data UnifyResult
= Win
| Lose
| Draw CoreSubst
instance Outputable UnifyResult where
ppr Win = text "Win"
ppr (Draw subst) = text "Draw" <+> ppr subst
ppr Lose = text "Lose"
unifyNats :: Ct -> CoreSOP -> CoreSOP -> TcPluginM UnifyResult
unifyNats ct u v = do
tcPluginTrace "unifyNats" (ppr ct $$ ppr u $$ ppr v)
return (unifyNats' ct u v)
unifyNats' :: Ct -> CoreSOP -> CoreSOP -> UnifyResult
unifyNats' ct u v
| eqFV u v = if u == v then Win else Lose
| otherwise = Draw (unifiers ct u v)
unifiers :: Ct -> CoreSOP -> CoreSOP -> CoreSubst
unifiers ct (S [P [V x]]) s
| isGiven (ctEvidence ct) = [SubstItem x s ct]
| otherwise = []
unifiers ct s (S [P [V x]])
| isGiven (ctEvidence ct) = [SubstItem x s ct]
| otherwise = []
unifiers ct u v = unifiers' ct u v
unifiers' :: Ct -> CoreSOP -> CoreSOP -> CoreSubst
unifiers' ct (S [P [V x]]) (S []) = [SubstItem x (S [P [I 0]]) ct]
unifiers' ct (S []) (S [P [V x]]) = [SubstItem x (S [P [I 0]]) ct]
unifiers' ct (S [P [V x]]) s = [SubstItem x s ct]
unifiers' ct s (S [P [V x]]) = [SubstItem x s ct]
unifiers' ct (S [P (p:ps1)]) (S [P (p':ps2)])
| p == p' = unifiers' ct (S [P ps1]) (S [P ps2])
| otherwise = []
unifiers' ct (S ps1) (S ps2)
| null psx = []
| otherwise = unifiers' ct (S (ps1 \\ psx)) (S (ps2 \\ psx))
where
psx = intersect ps1 ps2
fvSOP :: CoreSOP -> UniqSet TyVar
fvSOP = unionManyUniqSets . map fvProduct . unS
fvProduct :: CoreProduct -> UniqSet TyVar
fvProduct = unionManyUniqSets . map fvSymbol . unP
fvSymbol :: CoreSymbol -> UniqSet TyVar
fvSymbol (I _) = emptyUniqSet
fvSymbol (C _) = emptyUniqSet
fvSymbol (V v) = unitUniqSet v
fvSymbol (E s p) = fvSOP s `unionUniqSets` fvProduct p
eqFV :: CoreSOP -> CoreSOP -> Bool
eqFV = (==) `on` fvSOP