Tactics support for using given constraints

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