Wingman: Fix fundeps (#2611)

* Fix fundeps

* One unifier to rule them all

* Fix imports

Co-authored-by: Pepe Iborra <[email protected]>

Author: isovector

Diff for: plugins/hls-tactics-plugin/src/Wingman/GHC.hs

@@ -6,6 +6,7 @@ module Wingman.GHC where
 import           Control.Monad.State
 import           Control.Monad.Trans.Maybe (MaybeT(..))
 import           Data.Bool (bool)
+import           Data.Coerce (coerce)
 import           Data.Function (on)
 import           Data.Functor ((<&>))
 import           Data.List (isPrefixOf)
@@ -359,12 +360,17 @@ expandTyFam ctx = snd . normaliseType  (ctxFamInstEnvs ctx) Nominal
 -- | Like 'tcUnifyTy', but takes a list of skolems to prevent unification of.
 tryUnifyUnivarsButNotSkolems :: Set TyVar -> CType -> CType -> Maybe TCvSubst
 tryUnifyUnivarsButNotSkolems skolems goal inst =
-  case tcUnifyTysFG
-         (bool BindMe Skolem . flip S.member skolems)
-         [unCType inst]
-         [unCType goal] of
-    Unifiable subst -> pure subst
-    _               -> Nothing
+  tryUnifyUnivarsButNotSkolemsMany skolems $ coerce [(goal, inst)]
+
+------------------------------------------------------------------------------
+-- | Like 'tryUnifyUnivarsButNotSkolems', but takes a list
+-- of pairs of types to unify.
+tryUnifyUnivarsButNotSkolemsMany :: Set TyVar -> [(Type, Type)] -> Maybe TCvSubst
+tryUnifyUnivarsButNotSkolemsMany skolems (unzip -> (goal, inst)) =
+  tcUnifyTys
+    (bool BindMe Skolem . flip S.member skolems)
+    inst
+    goal
 
 
 updateSubst :: TCvSubst -> TacticState -> TacticState

Diff for: plugins/hls-tactics-plugin/src/Wingman/Machinery.hs

@@ -1,5 +1,6 @@
-{-# LANGUAGE TupleSections   #-}
-{-# LANGUAGE RankNTypes #-}
+{-# LANGUAGE CPP           #-}
+{-# LANGUAGE RankNTypes    #-}
+{-# LANGUAGE TupleSections #-}
 
 module Wingman.Machinery where
 
@@ -30,11 +31,19 @@ import           Refinery.Tactic
 import           Refinery.Tactic.Internal
 import           System.Timeout (timeout)
 import           Wingman.Context (getInstance)
-import           Wingman.GHC (tryUnifyUnivarsButNotSkolems, updateSubst, tacticsGetDataCons, freshTyvars)
+import           Wingman.GHC (tryUnifyUnivarsButNotSkolems, updateSubst, tacticsGetDataCons, freshTyvars, tryUnifyUnivarsButNotSkolemsMany)
 import           Wingman.Judgements
 import           Wingman.Simplify (simplify)
 import           Wingman.Types
 
+#if __GLASGOW_HASKELL__ < 900
+import FunDeps (fd_eqs, improveFromInstEnv)
+import Pair (unPair)
+#else
+import GHC.Tc.Instance.FunDeps (fd_eqs, improveFromInstEnv)
+import GHC.Data.Pair (unPair)
+#endif
+
 
 substCTy :: TCvSubst -> CType -> CType
 substCTy subst = coerce . substTy subst . coerce
@@ -245,6 +254,23 @@ unify goal inst = do
       modify $ updateSubst subst
     Nothing -> cut
 
+------------------------------------------------------------------------------
+-- | Get a substition out of a theta's fundeps
+learnFromFundeps
+    :: ThetaType
+    -> RuleM ()
+learnFromFundeps theta = do
+  inst_envs <- asks ctxInstEnvs
+  skolems <- gets ts_skolems
+  subst <- gets ts_unifier
+  let theta' = substTheta subst theta
+      fundeps = foldMap (foldMap fd_eqs . improveFromInstEnv inst_envs (\_ _ -> ())) theta'
+  case tryUnifyUnivarsButNotSkolemsMany skolems $ fmap unPair fundeps of
+    Just subst ->
+      modify $ updateSubst subst
+    Nothing -> cut
+
+
 cut :: RuleT jdg ext err s m a
 cut = RuleT Empty
 

Diff for: plugins/hls-tactics-plugin/src/Wingman/Tactics.hs

@@ -280,11 +280,12 @@ apply (Unsaturated n) hi = tracing ("apply' " <> show (hi_name hi)) $ do
       ty = unCType $ hi_type hi
       func = hi_name hi
   ty' <- freshTyvars ty
-  let (_, _, all_args, ret) = tacticsSplitFunTy ty'
+  let (_, theta, all_args, ret) = tacticsSplitFunTy ty'
       saturated_args = dropEnd n all_args
       unsaturated_args = takeEnd n all_args
   rule $ \jdg -> do
     unify g (CType $ mkVisFunTys unsaturated_args ret)
+    learnFromFundeps theta
     ext
         <- fmap unzipTrace
         $ traverse ( newSubgoal

Diff for: plugins/hls-tactics-plugin/test/CodeAction/RunMetaprogramSpec.hs

@@ -43,5 +43,7 @@ spec = do
     metaTest  5  9 "MetaIdiom"
     metaTest  7  9 "MetaIdiomRecord"
 
+    metaTest 14 10 "MetaFundeps"
+
     metaTest  2 12 "IntrosTooMany"
 

Diff for: plugins/hls-tactics-plugin/test/golden/MetaFundeps.expected.hs

@@ -0,0 +1,16 @@
+{-# LANGUAGE FunctionalDependencies #-}
+{-# LANGUAGE MultiParamTypeClasses  #-}
+
+class Blah a b | a -> b, b -> a
+instance Blah Int Bool
+
+foo :: Int
+foo = 10
+
+bar :: Blah a b => a -> b
+bar = undefined
+
+qux :: Bool
+qux = bar foo
+
+

Diff for: plugins/hls-tactics-plugin/test/golden/MetaFundeps.hs

@@ -0,0 +1,16 @@
+{-# LANGUAGE FunctionalDependencies #-}
+{-# LANGUAGE MultiParamTypeClasses  #-}
+
+class Blah a b | a -> b, b -> a
+instance Blah Int Bool
+
+foo :: Int
+foo = 10
+
+bar :: Blah a b => a -> b
+bar = undefined
+
+qux :: Bool
+qux = [wingman| use bar, use foo |]
+
+