Tactics support for using given constraints (#534)

This PR allows tactics to use methods from given constraints, meaning it can solve holes like this:

```haskell
showMe :: Show a => a -> String
showMe = _
```

It will not, however, discover instances. So this one *won't* solve:


```haskell
showMe :: Int -> String
showMe = _
```

There's quite a lot of finicky details going on in order to support this. The primary challenge is that our types are running after the typechecker has finished, meaning it's already solved the constraints and inlined their evidence. Our solution is to look up the written polymorphic type and unify it with the final, typechecked type. We can use the polymorphic type to find the theta context, and instantiate every class method in the theta.

In addition, this PR fixes a subtle bug in our unification code, which could cause skolems to unify in some circumstances.

Furthermore, it adds a tie-breaker to the scoring metric to prefer shorter programs.

Author: isovector

plugins/tactics/hls-tactics-plugin.cabal

@@ -84,6 +84,7 @@ test-suite tests
   main-is: Main.hs
   other-modules:
     AutoTupleSpec
+    UnificationSpec
   hs-source-dirs:
       test
   ghc-options: -Wall -Wredundant-constraints -threaded -rtsopts -with-rtsopts=-N

plugins/tactics/src/Ide/Plugin/Tactic.hs

@@ -265,9 +265,11 @@ judgementForHole state nfp range = do
                 $ getDefiningBindings binds rss)
               tcg
       hyps = hypothesisFromBindings rss binds
+      ambient = M.fromList $ contextMethodHypothesis ctx
   pure ( resulting_range
        , mkFirstJudgement
            hyps
+           ambient
            (isRhsHole rss tcs)
            (maybe
               mempty

plugins/tactics/src/Ide/Plugin/Tactic/CodeGen.hs

@@ -29,7 +29,8 @@ import           Type hiding (Var)
 
 useOccName :: MonadState TacticState m => Judgement -> OccName -> m ()
 useOccName jdg name =
-  case M.lookup name $ jHypothesis jdg of
+  -- Only score points if this is in the local hypothesis
+  case M.lookup name $ jLocalHypothesis jdg of
     Just{}  -> modify $ withUsedVals $ S.insert name
     Nothing -> pure ()
 

plugins/tactics/src/Ide/Plugin/Tactic/Context.hs

@@ -10,16 +10,50 @@ import Development.IDE.GHC.Compat
 import Ide.Plugin.Tactic.Types
 import OccName
 import TcRnTypes
+import Ide.Plugin.Tactic.GHC (tacticsThetaTy)
+import Ide.Plugin.Tactic.Machinery (methodHypothesis)
+import Data.Maybe (mapMaybe)
+import Data.List
+import TcType (substTy, tcSplitSigmaTy)
+import Unify (tcUnifyTy)
 
 
 mkContext :: [(OccName, CType)] -> TcGblEnv -> Context
-mkContext locals
-  = Context locals
-  . fmap splitId
-  . (getFunBindId =<<)
-  . fmap unLoc
-  . bagToList
-  . tcg_binds
+mkContext locals tcg = Context
+  { ctxDefiningFuncs = locals
+  , ctxModuleFuncs = fmap splitId
+                   . (getFunBindId =<<)
+                   . fmap unLoc
+                   . bagToList
+                   $ tcg_binds tcg
+  }
+
+
+------------------------------------------------------------------------------
+-- | Find all of the class methods that exist from the givens in the context.
+contextMethodHypothesis :: Context -> [(OccName, CType)]
+contextMethodHypothesis ctx
+  = join
+  . concatMap
+      ( mapMaybe methodHypothesis
+      . tacticsThetaTy
+      . unCType
+      )
+  . mapMaybe (definedThetaType ctx)
+  . fmap fst
+  $ ctxDefiningFuncs ctx
+
+
+------------------------------------------------------------------------------
+-- | Given the name of a function that exists in 'ctxDefiningFuncs', get its
+-- theta type.
+definedThetaType :: Context -> OccName -> Maybe CType
+definedThetaType ctx name = do
+  (_, CType mono) <- find ((== name) . fst) $ ctxDefiningFuncs ctx
+  (_, CType poly) <- find ((== name) . fst) $ ctxModuleFuncs ctx
+  let (_, _, poly') = tcSplitSigmaTy poly
+  subst <- tcUnifyTy poly' mono
+  pure $ CType $ substTy subst $ snd $ splitForAllTys poly
 
 
 splitId :: Id -> (OccName, CType)

plugins/tactics/src/Ide/Plugin/Tactic/GHC.hs

@@ -1,18 +1,24 @@
-{-# LANGUAGE CPP             #-}
-{-# LANGUAGE PatternSynonyms #-}
-{-# LANGUAGE ViewPatterns    #-}
+{-# LANGUAGE CPP              #-}
+{-# LANGUAGE FlexibleContexts #-}
+{-# LANGUAGE PatternSynonyms  #-}
+{-# LANGUAGE ViewPatterns     #-}
 
 module Ide.Plugin.Tactic.GHC where
 
-import Data.Maybe (isJust)
-import Development.IDE.GHC.Compat
-import OccName
-import TcType
-import TyCoRep
-import Type
-import TysWiredIn (intTyCon, floatTyCon, doubleTyCon, charTyCon)
-import Unique
-import Var
+import           Control.Monad.State
+import qualified Data.Map as M
+import           Data.Maybe (isJust)
+import           Data.Traversable
+import           Development.IDE.GHC.Compat
+import           Generics.SYB (mkT, everywhere)
+import           Ide.Plugin.Tactic.Types
+import           OccName
+import           TcType
+import           TyCoRep
+import           Type
+import           TysWiredIn (intTyCon, floatTyCon, doubleTyCon, charTyCon)
+import           Unique
+import           Var
 
 
 tcTyVar_maybe :: Type -> Maybe Var
@@ -43,8 +49,44 @@ cloneTyVar t =
 ------------------------------------------------------------------------------
 -- | Is this a function type?
 isFunction :: Type -> Bool
-isFunction (tcSplitFunTys -> ((_:_), _)) = True
-isFunction _ = False
+isFunction (tacticsSplitFunTy -> (_, _, [], _)) = False
+isFunction _ = True
+
+
+------------------------------------------------------------------------------
+-- | Split a function, also splitting out its quantified variables and theta
+-- context.
+tacticsSplitFunTy :: Type -> ([TyVar], ThetaType, [Type], Type)
+tacticsSplitFunTy t
+  = let (vars, theta, t') = tcSplitSigmaTy t
+        (args, res) = tcSplitFunTys t'
+     in (vars, theta, args, res)
+
+
+------------------------------------------------------------------------------
+-- | Rip the theta context out of a regular type.
+tacticsThetaTy :: Type -> ThetaType
+tacticsThetaTy (tcSplitSigmaTy -> (_, theta,  _)) = theta
+
+
+------------------------------------------------------------------------------
+-- | Instantiate all of the quantified type variables in a type with fresh
+-- skolems.
+freshTyvars :: MonadState TacticState m => Type -> m Type
+freshTyvars t = do
+  let (tvs, _, _, _) = tacticsSplitFunTy t
+  reps <- fmap M.fromList
+        $ for tvs $ \tv -> do
+            uniq <- freshUnique
+            pure $ (tv, setTyVarUnique tv uniq)
+  pure $
+    everywhere
+      (mkT $ \tv ->
+        case M.lookup tv reps of
+          Just tv' -> tv'
+          Nothing -> tv
+      ) t
+
 
 ------------------------------------------------------------------------------
 -- | Is this an algebraic type?

plugins/tactics/src/Ide/Plugin/Tactic/Judgements.hs

@@ -162,8 +162,17 @@ disallowing ns =
   field @"_jHypothesis" %~ flip M.withoutKeys (S.fromList ns)
 
 
+------------------------------------------------------------------------------
+-- | The hypothesis, consisting of local terms and the ambient environment
+-- (includes and class methods.)
 jHypothesis :: Judgement' a -> Map OccName a
-jHypothesis = _jHypothesis
+jHypothesis = _jHypothesis <> _jAmbientHypothesis
+
+
+------------------------------------------------------------------------------
+-- | Just the local hypothesis.
+jLocalHypothesis :: Judgement' a -> Map OccName a
+jLocalHypothesis = _jHypothesis
 
 
 isPatVal :: Judgement' a -> OccName -> Bool
@@ -191,13 +200,15 @@ substJdg :: TCvSubst -> Judgement -> Judgement
 substJdg subst = fmap $ coerce . substTy subst . coerce
 
 mkFirstJudgement
-    :: M.Map OccName CType
+    :: M.Map OccName CType  -- ^ local hypothesis
+    -> M.Map OccName CType  -- ^ ambient hypothesis
     -> Bool  -- ^ are we in the top level rhs hole?
     -> M.Map OccName [[OccName]]  -- ^ existing pos vals
     -> Type
     -> Judgement' CType
-mkFirstJudgement hy top posvals goal = Judgement
+mkFirstJudgement hy ambient top posvals goal = Judgement
   { _jHypothesis        = hy
+  , _jAmbientHypothesis = ambient
   , _jDestructed        = mempty
   , _jPatternVals       = mempty
   , _jBlacklistDestruct = False

plugins/tactics/src/Ide/Plugin/Tactic/Machinery.hs

@@ -1,3 +1,4 @@
+{-# LANGUAGE ScopedTypeVariables #-}
 {-# LANGUAGE DeriveFunctor         #-}
 {-# LANGUAGE DeriveGeneric         #-}
 {-# LANGUAGE DerivingStrategies    #-}
@@ -17,6 +18,7 @@ module Ide.Plugin.Tactic.Machinery
   ( module Ide.Plugin.Tactic.Machinery
   ) where
 
+import           Class (Class(classTyVars))
 import           Control.Arrow
 import           Control.Monad.Error.Class
 import           Control.Monad.Reader
@@ -25,12 +27,15 @@ import           Control.Monad.State.Class (gets, modify)
 import           Control.Monad.State.Strict (StateT (..))
 import           Data.Coerce
 import           Data.Either
-import           Data.List (intercalate, sortBy)
+import           Data.Functor ((<&>))
+import           Data.Generics (mkQ, everything, gcount)
+import           Data.List (sortBy)
 import           Data.Ord (comparing, Down(..))
 import qualified Data.Set as S
 import           Development.IDE.GHC.Compat
 import           Ide.Plugin.Tactic.Judgements
 import           Ide.Plugin.Tactic.Types
+import           OccName (HasOccName(occName))
 import           Refinery.ProofState
 import           Refinery.Tactic
 import           Refinery.Tactic.Internal
@@ -74,7 +79,8 @@ runTactic ctx jdg t =
       (errs, []) -> Left $ take 50 $ errs
       (_, fmap assoc23 -> solns) -> do
         let sorted =
-              sortBy (comparing $ Down . uncurry scoreSolution . snd) solns
+              flip sortBy solns $ comparing $ \((_, ext), (jdg, holes)) ->
+                Down $ scoreSolution ext jdg holes
         case sorted of
           (((tr, ext), _) : _) ->
             Right
@@ -121,56 +127,97 @@ setRecursionFrameData b = do
 
 
 scoreSolution
-    :: TacticState
+    :: LHsExpr GhcPs
+    -> TacticState
     -> [Judgement]
     -> ( Penalize Int  -- number of holes
        , Reward Bool   -- all bindings used
        , Penalize Int  -- number of introduced bindings
        , Reward Int    -- number used bindings
+       , Penalize Int  -- size of extract
        )
-scoreSolution TacticState{..} holes
+scoreSolution ext TacticState{..} holes
   = ( Penalize $ length holes
-    , Reward $ S.null $ ts_intro_vals S.\\ ts_used_vals
+    , Reward   $ S.null $ ts_intro_vals S.\\ ts_used_vals
     , Penalize $ S.size ts_intro_vals
-    , Reward $ S.size ts_used_vals
+    , Reward   $ S.size ts_used_vals
+    , Penalize $ solutionSize ext
     )
 
 
+------------------------------------------------------------------------------
+-- | Compute the number of 'LHsExpr' nodes; used as a rough metric for code
+-- size.
+solutionSize :: LHsExpr GhcPs -> Int
+solutionSize = everything (+) $ gcount $ mkQ False $ \case
+  (_ :: LHsExpr GhcPs) -> True
+
+
 newtype Penalize a = Penalize a
   deriving (Eq, Ord, Show) via (Down a)
 
 newtype Reward a = Reward a
   deriving (Eq, Ord, Show) via a
 
 
+------------------------------------------------------------------------------
+-- | Like 'tcUnifyTy', but takes a list of skolems to prevent unification of.
+tryUnifyUnivarsButNotSkolems :: [TyVar] -> CType -> CType -> Maybe TCvSubst
+tryUnifyUnivarsButNotSkolems skolems goal inst =
+  case tcUnifyTysFG (skolemsOf skolems) [unCType inst] [unCType goal] of
+    Unifiable subst -> pure subst
+    _ -> Nothing
+
 
 ------------------------------------------------------------------------------
--- | We need to make sure that we don't try to unify any skolems.
--- To see why, consider the case:
---
--- uhh :: (Int -> Int) -> a
--- uhh f = _
---
--- If we were to apply 'f', then we would try to unify 'Int' and 'a'.
--- This is fine from the perspective of 'tcUnifyTy', but will cause obvious
--- type errors in our use case. Therefore, we need to ensure that our
--- 'TCvSubst' doesn't try to unify skolems.
-checkSkolemUnification :: CType -> CType -> TCvSubst -> RuleM ()
-checkSkolemUnification t1 t2 subst = do
-    skolems <- gets ts_skolems
-    unless (all (flip notElemTCvSubst subst) skolems) $
-      throwError (UnificationError t1 t2)
+-- | Helper method for 'tryUnifyUnivarsButNotSkolems'
+skolemsOf :: [TyVar] -> TyVar -> BindFlag
+skolemsOf tvs tv =
+  case elem tv tvs of
+    True  -> Skolem
+    False -> BindMe
 
 
 ------------------------------------------------------------------------------
 -- | Attempt to unify two types.
 unify :: CType -- ^ The goal type
       -> CType -- ^ The type we are trying unify the goal type with
       -> RuleM ()
-unify goal inst =
-    case tcUnifyTy (unCType inst) (unCType goal) of
-      Just subst -> do
-          checkSkolemUnification inst goal subst
-          modify (\s -> s { ts_unifier = unionTCvSubst subst (ts_unifier s) })
-      Nothing -> throwError (UnificationError inst goal)
+unify goal inst = do
+  skolems <- gets ts_skolems
+  case tryUnifyUnivarsButNotSkolems skolems goal inst of
+    Just subst ->
+      modify (\s -> s { ts_unifier = unionTCvSubst subst (ts_unifier s) })
+    Nothing -> throwError (UnificationError inst goal)
+
+
+------------------------------------------------------------------------------
+-- | Get the class methods of a 'PredType', correctly dealing with
+-- instantiation of quantified class types.
+methodHypothesis :: PredType -> Maybe [(OccName, CType)]
+methodHypothesis ty = do
+  (tc, apps) <- splitTyConApp_maybe ty
+  cls <- tyConClass_maybe tc
+  let methods = classMethods cls
+      tvs     = classTyVars cls
+      subst   = zipTvSubst tvs apps
+  sc_methods <- fmap join
+              $ traverse (methodHypothesis . substTy subst)
+              $ classSCTheta cls
+  pure $ mappend sc_methods $ methods <&> \method ->
+    let (_, _, ty) = tcSplitSigmaTy $ idType method
+    in (occName method,  CType $ substTy subst ty)
+
+
+------------------------------------------------------------------------------
+-- | Run the given tactic iff the current hole contains no univars. Skolems and
+-- already decided univars are OK though.
+requireConcreteHole :: TacticsM a -> TacticsM a
+requireConcreteHole m = do
+  jdg     <- goal
+  skolems <- gets $ S.fromList . ts_skolems
+  let vars = S.fromList $ tyCoVarsOfTypeWellScoped $ unCType $ jGoal jdg
+  case S.size $ vars S.\\ skolems of
+    0 -> m
+    _ -> throwError TooPolymorphic