Properly track skolems for plugin unification (#423)

* Properly track skolems for plugin unification

* Update haddock

* Cleanup imports

* Guard the import of unify

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

@@ -111,8 +111,10 @@ findMatchingEffectIfSingular (FindConstraint _ eff_name wanted r) ts =
     FindConstraint _ eff_name' eff' r' <- ts
     guard $ eqType eff_name eff_name'
     guard $ eqType r r'
-    guard $ canUnifyRecursive FunctionDef wanted eff'
+    guard $ canUnify (FunctionDef skolems) wanted eff'
     pure eff'
+  where
+    skolems = S.fromList $ foldMap (tyCoVarsOfTypeWellScoped . fcEffect) ts
 
 
 ------------------------------------------------------------------------------
@@ -148,7 +150,7 @@ mkWanted
     -> Type  -- ^ The given effect.
     -> TcPluginM (Maybe (Unification, Ct))
 mkWanted fc solve_ctx given =
-  whenA (not (mustUnify solve_ctx) || canUnifyRecursive solve_ctx wanted given) $
+  whenA (not (mustUnify solve_ctx) || canUnify solve_ctx wanted given) $
     mkWantedForce fc given
   where
     wanted = fcEffect fc
@@ -250,3 +252,4 @@ solveFundep (ref, stuff) given _ wanted = do
   tcPluginIO $ modifyIORef ref $ S.union $ S.fromList unifications
 
   pure $ TcPluginOk (solveBogusError stuff wanted) new_wanteds
+

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

@@ -1,9 +1,11 @@
-{-# LANGUAGE CPP                       #-}
+{-# LANGUAGE CPP #-}
 
 module Polysemy.Plugin.Fundep.Unification where
 
 import           Data.Bool
 import           Data.Function (on)
+import           Data.Maybe (isJust)
+import           Data.Set (Set)
 import qualified Data.Set as S
 #if __GLASGOW_HASKELL__ >= 900
 import           GHC.Tc.Types.Constraint
@@ -15,8 +17,10 @@ import           TcRnTypes
 
 #if __GLASGOW_HASKELL__ >= 900
 import           GHC.Core.Type
+import           GHC.Core.Unify
 #else
 import           Type
+import           Unify
 #endif
 
 
@@ -24,13 +28,14 @@ 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 a function definition. The @Set TyVar@ is all of the
+    -- skolems that exist in the [G] constraints for this function.
+    FunctionDef (Set TyVar)
     -- | 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)
+  deriving (Eq, Ord)
 
 
 ------------------------------------------------------------------------------
@@ -39,79 +44,39 @@ data SolveContext
 -- 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 (FunctionDef _) = True
 mustUnify (InterpreterUse b) = b
 
 
 ------------------------------------------------------------------------------
--- | Determine whether or not two effects are unifiable. This is nuanced.
+-- | Determine whether or not two effects are unifiable.
 --
--- 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
+-- All free variables in [W] constraints are considered skolems, and thus are
+-- not allowed to unify with anything but themselves. This properly handles all
+-- cases in which we are unifying ambiguous [W] constraints (which are true
+-- type variables) against [G] constraints.
+canUnify
     :: SolveContext
     -> Type  -- ^ wanted
     -> Type  -- ^ given
     -> Bool
-canUnifyRecursive solve_ctx = go True
+canUnify solve_ctx = (isJust .) . tryUnifyUnivarsButNotSkolems skolems
   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 =
+    skolems :: Set TyVar
+    skolems =
       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
+        InterpreterUse _ -> mempty
+        FunctionDef s    -> s
 
 
-------------------------------------------------------------------------------
--- | 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
-     ]
+tryUnifyUnivarsButNotSkolems :: Set TyVar -> Type -> Type -> Maybe TCvSubst
+tryUnifyUnivarsButNotSkolems skolems goal inst =
+  case tcUnifyTysFG
+         (bool BindMe Skolem . flip S.member skolems)
+         [inst]
+         [goal] of
+    Unifiable subst -> pure subst
+    _               -> Nothing
 
 
 ------------------------------------------------------------------------------

polysemy-plugin/test/ExampleSpec.hs

@@ -1,5 +1,4 @@
 {-# LANGUAGE TemplateHaskell #-}
-{-# OPTIONS_GHC -fplugin=Polysemy.Plugin #-}
 
 module ExampleSpec where