From ed3b1402e05b57d133495ae11fb606e082ca2a67 Mon Sep 17 00:00:00 2001
From: Dan Doel <dan.doel@gmail.com>
Date: Thu, 19 Aug 2021 15:49:52 -0400
Subject: [PATCH] Only elide covariant variables in removePureEffects

---
 unison-core/src/Unison/Type.hs | 27 +++++++++++++++++++++------
 1 file changed, 21 insertions(+), 6 deletions(-)

diff --git a/unison-core/src/Unison/Type.hs b/unison-core/src/Unison/Type.hs
index b2e9de56b..6e8862f55 100644
--- a/unison-core/src/Unison/Type.hs
+++ b/unison-core/src/Unison/Type.hs
@@ -552,14 +552,29 @@ removeAllEffectVars t = let
 removePureEffects :: ABT.Var v => Type v a -> Type v a
 removePureEffects t | not Settings.removePureEffects = t
                     | otherwise =
-  generalize vs $ removeEffectVars (Set.filter isPure fvs) tu
+  generalize vs $ removeEffectVars fvs tu
   where
     (vs, tu) = unforall' t
-    fvs = freeEffectVars tu `Set.difference` ABT.freeVars t
-    -- If an effect variable is mentioned only once, it is on
-    -- an arrow `a ->{e} b`. Generalizing this to
-    -- `∀ e . a ->{e} b` gives us the pure arrow `a -> b`.
-    isPure v = ABT.occurrences v tu <= 1
+    vss = Set.fromList vs
+    fvs = freeEffectVars tu `Set.difference` keep
+
+    keep = keepVarsT True tu
+
+    keepVarsT pos (Arrow' i o)
+      = keepVarsT (not pos) i <> keepVarsT pos o
+    keepVarsT pos (Effect1' e o)
+      = keepVarsT pos e <> keepVarsT pos o
+    keepVarsT pos (Effects' es) = foldMap (keepVarsE pos) es
+    keepVarsT pos (ForallNamed' _ t) = keepVarsT pos t
+    keepVarsT pos (IntroOuterNamed' _ t) = keepVarsT pos t
+    keepVarsT _ t = freeVars t
+
+    -- Note, this only allows removal if the variable was quantified,
+    -- so variables that were free in `t` will not be removed.
+    keepVarsE pos (Var' v)
+      | pos, v `Set.member` vss = mempty
+      | otherwise = Set.singleton v
+    keepVarsE pos e = keepVarsT pos e
 
 editFunctionResult
   :: forall v a