Cleanup and document fundep plugin (#155)

I made the code style agree with the rest of the codebase, and I wrote down everything I know about how this works.

polysemy-plugin/polysemy-plugin.cabal

@@ -4,7 +4,7 @@ cabal-version: 1.12
 --
 -- see: https://github.com/sol/hpack
 --
--- hash: 63b03550f0fe1e13941505a11116fac8e6eef8a871dff51b30d2bca2263b2a75
+-- hash: b8d6dd19e90295689617adfecbd3bb83127b112840f4f304d956a4d5b33bf821
 
 name:           polysemy-plugin
 version:        0.2.2.0
@@ -31,6 +31,9 @@ library
   exposed-modules:
       Polysemy.Plugin
       Polysemy.Plugin.Fundep
+      Polysemy.Plugin.Fundep.Stuff
+      Polysemy.Plugin.Fundep.Unification
+      Polysemy.Plugin.Fundep.Utils
       Polysemy.Plugin.Phases
   other-modules:
       Paths_polysemy_plugin

polysemy-plugin/src/Polysemy/Plugin/Fundep.hs

@@ -4,7 +4,7 @@
 ------------------------------------------------------------------------------
 -- The MIT License (MIT)
 --
--- Copyright (c) 2017 Luka Horvat
+-- Copyright (c) 2017 Luka Horvat, 2019 Sandy Maguire
 --
 -- Permission is hereby granted, free of charge, to any person obtaining a copy
 -- of this software and associated documentation files (the "Software"), to
@@ -26,218 +26,106 @@
 --
 ------------------------------------------------------------------------------
 --
--- This module is heavily based on 'Control.Effects.Plugin' from the
+-- This module was originally based on 'Control.Effects.Plugin' from the
--- 'simple-effects' package, originally by Luka Horvat.
+-- 'simple-effects' package, by Luka Horvat.
 --
 -- https://gitlab.com/LukaHorvat/simple-effects/commit/966ce80b8b5777a4bd8f87ffd443f5fa80cc8845#f51c1641c95dfaa4827f641013f8017e8cd02aab
 
 module Polysemy.Plugin.Fundep (fundepPlugin) where
 
-import           Class
-import           CoAxiom
-import           Control.Applicative
 import           Control.Monad
 import           Data.Bifunctor
-import           Data.Bool
 import           Data.Coerce
-import           Data.Function (on)
 import           Data.IORef
-import qualified Data.Kind as K
+import qualified Data.Map as M
-import           Data.List
 import           Data.Maybe
 import qualified Data.Set as S
-import           FastString (fsLit)
+import           Polysemy.Plugin.Fundep.Stuff
-import           GHC (TyCon, Name)
+import           Polysemy.Plugin.Fundep.Unification
-import           GHC.TcPluginM.Extra (lookupModule, lookupName)
+import           Polysemy.Plugin.Fundep.Utils
-import           Module (mkModuleName)
-import           OccName (mkTcOcc)
 import           TcEvidence
-import           TcPluginM (TcPluginM, tcLookupClass, tcLookupTyCon, tcPluginIO)
+import           TcPluginM (TcPluginM, tcPluginIO)
 import           TcRnTypes
 import           TcSMonad hiding (tcLookupClass)
-import           TyCoRep (Type (..))
 import           Type
 
 
-data LookupState
-  = Locations
-  | Things
-
-
-type family ThingOf (l :: LookupState) (a :: K.Type) :: K.Type where
-  ThingOf 'Locations _ = (String, String)
-  ThingOf 'Things a = a
-
-
-data PolysemyStuff (l :: LookupState) = PolysemyStuff
-  { findClass    :: ThingOf l Class
-  , semTyCon     :: ThingOf l TyCon
-  , ifStuckTyCon :: ThingOf l TyCon
-  , indexOfTyCon :: ThingOf l TyCon
-  }
-
-
-class CanLookup a where
-  lookupStrategy :: Name -> TcPluginM a
-
-instance CanLookup Class where
-  lookupStrategy = tcLookupClass
-
-instance CanLookup TyCon where
-  lookupStrategy = tcLookupTyCon
-
-
-doLookup :: CanLookup a => ThingOf 'Locations a -> TcPluginM (ThingOf 'Things a)
-doLookup (mdname, name) = do
-  md  <- lookupModule (mkModuleName mdname) $ fsLit "polysemy"
-  nm <- lookupName md $ mkTcOcc name
-  lookupStrategy nm
-
-
-lookupEverything :: PolysemyStuff 'Locations -> TcPluginM (PolysemyStuff 'Things)
-lookupEverything (PolysemyStuff a b c d) =
-  PolysemyStuff <$> doLookup a
-                <*> doLookup b
-                <*> doLookup c
-                <*> doLookup d
-
-
-polysemyStuffLocations :: PolysemyStuff 'Locations
-polysemyStuffLocations = PolysemyStuff
-  { findClass    = ("Polysemy.Internal.Union",                  "Find")
-  , semTyCon     = ("Polysemy.Internal",                        "Sem")
-  , ifStuckTyCon = ("Polysemy.Internal.CustomErrors.Redefined", "IfStuck")
-  , indexOfTyCon = ("Polysemy.Internal.Union",                  "IndexOf")
-  }
-
 
 fundepPlugin :: TcPlugin
 fundepPlugin = TcPlugin
-    { tcPluginInit =
+  { tcPluginInit =
-        (,) <$> tcPluginIO (newIORef S.empty)
+      (,) <$> tcPluginIO (newIORef S.empty)
-            <*> lookupEverything polysemyStuffLocations
+          <*> polysemyStuff
-    , tcPluginSolve = solveFundep
+  , tcPluginSolve = solveFundep
-    , tcPluginStop = const (return ()) }
+  , tcPluginStop = const $ pure ()
+  }
 
-allMonadEffectConstraints :: PolysemyStuff 'Things -> [Ct] -> [(CtLoc, (Type, Type, Type))]
-allMonadEffectConstraints (findClass -> cls) cts =
-    [ (ctLoc cd, (effName, eff, r))
-    | cd@CDictCan{cc_class = cls', cc_tyargs = [_, r, eff]} <- cts
-    , cls == cls'
-    , let effName = getEffName eff
-    ]
 
-singleListToJust :: [a] -> Maybe a
+------------------------------------------------------------------------------
-singleListToJust [a] = Just a
+-- | Corresponds to a 'Polysemy.Internal.Union.Find' constraint. For example,
-singleListToJust _ = Nothing
+-- given @Member (State s) r@, we would get:
+data FindConstraint = FindConstraint
+  { fcLoc        :: CtLoc
+  , fcEffectName :: Type  -- ^ @State@
+  , fcEffect     :: Type  -- ^ @State s@
+  , fcRow        :: Type  -- ^ @r@
+  }
 
-findMatchingEffectIfSingular :: (Type, Type, Type) -> [(Type, Type, Type)] -> Maybe Type
-findMatchingEffectIfSingular (effName, _, mon) ts = singleListToJust
-    [ eff'
-        | (effName', eff', mon') <- ts
-        , eqType effName effName'
-        , eqType mon mon' ]
 
+------------------------------------------------------------------------------
+-- | Given a list of constraints, filter out the 'FindConstraint's.
+getFindConstraints :: PolysemyStuff 'Things -> [Ct] -> [FindConstraint]
+getFindConstraints (findClass -> cls) cts = do
+  cd@CDictCan{cc_class = cls', cc_tyargs = [_, r, eff]} <- cts
+  guard $ cls == cls'
+  pure $ FindConstraint
+    { fcLoc = ctLoc cd
+    , fcEffectName = getEffName eff
+    , fcEffect = eff
+    , fcRow = r
+    }
+
+
+------------------------------------------------------------------------------
+-- | If there's only a single @Member@ in the same @r@ whose effect name
+-- matches, return its effect (including tyvars.)
+findMatchingEffectIfSingular
+    :: FindConstraint
+    -> [FindConstraint]
+    -> Maybe Type
+findMatchingEffectIfSingular (FindConstraint _ eff_name _ r) ts =
+  singleListToJust $ do
+    FindConstraint _ eff_name' eff' r' <- ts
+    guard $ eqType eff_name eff_name'
+    guard $ eqType r r'
+    pure eff'
+
+
+------------------------------------------------------------------------------
+-- | Given an effect, compute its effect name.
 getEffName :: Type -> Type
 getEffName t = fst $ splitAppTys t
 
 
-canUnifyRecursive :: SolveContext -> Type -> Type -> Bool
-canUnifyRecursive solve_ctx = go True
-  where
-    -- It's only OK to solve a polymorphic "given" if we're in the context of
-    -- an interpreter, because it's not really a given!
-    poly_given_ok :: Bool
-    poly_given_ok =
-      case solve_ctx of
-        InterpreterUse _ -> True
-        FunctionDef      -> False
-
-    -- On the first go around, we don't want to unify effects with tyvars, but
-    -- we _do_ want to unify their arguments, thus 'is_first'.
-    go :: Bool -> Type -> Type -> Bool
-    go is_first wanted given =
-      let (w, ws) = splitAppTys wanted
-          (g, gs) = splitAppTys given
-       in (&& bool (canUnify poly_given_ok) eqType is_first w g)
-        . flip all (zip ws gs)
-        $ \(wt, gt) -> canUnify poly_given_ok wt gt || go False wt gt
-
-
-canUnify :: Bool -> Type -> Type -> Bool
-canUnify poly_given_ok wt gt =
-  or [ isTyVarTy wt
-     , isTyVarTy gt && poly_given_ok
-     , eqType wt gt
-     ]
-
-
 ------------------------------------------------------------------------------
--- | Like 'Control.Monad.when', but in the context of an 'Alternative'.
+-- | Generate a wanted unification for the effect described by the
-whenA
+-- 'FindConstraint' and the given effect --- if they can be unified in this
-    :: (Monad m, Alternative z)
+-- context.
-    => Bool
-    -> m a
-    -> m (z a)
-whenA False _ = pure empty
-whenA True ma = fmap pure ma
-
-
 mkWanted
-    :: SolveContext
+    :: FindConstraint
-    -> CtLoc
+    -> SolveContext
-    -> Type
+    -> Type  -- ^ The given effect.
-    -> Type
+    -> TcPluginM (Maybe (Unification, Ct))
-    -> TcPluginM (Maybe ( (OrdType, OrdType)  -- the types we want to unify
+mkWanted fc solve_ctx given =
-                        , Ct                  -- the constraint
-                        ))
-mkWanted solve_ctx loc wanted given =
   whenA (not (mustUnify solve_ctx) || canUnifyRecursive solve_ctx wanted given) $ do
     (ev, _) <- unsafeTcPluginTcM
              . runTcSDeriveds
-             $ newWantedEq loc Nominal wanted given
+             $ newWantedEq (fcLoc fc) Nominal wanted given
-    pure ( (OrdType wanted, OrdType given)
+    pure ( Unification (OrdType wanted) (OrdType given)
          , CNonCanonical ev
          )
-
+  where
-thd :: (a, b, c) -> c
+    wanted = fcEffect fc
-thd (_, _, c) = c
-
-countLength :: (a -> a -> Bool) -> [a] -> [(a, Int)]
-countLength eq as =
-  let grouped = groupBy eq as
-   in zipWith (curry $ bimap head length) grouped grouped
-
-
-------------------------------------------------------------------------------
--- | 'Type's don't have 'Eq' or 'Ord' instances by default, even though there
--- are functions in GHC that implement these operations. This newtype gives us
--- those instances.
-newtype OrdType = OrdType
-  { getOrdType :: Type
-  }
-
-instance Eq OrdType where
-  (==) = eqType `on` getOrdType
-
-instance Ord OrdType where
-  compare = nonDetCmpType `on` getOrdType
-
-
-------------------------------------------------------------------------------
--- | The context in which we're attempting to solve a constraint.
-data SolveContext
-  = -- | In the context of a function definition.
-    FunctionDef
-    -- | In the context of running an interpreter. The 'Bool' corresponds to
-    -- whether we are only trying to solve a single 'Member' constraint right
-    -- now. If so, we *must* produce a unification wanted.
-  | InterpreterUse Bool
-  deriving (Eq, Ord, Show)
-
-mustUnify :: SolveContext -> Bool
-mustUnify FunctionDef = True
-mustUnify (InterpreterUse b) = b
 
 
 ------------------------------------------------------------------------------
@@ -246,68 +134,93 @@ mustUnify (InterpreterUse b) = b
 getBogusRs :: PolysemyStuff 'Things -> [Ct] -> [Type]
 getBogusRs stuff wanteds = do
   CIrredCan ct _ <- wanteds
-  case splitAppTys $ ctev_pred ct of
+  (_, [_, _, a, b]) <- pure . splitAppTys $ ctev_pred ct
-    (_, [_, _, a, b]) ->
+  maybeToList (extractRowFromSem stuff a)
-      maybeToList (getRIfSem stuff a) ++ maybeToList (getRIfSem stuff b)
+    ++ maybeToList (extractRowFromSem stuff b)
-    (_, _) -> []
 
 
 ------------------------------------------------------------------------------
 -- | Take the @r@ out of @Sem r a@.
-getRIfSem :: PolysemyStuff 'Things -> Type -> Maybe Type
+extractRowFromSem :: PolysemyStuff 'Things -> Type -> Maybe Type
-getRIfSem (semTyCon -> sem) ty =
+extractRowFromSem (semTyCon -> sem) ty = do
-  case splitTyConApp_maybe ty of
+  (tycon, [r, _]) <- splitTyConApp_maybe ty
-    Just (tycon, [r, _]) | tycon == sem -> pure r
+  guard $ tycon == sem
-    _                                   -> Nothing
+  pure r
 
 
 ------------------------------------------------------------------------------
 -- | Given a list of bogus @r@s, and the wanted constraints, produce bogus
 -- evidence terms that will prevent @IfStuck (IndexOf r _) _ _@ error messsages.
-solveBogusError :: PolysemyStuff 'Things -> [Type] -> [Ct] -> [(EvTerm, Ct)]
+solveBogusError :: PolysemyStuff 'Things -> [Ct] -> [(EvTerm, Ct)]
-solveBogusError stuff bogus wanteds = do
+solveBogusError stuff wanteds = do
+  let splitTyConApp_list = maybeToList  . splitTyConApp_maybe
+
+  let bogus = getBogusRs stuff wanteds
   ct@(CIrredCan ce _) <- wanteds
-  case splitTyConApp_maybe $ ctev_pred ce of
+  (stuck, [_, _, expr, _, _]) <- splitTyConApp_list $ ctev_pred ce
-    Just (stuck, [_, _, expr, _, _]) | stuck == ifStuckTyCon stuff -> do
+  guard $ stuck == ifStuckTyCon stuff
-      case splitTyConApp_maybe expr of
+  (idx, [_, r, _]) <- splitTyConApp_list expr
-        Just (idx, [_, r, _]) | idx == indexOfTyCon stuff -> do
+  guard $ idx == indexOfTyCon stuff
-          case elem @[] (OrdType r) $ coerce bogus of
+  guard $ elem @[] (OrdType r) $ coerce bogus
-            True -> pure (error "bogus proof for stuck type family", ct)
+  pure (error "bogus proof for stuck type family", ct)
-            False -> []
+
-        _ -> []
+
-    _ -> []
+------------------------------------------------------------------------------
+-- | Determine if there is exactly one wanted find for the @r@ in question.
+exactlyOneWantedForR
+    :: [FindConstraint]  -- ^ Wanted finds
+    -> Type              -- ^ Effect row
+    -> Bool
+exactlyOneWantedForR wanteds
+    = fromMaybe False
+    . flip M.lookup singular_r
+    . OrdType
+  where
+    singular_r = M.fromList
+               -- TODO(sandy): Nothing fails if this is just @second (const
+               -- True)@. Why not? Incomplete test suite, or doing too much
+               -- work?
+               . fmap (second (/= 1))
+               . countLength
+               $ fmap (OrdType . fcRow) wanteds
 
 
 solveFundep
-    :: (IORef (S.Set (OrdType, OrdType)), PolysemyStuff 'Things)
+    :: ( IORef (S.Set Unification)
+       , PolysemyStuff 'Things
+       )
     -> [Ct]
     -> [Ct]
     -> [Ct]
     -> TcPluginM TcPluginResult
 solveFundep _ _ _ [] = pure $ TcPluginOk [] []
-solveFundep (ref, stuff) giv _ want = do
+solveFundep (ref, stuff) given _ wanted = do
-    let bogus = getBogusRs stuff want
+  let wanted_finds = getFindConstraints stuff wanted
-        solved_bogus = solveBogusError stuff bogus want
+      given_finds  = getFindConstraints stuff given
 
-    let wantedEffs = allMonadEffectConstraints stuff want
+  eqs <- forM wanted_finds $ \fc -> do
-        givenEffs = snd <$> allMonadEffectConstraints stuff giv
+    let r  = fcRow fc
-        num_wanteds_by_r = countLength eqType $ fmap (thd . snd) wantedEffs
+    case findMatchingEffectIfSingular fc given_finds of
-        must_unify r =
+      -- We found a real given, therefore we are in the context of a function
-          let Just num_wanted = find (eqType r . fst) num_wanteds_by_r
+      -- with an explicit @Member e r@ constraint.
-           in snd num_wanted /= 1
+      Just eff' -> mkWanted fc FunctionDef eff'
 
-    eqs <- forM wantedEffs $ \(loc, e@(_, eff, r)) -> do
+      -- Otherwise, check to see if @r ~ (e ': r')@. If so, pretend we're
-      case findMatchingEffectIfSingular e givenEffs of
+      -- trying to solve a given @Member e r@. But this can only happen in the
-        Nothing -> do
+      -- context of an interpreter!
-          case splitAppTys r of
+      Nothing ->
-            (_, [_, eff', _]) -> mkWanted (InterpreterUse $ must_unify r) loc eff eff'
+        case splitAppTys r of
-            _                 -> pure Nothing
+          (_, [_, eff', _]) ->
-        Just eff' -> mkWanted FunctionDef loc eff eff'
+            mkWanted fc
+                     (InterpreterUse $ exactlyOneWantedForR wanted_finds r)
+                     eff'
+          _ -> pure Nothing
 
-    already_emitted <- tcPluginIO $ readIORef ref
+  -- We only want to emit a unification wanted once, otherwise a type error can
-    let new_wanteds = filter (not . flip S.member already_emitted . fst)
+  -- force the type checker to loop forever.
-                    $ catMaybes eqs
+  already_emitted <- tcPluginIO $ readIORef ref
+  let (unifications, new_wanteds) = unzipNewWanteds already_emitted $ catMaybes eqs
+  tcPluginIO $ modifyIORef ref $ S.union $ S.fromList unifications
 
-    tcPluginIO $ modifyIORef ref $ S.union $ S.fromList $ fmap fst new_wanteds
+  pure $ TcPluginOk (solveBogusError stuff wanted) new_wanteds
-    pure . TcPluginOk solved_bogus $ fmap snd new_wanteds
 

polysemy-plugin/src/Polysemy/Plugin/Fundep/Stuff.hs

@@ -0,0 +1,83 @@
+module Polysemy.Plugin.Fundep.Stuff
+  ( PolysemyStuff (..)
+  , LookupState (..)
+  , polysemyStuff
+  ) where
+
+import Data.Kind (Type)
+import FastString (fsLit)
+import GHC (Name, Class, TyCon, mkModuleName)
+import GHC.TcPluginM.Extra (lookupModule, lookupName)
+import OccName (mkTcOcc)
+import TcPluginM (TcPluginM, tcLookupClass, tcLookupTyCon)
+
+
+
+------------------------------------------------------------------------------
+-- | All of the things from "polysemy" that we need access to in the plugin.
+-- When @l ~ 'Locations@, each of these is just a pair of strings. When @l
+-- ~ 'Things@, it's actually references to the stuff.
+data PolysemyStuff (l :: LookupState) = PolysemyStuff
+  { findClass    :: ThingOf l Class
+  , semTyCon     :: ThingOf l TyCon
+  , ifStuckTyCon :: ThingOf l TyCon
+  , indexOfTyCon :: ThingOf l TyCon
+  }
+
+
+------------------------------------------------------------------------------
+-- | All of the things we need to lookup.
+polysemyStuffLocations :: PolysemyStuff 'Locations
+polysemyStuffLocations = PolysemyStuff
+  { findClass    = ("Polysemy.Internal.Union",                  "Find")
+  , semTyCon     = ("Polysemy.Internal",                        "Sem")
+  , ifStuckTyCon = ("Polysemy.Internal.CustomErrors.Redefined", "IfStuck")
+  , indexOfTyCon = ("Polysemy.Internal.Union",                  "IndexOf")
+  }
+
+
+------------------------------------------------------------------------------
+-- | Lookup all of the 'PolysemyStuff'.
+polysemyStuff :: TcPluginM (PolysemyStuff 'Things)
+polysemyStuff =
+  let PolysemyStuff a b c d = polysemyStuffLocations
+   in PolysemyStuff <$> doLookup a
+                    <*> doLookup b
+                    <*> doLookup c
+                    <*> doLookup d
+
+
+------------------------------------------------------------------------------
+-- | Data kind for 'ThingOf'.
+data LookupState
+  = Locations
+  | Things
+
+
+------------------------------------------------------------------------------
+-- | HKD indexed by the 'LookupState'; used by 'PolysemyStuff'.
+type family ThingOf (l :: LookupState) (a :: Type) :: Type where
+  ThingOf 'Locations _ = (String, String)
+  ThingOf 'Things    a = a
+
+
+------------------------------------------------------------------------------
+-- | Things that can be found in a 'TcPluginM' environment.
+class CanLookup a where
+  lookupStrategy :: Name -> TcPluginM a
+
+instance CanLookup Class where
+  lookupStrategy = tcLookupClass
+
+instance CanLookup TyCon where
+  lookupStrategy = tcLookupTyCon
+
+
+------------------------------------------------------------------------------
+-- | Transform a @'ThingOf' 'Locations@ into a @'ThingOf' 'Things@.
+doLookup :: CanLookup a => ThingOf 'Locations a -> TcPluginM (ThingOf 'Things a)
+doLookup (mdname, name) = do
+  md <- lookupModule (mkModuleName mdname) $ fsLit "polysemy"
+  nm <- lookupName md $ mkTcOcc name
+  lookupStrategy nm
+

polysemy-plugin/src/Polysemy/Plugin/Fundep/Unification.hs

@@ -0,0 +1,136 @@
+module Polysemy.Plugin.Fundep.Unification where
+
+import           Data.Bool
+import           Data.Function (on)
+import qualified Data.Set as S
+import           TcRnTypes
+import           Type
+
+
+------------------------------------------------------------------------------
+-- | The context in which we're attempting to solve a constraint.
+data SolveContext
+  = -- | In the context of a function definition.
+    FunctionDef
+    -- | In the context of running an interpreter. The 'Bool' corresponds to
+    -- whether we are only trying to solve a single 'Member' constraint right
+    -- now. If so, we *must* produce a unification wanted.
+  | InterpreterUse Bool
+  deriving (Eq, Ord, Show)
+
+
+------------------------------------------------------------------------------
+-- | Depending on the context in which we're solving a constraint, we may or
+-- may not want to force a unification of effects. For example, when defining
+-- user code whose type is @Member (State Int) r => ...@, if we see @get :: Sem
+-- r s@, we should unify @s ~ Int@.
+mustUnify :: SolveContext -> Bool
+mustUnify FunctionDef = True
+mustUnify (InterpreterUse b) = b
+
+
+------------------------------------------------------------------------------
+-- | Determine whether or not two effects are unifiable. This is nuanced.
+--
+-- There are several cases:
+--
+-- 1. [W] ∀ e1. e1   [G] ∀ e2. e2
+--    Always fails, because we never want to unify two effects if effect names
+--    are polymorphic.
+--
+-- 2. [W] State s    [G] State Int
+--    Always succeeds. It's safe to take our given as a fundep annotation.
+--
+-- 3. [W] State Int  [G] State s
+--        (when the [G] is a given that comes from a type signature)
+--
+--    This should fail, because it means we wrote the type signature @Member
+--    (State s) r => ...@, but are trying to use @s@ as an @Int@. Clearly
+--    bogus!
+--
+-- 4. [W] State Int  [G] State s
+--        (when the [G] was generated by running an interpreter)
+--
+--    Sometimes OK, but only if the [G] is the only thing we're trying to solve
+--    right now. Consider the case:
+--
+--      runState 5 $ pure @(Sem (State Int ': r)) ()
+--
+--    Here we have  [G] forall a. Num a => State a  and  [W] State Int. Clearly
+--    the typechecking should flow "backwards" here, out of the row and into
+--    the type of 'runState'.
+--
+--    What happens if there are multiple [G]s in scope for the same @r@? Then
+--    we'd emit multiple unification constraints for the same effect but with
+--    different polymorphic variables, which would unify a bunch of effects
+--    that shouldn't be!
+canUnifyRecursive
+    :: SolveContext
+    -> Type  -- ^ wanted
+    -> Type  -- ^ given
+    -> Bool
+canUnifyRecursive solve_ctx = go True
+  where
+    -- It's only OK to solve a polymorphic "given" if we're in the context of
+    -- an interpreter, because it's not really a given!
+    poly_given_ok :: Bool
+    poly_given_ok =
+      case solve_ctx of
+        InterpreterUse _ -> True
+        FunctionDef      -> False
+
+    -- On the first go around, we don't want to unify effects with tyvars, but
+    -- we _do_ want to unify their arguments, thus 'is_first'.
+    go :: Bool -> Type -> Type -> Bool
+    go is_first wanted given =
+      let (w, ws) = splitAppTys wanted
+          (g, gs) = splitAppTys given
+       in (&& bool (canUnify poly_given_ok) eqType is_first w g)
+        . flip all (zip ws gs)
+        $ \(wt, gt) -> canUnify poly_given_ok wt gt || go False wt gt
+
+
+------------------------------------------------------------------------------
+-- | A non-recursive version of 'canUnifyRecursive'.
+canUnify :: Bool -> Type -> Type -> Bool
+canUnify poly_given_ok wt gt =
+  or [ isTyVarTy wt
+     , isTyVarTy gt && poly_given_ok
+     , eqType wt gt
+     ]
+
+
+------------------------------------------------------------------------------
+-- | A wrapper for two types that we want to say have been unified.
+data Unification = Unification
+  { _unifyLHS :: OrdType
+  , _unifyRHS :: OrdType
+  }
+  deriving (Eq, Ord)
+
+
+------------------------------------------------------------------------------
+-- | 'Type's don't have 'Eq' or 'Ord' instances by default, even though there
+-- are functions in GHC that implement these operations. This newtype gives us
+-- those instances.
+newtype OrdType = OrdType
+  { getOrdType :: Type
+  }
+
+instance Eq OrdType where
+  (==) = eqType `on` getOrdType
+
+instance Ord OrdType where
+  compare = nonDetCmpType `on` getOrdType
+
+
+------------------------------------------------------------------------------
+-- | Filter out the unifications we've already emitted, and then give back the
+-- things we should put into the @S.Set Unification@, and the new constraints
+-- we should emit.
+unzipNewWanteds
+    :: S.Set Unification
+    -> [(Unification, Ct)]
+    -> ([Unification], [Ct])
+unzipNewWanteds old = unzip . filter (not . flip S.member old . fst)
+

polysemy-plugin/src/Polysemy/Plugin/Fundep/Utils.hs

@@ -0,0 +1,33 @@
+module Polysemy.Plugin.Fundep.Utils  where
+
+import Control.Applicative
+import Data.Bifunctor
+import Data.List
+
+
+
+------------------------------------------------------------------------------
+-- | Returns the head of the list iff there is exactly one element.
+singleListToJust :: [a] -> Maybe a
+singleListToJust [a] = Just a
+singleListToJust _ = Nothing
+
+
+------------------------------------------------------------------------------
+-- | Like 'Control.Monad.when', but in the context of an 'Alternative'.
+whenA
+    :: (Monad m, Alternative z)
+    => Bool
+    -> m a
+    -> m (z a)
+whenA False _ = pure empty
+whenA True ma = fmap pure ma
+
+
+------------------------------------------------------------------------------
+-- | Count the number of times 'a' is present in the list.
+countLength ::  Eq a => [a] -> [(a, Int)]
+countLength as =
+  let grouped = group as
+   in zipWith (curry $ bimap head length) grouped grouped
+