Improve how wingman uses evidence

plugins/hls-tactics-plugin/src/Wingman/CodeGen.hs

@@ -11,8 +11,11 @@ module Wingman.CodeGen
 
 import           Control.Lens ((%~), (<>~), (&))
 import           Control.Monad.Except
+import           Control.Monad.State
 import           Data.Generics.Labels ()
 import           Data.List
+import           Data.Maybe (mapMaybe)
+import           Data.Monoid (Endo(..))
 import qualified Data.Set as S
 import           Data.Traversable
 import           DataCon
@@ -26,6 +29,7 @@ import           Type hiding (Var)
 import           Wingman.CodeGen.Utils
 import           Wingman.GHC
 import           Wingman.Judgements
+import           Wingman.Judgements.Theta
 import           Wingman.Machinery
 import           Wingman.Naming
 import           Wingman.Types
@@ -50,12 +54,20 @@ destructMatches f scrut t jdg = do
       case dcs of
         [] -> throwError $ GoalMismatch "destruct" g
         _ -> fmap unzipTrace $ for dcs $ \dc -> do
-          let args = dataConInstOrigArgTys' dc apps
+          let ev = mapMaybe mkEvidence $ dataConInstArgTys dc apps
+              -- We explicitly do not need to add the method hypothesis to
+              -- #syn_scoped
+              method_hy = foldMap evidenceToHypothesis ev
+              args = dataConInstOrigArgTys' dc apps
+          modify $ appEndo $ foldMap (Endo . evidenceToSubst) ev
+          subst <- gets ts_unifier
           names <- mkManyGoodNames (hyNamesInScope hy) args
           let hy' = patternHypothesis scrut dc jdg
                   $ zip names
                   $ coerce args
-              j = introduce hy'
+              j = fmap (CType . substTyAddInScope subst . unCType)
+                $ introduce hy'
+                $ introduce method_hy
                 $ withNewGoal g jdg
           ext <- f dc j
           pure $ ext

plugins/hls-tactics-plugin/src/Wingman/Judgements.hs

@@ -369,6 +369,13 @@ isTopLevel TopLevelArgPrv{} = True
 isTopLevel _                = False
 
 
+------------------------------------------------------------------------------
+-- | Was this term defined by the user?
+isUserProv :: Provenance -> Bool
+isUserProv UserPrv{} = True
+isUserProv _         = False
+
+
 ------------------------------------------------------------------------------
 -- | Is this a local function argument, pattern match or user val?
 isLocalHypothesis :: Provenance -> Bool

plugins/hls-tactics-plugin/src/Wingman/Judgements/Theta.hs

@@ -1,30 +1,88 @@
 {-# LANGUAGE ViewPatterns #-}
 
 module Wingman.Judgements.Theta
-  ( getMethodHypothesisAtHole
+  ( Evidence
+  , getEvidenceAtHole
+  , mkEvidence
+  , evidenceToSubst
+  , evidenceToHypothesis
   ) where
 
-import           Data.Maybe (fromMaybe)
+import           Data.Maybe (fromMaybe, mapMaybe)
 import           Data.Set (Set)
 import qualified Data.Set as S
 import           Development.IDE.GHC.Compat
 import           Generics.SYB
-import           GhcPlugins (EvVar, mkVarOcc)
+import           GhcPlugins (mkVarOcc, splitTyConApp_maybe, eqTyCon, getTyVar_maybe)
+import           TcEvidence
+import           TcType (tcTyConAppTyCon_maybe)
+import           TysPrim (eqPrimTyCon)
 import           Wingman.Machinery
 import           Wingman.Types
 
 
 ------------------------------------------------------------------------------
--- | Create a 'Hypothesis' containing 'ClassMethodPrv' provenance. For every
--- dictionary that is in scope at the given 'SrcSpan', find every method and
--- superclass method available.
-getMethodHypothesisAtHole :: SrcSpan -> LHsBinds GhcTc -> Hypothesis CType
-getMethodHypothesisAtHole dst
-  = Hypothesis
-  . excludeForbiddenMethods
-  . fromMaybe []
-  . foldMap methodHypothesis
-  . (everything (<>) $ mkQ mempty $ evbinds dst)
+-- | Something we've learned about the type environment.
+data Evidence
+    -- | The two types are equal, via a @a ~ b@ relationship
+  = EqualityOfTypes Type Type
+    -- | We have an instance in scope
+  | HasInstance PredType
+  deriving (Show)
+
+
+------------------------------------------------------------------------------
+-- | Given a 'PredType', pull an 'Evidence' out of it.
+mkEvidence :: PredType -> Maybe Evidence
+mkEvidence (getEqualityTheta -> Just (a, b))
+  = Just $ EqualityOfTypes a b
+mkEvidence inst@(tcTyConAppTyCon_maybe -> Just (isClassTyCon -> True))
+  = Just $ HasInstance inst
+mkEvidence _ = Nothing
+
+
+------------------------------------------------------------------------------
+-- | Compute all the 'Evidence' implicitly bound at the given 'SrcSpan'.
+getEvidenceAtHole :: SrcSpan -> LHsBinds GhcTc -> [Evidence]
+getEvidenceAtHole dst
+  = mapMaybe mkEvidence
+  . (everything (<>) $
+        mkQ mempty (absBinds dst) `extQ` wrapperBinds dst `extQ` matchBinds dst)
+
+
+------------------------------------------------------------------------------
+-- | Update our knowledge of which types are equal.
+evidenceToSubst :: Evidence -> TacticState -> TacticState
+evidenceToSubst (EqualityOfTypes a b) ts =
+  let tyvars = S.fromList $ mapMaybe getTyVar_maybe [a, b]
+      -- If we can unify our skolems, at least one is no longer a skolem.
+      -- Removing them from this set ensures we can get a subtitution between
+      -- the two. But it's okay to leave them in 'ts_skolems' in general, since
+      -- they won't exist after running this substitution.
+      skolems = ts_skolems ts S.\\ tyvars
+   in
+  case tryUnifyUnivarsButNotSkolems skolems (CType a) (CType b) of
+    Just subst -> updateSubst subst ts
+    Nothing -> ts
+evidenceToSubst HasInstance{} ts = ts
+
+
+------------------------------------------------------------------------------
+-- | Get all of the methods that are in scope from this piece of 'Evidence'.
+evidenceToHypothesis :: Evidence -> Hypothesis CType
+evidenceToHypothesis EqualityOfTypes{} = mempty
+evidenceToHypothesis (HasInstance t) =
+  Hypothesis . excludeForbiddenMethods . fromMaybe [] $ methodHypothesis t
+
+
+------------------------------------------------------------------------------
+-- | Given @a ~ b@ or @a ~# b@, returns @Just (a, b)@, otherwise @Nothing@.
+getEqualityTheta :: PredType -> Maybe (Type, Type)
+getEqualityTheta (splitTyConApp_maybe -> Just (tc, [_k, a, b]))
+  | tc == eqTyCon = Just (a, b)
+getEqualityTheta (splitTyConApp_maybe -> Just (tc, [_k1, _k2, a, b]))
+  | tc == eqPrimTyCon = Just (a, b)
+getEqualityTheta _ = Nothing
 
 
 ------------------------------------------------------------------------------
@@ -38,13 +96,48 @@ excludeForbiddenMethods = filter (not . flip S.member forbiddenMethods . hi_name
     forbiddenMethods = S.map mkVarOcc $ S.fromList
       [ -- monadfail
         "fail"
+        -- show
+      , "showsPrec"
+      , "showList"
       ]
 
 
 ------------------------------------------------------------------------------
--- | Extract the types of the evidence bindings in scope.
-evbinds ::  SrcSpan -> LHsBindLR GhcTc GhcTc -> [PredType]
-evbinds dst (L src (AbsBinds _ _ h _ _ _ _))
+-- | Extract evidence from 'AbsBinds' in scope.
+absBinds ::  SrcSpan -> LHsBindLR GhcTc GhcTc -> [PredType]
+absBinds dst (L src (AbsBinds _ _ h _ _ _ _))
   | dst `isSubspanOf` src = fmap idType h
-evbinds _ _ = []
+absBinds _ _ = []
+
+
+------------------------------------------------------------------------------
+-- | Extract evidence from 'HsWrapper's in scope
+wrapperBinds ::  SrcSpan -> LHsExpr GhcTc -> [PredType]
+wrapperBinds dst (L src (HsWrap _ h _))
+  | dst `isSubspanOf` src = wrapper h
+wrapperBinds _ _ = []
+
+
+------------------------------------------------------------------------------
+-- | Extract evidence from the 'ConPatOut's bound in this 'Match'.
+matchBinds :: SrcSpan -> LMatch GhcTc (LHsExpr GhcTc) -> [PredType]
+matchBinds dst (L src (Match _ _ pats _))
+  | dst `isSubspanOf` src = everything (<>) (mkQ mempty patBinds) pats
+matchBinds _ _ = []
+
+
+------------------------------------------------------------------------------
+-- | Extract evidence from a 'ConPatOut'.
+patBinds ::  Pat GhcTc -> [PredType]
+patBinds (ConPatOut { pat_dicts = dicts })
+  = fmap idType dicts
+patBinds _ = []
+
+
+------------------------------------------------------------------------------
+-- | Extract the types of the evidence bindings in scope.
+wrapper ::  HsWrapper -> [PredType]
+wrapper (WpCompose h h2) = wrapper h <> wrapper h2
+wrapper (WpEvLam v) = [idType v]
+wrapper _ = []
 

plugins/hls-tactics-plugin/src/Wingman/LanguageServer.hs

@@ -31,7 +31,7 @@ import           Development.IDE.Spans.LocalBindings (Bindings, getDefiningBindi
 import           Development.Shake (Action, RuleResult)
 import           Development.Shake.Classes (Typeable, Binary, Hashable, NFData)
 import qualified FastString
-import           GhcPlugins (mkAppTys, tupleDataCon, consDataCon)
+import           GhcPlugins (mkAppTys, tupleDataCon, consDataCon, substTyAddInScope)
 import           Ide.Plugin.Config (PluginConfig (plcConfig))
 import qualified Ide.Plugin.Config as Plugin
 import           Language.LSP.Server (MonadLsp, sendNotification)
@@ -44,7 +44,7 @@ import           Wingman.Context
 import           Wingman.FeatureSet
 import           Wingman.GHC
 import           Wingman.Judgements
-import           Wingman.Judgements.Theta (getMethodHypothesisAtHole)
+import           Wingman.Judgements.Theta
 import           Wingman.Range
 import           Wingman.Types
 
@@ -143,8 +143,10 @@ mkJudgementAndContext features g binds rss tcmod = do
           top_provs = getRhsPosVals rss tcs
           local_hy = spliceProvenance top_provs
                    $ hypothesisFromBindings rss binds
-          cls_hy = getMethodHypothesisAtHole (RealSrcSpan rss) tcs
-       in ( mkFirstJudgement
+          evidence = getEvidenceAtHole (RealSrcSpan rss) tcs
+          cls_hy = foldMap evidenceToHypothesis evidence
+          subst = ts_unifier $ appEndo (foldMap (Endo . evidenceToSubst) evidence) defaultTacticState
+       in ( fmap (CType . substTyAddInScope subst . unCType) $ mkFirstJudgement
               (local_hy <> cls_hy)
               (isRhsHole rss tcs)
               g

plugins/hls-tactics-plugin/src/Wingman/Machinery.hs

@@ -17,6 +17,7 @@ import           Data.Generics (everything, gcount, mkQ)
 import           Data.Generics.Product (field')
 import           Data.List (sortBy)
 import qualified Data.Map as M
+import           Data.Maybe (mapMaybe)
 import           Data.Monoid (getSum)
 import           Data.Ord (Down (..), comparing)
 import           Data.Set (Set)
@@ -79,7 +80,7 @@ runTactic ctx jdg t =
       (_, fmap assoc23 -> solns) -> do
         let sorted =
               flip sortBy solns $ comparing $ \(ext, (_, holes)) ->
-                Down $ scoreSolution ext holes
+                Down $ scoreSolution ext jdg holes
         case sorted of
           ((syn, _) : _) ->
             Right $
@@ -142,6 +143,7 @@ mappingExtract f (TacticT m)
 -- to produce the right test results.
 scoreSolution
     :: Synthesized (LHsExpr GhcPs)
+    -> Judgement
     -> [Judgement]
     -> ( Penalize Int  -- number of holes
        , Reward Bool   -- all bindings used
@@ -151,18 +153,23 @@ scoreSolution
        , Penalize Int  -- number of recursive calls
        , Penalize Int  -- size of extract
        )
-scoreSolution ext holes
+scoreSolution ext goal holes
   = ( Penalize $ length holes
     , Reward   $ S.null $ intro_vals S.\\ used_vals
     , Penalize $ S.size unused_top_vals
     , Penalize $ S.size intro_vals
-    , Reward   $ S.size used_vals
+    , Reward   $ S.size used_vals + length used_user_vals
     , Penalize $ getSum $ syn_recursion_count ext
     , Penalize $ solutionSize $ syn_val ext
     )
   where
+    initial_scope = hyByName $ jEntireHypothesis goal
     intro_vals = M.keysSet $ hyByName $ syn_scoped ext
     used_vals = S.intersection intro_vals $ syn_used_vals ext
+    used_user_vals = filter (isUserProv . hi_provenance)
+                   $ mapMaybe (flip M.lookup initial_scope)
+                   $ S.toList
+                   $ syn_used_vals ext
     top_vals = S.fromList
              . fmap hi_name
              . filter (isTopLevel . hi_provenance)
@@ -198,6 +205,10 @@ tryUnifyUnivarsButNotSkolems skolems goal inst =
     _               -> Nothing
 
 
+updateSubst :: TCvSubst -> TacticState -> TacticState
+updateSubst subst s = s { ts_unifier = unionTCvSubst subst (ts_unifier s) }
+
+
 
 ------------------------------------------------------------------------------
 -- | Attempt to unify two types.
@@ -208,7 +219,7 @@ 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) })
+      modify $ updateSubst subst
     Nothing -> throwError (UnificationError inst goal)
 
 

plugins/hls-tactics-plugin/src/Wingman/Plugin.hs

@@ -31,6 +31,7 @@ import           Wingman.CaseSplit
 import           Wingman.GHC
 import           Wingman.LanguageServer
 import           Wingman.LanguageServer.TacticProviders
+import           Wingman.Machinery (scoreSolution)
 import           Wingman.Range
 import           Wingman.Tactics
 import           Wingman.Types
@@ -136,7 +137,9 @@ mkWorkspaceEdits
     -> RunTacticResults
     -> Either UserFacingMessage WorkspaceEdit
 mkWorkspaceEdits span dflags ccs uri pm rtr = do
-  for_ (rtr_other_solns rtr) $ traceMX "other solution"
+  for_ (rtr_other_solns rtr) $ \soln -> do
+    traceMX "other solution" soln
+    traceMX "with score" $ scoreSolution soln (rtr_jdg rtr) []
   traceMX "solution" $ rtr_extract rtr
   let g = graftHole (RealSrcSpan span) rtr
       response = transform dflags ccs uri g pm

plugins/hls-tactics-plugin/src/Wingman/Types.hs

@@ -144,6 +144,9 @@ instance Show (Pat GhcPs) where
 instance Show (LHsSigType GhcPs) where
   show  = unsafeRender
 
+instance Show TyCon where
+  show  = unsafeRender
+
 
 ------------------------------------------------------------------------------
 -- | The state that should be shared between subgoals. Extracts move towards

plugins/hls-tactics-plugin/test/CodeAction/AutoSpec.hs

@@ -47,14 +47,24 @@ spec = do
     autoTest  2  9 "Fgmap.hs"
     autoTest  4 19 "FmapJoinInLet.hs"
     autoTest  9 12 "AutoEndo.hs"
-    autoTest 12 10 "AutoThetaFix.hs"
 
     failing "flaky in CI" $
       autoTest 2 11 "GoldenApplicativeThen.hs"
 
     failing "not enough auto gas" $
       autoTest 5 18 "GoldenFish.hs"
 
+  describe "theta" $ do
+    autoTest 12 10 "AutoThetaFix.hs"
+    autoTest  7 20 "AutoThetaRankN.hs"
+    autoTest  6 10 "AutoThetaGADT.hs"
+    autoTest  6  8 "AutoThetaGADTDestruct.hs"
+    autoTest  4  8 "AutoThetaEqCtx.hs"
+    autoTest  6 10 "AutoThetaEqGADT.hs"
+    autoTest  6  8 "AutoThetaEqGADTDestruct.hs"
+    autoTest  6 10 "AutoThetaRefl.hs"
+    autoTest  6  8 "AutoThetaReflDestruct.hs"
+
 
   describe "messages" $ do
     mkShowMessageTest allFeatures Auto "" 2 8 "MessageForallA.hs" TacticErrors

plugins/hls-tactics-plugin/test/golden/AutoThetaEqCtx.hs

@@ -0,0 +1,5 @@
+{-# LANGUAGE GADTs #-}
+
+fun2 :: (a ~ b) => a -> b
+fun2 = _ -- id
+

plugins/hls-tactics-plugin/test/golden/AutoThetaEqCtx.hs.expected

@@ -0,0 +1,5 @@
+{-# LANGUAGE GADTs #-}
+
+fun2 :: (a ~ b) => a -> b
+fun2 = id -- id
+

plugins/hls-tactics-plugin/test/golden/AutoThetaEqGADT.hs

@@ -0,0 +1,7 @@
+{-# LANGUAGE GADTs #-}
+
+data Y a b = a ~ b => Y
+
+fun3 :: Y a b -> a -> b
+fun3 Y = _
+

plugins/hls-tactics-plugin/test/golden/AutoThetaEqGADT.hs.expected

@@ -0,0 +1,7 @@
+{-# LANGUAGE GADTs #-}
+
+data Y a b = a ~ b => Y
+
+fun3 :: Y a b -> a -> b
+fun3 Y = id
+