ghc-tcplugin-api 0.16.2.0 → 0.19.0.0
raw patch · 6 files changed
Files
- changelog.md +60/−1
- ghc-tcplugin-api.cabal +11/−7
- src/GHC/TcPlugin/API.hs +108/−4
- src/GHC/TcPlugin/API/Internal.hs +135/−46
- src/GHC/TcPlugin/API/Internal/Shim.hs +137/−14
- src/GHC/TcPlugin/API/TyConSubst.hs +171/−81
changelog.md view
@@ -1,4 +1,63 @@-# Version 0.16.1.0 (2025-08-22) + +# Version 0.19.0.0 (2026-05-12) + +- Starting from GHC 10.0, typechecker plugins are kept running during the + desugaring phase, in order to improve the accuracy of pattern-match warnings. + + For this reason, the `tcPluginStop` function has been split into two: + - `tcPluginPostTc`, for actions to run at the end of typechecking. + - `tcPluginShutdown`, to shut down the plugin. This action runs in `IO` + instead of a `TcM`-like monad. + + On GHC versions prior to 10.0, both of these actions will run at the end of + typechecking, and typechecker plugins aren't invoked during pattern-match + checking. + +# Version 0.18.2.1 (2026-04-28) + +- Add preliminary support for GHC 10.0. + +# Version 0.18.2.0 (2025-01-14) + +- Add a new function `ctsSubst` which computes a fix-point substitution from + a collection of constraints which define a terminating generalised + substitution. Mainly for use with the set of inert Given constraints. + +- Re-export several functions to deal with substitutions: `substTy`, + `substTys`, `extendTvSubst` and `mkTvSubstPrs`. Also re-export the `Subst` + type (alias for `TCvSubst` on GHC <= 9.6). + +# Version 0.18.1.0 (2025-10-09) + +- Bugfix for v0.18.0.0: deal with constraints generated by unflattening + in `splitTyConApp_upTo`. Fixes #19. + +# Version 0.18.0.0 (2025-09-15) + +- On GHC 9.0 and below, `ghc-tcplugin-api` will now automatically unflatten all + Given constraints. That is, all flattening skolems `[G] fsk ~ F tys` will be + substituted away, both in Givens and in Wanteds/Deriveds. + + No change for GHC 9.2 and above, as GHC stopped producing flattening variables + from 9.2 onwards. + +# Version 0.17.2.0 (2025-09-08) + +- Fix the package failing to build on GHC 9.6.1 through 9.6.6 and + on GHC 9.8.1 through 9.8.3. GHC versions 9.6.7 and 9.8.4 are unaffected. + +# Version 0.17.1.0 (2025-08-27) + +- Fix a regression, introduced in `0.17.0.0`, in which `splitTyConApp_upTo` would + fail to take into account equalities of the form `tv1 ~ tv2`. + +# Version 0.17.0.0 (2025-08-25) + +- `splitTyConApp_upTo` now additionally returns a `[Coercion]` for tracking + Given dependencies (which should be passed to functions such as `mkPluginUnivCo`, + `mkPluginUnivEvTerm` and `mkTyFamAppReduction`). + +# Version 0.16.2.0 (2025-08-22) - `splitTyConApp_upTo` now correctly splits apart type families. This ensures it is a valid drop-in replacement for `splitTyConApp_maybe` (fixes issue #13).
ghc-tcplugin-api.cabal view
@@ -1,6 +1,6 @@ cabal-version: 3.0 name: ghc-tcplugin-api -version: 0.16.2.0 +version: 0.19.0.0 synopsis: An API for type-checker plugins. license: BSD-3-Clause build-type: Simple @@ -15,7 +15,7 @@ for writing GHC type-checking plugins. Each stage in a type-checking plugin (initialisation, solving, rewriting, - shutdown) has a corresponding monad, preventing operations that are only + post-tc) has a corresponding monad, preventing operations that are only allowed in some stages to be used in the other stages. Operations that work across multiple stages are overloaded across monads using MTL-like typeclasses. @@ -33,15 +33,17 @@ build-depends: base - >= 4.13.0 && < 4.23, + >= 4.13.0 && < 5, + array + >= 0.5.3.0 && < 0.6, containers - >= 0.6 && < 0.9, + >= 0.6 && < 0.9, ghc - >= 8.8 && < 9.16, + >= 8.8 && < 10.1, transformers - >= 0.5 && < 0.7, + >= 0.5 && < 0.7, template-haskell - >= 2.15 && < 2.26, + >= 2.15 && < 2.26, default-language: Haskell2010 @@ -151,6 +153,7 @@ , Module as GHC.Unit.Module.Name , Module as GHC.Unit.Types + , FV as GHC.Utils.FV , Util as GHC.Utils.Misc , TcRnMonad as GHC.Utils.Monad , Outputable as GHC.Utils.Outputable @@ -165,6 +168,7 @@ , Constraint as GHC.Tc.Types.Constraint , TcOrigin as GHC.Tc.Types.Origin , GHC.ThToHs as GHC.ThToHs + -- I forget what this last line does, but it is needed on GHC 9.10 ) else
src/GHC/TcPlugin/API.hs view
@@ -450,6 +450,12 @@ , pattern OneTy, pattern ManyTy #endif + -- ** Substitution + , Subst + , substTy, substTys + , mkTvSubstPrs, extendTvSubst + , ctsSubst + -- ** Zonking -- | Zonking is the operation in which GHC actually switches out mutable unification variables @@ -539,10 +545,16 @@ -- base import Prelude hiding ( cos ) +#if !MIN_VERSION_base(4,20,0) +import Data.Foldable + ( foldl' ) +#endif #if !MIN_VERSION_ghc(8,11,0) import Data.List.NonEmpty ( NonEmpty(..) ) #endif +import Data.Maybe + ( mapMaybe ) -- ghc import GHC @@ -662,6 +674,12 @@ #elif MIN_VERSION_ghc(9,0,0) , pattern One, pattern Many #endif + , TvSubstEnv +#if MIN_VERSION_ghc(9,6,0) + , Subst +#else + , TCvSubst +#endif ) import GHC.Data.FastString ( FastString, fsLit, unpackFS ) @@ -748,6 +766,15 @@ ( TcType, TcLevel, MetaDetails, MetaInfo , isSkolemTyVar, isMetaTyVar , nonDetCmpType + , mkTvSubst, extendTvSubst + , substTy, substTys, mkTvSubstPrs + , tyCoVarsOfTypes + , scopedSort +#if MIN_VERSION_ghc(9,15,0) + , tyCoVarsOfTypes +#else + , tyCoFVsOfTypes +#endif ) import GHC.Tc.Utils.TcMType ( isFilledMetaTyVar_maybe, writeMetaTyVar ) @@ -793,21 +820,31 @@ ) import GHC.Types.Unique ( Unique ) +-- UniqFM didn't have the first argument in GHC < 9.0. +-- This adds a little compatibility shim for that. #if MIN_VERSION_ghc(9,0,0) import GHC.Types.Unique.FM as UniqFM - ( UniqFM, emptyUFM, listToUFM ) + ( UniqFM ) #else import qualified GHC.Types.Unique.FM as GHC ( UniqFM ) -import GHC.Types.Unique.FM as UniqFM - ( emptyUFM, listToUFM ) #endif +import GHC.Types.Unique.FM + ( emptyUFM, listToUFM, nonDetEltsUFM ) import GHC.Types.Unique.DFM ( UniqDFM, lookupUDFM, lookupUDFM_Directly, elemUDFM ) import GHC.Types.Var ( TyVar, CoVar, TcTyVar, EvVar - , mkTyVar, DFunId + , mkTyVar, DFunId, isTyVar, updateTyVarKind ) +import GHC.Types.Var.Env + ( InScopeSet, mkInScopeSet, mapVarEnv, elemVarEnv ) +import GHC.Types.Var.Set + ( mkVarSet, unionVarSet +#if MIN_VERSION_ghc(9,15,0) + , nonDetVarSetElems +#endif + ) import GHC.Utils.Outputable ( Outputable(..), SDoc, text ) #if MIN_VERSION_ghc(9,2,0) @@ -839,6 +876,13 @@ import GHC.Utils.Misc ( HasDebugCallStack ) #endif +#if MIN_VERSION_ghc(9,15,0) +#else +import GHC.Utils.FV + ( fvVarList ) +#endif +import GHC.Utils.Misc + ( filterOut ) -- ghc-tcplugin-api import GHC.TcPlugin.API.Internal @@ -981,6 +1025,66 @@ -- | Lookup an identifier, such as a type variable. tcLookupId :: MonadTcPlugin m => Name -> m Id tcLookupId = liftTcPluginM . GHC.tcLookupId + +-------------------------------------------------------------------------------- + +#if !MIN_VERSION_ghc(9,6,0) +type Subst = TCvSubst +#endif + +-- | Compute an idempotent substitution from a collection of constraints, in +-- particular from inert Givens. +-- +-- NB: don't pass any set of constraints to this function; they should form a +-- terminating generalised substitution as in Note [inert_eqs: the inert equalities] +-- in GHC.Tc.Solver.InertSet. +ctsSubst :: [Ct] -> Subst +ctsSubst cts = niFixSubst inScope $ listToUFM prs + where + toSubstPair ct = + case classifyPredType (ctEvPred (ctEvidence ct)) of + EqPred NomEq t1 t2 + | Just tv <- getTyVar_maybe t1 + -- NB: deliberately don't re-orient constraints: we are only + -- guaranteed a terminating substitution by keeping the orientation. + -> Just (tv, t2) + _ -> Nothing + prs = mapMaybe toSubstPair cts + inScope = + mkInScopeSet $ + mkVarSet (map fst prs) + `unionVarSet` + tyCoVarsOfTypes (map snd prs) + +-- NB: copied from GHC, as GHC doesn't export this function. +niFixSubst :: InScopeSet -> TvSubstEnv -> Subst +niFixSubst in_scope tenv + | not_fixpoint = niFixSubst in_scope (mapVarEnv (substTy subst) tenv) + | otherwise = tenv_subst + where + tenv_subst = mkTvSubst in_scope tenv + + range_tvs :: [TyVar] + range_tvs = +#if MIN_VERSION_ghc(9,15,0) + nonDetVarSetElems $ tyCoVarsOfTypes +#else + fvVarList $ tyCoFVsOfTypes +#endif + $ nonDetEltsUFM tenv + + not_fixpoint = any in_domain range_tvs + in_domain tv = tv `elemVarEnv` tenv + + free_tvs = scopedSort (filterOut in_domain range_tvs) + + subst = foldl' add_free_tv tenv_subst free_tvs + + add_free_tv subst0 tv + | isTyVar tv = extendTvSubst subst0 tv (mkTyVarTy tv') + | otherwise = subst0 + where + tv' = updateTyVarKind (substTy subst0) tv --------------------------------------------------------------------------------
src/GHC/TcPlugin/API/Internal.hs view
@@ -108,6 +108,9 @@ as GHC ( TcM, TcPlugin(..), TcPluginM , TcPluginSolver +#if !MIN_VERSION_ghc(9,1,0) + , TcPluginResult(..) +#endif #ifdef HAS_REWRITING , TcPluginRewriter #else @@ -137,6 +140,13 @@ #else import GHC.Driver.Types ( MonadThings(..) ) +import GHC.Tc.Types.Constraint + ( ctEvidence, ctEvId, isDerived ) +import GHC.Types.Var.Env + ( lookupVarEnv, mkVarEnv ) +import GHC.Utils.Outputable +import GHC.Tc.Plugin + ( tcPluginTrace ) #endif -- ghc-tcplugin-api @@ -145,6 +155,9 @@ ( TcPluginSolveResult, TcPluginRewriteResult(..) , RewriteEnv , shimRewriter +#if !MIN_VERSION_ghc(9,1,0) + , unflattenCts +#endif ) #endif @@ -153,10 +166,14 @@ -- | Stage of a type-checking plugin, used as a data kind. data TcPluginStage + -- | Plugin initialisation = Init + -- | Solve constraints | Solve + -- | Rewrite type-families | Rewrite - | Stop + -- | End of typechecking + | PostTc -- | The @solve@ function of a type-checking plugin takes in Given and Wanted -- constraints, and should return a 'GHC.Tc.Types.TcPluginSolveResult' @@ -200,22 +217,22 @@ -- > myTcPlugin args = ... data TcPlugin = forall s. TcPlugin { tcPluginInit :: TcPluginM Init s - -- ^ Initialise plugin, when entering type-checker. + -- ^ Initialize plugin (once per module), when starting the type-checker. , tcPluginSolve :: s -> TcPluginSolver -- ^ Solve some constraints. -- -- This function will be invoked at two points in the constraint solving - -- process: once to manipulate given constraints, and once to solve - -- wanted constraints. In the first case (and only in the first case), - -- no wanted constraints will be passed to the plugin. + -- process: once to simplify Given constraints, and once to solve + -- Wanted constraints. In the first case (and only in the first case), + -- no Wanted constraints will be passed to the plugin. -- -- The plugin can either return a contradiction, -- or specify that it has solved some constraints (with evidence), - -- and possibly emit additional wanted constraints. + -- and possibly emit additional constraints. These returned constraints + -- must be Givens in the first case, and Wanteds in the second. -- - -- Use @ \\ _ _ _ -> pure $ TcPluginOK [] [] @ if your plugin - -- does not provide this functionality. + -- Use @ \\ _ _ _ _ -> pure $ TcPluginOk [] [] @ if your plugin , tcPluginRewrite :: s -> GHC.UniqFM @@ -234,8 +251,18 @@ -- -- Use @ const emptyUFM @ if your plugin does not provide this functionality. - , tcPluginStop :: s -> TcPluginM Stop () - -- ^ Clean up after the plugin, when exiting the type-checker. + , tcPluginPostTc :: s -> TcPluginM PostTc () + -- ^ Action to run at the end of typechecking a module, e.g. to intercept + -- the final 'TcGblEnv'/'TcLclEnv' at the end of typechecking, possibly + -- modifying mutable fields. + -- + -- Should not terminate the plugin, as the plugin may continue to be invoked + -- when desugaring the module (as the pattern-match checker may invoke the + -- constraint solver). + + , tcPluginShutdown :: s -> IO () + -- ^ Clean up after the plugin, when GHC is done processing the given + -- module (e.g. after desugaring). } -- | The monad used for a type-checker plugin, parametrised by @@ -259,8 +286,8 @@ TcPluginRewriteM { tcPluginRewriteM :: BuiltinDefs -> RewriteEnv -> GHC.TcPluginM a } deriving ( Functor, Applicative, Monad ) via ( ReaderT BuiltinDefs ( ReaderT RewriteEnv GHC.TcPluginM ) ) -newtype instance TcPluginM Stop a = - TcPluginStopM { tcPluginStopM :: GHC.TcPluginM a } +newtype instance TcPluginM PostTc a = + TcPluginPostTcM { tcPluginPostTcM :: GHC.TcPluginM a } deriving newtype ( Functor, Applicative, Monad ) -- | Ask for the evidence currently gathered by the type-checker. @@ -371,14 +398,14 @@ #endif . ( \ f -> f builtinDefs rewriteEnv ) . tcPluginRewriteM ) -instance MonadTcPlugin ( TcPluginM Stop ) where - liftTcPluginM = TcPluginStopM +instance MonadTcPlugin ( TcPluginM PostTc ) where + liftTcPluginM = TcPluginPostTcM unsafeWithRunInTcM runInTcM = unsafeLiftTcM $ runInTcM #ifdef HAS_REWRITING - ( GHC.runTcPluginM . tcPluginStopM ) + ( GHC.runTcPluginM . tcPluginPostTcM ) #else - ( ( `GHC.runTcPluginM` ( error "tcPluginStop: cannot access EvBindsVar" ) ) . tcPluginStopM ) + ( ( `GHC.runTcPluginM` ( error "tcPluginPostTc: cannot access EvBindsVar" ) ) . tcPluginPostTcM ) #endif -- | Take a function whose argument and result types are both within the 'GHC.Tc.TcM' monad, @@ -396,19 +423,26 @@ { tcPluginInit = tcPluginInit :: TcPluginM Init userDefs , tcPluginSolve , tcPluginRewrite - , tcPluginStop + , tcPluginPostTc + , tcPluginShutdown } ) = GHC.TcPlugin - { GHC.tcPluginInit = adaptUserInit tcPluginInit + { GHC.tcPluginInit = adaptUserInit tcPluginInit #ifdef HAS_REWRITING - , GHC.tcPluginSolve = adaptUserSolve tcPluginSolve - , GHC.tcPluginRewrite = adaptUserRewrite tcPluginRewrite + , GHC.tcPluginSolve = adaptUserSolve tcPluginSolve + , GHC.tcPluginRewrite = adaptUserRewrite tcPluginRewrite #else - , GHC.tcPluginSolve = adaptUserSolveAndRewrite - tcPluginSolve tcPluginRewrite + , GHC.tcPluginSolve = adaptUserSolveAndRewrite + tcPluginSolve tcPluginRewrite #endif - , GHC.tcPluginStop = adaptUserStop tcPluginStop +#if MIN_VERSION_ghc(10,0,0) + , GHC.tcPluginPostTc = adaptUserPostTc tcPluginPostTc + , GHC.tcPluginShutdown = adaptUserShutdown tcPluginShutdown +#else + , GHC.tcPluginStop = adaptUserStop + tcPluginPostTc tcPluginShutdown +#endif } where adaptUserInit :: TcPluginM Init userDefs -> GHC.TcPluginM ( TcPluginDefs userDefs ) @@ -446,27 +480,82 @@ -> TcPluginDefs userDefs -> GHC.TcPluginSolver adaptUserSolveAndRewrite userSolve userRewrite ( TcPluginDefs { tcPluginUserDefs, tcPluginBuiltinDefs } ) - = \ givens deriveds wanteds -> do + = \ givens0 deriveds0 wanteds0 -> do +#if MIN_VERSION_ghc(9,1,0) + let givens = givens0 + deriveds = deriveds0 + wanteds = wanteds0 +#else + let ( givens, deriveds, wanteds ) = unflattenCts givens0 deriveds0 wanteds0 + tcPluginTrace "ghc-tcplugin-api: unflattening" $ + vcat [ text "givens:" <+> ppr givens0 + , text "deriveds:" <+> ppr deriveds0 + , text "wanteds:" <+> ppr wanteds0 + , text (replicate 80 '=') + , text "givens:" <+> ppr givens + , text "deriveds:" <+> ppr deriveds + , text "wanteds:" <+> ppr wanteds + ] +#endif evBindsVar <- GHC.getEvBindsTcPluginM - shimRewriter - givens deriveds wanteds - ( fmap - ( \ userRewriter rewriteEnv gs tys -> - tcPluginRewriteM ( userRewriter gs tys ) - tcPluginBuiltinDefs rewriteEnv - ) - ( userRewrite tcPluginUserDefs ) - ) - ( \ gs ds ws -> - tcPluginSolveM ( userSolve tcPluginUserDefs gs ws ) - tcPluginBuiltinDefs evBindsVar ds - ) + res <- + shimRewriter + givens deriveds wanteds + ( fmap + ( \ userRewriter rewriteEnv gs tys -> + tcPluginRewriteM ( userRewriter gs tys ) + tcPluginBuiltinDefs rewriteEnv + ) + ( userRewrite tcPluginUserDefs ) + ) + ( \ gs ds ws -> + tcPluginSolveM ( userSolve tcPluginUserDefs gs ws ) + tcPluginBuiltinDefs evBindsVar ds + ) +#if MIN_VERSION_ghc(9,1,0) + return res +#else + let origCts = + mkVarEnv + [ ( ctEvId ct, ct ) + | ct <- givens0 ++ wanteds0 ] + case res of + GHC.TcPluginOk solved new -> do + -- If we are solving a constraint that has been flattened, make + -- sure to solve the original constraint instead of the flattened + -- constraint, otherwise GHC gets very confused. + let + lookupOrigCt ct + | isDerived (ctEvidence ct) + = ct + | Just ct' <- lookupVarEnv origCts ( ctEvId ct ) + = ct' + | otherwise + = ct + solved' = map ( \ ( ev, ct ) -> ( ev, lookupOrigCt ct ) ) solved + return $ GHC.TcPluginOk solved' new + GHC.TcPluginContradiction {} -> + return res #endif +#endif - adaptUserStop :: ( userDefs -> TcPluginM Stop () ) -> TcPluginDefs userDefs -> GHC.TcPluginM () - adaptUserStop userStop ( TcPluginDefs { tcPluginUserDefs } ) = - tcPluginStopM $ userStop tcPluginUserDefs +#if MIN_VERSION_ghc(10,0,0) + adaptUserPostTc :: ( userDefs -> TcPluginM PostTc () ) -> TcPluginDefs userDefs -> GHC.TcPluginM () + adaptUserPostTc userPostTc ( TcPluginDefs { tcPluginUserDefs } ) = + tcPluginPostTcM $ userPostTc tcPluginUserDefs + adaptUserShutdown :: ( userDefs -> IO () ) -> TcPluginDefs userDefs -> IO () + adaptUserShutdown userShutdown ( TcPluginDefs { tcPluginUserDefs } ) = + userShutdown tcPluginUserDefs +#else + -- Prior to GHC 10.0, run the "post-tc" and "shutdown" actions in sequence, + -- at the end of typechecking. Typechecker plugins do not persist to desugaring. + adaptUserStop :: ( userDefs -> TcPluginM PostTc () ) -> ( userDefs -> IO () ) -> TcPluginDefs userDefs -> GHC.TcPluginM () + adaptUserStop userPostTc userShutdown ( TcPluginDefs { tcPluginUserDefs } ) = do + tcPluginPostTcM $ userPostTc tcPluginUserDefs + GHC.unsafeTcPluginTcM . liftIO $ userShutdown tcPluginUserDefs +#endif + -- | @since 0.15.0.0 instance ( Monad ( TcPluginM s ), MonadTcPlugin ( TcPluginM s ) ) => MonadIO ( TcPluginM s ) where liftIO = unsafeLiftTcM . liftIO @@ -476,7 +565,7 @@ -- -- These operations are supported by the monads that 'tcPluginSolve' -- and 'tcPluginRewrite' use; it is not possible to emit work or --- throw type errors in 'tcPluginInit' or 'tcPluginStop'. +-- throw type errors in 'tcPluginInit' or 'tcPluginPostTc'. -- -- See 'mkTcPluginErrorTy' and 'GHC.TcPlugin.API.emitWork' for functions -- which require this typeclass. @@ -498,9 +587,9 @@ instance TypeError ( 'Text "Cannot emit new work in 'tcPluginInit'." ) => MonadTcPluginWork ( TcPluginM Init ) where askBuiltins = error "Cannot emit new work in 'tcPluginInit'." -instance TypeError ( 'Text "Cannot emit new work in 'tcPluginStop'." ) - => MonadTcPluginWork ( TcPluginM Stop ) where - askBuiltins = error "Cannot emit new work in 'tcPluginStop'." +instance TypeError ( 'Text "Cannot emit new work in 'tcPluginPostTc'." ) + => MonadTcPluginWork ( TcPluginM PostTc ) where + askBuiltins = error "Cannot emit new work in 'tcPluginPostTc'." -- | Use this type like 'GHC.TypeLits.ErrorMessage' to write an error message. -- This error message can then be thrown at the type-level by the plugin, @@ -533,8 +622,8 @@ errorMsgTy = interpretErrorMessage builtinDefs msg pure $ GHC.mkTyConApp typeErrorTyCon [ GHC.constraintKind, errorMsgTy ] -instance ( Monad (TcPluginM s), MonadTcPlugin (TcPluginM s) ) - => MonadThings (TcPluginM s) where +instance ( Monad ( TcPluginM s ), MonadTcPlugin ( TcPluginM s ) ) + => MonadThings ( TcPluginM s ) where lookupThing = unsafeLiftTcM . lookupThing --------------------------------------------------------------------------------
src/GHC/TcPlugin/API/Internal/Shim.hs view
@@ -21,6 +21,9 @@ , TcPluginSolveResult(TcPluginContradiction, TcPluginOk, ..), TcPluginRewriteResult(..) , RewriteEnv(..) , shimRewriter +#if !MIN_VERSION_ghc(9,1,0) + , unflattenCts +#endif ) where @@ -31,11 +34,11 @@ ( forM, unless, when ) import Data.Foldable ( traverse_ -#if !MIN_VERSION_ghc(9,2,0) +#if !MIN_VERSION_ghc(9,1,0) , foldlM #endif ) -#if MIN_VERSION_ghc(9,2,0) +#if MIN_VERSION_ghc(9,1,0) import Data.List.NonEmpty ( NonEmpty((:|)) ) #endif @@ -55,6 +58,9 @@ , mkAppCos, mkNomReflCo, mkSubCo , mkTyConAppCo, tyConRolesX , tyConRolesRepresentational +#if !MIN_VERSION_ghc(9,1,0) + , CoercionHole(..), setCoHoleCoVar, coHoleCoVar +#endif ) import GHC.Core.Predicate ( EqRel(..), eqRelRole ) @@ -64,7 +70,7 @@ ) import GHC.Core.TyCon ( TyCon(..), TyConBinder, TyConBndrVis(..) -#if MIN_VERSION_ghc(9,2,0) +#if MIN_VERSION_ghc(9,1,0) , isForgetfulSynTyCon #endif , isFamFreeTyCon, isTypeSynonymTyCon @@ -80,7 +86,7 @@ #endif , coreView, tyVarKind ) -#if MIN_VERSION_ghc(9,2,0) +#if MIN_VERSION_ghc(9,1,0) import GHC.Data.Maybe ( firstJustsM ) #endif @@ -96,7 +102,7 @@ , runTcPluginTcS, runTcSWithEvBinds , traceTcS , setWantedEvTerm -#if MIN_VERSION_ghc(9,2,0) +#if MIN_VERSION_ghc(9,1,0) , lookupFamAppCache, lookupFamAppInert, extendFamAppCache , pattern EqualCtList #else @@ -112,8 +118,10 @@ import GHC.Tc.Types.Constraint ( Ct(..), CtEvidence(..) , CtLoc, CtFlavour(..), CtFlavourRole, ShadowInfo(..) -#if MIN_VERSION_ghc(9,2,0) +#if MIN_VERSION_ghc(9,1,0) , CanEqLHS(..) +#else + , QCInst(..), TcEvDest(..) #endif , ctLoc, ctFlavour, ctEvidence, ctEqRel, ctEvPred , ctEvExpr, ctEvCoercion, ctEvFlavour @@ -125,10 +133,14 @@ ) import GHC.Tc.Utils.TcType ( TcTyCoVarSet -#if MIN_VERSION_ghc(9,2,0) +#if MIN_VERSION_ghc(9,1,0) , tcSplitForAllTyVarBinders #else , tcSplitForAllVarBndrs + + -- Unflattening + , TCvSubst(..), substTy, substTys, extendTvSubst + , emptyTCvSubst, lookupTyVar #endif , tcSplitTyConApp_maybe , tcTypeKind @@ -138,8 +150,10 @@ ( UniqFM, lookupUFM, isNullUFM ) import GHC.Types.Var ( TcTyVar, VarBndr(..) -#if !MIN_VERSION_ghc(9,2,0) +#if !MIN_VERSION_ghc(9,1,0) , TyVarBinder + , updateTyVarKind + , setVarType #endif , updateTyVarKindM ) @@ -152,7 +166,7 @@ import GHC.Utils.Monad ( zipWith3M ) import GHC.Utils.Outputable - ( Outputable(..), SDoc, empty ) + hiding ( (<>) ) -- ghc-tcplugin-api import GHC.TcPlugin.API.Internal.Shim.Reduction @@ -532,7 +546,7 @@ liftTcS $ extendFamAppCache tc tys ( mkSymCoOnGHC92 final_co, final_xi ) -#if !MIN_VERSION_ghc(9,2,0) +#if !MIN_VERSION_ghc(9,1,0) flavour #endif return final_redn @@ -545,7 +559,7 @@ -- (Recall: this module is only used for GHC 9.2 and below.) mkSymCoOnGHC92 :: Coercion -> Coercion mkSymCoOnGHC92 co = -#if MIN_VERSION_ghc(9,2,0) +#if MIN_VERSION_ghc(9,1,0) mkSymCo co #else co @@ -615,7 +629,7 @@ rewriterView :: Type -> Maybe Type rewriterView ty@(TyConApp tc _) | ( isTypeSynonymTyCon tc && not (isFamFreeTyCon tc) ) -#if MIN_VERSION_ghc(9,2,0) +#if MIN_VERSION_ghc(9,1,0) || isForgetfulSynTyCon tc #endif = tcView ty @@ -651,7 +665,7 @@ rewrite_tyvar2 tv fr@(_, eq_rel) = do ieqs <- liftTcS $ getInertEqs case lookupDVarEnv ieqs tv of -#if MIN_VERSION_ghc(9,2,0) +#if MIN_VERSION_ghc(9,1,0) Just (EqualCtList (ct :| _)) | CEqCan { cc_ev = ctev, cc_lhs = TyVarLHS {} , cc_rhs = rhs_ty, cc_eq_rel = ct_eq_rel } <- ct @@ -942,7 +956,7 @@ !env' = env { rewriteEnv = renv' } thing_inside env' s -#if !MIN_VERSION_ghc(9,2,0) +#if !MIN_VERSION_ghc(9,1,0) -------------------------------------------------------------------------------- -- GHC 9.0 compatibility. @@ -971,5 +985,114 @@ tcSplitForAllTyVarBinders :: Type -> ([TyVarBinder], Type) tcSplitForAllTyVarBinders = tcSplitForAllVarBndrs + +#endif + +-------------------------------------------------------------------------------- +-- Unflattening constraints + +#if !MIN_VERSION_ghc(9,1,0) +unflattenCts :: [ Ct ] -> [ Ct ] -> [ Ct ] -> ( [ Ct ], [ Ct ], [ Ct ] ) +unflattenCts givens0 deriveds0 wanteds0 = + -- NB: on all GHC versions we support (GHC >= 8.8), + -- GHC passes zonked constraints to typechecker plugins. + -- So there's no need for us to zonk anything. + let ( subst, givens ) = add_fsks emptyTCvSubst [] givens0 + deriveds = substCts subst deriveds0 + wanteds = substCts subst wanteds0 + + in + ( givens, deriveds, wanteds ) + +-- Construct an unflattening substitution from Givens of the form +-- +-- [G] fsk ~ rhs +-- +-- and filter them out. +add_fsks :: TCvSubst -> [ Ct ] -> [ Ct ] -> ( TCvSubst, [ Ct ] ) +add_fsks subst rev_acc [] = ( subst, substCts subst $ reverse rev_acc ) +add_fsks subst rev_acc ( ct : cts ) = + case ct of + CFunEqCan { cc_fsk, cc_fun, cc_tyargs } -> + let subst' = extendSubst subst cc_fsk $ mkTyConApp cc_fun cc_tyargs + in add_fsks subst' rev_acc cts + _ -> add_fsks subst ( ct : rev_acc ) cts + +extendSubst :: TCvSubst -> TyVar -> Type -> TCvSubst +extendSubst subst@( TCvSubst in_scope tv_env cv_env ) tv ty = + -- Keep the substitution idempotent + let ty' = substTy subst ty + tv_subst = extendTvSubst emptyTCvSubst tv ty' + subst' = TCvSubst in_scope ( fmap ( substTy tv_subst ) tv_env ) cv_env + in extendTvSubst subst' tv ty' + +substCts :: TCvSubst -> [ Ct ] -> [ Ct ] +substCts subst = map $ substCt subst + +substCt :: TCvSubst -> Ct -> Ct +substCt subst ct = + case ct of + CIrredCan { cc_ev } -> + ct { cc_ev = substCtEv subst cc_ev } + CNonCanonical { cc_ev } -> + ct { cc_ev = substCtEv subst cc_ev } + CDictCan { cc_ev, cc_tyargs } -> + ct { cc_ev = substCtEv subst cc_ev + , cc_tyargs = substTys subst cc_tyargs + } + CFunEqCan {} -> + pprPanic "substCt: CFunEqCan" (ppr ct) + CTyEqCan { cc_ev, cc_tyvar, cc_rhs } -> + let ev' = substCtEv subst cc_ev + rhs' = substTy subst cc_rhs + in + case lookupTyVar subst cc_tyvar of + Nothing -> + ct { cc_ev = ev' + , cc_tyvar = updateTyVarKind ( substTy subst ) cc_tyvar + , cc_rhs = rhs' + } + Just ty' -> + case ty' of + TyVarTy tv' -> + ct { cc_ev = ev' + , cc_tyvar = tv' + , cc_rhs = rhs' + } + _ -> + CNonCanonical { cc_ev = ev' } + CQuantCan qci@( QCI { qci_ev, qci_tvs, qci_pred } ) -> + CQuantCan $ + qci { qci_ev = substCtEv subst qci_ev + , qci_tvs = map ( updateTyVarKind $ substTy subst ) qci_tvs + , qci_pred = substTy subst qci_pred + } +#if !MIN_VERSION_ghc(8,11,0) + CHoleCan { cc_ev } -> + let ev' = substCtEv subst cc_ev + in ct { cc_ev = ev' } +#endif + +substCtEv :: TCvSubst -> CtEvidence -> CtEvidence +substCtEv subst ctev = + setCtEvPredType ctev ( substTy subst ( ctev_pred ctev ) ) + +setCtEvPredType :: CtEvidence -> Type -> CtEvidence +setCtEvPredType old_ev@(CtGiven { ctev_evar = ev }) new_pred + = old_ev { ctev_pred = new_pred + , ctev_evar = setVarType ev new_pred } + +setCtEvPredType old_ev@(CtWanted { ctev_dest = dest }) new_pred + = old_ev { ctev_pred = new_pred + , ctev_dest = new_dest } + where + new_dest = case dest of + EvVarDest ev -> EvVarDest (setVarType ev new_pred) + HoleDest h -> HoleDest (setCoHoleType h new_pred) +setCtEvPredType old_ev@( CtDerived {} ) new_pred + = old_ev { ctev_pred = new_pred } + +setCoHoleType :: CoercionHole -> Type -> CoercionHole +setCoHoleType h t = setCoHoleCoVar h (setVarType (coHoleCoVar h) t) #endif
src/GHC/TcPlugin/API/TyConSubst.hs view
@@ -1,10 +1,13 @@ {-# LANGUAGE CPP #-} +{-# LANGUAGE BangPatterns #-} {-# LANGUAGE LambdaCase #-} +{-# LANGUAGE MultiWayIf #-} {-# LANGUAGE NamedFieldPuns #-} {-# LANGUAGE RecordWildCards #-} {-# LANGUAGE ScopedTypeVariables #-} {-# LANGUAGE TupleSections #-} +{-# LANGUAGE TypeApplications #-} {-# LANGUAGE ViewPatterns #-} {-| @@ -39,27 +42,51 @@ , splitTyConApp_upTo ) where +-- Word64Map is available: +-- +-- - from 9.10.1 onwards +-- - in 9.8.4 +-- - in 9.6.7 +-- +-- This is unusual: whether a module is exposed depends on the MINOR version +#define HAS_WORD64MAP \ + (MIN_VERSION_ghc(9,9,0) || \ + (MIN_VERSION_ghc(9,8,4) && !(MIN_VERSION_ghc(9,9,0))) || \ + (MIN_VERSION_ghc(9,6,7) && !(MIN_VERSION_ghc(9,7,0)))) + -- base import Data.Bifunctor + ( Bifunctor(first, second) ) import Data.Either ( partitionEithers ) import Data.Foldable ( toList, asum ) import Data.List.NonEmpty ( NonEmpty(..) ) +import Data.Maybe + ( fromJust ) +#if !MIN_VERSION_ghc(8,11,0) +import Unsafe.Coerce + ( unsafeCoerce ) +#endif -- containers -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 Data.Sequence + ( Seq ) +import qualified Data.Sequence as Seq -- ghc +#if HAS_WORD64MAP +import qualified GHC.Data.Word64Map.Strict as Word64Map +#else +import qualified Data.IntMap.Strict as IntMap +#endif +import GHC.Types.Unique + ( Uniquable, getUnique ) +import qualified GHC.Types.Unique.FM as UFM +import GHC.Types.Unique.Set + ( UniqSet ) +import qualified GHC.Types.Unique.Set as UniqSet import GHC.Utils.Outputable hiding ( (<>) ) @@ -76,8 +103,8 @@ -- | Substitution for recognizing 'TyCon' applications modulo nominal equalities. data TyConSubst = TyConSubst { - tyConSubstMap :: Map TcTyVar (NonEmpty (TyCon, [Type])) - , tyConSubstCanon :: Map TcTyVar TcTyVar + tyConSubstMap :: UniqFM TcTyVar (NonEmpty (TyCon, [Type], [Coercion])) + , tyConSubstCanon :: UniqFM TcTyVar (TcTyVar, [Coercion]) } -- During constraint solving the set of Given constraints includes so-called -- "canonical equalities": equalities of the form @@ -97,7 +124,7 @@ -- derive any additional equalities from them (i.e. if, say, we know that -- @x ~ T1@ and @x ~ T2@, we will not attempt to use the fact that this means -- that @T1 ~ T2@, nor any derived conclusions thereof). We /will/ however --- try to apply the canononical equalities as often as is necessary (e.g., +-- try to apply the canonical equalities as often as is necessary (e.g., -- first applying @x ~ T y@, then applying @y ~ T2@). -- -- We solve this problem by constructing a 'TyConSubst': a possibly @@ -139,31 +166,36 @@ -- -- The canonical variables map is established once when the initial substitution -- is generated and not updated thereafter. -tyConSubstEmpty :: Map TcTyVar TcTyVar -> TyConSubst +tyConSubstEmpty :: UniqFM TcTyVar (TcTyVar, [Coercion]) -> TyConSubst tyConSubstEmpty canon = TyConSubst { - tyConSubstMap = Map.empty + tyConSubstMap = UFM.emptyUFM , tyConSubstCanon = canon } -- | Lookup a variable in the substitution -tyConSubstLookup :: TcTyVar -> TyConSubst -> Maybe (NonEmpty (TyCon, [Type])) -tyConSubstLookup var TyConSubst{..} = Map.lookup var' tyConSubstMap +tyConSubstLookup :: TcTyVar -> TyConSubst -> Maybe (NonEmpty (TyCon, [Type], [Coercion])) +tyConSubstLookup var TyConSubst{..} = + fmap ( \ (tc, tys, deps) -> (tc, tys, deps ++ deps')) <$> UFM.lookupUFM tyConSubstMap var' where var' :: TcTyVar - var' = canonicalize tyConSubstCanon var + deps' :: [Coercion] + (var', deps') = canonicalize tyConSubstCanon var -- | Extend substitution with new bindings tyConSubstExtend :: - [(TcTyVar, (TyCon, [Type]))] + [(TcTyVar, (TyCon, [Type]), [Coercion])] -> TyConSubst -> TyConSubst tyConSubstExtend new subst@TyConSubst{..} = subst { - tyConSubstMap = Map.unionWith (<>) - (Map.fromList $ map (uncurry aux) new) + tyConSubstMap = UFM.plusUFM_C (<>) + (UFM.listToUFM $ map aux new) tyConSubstMap } where - aux :: TcTyVar -> (TyCon, [Type]) -> (TcTyVar, NonEmpty (TyCon, [Type])) - aux var s = (canonicalize tyConSubstCanon var, s :| []) + aux :: (TcTyVar, (TyCon, [Type]), [Coercion]) + -> (TcTyVar, NonEmpty (TyCon, [Type] , [Coercion])) + aux (var, (tc, args), deps) = + let (var', deps') = canonicalize tyConSubstCanon var + in (var', (tc, args, deps ++ deps') :| []) {------------------------------------------------------------------------------- Classification @@ -181,13 +213,13 @@ -- an application of a concrete type constructor (note that we only ever -- apply the substitution to the head @t@ of a type @t args@, never to the -- arguments). - classifiedProductive :: [(TcTyVar, (TyCon, [Type]))] + classifiedProductive :: [(TcTyVar, (TyCon, [Type]), [Coercion])] -- | Extend equivalence class of variables -- -- An equality @var1 := var2@ we will regard as extending the equivalence -- classes of variables (see 'constructEquivClasses'). - , classifiedExtendEquivClass :: [(TcTyVar, TcTyVar)] + , classifiedExtendEquivClass :: [(TcTyVar, TcTyVar, [Coercion])] -- | Substitutions we need to reconsider later -- @@ -195,7 +227,7 @@ -- arguments) is most problematic. Applying it /may/ allow us to make -- progress, but it may not (consider for example @var := var arg@). We -- will reconsider such equalities at the end (see 'process'). - , classifiedReconsider :: [(TcTyVar, (TcTyVar, NonEmpty Type))] + , classifiedReconsider :: [(TcTyVar, (TcTyVar, NonEmpty Type), [Coercion])] } instance Semigroup Classified where @@ -211,19 +243,19 @@ instance Monoid Classified where mempty = Classified [] [] [] -productive :: TcTyVar -> (TyCon, [Type]) -> Classified -productive var (tyCon, args) = mempty { - classifiedProductive = [(var, (tyCon, args))] +productive :: TcTyVar -> (TyCon, [Type]) -> [Coercion] -> Classified +productive var app deps = mempty { + classifiedProductive = [(var, app, deps)] } -extendEquivClass :: TcTyVar -> TcTyVar -> Classified -extendEquivClass var var' = mempty { - classifiedExtendEquivClass = [(var, var')] +extendEquivClass :: TcTyVar -> TcTyVar -> [Coercion] -> Classified +extendEquivClass var var' deps = mempty { + classifiedExtendEquivClass = [(var, var', deps)] } -reconsider :: TcTyVar -> (TcTyVar, NonEmpty Type) -> Classified -reconsider var (var', args) = mempty { - classifiedReconsider = [(var, (var', args))] +reconsider :: TcTyVar -> (TcTyVar, NonEmpty Type) -> [Coercion] -> Classified +reconsider var (var', args) deps = mempty { + classifiedReconsider = [(var, (var', args), deps)] } -- | Classify a set of Given constraints. @@ -235,14 +267,15 @@ go :: Classified -> [Ct] -> Classified go acc [] = acc go acc (c:cs) = + let deps = [ctEvCoercion (ctEvidence c)] in case isCanonicalVarEq c of Just (var, splitAppTys -> (fn, args), NomEq) | Just (tyCon, inner) <- splitTyConApp_maybe fn -> - go (productive var (tyCon, inner ++ args) <> acc) cs + go (productive var (tyCon, inner ++ args) deps <> acc) cs | Just var' <- getTyVar_maybe fn, null args -> - go (extendEquivClass var var' <> acc) cs + go (extendEquivClass var var' deps <> acc) cs | Just var' <- getTyVar_maybe fn, x:xs <- args -> - go (reconsider var (var', x :| xs) <> acc) cs + go (reconsider var (var', x :| xs) deps <> acc) cs _otherwise -> go acc cs @@ -269,7 +302,7 @@ -- o We keep doing this until we can make no more progress. -- -- The functions for working with 'TyConSubst' take the variable equivalence --- classes into acocunt, so we do not need to do that here. +-- classes into account, so we do not need to do that here. -- -- Two observations: -- @@ -302,7 +335,7 @@ $ tyConSubstEmpty (constructEquivClasses classifiedExtendEquivClass) go :: TyConSubst - -> [(TcTyVar, (TcTyVar, NonEmpty Type))] + -> [(TcTyVar, (TcTyVar, NonEmpty Type), [Coercion])] -> TyConSubst go acc rs = let (prod, rest) = tryApply makeProductive rs in @@ -311,13 +344,15 @@ else go (tyConSubstExtend prod acc) rest where makeProductive :: - (TcTyVar, (TcTyVar, NonEmpty Type)) - -> Maybe (NonEmpty (TcTyVar, (TyCon, [Type]))) - makeProductive (var, (var', args)) = - fmap (fmap (uncurry aux)) (tyConSubstLookup var' acc) + (TcTyVar, (TcTyVar, NonEmpty Type), [Coercion]) + -> Maybe (NonEmpty (TcTyVar, (TyCon, [Type]), [Coercion])) + makeProductive (var, (var', args), deps) = do + tcApp <- tyConSubstLookup var' acc + return $ fmap aux tcApp where - aux :: TyCon -> [Type] -> (TcTyVar, (TyCon, [Type])) - aux tyCon args' = (var, (tyCon, (args' ++ toList args))) + aux :: (TyCon, [Type], [Coercion]) -> (TcTyVar, (TyCon, [Type]), [Coercion]) + aux (tyCon, args', deps') = + (var, (tyCon, args' ++ toList args), deps ++ deps') -- | Construct a 'TyConSubst' from a collection of Given constraints. mkTyConSubst :: [Ct] -> TyConSubst @@ -329,16 +364,21 @@ -- | Like 'splitTyConApp_maybe', but taking Given constraints into account. -- +-- Alongside the @TyCon@ and its arguments, also returns a list of coercions +-- that embody the Givens that we depended on. +-- -- Looks through type synonyms, just like 'splitTyConApp_maybe' does. -splitTyConApp_upTo :: TyConSubst -> Type -> Maybe (NonEmpty (TyCon, [Type])) +splitTyConApp_upTo :: TyConSubst -> Type -> Maybe (NonEmpty (TyCon, [Type], [Coercion])) splitTyConApp_upTo subst typ = asum [ -- Direct match do (tyCon, inner) <- splitTyConApp_maybe fn - return ((tyCon, inner ++ args) :| []) + return ((tyCon, inner ++ args, []) :| []) -- Indirect match , do var <- getTyVar_maybe fn - fmap (fmap (second (++ args))) $ tyConSubstLookup var subst + tcApps <- tyConSubstLookup var subst + return $ + fmap (\ (tc, inner, deps) -> (tc, inner ++ args, deps)) tcApps ] where (fn, args) = splitAppTys typ @@ -364,7 +404,6 @@ 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 @@ -372,8 +411,6 @@ | TyFamLHS tyCon args <- cc_lhs , Just var <- getTyVar_maybe cc_rhs -> Just (var, mkTyConApp tyCon args, NomEq) - _otherwise - -> Nothing #else CEqCan eqCt | TyVarLHS var <- lhs @@ -385,9 +422,18 @@ lhs = eq_lhs eqCt rhs = eq_rhs eqCt rel = eq_eq_rel eqCt - _otherwise - -> Nothing #endif + -- Deal with CNonCanonical, which are produced by ghc-tcplugin-api + -- in GHC 9.0 and below due to 'unflattenCts'. + ct@(CNonCanonical {}) + | EqPred rel lhs rhs <- classifyPredType (ctPred ct) + -> if | Just tv <- getTyVar_maybe lhs + -> Just (tv, rhs, rel) + | Just tv <- getTyVar_maybe rhs + -> Just (tv, lhs, rel) + | otherwise + -> Nothing + _ -> Nothing {------------------------------------------------------------------------------- Internal auxiliary @@ -407,41 +453,85 @@ Equivalence classes -------------------------------------------------------------------------------} --- | Given a set of equivalent pairs, map every value to canonical value +-- | Given a set of labelled equivalent pairs, map every value to a canonical +-- value in the same equivalence class, with the shortest path (as a list of +-- labels) connecting the value to the canonical value. -- --- Example with two classes: +-- Example with two equivalence classes: -- --- >>> constructEquivClasses [(1, 2), (4, 5), (2, 3)] --- fromList [(1,1),(2,1),(3,1),(4,4),(5,4)] +-- >>> constructEquivClasses [(1, 2, ["12"]), (4, 5, ["45"]), (2, 3, ["23"])] +-- fromList [(1,(1,[])),(2,(1,["12"])),(3,(1,["23","12"])),(4,(4,[])),(5,(4,["45"]))] -- --- Adding one element that connects both classes: +-- Adding one element that connects both equivalence classes: -- --- >>> constructEquivClasses [(1, 2), (4, 5), (2, 3), (3, 4)] --- fromList [(1,1),(2,1),(3,1),(4,1),(5,1)] -constructEquivClasses :: forall a. Ord a => [(a, a)] -> Map a a -constructEquivClasses equivs = - Map.unions $ map (pickCanonical . map fromVertex . toList) $ - Graph.components graph +-- >>> constructEquivClasses [(1, 2, ["12"]), (4, 5, ["45"]), (2, 3, ["23"]), (3, 4, ["34"])] +-- fromList [(1,(1,[])),(2,(1,["12"])),(3,(1,["23","12"])),(4,(1,["34","23","12"])),(5,(1,["45","34","23","12"]))] +constructEquivClasses :: forall a l. (Uniquable a, Monoid l) => [(a, a, l)] -> UniqFM a (a, l) +constructEquivClasses equivs = canonicals where - allValues :: Set a - allValues = Set.fromList $ concatMap (\(x, y) -> [x, y]) equivs + neighbours :: UniqFM a (UniqFM a l) + neighbours = neighboursMap equivs - toVertex :: a -> Vertex - fromVertex :: Vertex -> a + allValues :: UniqSet a + allValues = UniqSet.mkUniqSet $ concat [ [x,y] | (x,y,_) <- equivs ] - toVertex a = Map.findWithDefault (error "toVertex: impossible") a $ - Map.fromList $ zip (Set.toList allValues) [1..] - fromVertex v = Map.findWithDefault (error "fromVertex: impossible") v $ - Map.fromList $ zip [1..] (Set.toList allValues) + canonicals :: UniqFM a (a, l) + canonicals = go UFM.emptyUFM allValues + where + go :: UniqFM a (a, l) -> UniqSet a -> UniqFM a (a, l) + go acc vs = + case minViewUniqSet vs of + Nothing -> acc + Just (v, vs') -> + let + !comp = doComp + ( UFM.unitUFM v mempty ) + ( Seq.singleton $ getUnique v ) + in + go + ( UFM.plusUFM acc ( ( v , ) <$> comp ) ) + ( vs' `UniqSet.uniqSetMinusUFM` comp ) - graph :: Graph - graph = Graph.buildG (1, Set.size allValues) $ - map (bimap toVertex toVertex) equivs + doComp :: UniqFM a l -> Seq Unique -> UniqFM a l + doComp !ds Seq.Empty = ds + doComp ds (v Seq.:<| vs) = + let + -- unvisited neighbours + !us = ( fromJust $ UFM.lookupUFM_Directly neighbours v ) `UFM.minusUFM` ds + !d = fromJust $ UFM.lookupUFM_Directly ds v - -- Given a previously established equivalence class, construct a mapping - -- that maps each value to an (arbitrary) canonical value. - pickCanonical :: [a] -> Map a a - pickCanonical cls = Map.fromList $ zip cls (repeat (minimum cls)) + !ds' = UFM.plusUFM ds ( UFM.mapUFM (<> d) us ) + !vs' = vs Seq.>< Seq.fromList ( UFM.nonDetKeysUFM us ) + in + doComp ds' vs' -canonicalize :: Ord a => Map a a -> a -> a -canonicalize canon x = Map.findWithDefault x x canon+minViewUniqSet :: forall a. UniqSet a -> Maybe (a, UniqSet a) +minViewUniqSet s = + let m = UFM.ufmToIntMap $ UniqSet.getUniqSet s + in second + ( UniqSet.unsafeUFMToUniqSet + . +#if MIN_VERSION_ghc(8,11,0) + UFM.unsafeIntMapToUFM +#else + ( unsafeCoerce :: IntMap.IntMap a -> UniqFM a a ) +#endif + ) <$> +#if HAS_WORD64MAP + Word64Map.minView +#else + IntMap.minView +#endif + m + +neighboursMap :: forall a l. Uniquable a => [(a, a, l)] -> UniqFM a (UniqFM a l) +neighboursMap edges = foldr addEdge UFM.emptyUFM edges + where + addEdge :: (a, a, l) -> UniqFM a (UniqFM a l) -> UniqFM a (UniqFM a l) + addEdge (u, v, l) m + = UFM.addToUFM_C UFM.plusUFM + (UFM.addToUFM_C UFM.plusUFM m v (UFM.unitUFM u l)) + u (UFM.unitUFM v l) + +canonicalize :: (Uniquable a, Monoid l) => UniqFM a (a, l) -> a -> (a, l) +canonicalize canon x = UFM.lookupWithDefaultUFM canon (x, mempty) x