Preliminary Codensity STM

Control/Concurrent/Speculation.hs

@@ -1,4 +1,4 @@
-{-# LANGUAGE BangPatterns #-}
+{-# LANGUAGE BangPatterns, DeriveDataTypeable #-}
 module Control.Concurrent.Speculation
     ( 
     -- * Speculative application
@@ -8,30 +8,39 @@ module Control.Concurrent.Speculation
     , specBy'
     , specOn
     , specOn'
-    -- * Detecting closure evaluation
-    , evaluated
+    -- * Speculative application with transactional rollback
+    , specSTM
+    , specSTM'
+    , specOnSTM
+    , specOnSTM'
+    , specBySTM
+    , specBySTM'
+    -- * Codensity STM speculation
+    , specCSTM
+    , specCSTM'
+    , specOnCSTM
+    , specOnCSTM'
+    , specByCSTM
+    , specByCSTM'
+    , CSTM
+    -- * Codensity
+    , Codensity(..)
+    , liftCodensity
+    , lowerCodensity
     ) where
 
+import Control.Concurrent.STM
+import Control.Concurrent.Speculation.Internal 
+    (Codensity(..), liftCodensity, lowerCodensity, evaluated)
+import Control.Exception (Exception, throw, fromException)
 import Control.Parallel (par)
+import Data.Typeable (Typeable)
 import Data.Function (on)
-import Data.Bits ((.&.))
-import Foreign (sizeOf)
-import Unsafe.Coerce (unsafeCoerce)
+import System.IO.Unsafe (unsafePerformIO)
 
-data Box a = Box a
+type CSTM = Codensity STM
 
--- | Inspect the dynamic pointer tagging bits of a closure. This is an impure function that
--- relies on GHC internals and will falsely return 0, but (hopefully) never give the wrong tag number if it returns 
--- a non-0 value.
-tag :: a -> Int
-tag a = unsafeCoerce (Box a) .&. (sizeOf (undefined :: Int) - 1)
-{-# INLINE tag #-}
-
--- | Returns a guess as to whether or not a value has been evaluated. This is an impure function
--- that relies on GHC internals and will return false negatives, but (hopefully) no false positives.
-evaluated :: a -> Bool
-evaluated a = tag a /= 0
-{-# INLINE evaluated #-}
+-- * Basic speculation
 
 -- | @'spec' g f a@ evaluates @f g@ while forcing @a@, if @g == a@ then @f g@ is returned. Otherwise @f a@ is evaluated.
 --
@@ -100,3 +109,140 @@ specOn = specBy . on (==)
 specOn' :: Eq c => (a -> c) -> a -> (a -> b) -> a -> b
 specOn' = specBy' . on (==)
 {-# INLINE specOn' #-}
+
+-- * STM-based speculation
+
+-- | @'specSTM' g f a@ evaluates @f g@ while forcing @a@, if @g == a@ then @f g@ is returned. Otherwise the side-effects 
+-- of the current STM transaction are rolled back and @f a@ is evaluated.
+--   
+-- If the argument @a@ is already evaluated, we don\'t bother to perform @f g@ at all.
+--
+-- If a good guess at the value of @a@ is available, this is one way to induce parallelism in an otherwise sequential task. 
+--
+-- However, if the guess isn\'t available more cheaply than the actual answer then this saves no work, and if the guess is
+-- wrong, you risk evaluating the function twice.
+--
+-- > specSTM a f a = f $! a
+--
+-- The best-case timeline looks like:
+--
+-- > [------ f g ------]
+-- >     [------- a -------]
+-- > [--- specSTM g f a ---]
+--
+-- The worst-case timeline looks like:
+--
+-- > [------ f g ------] 
+-- >     [------- a -------]
+-- >                       [-- rollback --]
+-- >                                      [------ f a ------]     
+-- > [------------------ spec g f a ------------------------]
+--
+-- Compare these to the timeline of @f $! a@:
+--
+-- > [------- a -------]
+-- >                   [------ f a ------]
+
+
+specSTM :: Eq a => a -> (a -> STM b) -> a -> STM b
+specSTM = specBySTM (==)
+{-# INLINE specSTM #-}
+
+-- | Unlike 'specSTM', 'specSTM'' doesn't check if the argument has already been evaluated.
+
+specSTM' :: Eq a => a -> (a -> STM b) -> a -> STM b
+specSTM' = specBySTM' (==)
+{-# INLINE specSTM' #-}
+
+-- | 'specSTM' using a user defined comparison function
+specBySTM :: (a -> a -> Bool) -> a -> (a -> STM b) -> a -> STM b
+specBySTM cmp g f a 
+    | evaluated a = f a 
+    | otherwise   = specBySTM' cmp g f a
+{-# INLINE specBySTM #-}
+
+-- | 'specSTM'' using a user defined comparison function
+specBySTM' :: (a -> a -> Bool) -> a -> (a -> STM b) -> a -> STM b
+specBySTM' cmp g f a = a `par` do
+    exn <- freshSpeculation
+    let 
+      try = do
+        result <- f g 
+        if cmp g a
+          then return result
+          else throw exn
+    try `catchSTM` \e -> case fromException e of
+        Just exn' | exn == exn' -> f a -- rerun with alternative inputs
+        _ -> throw e                   -- this is somebody else's problem
+
+{-# INLINE specBySTM' #-}
+
+-- | 'specBySTM' . 'on' (==)'
+specOnSTM :: Eq c => (a -> c) -> a -> (a -> STM b) -> a -> STM b
+specOnSTM = specBySTM . on (==)
+{-# INLINE specOnSTM #-}
+
+-- | 'specBySTM'' . 'on' (==)'
+specOnSTM' :: Eq c => (a -> c) -> a -> (a -> STM b) -> a -> STM b
+specOnSTM' = specBySTM' . on (==)
+{-# INLINE specOnSTM' #-}
+
+-- ** Codensity STM speculation
+
+specCSTM :: Eq a => a -> (a -> CSTM b) -> a -> CSTM b
+specCSTM = specByCSTM (==)
+{-# INLINE specCSTM #-}
+
+-- | Unlike 'specSTM', 'specSTM'' doesn't check if the argument has already been evaluated.
+
+specCSTM' :: Eq a => a -> (a -> CSTM b) -> a -> CSTM b
+specCSTM' = specByCSTM' (==)
+{-# INLINE specCSTM' #-}
+
+-- | 'specSTM' using a user defined comparison function
+specByCSTM :: (a -> a -> Bool) -> a -> (a -> CSTM b) -> a -> CSTM b
+specByCSTM cmp g f a 
+    | evaluated a = f a 
+    | otherwise   = specByCSTM' cmp g f a
+{-# INLINE specByCSTM #-}
+
+-- | 'specCSTM'' using a user defined comparison function
+specByCSTM' :: (a -> a -> Bool) -> a -> (a -> CSTM b) -> a -> CSTM b
+specByCSTM' cmp g f a = a `par` Codensity $ \k -> do
+    exn <- freshSpeculation
+    let 
+      try = do
+        result <- lowerCodensity (f g)
+        if cmp g a
+          then k result
+          else throw exn
+    try `catchSTM` \e -> case fromException e of
+        Just exn' | exn == exn' -> lowerCodensity (f a) >>= k -- rerun with alternative inputs
+        _ -> throw e                         -- this is somebody else's problem
+
+{-# INLINE specByCSTM' #-}
+
+-- | 'specByCSTM' . 'on' (==)'
+specOnCSTM :: Eq c => (a -> c) -> a -> (a -> CSTM b) -> a -> CSTM b
+specOnCSTM = specByCSTM . on (==)
+{-# INLINE specOnCSTM #-}
+
+-- | 'specByCSTM'' . 'on' (==)'
+specOnCSTM' :: Eq c => (a -> c) -> a -> (a -> CSTM b) -> a -> CSTM b
+specOnCSTM' = specByCSTM' . on (==)
+{-# INLINE specOnCSTM' #-}
+
+-- | TVar used to allocate speculation exceptions
+speculationSupply :: TVar Int
+speculationSupply = unsafePerformIO $ newTVarIO 0
+{-# NOINLINE speculationSupply #-}
+
+freshSpeculation :: STM Speculation
+freshSpeculation = do
+    n <- readTVar speculationSupply
+    writeTVar speculationSupply $! n + 1
+    return (Speculation n)
+{-# INLINE freshSpeculation #-}
+
+newtype Speculation = Speculation Int deriving (Show,Eq,Typeable)
+instance Exception Speculation

Control/Concurrent/Speculation/Internal.hs

@@ -0,0 +1,57 @@
+{-# LANGUAGE Rank2Types #-}
+module Control.Concurrent.Speculation.Internal 
+    ( 
+    -- * Determining if a closure is evaluated
+      evaluated
+    -- * Codensity monad
+    , Codensity(..)
+    , liftCodensity
+    , lowerCodensity
+--  , returning
+    ) where
+
+import Control.Applicative
+import Control.Monad
+
+import Data.Bits ((.&.))
+import Foreign (sizeOf)
+import Unsafe.Coerce (unsafeCoerce)
+
+
+-- | Used to inspect tag bits
+data Box a = Box a
+
+-- | Inspect the dynamic pointer tagging bits of a closure. This is an impure function that
+-- relies on GHC internals and will falsely return 0, but (hopefully) never give the wrong tag number if it returns a non-0 value.
+tag :: a -> Int
+tag a = unsafeCoerce (Box a) .&. (sizeOf (undefined :: Int) - 1)
+{-# INLINE tag #-}
+
+-- | Returns a guess as to whether or not a value has been evaluated. This is an impure function
+-- that relies on GHC internals and will return false negatives, but (hopefully) no false positives. This is unsafe as the value of this function will vary (from False to True) over the course of pure invocations!
+evaluated :: a -> Bool
+evaluated a = tag a /= 0
+{-# INLINE evaluated #-}
+
+-- returning :: Monad m => (a -> a -> b) -> a -> a -> m b
+-- returning f a b = return (f a b)
+-- {-# INLINE returning #-}
+
+newtype Codensity f a = Codensity { runCodensity :: forall r. (a -> f r) -> f r } 
+
+instance Functor (Codensity f) where
+    fmap f (Codensity m) = Codensity $ \k -> m (k . f)
+
+instance Applicative (Codensity f) where
+    pure = return
+    (<*>) = ap
+
+instance Monad (Codensity f) where
+    return x = Codensity (\k -> k x)
+    Codensity m >>= f = Codensity (\k -> m (\a -> runCodensity (f a) k))
+
+liftCodensity :: Monad m => m a -> Codensity m a 
+liftCodensity m = Codensity (m >>=)
+
+lowerCodensity :: Monad m => Codensity m a -> m a
+lowerCodensity a = runCodensity a return

Data/Foldable/Speculation.hs

@@ -40,8 +40,10 @@ import Data.Ix ()
 import Data.Function (on)
 import Data.Foldable (Foldable)
 import qualified Data.Foldable as Foldable
-import Control.Applicative
+-- import Control.Concurrent.STM
 import Control.Concurrent.Speculation
+-- import Control.Concurrent.Speculation.Internal (Codensity(..))
+import Control.Applicative
 import Control.Monad hiding (mapM_, msum, forM_, sequence_)
 
 -- | Given a valid estimate @g@, @'fold' g f xs@ yields the same answer as @'fold' f xs@.
@@ -92,6 +94,26 @@ foldrBy cmp g f z = extractAcc . Foldable.foldr mf (Acc 0 z)
     mf a (Acc n b) = let n' = n + 1 in Acc n' (specBy' cmp (g n') (f a) b)
 {-# INLINE foldrBy #-}
 
+{-
+foldrSTM :: (Foldable f, Eq b) => (Int -> STM b) -> (a -> b -> STM b) -> b -> f a -> STM b
+foldrSTM = foldrSTMBy (==)
+{-# INLINE foldrSTM #-}
+
+foldrSTMBy :: Foldable f => (b -> b -> Bool) -> (Int -> STM b) -> (a -> b -> STM b) -> b -> f a -> STM b
+foldrSTMBy = undefined
+{-# INLINE foldrSTMBy #-}
+-}
+
+{-
+foldrSTMBy cmp g f z xs = liftM extractAcc . Foldable.foldl mf return xs (Acc 0 z)
+  where
+    mf h t = do
+        Acc n t' <- t
+        let !n' = n + 1 
+        specSTMBy' cmp (g n') (flip f h >=> t) 
+        ...
+-}
+
 -- | Given a valid estimator @g@, @'foldl' g f z xs@ yields the same answer as @'foldl'' f z xs@.
 --
 -- @g n@ should supply an estimate of the value returned from folding over the first @n@ elements of the container.

README.markdown

@@ -10,6 +10,8 @@ This package provides speculative evaluation primitives for Haskell, very loosel
 
 ### speculative function application
 
+Various speculative function application combinators are provided. Two fairly canonical samples are described here.
+
 #### spec
 
     spec :: Eq a => a -> (a -> b) -> a -> b
@@ -42,24 +44,31 @@ Compare these to the timeline of @f $! a@:
 
 `specSTM` provides a similar compressed timeline for speculated STM actions, but also rolls back side-effects.
 
-### folds
+### speculative folds
 
-A number of speculative folds are also provided.
+A speculative version of Data.Foldable is provided as Data.Foldable.Speculation.
     
-These take an extra argument which is a function that guesses the result of of the fold up to a given point.
+Each combinator therein takes an extra argument that is used to speculate on the value of the list.
 
-#### specFoldr
+#### foldr
 
-    specFoldr :: (Foldable f, Eq b) => (Int -> b) -> (a -> b -> b) -> b -> f a -> b
+    foldr :: (Foldable f, Eq b) => (Int -> b) -> (a -> b -> b) -> b -> f a -> b
 
-Given a valid estimator `g`, `'specFoldr g f z xs` yields the same answer as `foldr' f z xs`.
+Given a valid estimator `g`, `'foldr g f z xs` yields the same answer as `Foldable.foldr' f z xs`.
 
 `g n` should supply an estimate of the value returned from folding over the /last/ `n` elements of the container.
 
 As with `spec`, if the guess `g n` is accurate a reasonable percentage of the time and faster to compute than the fold, then this can provide increased opportunities for parallelism.
 
-#### specFoldl
+#### foldl
 
-    specFoldl :: (Foldable f, Eq b) => (Int -> b) -> (b -> a -> b) -> b -> f a -> b
+    foldl :: (Foldable f, Eq b) => (Int -> b) -> (b -> a -> b) -> b -> f a -> b
 
-`specFoldl` works similarly to `foldl'`, except that `g n` should provide an estimate for the /first/ `n` elements.
+`foldl` works similarly to `Foldable.foldl'`, except that `g n` should provide an estimate for the /first/ `n` elements.
+
+contact information
+-------------------
+
+I can be reached through the user ekmett on github, as edwardk on irc.freenode.net #haskell channel, or by email to <ekmett@gmail.com>.
+
+-Edward Kmett

speculation.cabal

@@ -1,5 +1,5 @@
 name:           speculation
-version:        0.3.0
+version:        0.4.0
 license:        BSD3
 license-file:   LICENSE
 author:         Edward A. Kmett
@@ -61,5 +61,7 @@ library
 
   exposed-modules:
     Control.Concurrent.Speculation
-    Control.Concurrent.STM.Speculation
     Data.Foldable.Speculation
+
+  other-modules:
+    Control.Concurrent.Speculation.Internal