Let Wingman peek through type families (#1881)

* Let Wingman peek through type families

* Cleanup imports

* We can just use normalizejudgment in mkFirstJudgment

* Add support for data families too

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

Author: isovector

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

@@ -13,6 +13,7 @@ module Wingman.CodeGen
 import           ConLike
 import           Control.Lens ((%~), (<>~), (&))
 import           Control.Monad.Except
+import           Control.Monad.Reader (ask)
 import           Control.Monad.State
 import           Data.Bool (bool)
 import           Data.Functor ((<&>))
@@ -24,6 +25,7 @@ import           Data.Traversable
 import           DataCon
 import           Development.IDE.GHC.Compat
 import           GHC.Exts
+import           GHC.SourceGen (occNameToStr)
 import           GHC.SourceGen.Binds
 import           GHC.SourceGen.Expr
 import           GHC.SourceGen.Overloaded
@@ -39,7 +41,6 @@ import           Wingman.Judgements.Theta
 import           Wingman.Machinery
 import           Wingman.Naming
 import           Wingman.Types
-import GHC.SourceGen (occNameToStr)
 
 
 destructMatches
@@ -69,6 +70,7 @@ destructMatches use_field_puns f scrut t jdg = do
             args = conLikeInstOrigArgTys' con apps
         modify $ appEndo $ foldMap (Endo . evidenceToSubst) ev
         subst <- gets ts_unifier
+        ctx <- ask
 
         let names_in_scope = hyNamesInScope hy
             names = mkManyGoodNames (hyNamesInScope hy) args
@@ -79,8 +81,8 @@ destructMatches use_field_puns f scrut t jdg = do
                 $ zip names'
                 $ coerce args
             j = fmap (CType . substTyAddInScope subst . unCType)
-              $ introduce hy'
-              $ introduce method_hy
+              $ introduce ctx hy'
+              $ introduce ctx method_hy
               $ withNewGoal g jdg
         ext <- f con j
         pure $ ext
@@ -289,13 +291,14 @@ letForEach rename solve (unHypothesis -> hy) jdg = do
   case hy of
     [] -> newSubgoal jdg
     _ -> do
+      ctx <- ask
       let g = jGoal jdg
       terms <- fmap sequenceA $ for hy $ \hi -> do
         let name = rename $ hi_name hi
         res <- tacticToRule jdg $ solve hi
         pure $ fmap ((name,) . unLoc) res
       let hy' = fmap (g <$) $ syn_val terms
           matches = fmap (fmap (\(occ, expr) -> valBind (occNameToStr occ) expr)) terms
-      g <- fmap (fmap unLoc) $ newSubgoal $ introduce (userHypothesis hy') jdg
+      g <- fmap (fmap unLoc) $ newSubgoal $ introduce ctx (userHypothesis hy') jdg
       pure $ fmap noLoc $ let' <$> matches <*> g
 

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

@@ -3,16 +3,18 @@ module Wingman.Context where
 import           Bag
 import           Control.Arrow
 import           Control.Monad.Reader
+import           Data.Coerce (coerce)
 import           Data.Foldable.Extra (allM)
 import           Data.Maybe (fromMaybe, isJust, mapMaybe)
 import qualified Data.Set as S
 import           Development.IDE.GHC.Compat
-import           GhcPlugins (ExternalPackageState (eps_inst_env), piResultTys)
+import           GhcPlugins (ExternalPackageState (eps_inst_env), piResultTys, eps_fam_inst_env)
 import           InstEnv (lookupInstEnv, InstEnvs(..), is_dfun)
 import           OccName
 import           TcRnTypes
 import           TcType (tcSplitTyConApp, tcSplitPhiTy)
 import           TysPrim (alphaTys)
+import           Wingman.GHC (normalizeType)
 import           Wingman.Judgements.Theta
 import           Wingman.Types
 
@@ -25,24 +27,28 @@ mkContext
     -> KnownThings
     -> [Evidence]
     -> Context
-mkContext cfg locals tcg eps kt ev = Context
-  { ctxDefiningFuncs = locals
-  , ctxModuleFuncs
-      = fmap splitId
-      . mappend (locallyDefinedMethods tcg)
-      . (getFunBindId =<<)
-      . fmap unLoc
-      . bagToList
-      $ tcg_binds tcg
-  , ctxConfig = cfg
-  , ctxInstEnvs =
-      InstEnvs
-        (eps_inst_env eps)
-        (tcg_inst_env tcg)
-        (tcVisibleOrphanMods tcg)
-  , ctxKnownThings = kt
-  , ctxTheta = evidenceToThetaType ev
-  }
+mkContext cfg locals tcg eps kt ev = fix $ \ctx ->
+  Context
+    { ctxDefiningFuncs
+        = fmap (second $ coerce $ normalizeType ctx) locals
+    , ctxModuleFuncs
+        = fmap (second (coerce $ normalizeType ctx) . splitId)
+        . mappend (locallyDefinedMethods tcg)
+        . (getFunBindId =<<)
+        . fmap unLoc
+        . bagToList
+        $ tcg_binds tcg
+    , ctxConfig = cfg
+    , ctxFamInstEnvs =
+        (eps_fam_inst_env eps, tcg_fam_inst_env tcg)
+    , ctxInstEnvs =
+        InstEnvs
+          (eps_inst_env eps)
+          (tcg_inst_env tcg)
+          (tcVisibleOrphanMods tcg)
+    , ctxKnownThings = kt
+    , ctxTheta = evidenceToThetaType ev
+    }
 
 
 locallyDefinedMethods :: TcGblEnv -> [Id]

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

@@ -23,9 +23,11 @@ import           Development.IDE.Core.Compile (lookupName)
 import           Development.IDE.GHC.Compat hiding (exprType)
 import           DsExpr (dsExpr)
 import           DsMonad (initDs)
+import           FamInst (tcLookupDataFamInst_maybe)
+import           FamInstEnv (normaliseType)
 import           GHC.SourceGen (lambda)
 import           Generics.SYB (Data, everything, everywhere, listify, mkQ, mkT)
-import           GhcPlugins (extractModule, GlobalRdrElt (gre_name))
+import           GhcPlugins (extractModule, GlobalRdrElt (gre_name), Role (Nominal))
 import           OccName
 import           TcRnMonad
 import           TcType
@@ -374,3 +376,23 @@ mkFunTys' =
   mkVisFunTys
 #endif
 
+
+------------------------------------------------------------------------------
+-- | Expand type and data families
+normalizeType :: Context -> Type -> Type
+normalizeType ctx ty =
+  let ty' = expandTyFam ctx ty
+   in case tcSplitTyConApp_maybe ty' of
+        Just (tc, tys) ->
+          -- try to expand any data families
+          case tcLookupDataFamInst_maybe (ctxFamInstEnvs ctx) tc tys of
+            Just (dtc, dtys, _) -> mkAppTys (mkTyConTy dtc) dtys
+            Nothing -> ty'
+        Nothing -> ty'
+
+------------------------------------------------------------------------------
+-- | Expand type families
+expandTyFam :: Context -> Type -> Type
+expandTyFam ctx = snd . normaliseType  (ctxFamInstEnvs ctx) Nominal
+
+

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

@@ -17,7 +17,7 @@ import           Development.IDE.Spans.LocalBindings
 import           OccName
 import           SrcLoc
 import           Type
-import           Wingman.GHC (algebraicTyCon)
+import           Wingman.GHC (algebraicTyCon, normalizeType)
 import           Wingman.Types
 
 
@@ -69,11 +69,19 @@ withNewGoal :: a -> Judgement' a -> Judgement' a
 withNewGoal t = field @"_jGoal" .~ t
 
 
-introduce :: Hypothesis a -> Judgement' a -> Judgement' a
+normalizeHypothesis :: Functor f => Context -> f CType -> f CType
+normalizeHypothesis = fmap . coerce . normalizeType
+
+normalizeJudgement :: Functor f => Context -> f CType -> f CType
+normalizeJudgement = normalizeHypothesis
+
+
+introduce :: Context -> Hypothesis CType -> Judgement' CType -> Judgement' CType
 -- NOTE(sandy): It's important that we put the new hypothesis terms first,
 -- since 'jAcceptableDestructTargets' will never destruct a pattern that occurs
 -- after a previously-destructed term.
-introduce hy = field @"_jHypothesis" %~ mappend hy
+introduce ctx hy =
+  field @"_jHypothesis" %~ mappend (normalizeHypothesis ctx hy)
 
 
 ------------------------------------------------------------------------------
@@ -393,17 +401,20 @@ substJdg subst = fmap $ coerce . substTy subst . coerce
 
 
 mkFirstJudgement
-    :: Hypothesis CType
+    :: Context
+    -> Hypothesis CType
     -> Bool  -- ^ are we in the top level rhs hole?
     -> Type
     -> Judgement' CType
-mkFirstJudgement hy top goal = Judgement
-  { _jHypothesis        = hy
-  , _jBlacklistDestruct = False
-  , _jWhitelistSplit    = True
-  , _jIsTopHole         = top
-  , _jGoal              = CType goal
-  }
+mkFirstJudgement ctx hy top goal =
+  normalizeJudgement ctx $
+    Judgement
+      { _jHypothesis        = hy
+      , _jBlacklistDestruct = False
+      , _jWhitelistSplit    = True
+      , _jIsTopHole         = top
+      , _jGoal              = CType goal
+      }
 
 
 ------------------------------------------------------------------------------

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

@@ -275,7 +275,9 @@ mkJudgementAndContext cfg g (TrackedStale binds bmap) rss (TrackedStale tcg tcgm
       subst = ts_unifier $ appEndo (foldMap (Endo . evidenceToSubst) evidence) defaultTacticState
   pure $
     ( disallowing AlreadyDestructed already_destructed
-    $ fmap (CType . substTyAddInScope subst . unCType) $ mkFirstJudgement
+    $ fmap (CType . substTyAddInScope subst . unCType) $
+        mkFirstJudgement
+          ctx
           (local_hy <> cls_hy)
           (isRhsHole tcg_rss tcs)
           g

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

@@ -49,10 +49,13 @@ newSubgoal
     :: Judgement
     -> Rule
 newSubgoal j = do
-    unifier <- gets ts_unifier
-    subgoal
-      $ substJdg unifier
-      $ unsetIsTopHole j
+  ctx <- ask
+  unifier <- gets ts_unifier
+  subgoal
+    $ normalizeJudgement ctx
+    $ substJdg unifier
+    $ unsetIsTopHole
+    $ normalizeJudgement ctx j
 
 
 tacticToRule :: Judgement -> TacticsM () -> Rule

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

@@ -83,7 +83,8 @@ recursion = requireConcreteHole $ tracing "recursion" $ do
       unless (any (flip M.member pat_vals) $ syn_used_vals ext) empty
 
     let hy' = recursiveHypothesis defs
-    localTactic (apply $ HyInfo name RecursivePrv ty) (introduce hy')
+    ctx <- ask
+    localTactic (apply $ HyInfo name RecursivePrv ty) (introduce ctx hy')
       <@> fmap (localTactic assumption . filterPosition name) [0..]
 
 
@@ -110,15 +111,15 @@ intros'
     -> TacticsM ()
 intros' names = rule $ \jdg -> do
   let g  = jGoal jdg
-  ctx <- ask
   case tacticsSplitFunTy $ unCType g of
     (_, _, [], _) -> throwError $ GoalMismatch "intros" g
     (_, _, as, b) -> do
+      ctx <- ask
       let vs = fromMaybe (mkManyGoodNames (hyNamesInScope $ jEntireHypothesis jdg) as) names
           num_args = length vs
           top_hole = isTopHole ctx jdg
           hy' = lambdaHypothesis top_hole $ zip vs $ coerce as
-          jdg' = introduce hy'
+          jdg' = introduce ctx hy'
                $ withNewGoal (CType $ mkFunTys' (drop num_args as) b) jdg
       ext <- newSubgoal jdg'
       pure $

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

@@ -26,8 +26,11 @@ import           Data.Set (Set)
 import           Data.Text (Text)
 import qualified Data.Text as T
 import           Data.Tree
+import           Development.IDE (Range)
+import           Development.IDE.Core.UseStale
 import           Development.IDE.GHC.Compat hiding (Node)
 import           Development.IDE.GHC.Orphans ()
+import           FamInstEnv (FamInstEnvs)
 import           GHC.Generics
 import           GHC.SourceGen (var)
 import           InstEnv (InstEnvs(..))
@@ -39,8 +42,6 @@ import           UniqSupply (takeUniqFromSupply, mkSplitUniqSupply, UniqSupply)
 import           Unique (nonDetCmpUnique, Uniquable, getUnique, Unique)
 import           Wingman.Debug
 import           Wingman.FeatureSet
-import Development.IDE.Core.UseStale
-import Development.IDE (Range)
 
 
 ------------------------------------------------------------------------------
@@ -422,6 +423,7 @@ data Context = Context
   , ctxConfig        :: Config
   , ctxKnownThings   :: KnownThings
   , ctxInstEnvs      :: InstEnvs
+  , ctxFamInstEnvs   :: FamInstEnvs
   , ctxTheta         :: Set CType
   }
 
@@ -452,6 +454,7 @@ emptyContext
       , ctxModuleFuncs = mempty
       , ctxConfig = emptyConfig
       , ctxKnownThings = error "empty known things from emptyContext"
+      , ctxFamInstEnvs = mempty
       , ctxInstEnvs = InstEnvs mempty mempty mempty
       , ctxTheta = mempty
       }

plugins/hls-tactics-plugin/test/AutoTupleSpec.hs

@@ -36,6 +36,7 @@ spec = describe "auto for tuple" $ do
           runTactic
             emptyContext
             (mkFirstJudgement
+              emptyContext
               (Hypothesis $ pure $ HyInfo (mkVarOcc "x") UserPrv $ CType in_type)
               True
               out_type)

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

@@ -17,6 +17,9 @@ spec = do
     destructTest "a"     7 25 "SplitPattern"
     destructTest "a"     6 18 "DestructPun"
     destructTest "fp"   31 14 "DestructCthulhu"
+    destructTest "b"     7 10 "DestructTyFam"
+    destructTest "b"     7 10 "DestructDataFam"
+    destructTest "b"    17 10 "DestructTyToDataFam"
 
   describe "layout" $ do
     destructTest "b"  4  3 "LayoutBind"

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

@@ -0,0 +1,8 @@
+{-# LANGUAGE TypeFamilies #-}
+
+data family Yo
+data instance Yo = Heya Int
+
+test :: Yo -> Int
+test (Heya n) = _
+

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

@@ -0,0 +1,8 @@
+{-# LANGUAGE TypeFamilies #-}
+
+data family Yo
+data instance Yo = Heya Int
+
+test :: Yo -> Int
+test b = _
+

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

@@ -0,0 +1,9 @@
+{-# LANGUAGE TypeFamilies #-}
+
+type family Yo where
+  Yo = Bool
+
+test :: Yo -> Int
+test False = _
+test True = _
+

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

@@ -0,0 +1,8 @@
+{-# LANGUAGE TypeFamilies #-}
+
+type family Yo where
+  Yo = Bool
+
+test :: Yo -> Int
+test b = _
+

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

@@ -0,0 +1,18 @@
+{-# LANGUAGE TypeFamilies         #-}
+{-# LANGUAGE UndecidableInstances #-}
+
+type family T1 a where
+  T1 a = T2 Int
+
+type family T2 a
+type instance T2 Int = T3
+
+type family T3 where
+  T3 = Yo
+
+data family Yo
+data instance Yo = Heya Int
+
+test :: T1 Bool -> Int
+test (Heya n) = _
+