{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ViewPatterns #-}
module Data.Record.Anon.Internal.Plugin.TC.TyConSubst (
TyConSubst
, mkTyConSubst
, splitTyConApp_upTo
) where
import Data.Bifunctor
import Data.Either (partitionEithers)
import Data.Foldable (toList, asum)
import Data.List.NonEmpty (NonEmpty(..))
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Record.Anon.Internal.Plugin.TC.EquivClasses
import Data.Record.Anon.Internal.Plugin.TC.GhcTcPluginAPI hiding ((<>))
data TyConSubst = TyConSubst {
TyConSubst -> Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstMap :: Map TcTyVar (NonEmpty (TyCon, [Type]))
, TyConSubst -> Map TcTyVar TcTyVar
tyConSubstCanon :: Map TcTyVar TcTyVar
}
tyConSubstEmpty :: Map TcTyVar TcTyVar -> TyConSubst
tyConSubstEmpty :: Map TcTyVar TcTyVar -> TyConSubst
tyConSubstEmpty Map TcTyVar TcTyVar
canon = TyConSubst {
tyConSubstMap :: Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstMap = Map TcTyVar (NonEmpty (TyCon, [Type]))
forall k a. Map k a
Map.empty
, tyConSubstCanon :: Map TcTyVar TcTyVar
tyConSubstCanon = Map TcTyVar TcTyVar
canon
}
tyConSubstLookup :: TcTyVar -> TyConSubst -> Maybe (NonEmpty (TyCon, [Type]))
tyConSubstLookup :: TcTyVar -> TyConSubst -> Maybe (NonEmpty (TyCon, [Type]))
tyConSubstLookup TcTyVar
var TyConSubst{Map TcTyVar (NonEmpty (TyCon, [Type]))
Map TcTyVar TcTyVar
tyConSubstMap :: TyConSubst -> Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstCanon :: TyConSubst -> Map TcTyVar TcTyVar
tyConSubstMap :: Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstCanon :: Map TcTyVar TcTyVar
..} = TcTyVar
-> Map TcTyVar (NonEmpty (TyCon, [Type]))
-> Maybe (NonEmpty (TyCon, [Type]))
forall k a. Ord k => k -> Map k a -> Maybe a
Map.lookup TcTyVar
var' Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstMap
where
var' :: TcTyVar
var' :: TcTyVar
var' = Map TcTyVar TcTyVar -> TcTyVar -> TcTyVar
forall a. Ord a => Map a a -> a -> a
canonicalize Map TcTyVar TcTyVar
tyConSubstCanon TcTyVar
var
tyConSubstExtend ::
[(TcTyVar, (TyCon, [Type]))]
-> TyConSubst -> TyConSubst
tyConSubstExtend :: [(TcTyVar, (TyCon, [Type]))] -> TyConSubst -> TyConSubst
tyConSubstExtend [(TcTyVar, (TyCon, [Type]))]
new subst :: TyConSubst
subst@TyConSubst{Map TcTyVar (NonEmpty (TyCon, [Type]))
Map TcTyVar TcTyVar
tyConSubstMap :: TyConSubst -> Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstCanon :: TyConSubst -> Map TcTyVar TcTyVar
tyConSubstMap :: Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstCanon :: Map TcTyVar TcTyVar
..} = TyConSubst
subst {
tyConSubstMap = Map.unionWith (<>)
(Map.fromList $ map (uncurry aux) new)
tyConSubstMap
}
where
aux :: TcTyVar -> (TyCon, [Type]) -> (TcTyVar, NonEmpty (TyCon, [Type]))
aux :: TcTyVar -> (TyCon, [Type]) -> (TcTyVar, NonEmpty (TyCon, [Type]))
aux TcTyVar
var (TyCon, [Type])
s = (Map TcTyVar TcTyVar -> TcTyVar -> TcTyVar
forall a. Ord a => Map a a -> a -> a
canonicalize Map TcTyVar TcTyVar
tyConSubstCanon TcTyVar
var, (TyCon, [Type])
s (TyCon, [Type]) -> [(TyCon, [Type])] -> NonEmpty (TyCon, [Type])
forall a. a -> [a] -> NonEmpty a
:| [])
data Classified = Classified {
Classified -> [(TcTyVar, (TyCon, [Type]))]
classifiedProductive :: [(TcTyVar, (TyCon, [Type]))]
, Classified -> [(TcTyVar, TcTyVar)]
classifiedExtendEquivClass :: [(TcTyVar, TcTyVar)]
, Classified -> [(TcTyVar, (TcTyVar, NonEmpty Type))]
classifiedReconsider :: [(TcTyVar, (TcTyVar, NonEmpty Type))]
}
instance Semigroup Classified where
Classified
c1 <> :: Classified -> Classified -> Classified
<> Classified
c2 = Classified {
classifiedProductive :: [(TcTyVar, (TyCon, [Type]))]
classifiedProductive = (Classified -> [(TcTyVar, (TyCon, [Type]))])
-> [(TcTyVar, (TyCon, [Type]))]
forall a. (Classified -> [a]) -> [a]
combine Classified -> [(TcTyVar, (TyCon, [Type]))]
classifiedProductive
, classifiedExtendEquivClass :: [(TcTyVar, TcTyVar)]
classifiedExtendEquivClass = (Classified -> [(TcTyVar, TcTyVar)]) -> [(TcTyVar, TcTyVar)]
forall a. (Classified -> [a]) -> [a]
combine Classified -> [(TcTyVar, TcTyVar)]
classifiedExtendEquivClass
, classifiedReconsider :: [(TcTyVar, (TcTyVar, NonEmpty Type))]
classifiedReconsider = (Classified -> [(TcTyVar, (TcTyVar, NonEmpty Type))])
-> [(TcTyVar, (TcTyVar, NonEmpty Type))]
forall a. (Classified -> [a]) -> [a]
combine Classified -> [(TcTyVar, (TcTyVar, NonEmpty Type))]
classifiedReconsider
}
where
combine :: (Classified -> [a]) -> [a]
combine :: forall a. (Classified -> [a]) -> [a]
combine Classified -> [a]
f = Classified -> [a]
f Classified
c1 [a] -> [a] -> [a]
forall a. [a] -> [a] -> [a]
++ Classified -> [a]
f Classified
c2
instance Monoid Classified where
mempty :: Classified
mempty = [(TcTyVar, (TyCon, [Type]))]
-> [(TcTyVar, TcTyVar)]
-> [(TcTyVar, (TcTyVar, NonEmpty Type))]
-> Classified
Classified [] [] []
productive :: TcTyVar -> (TyCon, [Type]) -> Classified
productive :: TcTyVar -> (TyCon, [Type]) -> Classified
productive TcTyVar
var (TyCon
tyCon, [Type]
args) = Classified
forall a. Monoid a => a
mempty {
classifiedProductive = [(var, (tyCon, args))]
}
extendEquivClass :: TcTyVar -> TcTyVar -> Classified
extendEquivClass :: TcTyVar -> TcTyVar -> Classified
extendEquivClass TcTyVar
var TcTyVar
var' = Classified
forall a. Monoid a => a
mempty {
classifiedExtendEquivClass = [(var, var')]
}
reconsider :: TcTyVar -> (TcTyVar, NonEmpty Type) -> Classified
reconsider :: TcTyVar -> (TcTyVar, NonEmpty Type) -> Classified
reconsider TcTyVar
var (TcTyVar
var', NonEmpty Type
args) = Classified
forall a. Monoid a => a
mempty {
classifiedReconsider = [(var, (var', args))]
}
classify :: [Ct] -> Classified
classify :: [Ct] -> Classified
classify = Classified -> [Ct] -> Classified
go Classified
forall a. Monoid a => a
mempty
where
go :: Classified -> [Ct] -> Classified
go :: Classified -> [Ct] -> Classified
go Classified
acc [] = Classified
acc
go Classified
acc (Ct
c:[Ct]
cs) =
case Ct -> Maybe (TcTyVar, Type)
isCanonicalVarEq Ct
c of
Just (TcTyVar
var, Type -> (Type, [Type])
splitAppTys -> (Type
fn, [Type]
args))
| Just TyCon
tyCon <- Type -> Maybe TyCon
tyConAppTyCon_maybe Type
fn ->
Classified -> [Ct] -> Classified
go (TcTyVar -> (TyCon, [Type]) -> Classified
productive TcTyVar
var (TyCon
tyCon, [Type]
args) Classified -> Classified -> Classified
forall a. Semigroup a => a -> a -> a
<> Classified
acc) [Ct]
cs
| Just TcTyVar
var' <- Type -> Maybe TcTyVar
getTyVar_maybe Type
fn, [Type] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [Type]
args ->
Classified -> [Ct] -> Classified
go (TcTyVar -> TcTyVar -> Classified
extendEquivClass TcTyVar
var TcTyVar
var' Classified -> Classified -> Classified
forall a. Semigroup a => a -> a -> a
<> Classified
acc) [Ct]
cs
| Just TcTyVar
var' <- Type -> Maybe TcTyVar
getTyVar_maybe Type
fn, Type
x:[Type]
xs <- [Type]
args ->
Classified -> [Ct] -> Classified
go (TcTyVar -> (TcTyVar, NonEmpty Type) -> Classified
reconsider TcTyVar
var (TcTyVar
var', Type
x Type -> [Type] -> NonEmpty Type
forall a. a -> [a] -> NonEmpty a
:| [Type]
xs) Classified -> Classified -> Classified
forall a. Semigroup a => a -> a -> a
<> Classified
acc) [Ct]
cs
Maybe (TcTyVar, Type)
_otherwise ->
Classified -> [Ct] -> Classified
go Classified
acc [Ct]
cs
process :: Classified -> TyConSubst
process :: Classified -> TyConSubst
process Classified{[(TcTyVar, (TcTyVar, NonEmpty Type))]
[(TcTyVar, (TyCon, [Type]))]
[(TcTyVar, TcTyVar)]
classifiedProductive :: Classified -> [(TcTyVar, (TyCon, [Type]))]
classifiedExtendEquivClass :: Classified -> [(TcTyVar, TcTyVar)]
classifiedReconsider :: Classified -> [(TcTyVar, (TcTyVar, NonEmpty Type))]
classifiedProductive :: [(TcTyVar, (TyCon, [Type]))]
classifiedExtendEquivClass :: [(TcTyVar, TcTyVar)]
classifiedReconsider :: [(TcTyVar, (TcTyVar, NonEmpty Type))]
..} =
TyConSubst -> [(TcTyVar, (TcTyVar, NonEmpty Type))] -> TyConSubst
go TyConSubst
initSubst [(TcTyVar, (TcTyVar, NonEmpty Type))]
classifiedReconsider
where
initSubst :: TyConSubst
initSubst :: TyConSubst
initSubst =
[(TcTyVar, (TyCon, [Type]))] -> TyConSubst -> TyConSubst
tyConSubstExtend [(TcTyVar, (TyCon, [Type]))]
classifiedProductive
(TyConSubst -> TyConSubst) -> TyConSubst -> TyConSubst
forall a b. (a -> b) -> a -> b
$ Map TcTyVar TcTyVar -> TyConSubst
tyConSubstEmpty ([(TcTyVar, TcTyVar)] -> Map TcTyVar TcTyVar
forall a. Ord a => [(a, a)] -> Map a a
constructEquivClasses [(TcTyVar, TcTyVar)]
classifiedExtendEquivClass)
go :: TyConSubst
-> [(TcTyVar, (TcTyVar, NonEmpty Type))]
-> TyConSubst
go :: TyConSubst -> [(TcTyVar, (TcTyVar, NonEmpty Type))] -> TyConSubst
go TyConSubst
acc [(TcTyVar, (TcTyVar, NonEmpty Type))]
rs =
let ([(TcTyVar, (TyCon, [Type]))]
prod, [(TcTyVar, (TcTyVar, NonEmpty Type))]
rest) = ((TcTyVar, (TcTyVar, NonEmpty Type))
-> Maybe (NonEmpty (TcTyVar, (TyCon, [Type]))))
-> [(TcTyVar, (TcTyVar, NonEmpty Type))]
-> ([(TcTyVar, (TyCon, [Type]))],
[(TcTyVar, (TcTyVar, NonEmpty Type))])
forall a b. (a -> Maybe (NonEmpty b)) -> [a] -> ([b], [a])
tryApply (TcTyVar, (TcTyVar, NonEmpty Type))
-> Maybe (NonEmpty (TcTyVar, (TyCon, [Type])))
makeProductive [(TcTyVar, (TcTyVar, NonEmpty Type))]
rs in
if [(TcTyVar, (TyCon, [Type]))] -> Bool
forall a. [a] -> Bool
forall (t :: * -> *) a. Foldable t => t a -> Bool
null [(TcTyVar, (TyCon, [Type]))]
prod
then TyConSubst
acc
else TyConSubst -> [(TcTyVar, (TcTyVar, NonEmpty Type))] -> TyConSubst
go ([(TcTyVar, (TyCon, [Type]))] -> TyConSubst -> TyConSubst
tyConSubstExtend [(TcTyVar, (TyCon, [Type]))]
prod TyConSubst
acc) [(TcTyVar, (TcTyVar, NonEmpty Type))]
rest
where
makeProductive ::
(TcTyVar, (TcTyVar, NonEmpty Type))
-> Maybe (NonEmpty (TcTyVar, (TyCon, [Type])))
makeProductive :: (TcTyVar, (TcTyVar, NonEmpty Type))
-> Maybe (NonEmpty (TcTyVar, (TyCon, [Type])))
makeProductive (TcTyVar
var, (TcTyVar
var', NonEmpty Type
args)) =
(NonEmpty (TyCon, [Type]) -> NonEmpty (TcTyVar, (TyCon, [Type])))
-> Maybe (NonEmpty (TyCon, [Type]))
-> Maybe (NonEmpty (TcTyVar, (TyCon, [Type])))
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (((TyCon, [Type]) -> (TcTyVar, (TyCon, [Type])))
-> NonEmpty (TyCon, [Type]) -> NonEmpty (TcTyVar, (TyCon, [Type]))
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap ((TyCon -> [Type] -> (TcTyVar, (TyCon, [Type])))
-> (TyCon, [Type]) -> (TcTyVar, (TyCon, [Type]))
forall a b c. (a -> b -> c) -> (a, b) -> c
uncurry TyCon -> [Type] -> (TcTyVar, (TyCon, [Type]))
aux)) (TcTyVar -> TyConSubst -> Maybe (NonEmpty (TyCon, [Type]))
tyConSubstLookup TcTyVar
var' TyConSubst
acc)
where
aux :: TyCon -> [Type] -> (TcTyVar, (TyCon, [Type]))
aux :: TyCon -> [Type] -> (TcTyVar, (TyCon, [Type]))
aux TyCon
tyCon [Type]
args' = (TcTyVar
var, (TyCon
tyCon, ([Type]
args' [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ NonEmpty Type -> [Type]
forall a. NonEmpty a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList NonEmpty Type
args)))
mkTyConSubst :: [Ct] -> TyConSubst
mkTyConSubst :: [Ct] -> TyConSubst
mkTyConSubst = Classified -> TyConSubst
process (Classified -> TyConSubst)
-> ([Ct] -> Classified) -> [Ct] -> TyConSubst
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Ct] -> Classified
classify
splitTyConApp_upTo :: TyConSubst -> Type -> Maybe (NonEmpty (TyCon, [Type]))
splitTyConApp_upTo :: TyConSubst -> Type -> Maybe (NonEmpty (TyCon, [Type]))
splitTyConApp_upTo TyConSubst
subst Type
typ = [Maybe (NonEmpty (TyCon, [Type]))]
-> Maybe (NonEmpty (TyCon, [Type]))
forall (t :: * -> *) (f :: * -> *) a.
(Foldable t, Alternative f) =>
t (f a) -> f a
asum [
do TyCon
tyCon <- Type -> Maybe TyCon
tyConAppTyCon_maybe Type
fn
NonEmpty (TyCon, [Type]) -> Maybe (NonEmpty (TyCon, [Type]))
forall a. a -> Maybe a
forall (m :: * -> *) a. Monad m => a -> m a
return ((TyCon
tyCon, [Type]
args) (TyCon, [Type]) -> [(TyCon, [Type])] -> NonEmpty (TyCon, [Type])
forall a. a -> [a] -> NonEmpty a
:| [])
, do TcTyVar
var <- Type -> Maybe TcTyVar
getTyVar_maybe Type
fn
(NonEmpty (TyCon, [Type]) -> NonEmpty (TyCon, [Type]))
-> Maybe (NonEmpty (TyCon, [Type]))
-> Maybe (NonEmpty (TyCon, [Type]))
forall a b. (a -> b) -> Maybe a -> Maybe b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (((TyCon, [Type]) -> (TyCon, [Type]))
-> NonEmpty (TyCon, [Type]) -> NonEmpty (TyCon, [Type])
forall a b. (a -> b) -> NonEmpty a -> NonEmpty b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (([Type] -> [Type]) -> (TyCon, [Type]) -> (TyCon, [Type])
forall b c a. (b -> c) -> (a, b) -> (a, c)
forall (p :: * -> * -> *) b c a.
Bifunctor p =>
(b -> c) -> p a b -> p a c
second ([Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [Type]
args))) (Maybe (NonEmpty (TyCon, [Type]))
-> Maybe (NonEmpty (TyCon, [Type])))
-> Maybe (NonEmpty (TyCon, [Type]))
-> Maybe (NonEmpty (TyCon, [Type]))
forall a b. (a -> b) -> a -> b
$ TcTyVar -> TyConSubst -> Maybe (NonEmpty (TyCon, [Type]))
tyConSubstLookup TcTyVar
var TyConSubst
subst
]
where
(Type
fn, [Type]
args) = Type -> (Type, [Type])
splitAppTys Type
typ
instance Outputable TyConSubst where
ppr :: TyConSubst -> SDoc
ppr TyConSubst{Map TcTyVar (NonEmpty (TyCon, [Type]))
Map TcTyVar TcTyVar
tyConSubstMap :: TyConSubst -> Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstCanon :: TyConSubst -> Map TcTyVar TcTyVar
tyConSubstMap :: Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstCanon :: Map TcTyVar TcTyVar
..} = SDoc -> SDoc
forall doc. IsLine doc => doc -> doc
parens (SDoc -> SDoc) -> SDoc -> SDoc
forall a b. (a -> b) -> a -> b
$
String -> SDoc
forall doc. IsLine doc => String -> doc
text String
"TyConSubst"
SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<+> Map TcTyVar (NonEmpty (TyCon, [Type])) -> SDoc
forall a. Outputable a => a -> SDoc
ppr Map TcTyVar (NonEmpty (TyCon, [Type]))
tyConSubstMap
SDoc -> SDoc -> SDoc
forall doc. IsLine doc => doc -> doc -> doc
<+> Map TcTyVar TcTyVar -> SDoc
forall a. Outputable a => a -> SDoc
ppr Map TcTyVar TcTyVar
tyConSubstCanon
tryApply :: forall a b. (a -> Maybe (NonEmpty b)) -> [a] -> ([b], [a])
tryApply :: forall a b. (a -> Maybe (NonEmpty b)) -> [a] -> ([b], [a])
tryApply a -> Maybe (NonEmpty b)
f = ([NonEmpty b] -> [b]) -> ([NonEmpty b], [a]) -> ([b], [a])
forall a b c. (a -> b) -> (a, c) -> (b, c)
forall (p :: * -> * -> *) a b c.
Bifunctor p =>
(a -> b) -> p a c -> p b c
first ([[b]] -> [b]
forall (t :: * -> *) a. Foldable t => t [a] -> [a]
concat ([[b]] -> [b]) -> ([NonEmpty b] -> [[b]]) -> [NonEmpty b] -> [b]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (NonEmpty b -> [b]) -> [NonEmpty b] -> [[b]]
forall a b. (a -> b) -> [a] -> [b]
map NonEmpty b -> [b]
forall a. NonEmpty a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) (([NonEmpty b], [a]) -> ([b], [a]))
-> ([a] -> ([NonEmpty b], [a])) -> [a] -> ([b], [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. [Either (NonEmpty b) a] -> ([NonEmpty b], [a])
forall a b. [Either a b] -> ([a], [b])
partitionEithers ([Either (NonEmpty b) a] -> ([NonEmpty b], [a]))
-> ([a] -> [Either (NonEmpty b) a]) -> [a] -> ([NonEmpty b], [a])
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (a -> Either (NonEmpty b) a) -> [a] -> [Either (NonEmpty b) a]
forall a b. (a -> b) -> [a] -> [b]
map a -> Either (NonEmpty b) a
f'
where
f' :: a -> Either (NonEmpty b) a
f' :: a -> Either (NonEmpty b) a
f' a
a = Either (NonEmpty b) a
-> (NonEmpty b -> Either (NonEmpty b) a)
-> Maybe (NonEmpty b)
-> Either (NonEmpty b) a
forall b a. b -> (a -> b) -> Maybe a -> b
maybe (a -> Either (NonEmpty b) a
forall a b. b -> Either a b
Right a
a) NonEmpty b -> Either (NonEmpty b) a
forall a b. a -> Either a b
Left (Maybe (NonEmpty b) -> Either (NonEmpty b) a)
-> Maybe (NonEmpty b) -> Either (NonEmpty b) a
forall a b. (a -> b) -> a -> b
$ a -> Maybe (NonEmpty b)
f a
a