{-# LANGUAGE CPP #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE NamedFieldPuns #-}
{-# LANGUAGE RecordWildCards #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ViewPatterns #-}
module GHC.TcPlugin.API.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.Graph
( Graph, Vertex )
import qualified Data.Graph as Graph
import Data.Map
( Map )
import qualified Data.Map as Map
import Data.Set
( Set )
import qualified Data.Set as Set
import GHC.Utils.Outputable
hiding ( (<>) )
import GHC.TcPlugin.API
import GHC.Tc.Types.Constraint
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, EqRel)
isCanonicalVarEq Ct
c of
Just (TcTyVar
var, Type -> (Type, [Type])
splitAppTys -> (Type
fn, [Type]
args), EqRel
NomEq)
| Just (TyCon
tyCon, [Type]
inner) <- HasDebugCallStack => Type -> Maybe (TyCon, [Type])
Type -> Maybe (TyCon, [Type])
splitTyConApp_maybe Type
fn ->
Classified -> [Ct] -> Classified
go (TcTyVar -> (TyCon, [Type]) -> Classified
productive TcTyVar
var (TyCon
tyCon, [Type]
inner [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [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, EqRel)
_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]
inner) <- HasDebugCallStack => Type -> Maybe (TyCon, [Type])
Type -> Maybe (TyCon, [Type])
splitTyConApp_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]
inner [Type] -> [Type] -> [Type]
forall a. [a] -> [a] -> [a]
++ [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
isCanonicalVarEq :: Ct -> Maybe (TcTyVar, Type, EqRel)
isCanonicalVarEq :: Ct -> Maybe (TcTyVar, Type, EqRel)
isCanonicalVarEq = \case
#if __GLASGOW_HASKELL__ < 902
CTyEqCan { cc_tyvar, cc_rhs, cc_eq_rel } ->
Just (cc_tyvar, cc_rhs, cc_eq_rel)
CFunEqCan { cc_fsk, cc_fun, cc_tyargs } ->
Just (cc_fsk, mkTyConApp cc_fun cc_tyargs, NomEq)
_otherwise -> Nothing
#elif __GLASGOW_HASKELL__ < 907
CEqCan { cc_lhs, cc_rhs, cc_eq_rel }
| TyVarLHS var <- cc_lhs
-> Just (var, cc_rhs, cc_eq_rel)
| TyFamLHS tyCon args <- cc_lhs
, Just var <- getTyVar_maybe cc_rhs
-> Just (var, mkTyConApp tyCon args, NomEq)
_otherwise
-> Nothing
#else
CEqCan EqCt
eqCt
| TyVarLHS TcTyVar
var <- CanEqLHS
lhs
-> (TcTyVar, Type, EqRel) -> Maybe (TcTyVar, Type, EqRel)
forall a. a -> Maybe a
Just (TcTyVar
var, Type
rhs, EqRel
rel)
| TyFamLHS TyCon
tyCon [Type]
args <- CanEqLHS
lhs
, Just TcTyVar
var <- Type -> Maybe TcTyVar
getTyVar_maybe Type
rhs
-> (TcTyVar, Type, EqRel) -> Maybe (TcTyVar, Type, EqRel)
forall a. a -> Maybe a
Just (TcTyVar
var, TyCon -> [Type] -> Type
mkTyConApp TyCon
tyCon [Type]
args, EqRel
rel)
where
lhs :: CanEqLHS
lhs = EqCt -> CanEqLHS
eq_lhs EqCt
eqCt
rhs :: Type
rhs = EqCt -> Type
eq_rhs EqCt
eqCt
rel :: EqRel
rel = EqCt -> EqRel
eq_eq_rel EqCt
eqCt
Ct
_otherwise
-> Maybe (TcTyVar, Type, EqRel)
forall a. Maybe a
Nothing
#endif
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
constructEquivClasses :: forall a. Ord a => [(a, a)] -> Map a a
constructEquivClasses :: forall a. Ord a => [(a, a)] -> Map a a
constructEquivClasses [(a, a)]
equivs =
[Map a a] -> Map a a
forall (f :: * -> *) k a.
(Foldable f, Ord k) =>
f (Map k a) -> Map k a
Map.unions ([Map a a] -> Map a a) -> [Map a a] -> Map a a
forall a b. (a -> b) -> a -> b
$ (Tree Vertex -> Map a a) -> [Tree Vertex] -> [Map a a]
forall a b. (a -> b) -> [a] -> [b]
map ([a] -> Map a a
pickCanonical ([a] -> Map a a) -> (Tree Vertex -> [a]) -> Tree Vertex -> Map a a
forall b c a. (b -> c) -> (a -> b) -> a -> c
. (Vertex -> a) -> [Vertex] -> [a]
forall a b. (a -> b) -> [a] -> [b]
map Vertex -> a
fromVertex ([Vertex] -> [a])
-> (Tree Vertex -> [Vertex]) -> Tree Vertex -> [a]
forall b c a. (b -> c) -> (a -> b) -> a -> c
. Tree Vertex -> [Vertex]
forall a. Tree a -> [a]
forall (t :: * -> *) a. Foldable t => t a -> [a]
toList) ([Tree Vertex] -> [Map a a]) -> [Tree Vertex] -> [Map a a]
forall a b. (a -> b) -> a -> b
$
Graph -> [Tree Vertex]
Graph.components Graph
graph
where
allValues :: Set a
allValues :: Set a
allValues = [a] -> Set a
forall a. Ord a => [a] -> Set a
Set.fromList ([a] -> Set a) -> [a] -> Set a
forall a b. (a -> b) -> a -> b
$ ((a, a) -> [a]) -> [(a, a)] -> [a]
forall (t :: * -> *) a b. Foldable t => (a -> [b]) -> t a -> [b]
concatMap (\(a
x, a
y) -> [a
x, a
y]) [(a, a)]
equivs
toVertex :: a -> Vertex
fromVertex :: Vertex -> a
toVertex :: a -> Vertex
toVertex a
a = Vertex -> a -> Map a Vertex -> Vertex
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault (String -> Vertex
forall a. HasCallStack => String -> a
error String
"toVertex: impossible") a
a (Map a Vertex -> Vertex) -> Map a Vertex -> Vertex
forall a b. (a -> b) -> a -> b
$
[(a, Vertex)] -> Map a Vertex
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(a, Vertex)] -> Map a Vertex) -> [(a, Vertex)] -> Map a Vertex
forall a b. (a -> b) -> a -> b
$ [a] -> [Vertex] -> [(a, Vertex)]
forall a b. [a] -> [b] -> [(a, b)]
zip (Set a -> [a]
forall a. Set a -> [a]
Set.toList Set a
allValues) [Vertex
1..]
fromVertex :: Vertex -> a
fromVertex Vertex
v = a -> Vertex -> Map Vertex a -> a
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault (String -> a
forall a. HasCallStack => String -> a
error String
"fromVertex: impossible") Vertex
v (Map Vertex a -> a) -> Map Vertex a -> a
forall a b. (a -> b) -> a -> b
$
[(Vertex, a)] -> Map Vertex a
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(Vertex, a)] -> Map Vertex a) -> [(Vertex, a)] -> Map Vertex a
forall a b. (a -> b) -> a -> b
$ [Vertex] -> [a] -> [(Vertex, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [Vertex
1..] (Set a -> [a]
forall a. Set a -> [a]
Set.toList Set a
allValues)
graph :: Graph
graph :: Graph
graph = Bounds -> [Bounds] -> Graph
Graph.buildG (Vertex
1, Set a -> Vertex
forall a. Set a -> Vertex
Set.size Set a
allValues) ([Bounds] -> Graph) -> [Bounds] -> Graph
forall a b. (a -> b) -> a -> b
$
((a, a) -> Bounds) -> [(a, a)] -> [Bounds]
forall a b. (a -> b) -> [a] -> [b]
map ((a -> Vertex) -> (a -> Vertex) -> (a, a) -> Bounds
forall a b c d. (a -> b) -> (c -> d) -> (a, c) -> (b, d)
forall (p :: * -> * -> *) a b c d.
Bifunctor p =>
(a -> b) -> (c -> d) -> p a c -> p b d
bimap a -> Vertex
toVertex a -> Vertex
toVertex) [(a, a)]
equivs
pickCanonical :: [a] -> Map a a
pickCanonical :: [a] -> Map a a
pickCanonical [a]
cls = [(a, a)] -> Map a a
forall k a. Ord k => [(k, a)] -> Map k a
Map.fromList ([(a, a)] -> Map a a) -> [(a, a)] -> Map a a
forall a b. (a -> b) -> a -> b
$ [a] -> [a] -> [(a, a)]
forall a b. [a] -> [b] -> [(a, b)]
zip [a]
cls (a -> [a]
forall a. a -> [a]
repeat ([a] -> a
forall a. Ord a => [a] -> a
forall (t :: * -> *) a. (Foldable t, Ord a) => t a -> a
minimum [a]
cls))
canonicalize :: Ord a => Map a a -> a -> a
canonicalize :: forall a. Ord a => Map a a -> a -> a
canonicalize Map a a
canon a
x = a -> a -> Map a a -> a
forall k a. Ord k => a -> k -> Map k a -> a
Map.findWithDefault a
x a
x Map a a
canon