Allow module-local and imported functions in Wingman metaprograms (#1856)

* Perform lookups of terms in scope and context

* Fix imports

* Add tests

* Cleanup and haddock

* mkFunTys'

* Wrangle some imports for GHC 8.6

Co-authored-by: mergify[bot] <37929162+mergify[bot]@users.noreply.github.com>

Author: isovector

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

@@ -4,7 +4,7 @@ import           Bag
 import           Control.Arrow
 import           Control.Monad.Reader
 import           Data.Foldable.Extra (allM)
-import           Data.Maybe (fromMaybe, isJust)
+import           Data.Maybe (fromMaybe, isJust, mapMaybe)
 import qualified Data.Set as S
 import           Development.IDE.GHC.Compat
 import           GhcPlugins (ExternalPackageState (eps_inst_env), piResultTys)
@@ -27,11 +27,13 @@ mkContext
     -> Context
 mkContext cfg locals tcg eps kt ev = Context
   { ctxDefiningFuncs = locals
-  , ctxModuleFuncs = fmap splitId
-                   . (getFunBindId =<<)
-                   . fmap unLoc
-                   . bagToList
-                   $ tcg_binds tcg
+  , ctxModuleFuncs
+      = fmap splitId
+      . mappend (locallyDefinedMethods tcg)
+      . (getFunBindId =<<)
+      . fmap unLoc
+      . bagToList
+      $ tcg_binds tcg
   , ctxConfig = cfg
   , ctxInstEnvs =
       InstEnvs
@@ -43,6 +45,14 @@ mkContext cfg locals tcg eps kt ev = Context
   }
 
 
+locallyDefinedMethods :: TcGblEnv -> [Id]
+locallyDefinedMethods
+  = foldMap classMethods
+  . mapMaybe tyConClass_maybe
+  . tcg_tcs
+
+
+
 splitId :: Id -> (OccName, CType)
 splitId = occName &&& CType . idType
 

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

@@ -73,7 +73,7 @@ isFunction _                                    = True
 -- context.
 tacticsSplitFunTy :: Type -> ([TyVar], ThetaType, [Type], Type)
 tacticsSplitFunTy t
-  = let (vars, theta, t') = tcSplitSigmaTy t
+  = let (vars, theta, t') = tcSplitNestedSigmaTys t
         (args, res) = tcSplitFunTys t'
      in (vars, theta, args, res)
 

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

@@ -229,6 +229,7 @@ provAncestryOf (PatternMatchPrv (PatVal mo so _ _)) =
 provAncestryOf (ClassMethodPrv _) = mempty
 provAncestryOf UserPrv = mempty
 provAncestryOf RecursivePrv = mempty
+provAncestryOf ImportPrv = mempty
 provAncestryOf (DisallowedPrv _ p2) = provAncestryOf p2
 
 

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

@@ -1,6 +1,7 @@
 {-# LANGUAGE CPP               #-}
 {-# LANGUAGE OverloadedStrings #-}
 {-# LANGUAGE RankNTypes        #-}
+{-# LANGUAGE TupleSections     #-}
 {-# LANGUAGE TypeFamilies      #-}
 
 {-# LANGUAGE NoMonoLocalBinds  #-}
@@ -43,6 +44,7 @@ import           Development.IDE.Spans.LocalBindings (Bindings, getDefiningBindi
 import qualified FastString
 import           GHC.Generics (Generic)
 import           Generics.SYB hiding (Generic)
+import           GhcPlugins (extractModule)
 import           GhcPlugins (tupleDataCon, consDataCon, substTyAddInScope, ExternalPackageState, HscEnv (hsc_EPS), unpackFS)
 import qualified Ide.Plugin.Config as Plugin
 import           Ide.Plugin.Properties
@@ -57,13 +59,14 @@ import           OccName
 import           Prelude hiding (span)
 import           Retrie (transformA)
 import           SrcLoc (containsSpan)
-import           TcRnTypes (tcg_binds, TcGblEnv)
+import           TcRnTypes (tcg_binds, TcGblEnv (tcg_rdr_env))
 import           Wingman.Context
 import           Wingman.FeatureSet
 import           Wingman.GHC
 import           Wingman.Judgements
 import           Wingman.Judgements.SYB (everythingContaining, metaprogramQ)
 import           Wingman.Judgements.Theta
+import           Wingman.Metaprogramming.Lexer (ParserContext(..))
 import           Wingman.Range
 import           Wingman.StaticPlugin (pattern WingmanMetaprogram, pattern MetaprogramSyntax)
 import           Wingman.Types
@@ -576,6 +579,10 @@ mkWorkspaceEdits dflags ccs uri pm g = do
    in first (InfrastructureError . T.pack) response
 
 
+------------------------------------------------------------------------------
+-- | Add ExactPrint annotations to every metaprogram in the source tree.
+-- Usually the ExactPrint module can do this for us, but we've enabled
+-- QuasiQuotes, so the round-trip print/parse journey will crash.
 annotateMetaprograms :: Data a => a -> Transform a
 annotateMetaprograms = everywhereM $ mkM $ \case
   L ss (WingmanMetaprogram mp) -> do
@@ -585,6 +592,9 @@ annotateMetaprograms = everywhereM $ mkM $ \case
     pure x
   (x :: LHsExpr GhcPs) -> pure x
 
+
+------------------------------------------------------------------------------
+-- | Find the source of a tactic metaprogram at the given span.
 getMetaprogramAtSpan
     :: Tracked age SrcSpan
     -> Tracked age TcGblEnv
@@ -596,3 +606,25 @@ getMetaprogramAtSpan (unTrack -> ss)
   . tcg_binds
   . unTrack
 
+
+------------------------------------------------------------------------------
+-- | The metaprogram parser needs the ability to lookup terms from the module
+-- and imports. The 'ParserContext' contains everything we need to find that
+-- stuff.
+getParserState
+    :: IdeState
+    -> NormalizedFilePath
+    -> Context
+    -> MaybeT IO ParserContext
+getParserState state nfp ctx = do
+  let stale a = runStaleIde "getParserState" state nfp a
+
+  TrackedStale (unTrack -> tcmod) _  <- stale TypeCheck
+  TrackedStale (unTrack -> hscenv) _ <- stale GhcSessionDeps
+
+  let tcgblenv = tmrTypechecked tcmod
+      modul = extractModule tcgblenv
+      rdrenv = tcg_rdr_env tcgblenv
+
+  pure $ ParserContext (hscEnv hscenv) rdrenv modul ctx
+

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

@@ -10,6 +10,7 @@ module Wingman.LanguageServer.Metaprogram
 
 import           Control.Applicative (empty)
 import           Control.Monad
+import           Control.Monad.Reader (runReaderT)
 import           Control.Monad.Trans
 import           Control.Monad.Trans.Maybe
 import           Data.List (find)
@@ -30,8 +31,8 @@ import           Prelude hiding (span)
 import           TcRnTypes (tcg_binds)
 import           Wingman.GHC
 import           Wingman.Judgements.SYB (metaprogramQ)
-import           Wingman.Metaprogramming.Parser (attempt_it)
 import           Wingman.LanguageServer
+import           Wingman.Metaprogramming.Parser (attempt_it)
 import           Wingman.Types
 
 
@@ -53,12 +54,12 @@ hoverProvider state plId (HoverParams (TextDocumentIdentifier uri) (unsafeMkCurr
             Just (trss, program) -> do
               let tr_range = fmap realSrcSpanToRange trss
               HoleJudgment{hj_jdg=jdg, hj_ctx=ctx} <- judgementForHole state nfp tr_range cfg
+              ps <- getParserState state nfp ctx
+              z <- liftIO $ flip runReaderT ps $ attempt_it ctx jdg $ T.unpack program
               pure $ Hover
                 { _contents = HoverContents
                             $ MarkupContent MkMarkdown
-                            $ either T.pack T.pack
-                            $ attempt_it ctx jdg
-                            $ T.unpack program
+                            $ either T.pack T.pack z
                 , _range = Just $ unTrack tr_range
                 }
             Nothing -> empty

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

@@ -34,26 +34,28 @@ import           Wingman.FeatureSet
 import           Wingman.GHC
 import           Wingman.Judgements
 import           Wingman.Machinery (useNameFromHypothesis)
+import           Wingman.Metaprogramming.Lexer (ParserContext)
 import           Wingman.Metaprogramming.Parser (parseMetaprogram)
 import           Wingman.Tactics
 import           Wingman.Types
+import Control.Monad.Reader (runReaderT)
 
 
 ------------------------------------------------------------------------------
 -- | A mapping from tactic commands to actual tactics for refinery.
-commandTactic :: TacticCommand -> T.Text -> TacticsM ()
-commandTactic Auto                   = const auto
-commandTactic Intros                 = const intros
-commandTactic Destruct               = useNameFromHypothesis destruct . mkVarOcc . T.unpack
-commandTactic DestructPun            = useNameFromHypothesis destructPun . mkVarOcc . T.unpack
-commandTactic Homomorphism           = useNameFromHypothesis homo . mkVarOcc . T.unpack
-commandTactic DestructLambdaCase     = const destructLambdaCase
-commandTactic HomomorphismLambdaCase = const homoLambdaCase
-commandTactic DestructAll            = const destructAll
-commandTactic UseDataCon             = userSplit . mkVarOcc . T.unpack
-commandTactic Refine                 = const refine
-commandTactic BeginMetaprogram       = const metaprogram
-commandTactic RunMetaprogram         = parseMetaprogram
+commandTactic :: ParserContext -> TacticCommand -> T.Text -> IO (TacticsM ())
+commandTactic _ Auto                   = pure . const auto
+commandTactic _ Intros                 = pure . const intros
+commandTactic _ Destruct               = pure . useNameFromHypothesis destruct . mkVarOcc . T.unpack
+commandTactic _ DestructPun            = pure . useNameFromHypothesis destructPun . mkVarOcc . T.unpack
+commandTactic _ Homomorphism           = pure . useNameFromHypothesis homo . mkVarOcc . T.unpack
+commandTactic _ DestructLambdaCase     = pure . const destructLambdaCase
+commandTactic _ HomomorphismLambdaCase = pure . const homoLambdaCase
+commandTactic _ DestructAll            = pure . const destructAll
+commandTactic _ UseDataCon             = pure . userSplit . mkVarOcc . T.unpack
+commandTactic _ Refine                 = pure . const refine
+commandTactic _ BeginMetaprogram       = pure . const metaprogram
+commandTactic c RunMetaprogram         = flip runReaderT c . parseMetaprogram
 
 
 ------------------------------------------------------------------------------

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

@@ -3,6 +3,7 @@
 module Wingman.Machinery where
 
 import           Class (Class (classTyVars))
+import           Control.Applicative (empty)
 import           Control.Lens ((<>~))
 import           Control.Monad.Error.Class
 import           Control.Monad.Reader
@@ -22,13 +23,15 @@ import           Data.Monoid (getSum)
 import           Data.Ord (Down (..), comparing)
 import           Data.Set (Set)
 import qualified Data.Set as S
+import           Development.IDE.Core.Compile (lookupName)
 import           Development.IDE.GHC.Compat
-import           OccName (HasOccName (occName))
+import           GhcPlugins (GlobalRdrElt (gre_name), lookupOccEnv, varType)
+import           OccName (HasOccName (occName), OccEnv)
 import           Refinery.ProofState
 import           Refinery.Tactic
 import           Refinery.Tactic.Internal
 import           TcType
-import           Type
+import           Type (tyCoVarsOfTypeWellScoped, splitTyConApp_maybe)
 import           Unify
 import           Wingman.Judgements
 import           Wingman.Simplify (simplify)
@@ -315,3 +318,52 @@ useNameFromHypothesis f name = do
     Just hi -> f hi
     Nothing -> throwError $ NotInScope name
 
+------------------------------------------------------------------------------
+-- | Lift a function over 'HyInfo's to one that takes an 'OccName' and tries to
+-- look it up in the hypothesis.
+useNameFromContext :: (HyInfo CType -> TacticsM a) -> OccName -> TacticsM a
+useNameFromContext f name = do
+  lookupNameInContext name >>= \case
+    Just ty -> f $ createImportedHyInfo name ty
+    Nothing -> throwError $ NotInScope name
+
+
+------------------------------------------------------------------------------
+-- | Find the type of an 'OccName' that is defined in the current module.
+lookupNameInContext :: MonadReader Context m => OccName -> m (Maybe CType)
+lookupNameInContext name = do
+  ctx <- asks ctxModuleFuncs
+  pure $ case find ((== name) . fst) ctx of
+    Just (_, ty) -> pure ty
+    Nothing      -> empty
+
+
+------------------------------------------------------------------------------
+-- | Build a 'HyInfo' for an imported term.
+createImportedHyInfo :: OccName -> CType -> HyInfo CType
+createImportedHyInfo on ty = HyInfo
+  { hi_name = on
+  , hi_provenance = ImportPrv
+  , hi_type = ty
+  }
+
+
+------------------------------------------------------------------------------
+-- | Lookup the type of any 'OccName' that was imported. Necessarily done in
+-- IO, so we only expose this functionality to the parser. Internal Haskell
+-- code that wants to lookup terms should do it via 'KnownThings'.
+getOccNameType
+    :: HscEnv
+    -> OccEnv [GlobalRdrElt]
+    -> Module
+    -> OccName
+    -> IO (Maybe Type)
+getOccNameType hscenv rdrenv modul occ =
+  case lookupOccEnv rdrenv occ of
+    Just (elt : _) -> do
+      mvar <- lookupName hscenv modul $ gre_name elt
+      pure $ case mvar of
+        Just (AnId v) -> pure $ varType v
+        _ -> Nothing
+    _ -> pure Nothing
+

plugins/hls-tactics-plugin/src/Wingman/Metaprogramming/Lexer.hs

@@ -7,16 +7,31 @@ module Wingman.Metaprogramming.Lexer where
 
 import           Control.Applicative
 import           Control.Monad
+import           Control.Monad.Reader (ReaderT)
 import           Data.Text (Text)
 import qualified Data.Text as T
 import           Data.Void
+import           Development.IDE.GHC.Compat (HscEnv, Module)
+import           GhcPlugins (GlobalRdrElt)
 import           Name
 import qualified Text.Megaparsec as P
 import qualified Text.Megaparsec.Char as P
 import qualified Text.Megaparsec.Char.Lexer as L
+import Wingman.Types (Context)
+
+
+------------------------------------------------------------------------------
+-- | Everything we need in order to call 'Wingman.Machinery.getOccNameType'.
+data ParserContext = ParserContext
+  { ps_hscEnv :: HscEnv
+  , ps_occEnv :: OccEnv [GlobalRdrElt]
+  , ps_module :: Module
+  , ps_context :: Context
+  }
+
+type Parser = P.ParsecT Void Text (ReaderT ParserContext IO)
 
 
-type Parser = P.Parsec Void Text
 
 lineComment :: Parser ()
 lineComment = L.skipLineComment "--"