Add -foptimal-applicative-do

Summary:
The algorithm for ApplicativeDo rearrangement is based on a heuristic
that runs in O(n^2).  This patch adds the optimal algorithm, which is
O(n^3), selected by a flag (-foptimal-applicative-do).  It finds better
solutions in a small number of cases (about 2% of the cases where
ApplicativeDo makes a difference), but it can be very slow for large do
expressions.  I'm mainly adding it for experimental reasons.

ToDo: user guide docs

Test Plan: validate

Reviewers: simonpj, bgamari, austin, niteria, erikd

Subscribers: thomie

Differential Revision: https://phabricator.haskell.org/D1969

5 files changed, 76 insertions, 0 deletions

diff --git a/testsuite/tests/ado/ado-optimal.hs b/testsuite/tests/ado/ado-optimal.hs
new file mode 100644
index 0000000000..aab8d5397f
--- /dev/null
+++ b/testsuite/tests/ado/ado-optimal.hs
@@ -0,0 +1,59 @@
+{-# LANGUAGE ScopedTypeVariables, ExistentialQuantification, ApplicativeDo #-}
+{-# OPTIONS_GHC -foptimal-applicative-do #-}
+module Main where
+
+import Control.Applicative
+import Text.PrettyPrint
+
+(a:b:c:d:e:f:g:h:_) = map (\c -> doc [c]) ['a'..]
+
+-- This one requires -foptimal-applicative-do to find the best solution
+-- ((a; b) | (c; d)); e
+test1 :: M ()
+test1 = do
+  x1 <- a
+  x2 <- const b x1
+  x3 <- c
+  x4 <- const d x3
+  x5 <- const e (x1,x4)
+  return (const () x5)
+
+main = mapM_ run
+ [ test1
+ ]
+
+-- Testing code, prints out the structure of a monad/applicative expression
+
+newtype M a = M (Bool -> (Maybe Doc, a))
+
+maybeParen True d = parens d
+maybeParen _ d = d
+
+run :: M a -> IO ()
+run (M m) = print d where (Just d,_) = m False
+
+instance Functor M where
+  fmap f m = m >>= return . f
+
+instance Applicative M where
+  pure a = M $ \_ -> (Nothing, a)
+  M f <*> M a = M $ \p ->
+    let (Just d1, f') = f True
+        (Just d2, a') = a True
+    in
+        (Just (maybeParen p (d1 <+> char '|' <+> d2)), f' a')
+
+instance Monad M where
+  return = pure
+  M m >>= k = M $ \p ->
+    let (d1, a) = m True
+        (d2, b) = case k a of M f -> f True
+    in
+    case (d1,d2) of
+      (Nothing,Nothing) -> (Nothing, b)
+      (Just d, Nothing) -> (Just d, b)
+      (Nothing, Just d) -> (Just d, b)
+      (Just d1, Just d2) -> (Just (maybeParen p (d1 <> semi <+> d2)), b)
+
+doc :: String -> M ()
+doc d = M $ \_ -> (Just (text d), ())
diff --git a/testsuite/tests/ado/ado-optimal.stdout b/testsuite/tests/ado/ado-optimal.stdout
new file mode 100644
index 0000000000..29f9856d5f
--- /dev/null
+++ b/testsuite/tests/ado/ado-optimal.stdout
@@ -0,0 +1 @@
+((a; b) | (c; d)); e
diff --git a/testsuite/tests/ado/ado004.hs b/testsuite/tests/ado/ado004.hs
index 67e04c117a..6ddc8395ea 100644
--- a/testsuite/tests/ado/ado004.hs
+++ b/testsuite/tests/ado/ado004.hs
@@ -15,6 +15,15 @@ test2 f = do
   y <- f 4
   return (x + y)
 
+-- Test we can also infer the Functor version of the type
+test2a f = do
+  x <- f 3
+  return (x + 1)
+
+-- Test for just one statement
+test2b f = do
+  return (f 3)
+
 -- This one will use join
 test3 f g = do
   x <- f 3
diff --git a/testsuite/tests/ado/ado004.stderr b/testsuite/tests/ado/ado004.stderr
index 6e877617dc..c6c5e3544d 100644
--- a/testsuite/tests/ado/ado004.stderr
+++ b/testsuite/tests/ado/ado004.stderr
@@ -5,6 +5,12 @@ TYPE SIGNATURES
     forall (f :: * -> *) b t.
     (Applicative f, Num t, Num b) =>
     (t -> f b) -> f b
+  test2a ::
+    forall (f :: * -> *) b t.
+    (Num t, Num b, Functor f) =>
+    (t -> f b) -> f b
+  test2b ::
+    forall (m :: * -> *) a t. (Num t, Monad m) => (t -> a) -> m a
   test3 ::
     forall (m :: * -> *) a t t1.
     (Num t, Monad m) =>
diff --git a/testsuite/tests/ado/all.T b/testsuite/tests/ado/all.T
index e1efdf221c..06cdbf993d 100644
--- a/testsuite/tests/ado/all.T
+++ b/testsuite/tests/ado/all.T
@@ -6,3 +6,4 @@ test('ado005', normal, compile_fail, [''])
 test('ado006', normal, compile, [''])
 test('ado007', normal, compile, [''])
 test('T11607', normal, compile_and_run, [''])
+test('ado-optimal', normal, compile_and_run, [''])