From 56cc237e26da5ec81b7c17e444082e8c318c2d19 Mon Sep 17 00:00:00 2001
From: Rob Everest <rob.everest@gmail.com>
Date: Sun, 2 Nov 2014 00:06:34 +1100
Subject: [PATCH] Support where clauses, let bindings and case expressions in
 case splitting (fixes #395)

---
 Language/Haskell/GhcMod/CaseSplit.hs | 17 +++++++++++++----
 ghc-mod.cabal                        |  1 +
 test/CaseSplitSpec.hs                | 21 +++++++++++++++++++++
 test/data/case-split/Vect.hs         | 15 +++++++++++++++
 4 files changed, 50 insertions(+), 4 deletions(-)

diff --git a/Language/Haskell/GhcMod/CaseSplit.hs b/Language/Haskell/GhcMod/CaseSplit.hs
index 4d5a2ff..dabb67b 100644
--- a/Language/Haskell/GhcMod/CaseSplit.hs
+++ b/Language/Haskell/GhcMod/CaseSplit.hs
@@ -71,8 +71,8 @@ getSrcSpanTypeForFnSplit :: GhcMonad m => G.ModSummary -> Int -> Int -> m (Maybe
 getSrcSpanTypeForFnSplit modSum lineNo colNo = do
     p@ParsedModule{pm_parsed_source = _pms} <- G.parseModule modSum
     tcm@TypecheckedModule{tm_typechecked_source = tcs} <- G.typecheckModule p
-    let varPat  = find isPatternVar $ listifySpans tcs (lineNo, colNo) :: Maybe (LPat Id)
-        match:_ = listifySpans tcs (lineNo, colNo) :: [Gap.GLMatchI]
+    let varPat = find isPatternVar $ listifySpans tcs (lineNo, colNo) :: Maybe (LPat Id)
+        match  = last $ listifySpans tcs (lineNo, colNo) :: Gap.GLMatchI
     case varPat of
       Nothing  -> return Nothing
       Just varPat' -> do
@@ -188,8 +188,11 @@ getCaseSplitText text (SplitToTextInfo { sVarName = sVN, sBindingSpan = sBS
                                        , sVarSpan = sVS, sTycons = sT })  =
   let bindingText = getBindingText text sBS
       difference  = srcSpanDifference sBS sVS
-      replaced    = concatMap (replaceVarWithTyCon bindingText difference sVN) sT
-   in T.unpack $ T.intercalate (T.pack "\n") replaced
+      replaced    = map (replaceVarWithTyCon bindingText difference sVN) sT
+      -- The newly generated bindings need to be indented to align with the
+      -- original binding.
+      replaced'   = head replaced : map (indentBindingTo sBS) (tail replaced)
+   in T.unpack $ T.intercalate (T.pack "\n") (concat replaced')
 
 getBindingText :: [T.Text] -> SrcSpan -> [T.Text]
 getBindingText text srcSpan =
@@ -220,3 +223,9 @@ replaceVarWithTyCon text (vsl,vsc,_,vec) varname tycon =
                                then T.take vsc line `T.append` tycon'' `T.append` T.drop vec line
                                else T.replicate spacesToAdd (T.pack " ") `T.append` line)
               [0 ..] text
+
+indentBindingTo :: SrcSpan -> [T.Text] -> [T.Text]
+indentBindingTo bndLoc binds =
+  let Just (_,sl,_,_) = Gap.getSrcSpan bndLoc
+      indent      = (T.replicate (sl - 1) (T.pack " ") `T.append`)
+   in indent (head binds) : tail binds
diff --git a/ghc-mod.cabal b/ghc-mod.cabal
index e167607..a6d87af 100644
--- a/ghc-mod.cabal
+++ b/ghc-mod.cabal
@@ -34,6 +34,7 @@ Extra-Source-Files:     ChangeLog
                         test/data/broken-cabal/cabal.sandbox.config.in
                         test/data/broken-sandbox/*.cabal
                         test/data/broken-sandbox/cabal.sandbox.config
+                        test/data/case-split/*.hs
                         test/data/cabal-flags/*.cabal
                         test/data/check-test-subdir/*.cabal
                         test/data/check-test-subdir/src/Check/Test/*.hs
diff --git a/test/CaseSplitSpec.hs b/test/CaseSplitSpec.hs
index 700cc68..5e5db3f 100644
--- a/test/CaseSplitSpec.hs
+++ b/test/CaseSplitSpec.hs
@@ -18,3 +18,24 @@ spec = do
                 res `shouldBe` "24 1 24 30"++
                         " \"mlAppend Nil y = _mlAppend_body\NUL"++
                            "mlAppend (Cons x1 x2) y = _mlAppend_body\"\n"
+
+        it "works with case expressions" $ do
+            withDirectory_ "test/data/case-split" $ do
+                res <- runD $ splits "Vect.hs" 28 20
+                res `shouldBe` "28 19 28 39"++
+                        " \"Nil -> _mlAppend_body\NUL"++
+                        "                  (Cons x'1 x'2) -> _mlAppend_body\"\n"
+
+        it "works with where clauses" $ do
+            withDirectory_ "test/data/case-split" $ do
+                res <- runD $ splits "Vect.hs" 34 17
+                res `shouldBe` "34 5 34 43"++
+                        " \"mlReverse' Nil accum = _mlReverse_body\NUL"++
+                        "    mlReverse' (Cons xs'1 xs'2) accum = _mlReverse_body\"\n"
+
+        it "works with let bindings" $ do
+            withDirectory_ "test/data/case-split" $ do
+                res <- runD $ splits "Vect.hs" 38 33
+                res `shouldBe` "38 21 38 59"++
+                        " \"mlReverse' Nil accum = _mlReverse_body\NUL"++
+                        "                    mlReverse' (Cons xs'1 xs'2) accum = _mlReverse_body\"\n"
diff --git a/test/data/case-split/Vect.hs b/test/data/case-split/Vect.hs
index cd16e4c..9d11ada 100644
--- a/test/data/case-split/Vect.hs
+++ b/test/data/case-split/Vect.hs
@@ -22,3 +22,18 @@ data MyList a = Nil | Cons a (MyList a)
 
 mlAppend :: MyList a -> MyList a -> MyList a
 mlAppend x y = _mlAppend_body
+
+mlAppend2 :: MyList a -> MyList a -> MyList a
+mlAppend2 x y = case x of
+                  x' -> _mlAppend_body
+
+mlReverse :: MyList a -> MyList a
+mlReverse xs = mlReverse' xs Nil
+  where
+    mlReverse' :: MyList a -> MyList a -> MyList a
+    mlReverse' xs' accum = _mlReverse_body
+
+mlReverse2 :: MyList a -> MyList a
+mlReverse2 xs = let mlReverse' :: MyList a -> MyList a -> MyList a
+                    mlReverse' xs' accum = _mlReverse_body
+                 in mlReverse' xs Nil