Make getCurrentDefinitions return polymorphic types (#1945)

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

Author: isovector

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

@@ -5,10 +5,9 @@ import           Control.Monad.Reader.Class (asks)
 import           Control.Monad.State (gets)
 import qualified Data.Set as S
 import           Refinery.Tactic
-import           Wingman.Context
 import           Wingman.Judgements
 import           Wingman.KnownStrategies
-import           Wingman.Machinery (tracing)
+import           Wingman.Machinery (tracing, getCurrentDefinitions)
 import           Wingman.Tactics
 import           Wingman.Types
 

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

@@ -73,11 +73,6 @@ getFunBindId (AbsBinds _ _ _ abes _ _ _)
 getFunBindId _ = []
 
 
-getCurrentDefinitions :: MonadReader Context m => m [(OccName, CType)]
-getCurrentDefinitions = asks ctxDefiningFuncs
-
-
-
 ------------------------------------------------------------------------------
 -- | Determine if there is an instance that exists for the given 'Class' at the
 -- specified types. Deeply checks contexts to ensure the instance is actually

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

@@ -4,10 +4,9 @@ import Control.Monad.Error.Class
 import Data.Foldable (for_)
 import OccName (mkVarOcc, mkClsOcc)
 import Refinery.Tactic
-import Wingman.Context (getCurrentDefinitions)
 import Wingman.Judgements (jGoal)
 import Wingman.KnownStrategies.QuickCheck (deriveArbitrary)
-import Wingman.Machinery (tracing, getKnownInstance)
+import Wingman.Machinery (tracing, getKnownInstance, getCurrentDefinitions)
 import Wingman.Tactics
 import Wingman.Types
 

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

@@ -1,4 +1,5 @@
-{-# LANGUAGE RecordWildCards       #-}
+{-# LANGUAGE RecordWildCards #-}
+{-# LANGUAGE TupleSections   #-}
 
 module Wingman.Machinery where
 
@@ -21,6 +22,7 @@ import           Data.Maybe (mapMaybe)
 import           Data.Monoid (getSum)
 import           Data.Ord (Down (..), comparing)
 import qualified Data.Set as S
+import           Data.Traversable (for)
 import           Development.IDE.Core.Compile (lookupName)
 import           Development.IDE.GHC.Compat
 import           GhcPlugins (GlobalRdrElt (gre_name), lookupOccEnv, varType)
@@ -403,3 +405,10 @@ getOccNameType occ = do
     Just (AnId v) -> pure $ varType v
     _ -> throwError $ NotInScope occ
 
+
+getCurrentDefinitions :: TacticsM [(OccName, CType)]
+getCurrentDefinitions = do
+  ctx_funcs <- asks ctxDefiningFuncs
+  for ctx_funcs $ \res@(occ, _) ->
+    pure . maybe res (occ,) =<< lookupNameInContext occ
+

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

@@ -15,8 +15,7 @@ import qualified Data.Text as T
 import qualified Refinery.Tactic as R
 import qualified Text.Megaparsec as P
 import           Wingman.Auto
-import           Wingman.Context (getCurrentDefinitions)
-import           Wingman.Machinery (useNameFromHypothesis, getOccNameType, createImportedHyInfo, useNameFromContext, lookupNameInContext)
+import           Wingman.Machinery (useNameFromHypothesis, getOccNameType, createImportedHyInfo, useNameFromContext, lookupNameInContext, getCurrentDefinitions)
 import           Wingman.Metaprogramming.Lexer
 import           Wingman.Metaprogramming.Parser.Documentation
 import           Wingman.Metaprogramming.ProofState (proofState, layout)

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

@@ -37,7 +37,6 @@ import           TcType
 import           Type hiding (Var)
 import           TysPrim (betaTy, alphaTy, betaTyVar, alphaTyVar)
 import           Wingman.CodeGen
-import           Wingman.Context
 import           Wingman.GHC
 import           Wingman.Judgements
 import           Wingman.Machinery
@@ -85,10 +84,6 @@ use sat occ = do
 recursion :: TacticsM ()
 -- TODO(sandy): This tactic doesn't fire for the @AutoThetaFix@ golden test,
 -- presumably due to running afoul of 'requireConcreteHole'. Look into this!
-
--- TODO(sandy): There's a bug here! This should use the polymorphic defining
--- types, not the ones available via 'getCurrentDefinitions'. As it is, this
--- tactic doesn't support polymorphic recursion.
 recursion = requireConcreteHole $ tracing "recursion" $ do
   defs <- getCurrentDefinitions
   attemptOn (const defs) $ \(name, ty) -> markRecursion $ do

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

@@ -47,6 +47,7 @@ spec = do
     autoTest  2 17 "AutoInfixApply"
     autoTest  2 19 "AutoInfixApplyMany"
     autoTest  2 25 "AutoInfixInfix"
+    autoTest 19 12 "AutoTypeLevel"
 
     failing "flaky in CI" $
       autoTest 2 11 "GoldenApplicativeThen"

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

@@ -0,0 +1,21 @@
+{-# LANGUAGE DataKinds      #-}
+{-# LANGUAGE GADTs          #-}
+{-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE TypeOperators  #-}
+
+import Data.Kind
+
+data Nat = Z | S Nat
+
+data HList (ls :: [Type]) where
+  HNil :: HList '[]
+  HCons :: t -> HList ts -> HList (t ': ts)
+
+data ElemAt (n :: Nat) t (ts :: [Type]) where
+  AtZ :: ElemAt 'Z t (t ': ts)
+  AtS :: ElemAt k t ts -> ElemAt ('S k) t (u ': ts)
+
+lookMeUp :: ElemAt i ty tys -> HList tys -> ty
+lookMeUp AtZ (HCons t _) = t
+lookMeUp (AtS ea') (HCons _ hl') = lookMeUp ea' hl'
+

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

@@ -0,0 +1,20 @@
+{-# LANGUAGE DataKinds      #-}
+{-# LANGUAGE GADTs          #-}
+{-# LANGUAGE KindSignatures #-}
+{-# LANGUAGE TypeOperators  #-}
+
+import Data.Kind
+
+data Nat = Z | S Nat
+
+data HList (ls :: [Type]) where
+  HNil :: HList '[]
+  HCons :: t -> HList ts -> HList (t ': ts)
+
+data ElemAt (n :: Nat) t (ts :: [Type]) where
+  AtZ :: ElemAt 'Z t (t ': ts)
+  AtS :: ElemAt k t ts -> ElemAt ('S k) t (u ': ts)
+
+lookMeUp :: ElemAt i ty tys -> HList tys -> ty
+lookMeUp = _
+