Start case splitting on type families (NOT working)

Language/Haskell/GhcMod/CaseSplit.hs

@@ -7,25 +7,30 @@ module Language.Haskell.GhcMod.CaseSplit (
 import CoreMonad (liftIO)
 import Data.Function (on)
 import Data.List (find, intercalate, sortBy)
+import Data.Maybe (isJust)
 import qualified Data.Text as T
 import qualified Data.Text.IO as T (readFile)
 import qualified DataCon as Ty
 import Exception (ghandle, SomeException(..))
-import GHC (GhcMonad, LHsBind, LPat, Id, ParsedModule(..), TypecheckedModule(..), DynFlags, SrcSpan, Type, GenLocated(L))
+import GHC (GhcMonad, LPat, Id, ParsedModule(..), TypecheckedModule(..), DynFlags, SrcSpan, Type, GenLocated(L))
 import qualified GHC as G
 import Language.Haskell.GhcMod.Convert
 import qualified Language.Haskell.GhcMod.Gap as Gap
 import Language.Haskell.GhcMod.Monad
 import Language.Haskell.GhcMod.SrcUtils
-import Outputable (PprStyle)
+import Outputable (ppr, PprStyle)
 import qualified TyCon as Ty
 import qualified Type as Ty
 
+import Debug.Trace
+import Language.Haskell.GhcMod.Doc
+
 ----------------------------------------------------------------
 -- CASE SPLITTING
 ----------------------------------------------------------------
 
-data SplitInfo = SplitInfo G.Name (SrcSpan,Type) (SrcSpan, Type) [SrcSpan]
+data SplitInfo = SplitInfo G.Name SrcSpan (SrcSpan, Type) [SrcSpan]
+               | TySplitInfo G.Name SrcSpan (SrcSpan, Ty.Kind)
 data SplitToTextInfo = SplitToTextInfo { sVarName     :: String
                                        , sBindingSpan :: SrcSpan
                                        , sVarSpan     :: SrcSpan
@@ -43,8 +48,13 @@ splits file lineNo colNo = ghandle handler body
     body = inModuleContext file $ \dflag style -> do
         opt <- options
         modSum <- Gap.fileModSummary file
-        whenFound' opt (getSrcSpanTypeForSplit modSum lineNo colNo) $
+        whenFound' opt (getSrcSpanTypeForSplit dflag style modSum lineNo colNo) $ \x -> case x of
-          \(SplitInfo varName (bndLoc,_) (varLoc,varT) _matches) -> do
+          (SplitInfo varName bndLoc (varLoc,varT) _matches) -> do
+             let varName' = showName dflag style varName  -- Convert name to string
+             text <- genCaseSplitTextFile file (SplitToTextInfo varName' bndLoc varLoc $
+                                                getTyCons dflag style varName varT)
+             return (fourInts bndLoc, text)
+          (TySplitInfo varName bndLoc (varLoc,varT)) -> do
              let varName' = showName dflag style varName  -- Convert name to string
              text <- genCaseSplitTextFile file (SplitToTextInfo varName' bndLoc varLoc $
                                                 getTyCons dflag style varName varT)
@@ -54,22 +64,28 @@ splits file lineNo colNo = ghandle handler body
 ----------------------------------------------------------------
 -- a. Code for getting the information of the variable
 
-getSrcSpanTypeForSplit :: GhcMonad m => G.ModSummary -> Int -> Int -> m (Maybe SplitInfo)
+getSrcSpanTypeForSplit :: GhcMonad m => DynFlags -> PprStyle -> G.ModSummary -> Int -> Int -> m (Maybe SplitInfo)
-getSrcSpanTypeForSplit modSum lineNo colNo = do
+getSrcSpanTypeForSplit dflag style modSum lineNo colNo = do
+  fn <- getSrcSpanTypeForFnSplit modSum lineNo colNo
+  if isJust fn
+     then return fn
+     else getSrcSpanTypeForTypeSplit dflag style modSum lineNo colNo
+
+-- Information for a function case split
+getSrcSpanTypeForFnSplit :: GhcMonad m => G.ModSummary -> Int -> Int -> m (Maybe SplitInfo)
+getSrcSpanTypeForFnSplit modSum lineNo colNo = do
     p@ParsedModule{pm_parsed_source = pms} <- G.parseModule modSum
     tcm@TypecheckedModule{tm_typechecked_source = tcs} <- G.typecheckModule p
-    let bs:_ = sortBy (cmp `on` G.getLoc) $ listifySpans tcs (lineNo, colNo) :: [LHsBind Id]
+    let varPat  = find isPatternVar $ listifySpans tcs (lineNo, colNo) :: Maybe (LPat Id)
-        varPat  = find isPatternVar $ listifySpans tcs (lineNo, colNo) :: Maybe (LPat Id)
         match:_ = listifyParsedSpans pms (lineNo, colNo) :: [Gap.GLMatch]
     case varPat of
       Nothing  -> return Nothing
       Just varPat' -> do
         varT <- Gap.getType tcm varPat'  -- Finally we get the type of the var
-        bsT  <- Gap.getType tcm bs
+        case varT of
-        case (varT, bsT) of
+          Just varT' ->
-          (Just varT', Just (_,bsT')) ->
             let (L matchL (G.Match _ _ (G.GRHSs rhsLs _))) = match
-            in return $ Just (SplitInfo (getPatternVarName varPat') (matchL,bsT') varT' (map G.getLoc rhsLs) )
+            in return $ Just (SplitInfo (getPatternVarName varPat') matchL varT' (map G.getLoc rhsLs) )
           _ -> return Nothing
 
 isPatternVar :: LPat Id -> Bool
@@ -78,7 +94,24 @@ isPatternVar _                  = False
 
 getPatternVarName :: LPat Id -> G.Name
 getPatternVarName (L _ (G.VarPat vName)) = G.getName vName
-getPatternVarName _                      = error "This should never happend"
+getPatternVarName _                      = error "This should never happened"
+
+-- Information for a type family case split
+getSrcSpanTypeForTypeSplit :: GhcMonad m => DynFlags -> PprStyle -> G.ModSummary -> Int -> Int -> m (Maybe SplitInfo)
+getSrcSpanTypeForTypeSplit dflag style modSum lineNo colNo = do
+    p <- G.parseModule modSum
+    TypecheckedModule{tm_typechecked_source = tcs} <- G.typecheckModule p
+    let lst = listifySpans tcs (lineNo, colNo) :: [G.LTyFamInstEqn Id]
+        (L eqL _):_ = trace (showOneLine dflag style $ ppr lst) $ sortBy (cmp `on` G.getLoc) lst
+        varPat = find isTypePatternVar $ listifySpans tcs (lineNo, colNo) :: Maybe (G.LHsType Id)
+    case trace (showPage dflag style (ppr tcs)) $ varPat of
+      Just (L varSpan (G.HsTyVar vName)) ->
+        return $ Just (TySplitInfo (G.getName vName) eqL (varSpan, G.idType vName))
+      _ -> return Nothing
+
+isTypePatternVar :: G.LHsType Id -> Bool
+isTypePatternVar (L _ (G.HsTyVar _)) = True
+isTypePatternVar _                   = False
 
 ----------------------------------------------------------------
 -- b. Code for getting the possible constructors
@@ -89,6 +122,12 @@ getTyCons dflag style name ty | Just (tyCon, _) <- Ty.splitTyConApp_maybe ty =
   in getTyCon dflag style name' tyCon
 getTyCons dflag style name _ = [showName dflag style name]
 
+getKiCons :: DynFlags -> PprStyle -> G.Name -> G.Kind ->  [String]
+getKiCons dflag style name ty | Just (tyCon, _) <- Ty.splitTyConApp_maybe ty =
+  let name' = showName dflag style name  -- Convert name to string
+  in getTyCon dflag style name' tyCon
+getKiCons dflag style name _ = [showName dflag style name]
+
 -- Write cases for one type
 getTyCon :: DynFlags -> PprStyle -> String -> Ty.TyCon -> [String]
 -- 1. Non-matcheable type constructors