Make NonDet Higher-Order (#174)

* Higher-Order NonDet and runNonDetMaybe

* Fixed docs

* Fixed a word in AlternativeSpec docs

* Fixed tests and wrong documentation

src/Polysemy/Internal.hs

@@ -211,10 +211,7 @@ instance Monad (Sem f) where
 instance (Member NonDet r) => Alternative (Sem r) where
   empty = send Empty
   {-# INLINE empty #-}
-  a <|> b = do
-    send (Choose id) >>= \case
-      False -> a
-      True  -> b
+  a <|> b = send (Choose a b)
   {-# INLINE (<|>) #-}
 
 -- | @since 0.2.1.0

src/Polysemy/Internal/NonDet.hs

@@ -6,12 +6,10 @@
 
 module Polysemy.Internal.NonDet where
 
-import Data.Kind
-
 
 ------------------------------------------------------------------------------
 -- | An effect corresponding to the 'Control.Applicative.Alternative' typeclass.
-data NonDet (m :: Type -> Type) a
+data NonDet m a
   = Empty
-  | Choose (Bool -> a)
+  | Choose (m a) (m a)
 

src/Polysemy/NonDet.hs

@@ -7,9 +7,11 @@ module Polysemy.NonDet
 
     -- * Interpretations
   , runNonDet
+  , runNonDetMaybe
   ) where
 
 import Control.Applicative
+import Control.Monad.Trans.Maybe
 import Data.Maybe
 import Polysemy.Internal
 import Polysemy.Internal.NonDet
@@ -55,17 +57,42 @@ instance Monad (NonDetC m) where
 ------------------------------------------------------------------------------
 -- | Run a 'NonDet' effect in terms of some underlying 'Alternative' @f@.
 runNonDet :: Alternative f => Sem (NonDet ': r) a -> Sem r (f a)
-runNonDet (Sem m) = Sem $ \k -> runNonDetC $ m $ \u ->
+runNonDet = runNonDetC . runNonDetInC
+{-# INLINE runNonDet #-}
+
+runNonDetInC :: Sem (NonDet ': r) a -> NonDetC (Sem r) a
+runNonDetInC = usingSem $ \u ->
   case decomp u of
     Left x  -> NonDetC $ \cons nil -> do
-      z <- k $ weave [()]
+      z <- liftSem $ weave [()]
                      (fmap concat . traverse runNonDet)
                      -- TODO(sandy): Is this the right semantics?
                      listToMaybe
                      x
       foldr cons nil z
     Right (Weaving Empty _ _ _ _) -> empty
-    Right (Weaving (Choose ek) s _ y _) -> do
-      z <- pure (ek False) <|> pure (ek True)
-      pure $ y $ z <$ s
+    Right (Weaving (Choose left right) s wv ex _) -> fmap ex $
+      runNonDetInC (wv (left <$ s)) <|> runNonDetInC (wv (right <$ s))
+{-# INLINE runNonDetInC #-}
 
+------------------------------------------------------------------------------
+-- | Run a 'NonDet' effect in terms of an underlying 'Maybe'
+--
+-- Unlike 'runNonDet', uses of '<|>' will not execute the
+-- second branch at all if the first option succeeds.
+runNonDetMaybe :: Sem (NonDet ': r) a -> Sem r (Maybe a)
+runNonDetMaybe (Sem sem) = Sem $ \k -> runMaybeT $ sem $ \u ->
+  case decomp u of
+    Right (Weaving e s wv ex _) ->
+      case e of
+        Empty -> empty
+        Choose left right ->
+          MaybeT $ usingSem k $ runMaybeT $ fmap ex $ do
+              MaybeT (runNonDetMaybe (wv (left <$ s)))
+          <|> MaybeT (runNonDetMaybe (wv (right <$ s)))
+    Left x -> MaybeT $
+      k $ weave (Just ())
+          (maybe (pure Nothing) runNonDetMaybe)
+          id
+          x
+{-# INLINE runNonDetMaybe #-}

test/AlternativeSpec.hs

@@ -4,16 +4,23 @@ import Polysemy
 import Polysemy.NonDet
 import Test.Hspec
 import Control.Applicative
+import Polysemy.Trace
 
 semFail :: Member NonDet r => Maybe Bool -> Sem r Bool
 semFail mb = do
   Just b <- pure mb
   pure b
 
-
 runAlt :: Alternative f => Sem '[NonDet] a -> f a
 runAlt = run . runNonDet
 
+failtrace :: (Member NonDet r, Member Trace r)
+          => Sem r ()
+failtrace = pure () <|> trace "trace"
+
+failtrace' :: (Member NonDet r, Member Trace r)
+           => Sem r ()
+failtrace' = trace "sim" *> empty <|> trace "salabim"
 
 spec :: Spec
 spec = parallel $ do
@@ -32,3 +39,15 @@ spec = parallel $ do
       runAlt (semFail $ Just True) `shouldBe` Just True
       runAlt (semFail $ Just False) `shouldBe` [False]
 
+  describe "runNonDetMaybe" $ do
+    it "should skip second branch if the first branch succeeds" $ do
+      (run . runNonDetMaybe . runTraceAsList) failtrace
+        `shouldBe` Just ([], ())
+      (run . runTraceAsList . runNonDetMaybe) failtrace
+        `shouldBe` ([], Just ())
+
+    it "should respect local/global state semantics" $ do
+      (run . runNonDetMaybe . runTraceAsList) failtrace'
+        `shouldBe` Just (["salabim"], ())
+      (run . runTraceAsList . runNonDetMaybe) failtrace'
+        `shouldBe` (["sim", "salabim"], Just ())