Playing around with explicit CPS in Haskell

Continuations.hs

{-# language ScopedTypeVariables, LambdaCase, TypeApplications #-}

module Continuations where

import Data.String
import Control.Monad


-- Short-circuiting evaluation

fastProduct :: forall n. (Eq n, Num n) => [n] -> n
fastProduct list = fastProduct' list id
  where
    fastProduct' [      ] k = k 1
    fastProduct' (x : xs) k
      | x == 0    = k 0
      | otherwise = fastProduct' xs (\total -> k (x * total))


-- Building data structures inside continuations

shuffle :: [a] -> [a]
shuffle list = shuffle' True list id
  where
    shuffle' :: Bool -> [a] -> ([a] -> [a]) -> [a]
    shuffle' _     [      ] k = k []
    shuffle' True  (x : xs) k = shuffle' False xs ((x :) . k)
    shuffle' False (x : xs) k = shuffle' True  xs (k . (x :))

quicksort :: forall a. Ord a => [a] -> [a]
quicksort list =
  partition list []
  where
    merge :: [a] -> a -> [a] -> ([a] -> [a])
    merge lt c gt =
      \rest -> partition lt (c : partition gt rest)

    partition :: [a] -> ([a] -> [a])
    partition [      ] = id
    partition (x : xs) =
      partitionAt x xs (\lt gt -> merge lt x gt)

    partitionAt :: a -> [a] -> ([a] -> [a] -> ([a] -> [a])) -> ([a] -> [a])
    partitionAt pivot [      ] k = k [] []
    partitionAt pivot (x : xs) k
      | x < pivot = partitionAt pivot xs (\lt gt -> k (x : lt)     gt)
      | otherwise = partitionAt pivot xs (\lt gt -> k      lt (x : gt))


-- Cooperative multi-threading

newtype Thread a r = T ([Thread a r] -> a -> r)

runThread :: Thread a r -> [Thread a r] -> a -> r
runThread (T t) = t


-- Round-robin scheduling

roundRobin :: ((a -> r) -> a -> r) -> Thread a r
roundRobin f =
  T (\threads a ->
       case threads of
         [] ->
           f (runThread (roundRobin f) []) a
         (t : threads) ->
           f (runThread t (threads ++ [roundRobin f])) a)

runRound :: [(a -> r) -> a -> r] -> a -> r
runRound [           ] = error "No threads to run!"
runRound (t : threads) = runThread (roundRobin t) (map roundRobin threads)


-- Alternating sequences

yieldOne :: Num n => (() -> [n]) -> () -> [n]
yieldOne continue () =
  1 : continue ()

yieldZero :: Num n => (() -> [n]) -> () -> [n]
yieldZero continue () =
  0 : continue ()

alternating :: Num n => [n]
alternating = runRound [yieldOne, yieldZero] ()


-- Some useful combinators

stopWhen :: (a -> Bool)
         -> (a -> r)
         -> ((a -> r) -> a -> r)
         -> ((a -> r) -> a -> r)
stopWhen predicate return k continue a
  | predicate a = return a
  | otherwise   = k continue a

applyWhen :: (a -> Bool)
          -> ((a -> r) -> a -> r)
          -> ((a -> r) -> a -> r)
applyWhen predicate k continue a
  | predicate a = k continue a
  | otherwise   = continue a


-- Collatz sequences

halve :: Integral n => (n -> [n]) -> n -> [n]
halve =
  stopWhen (== 1) (\x -> [x]) $
    applyWhen even $
      \continue n ->
        n : continue (n `div` 2)

triplePlusOne :: Integral n => (n -> [n]) -> n -> [n]
triplePlusOne =
  stopWhen (== 1) (\x -> [x]) $
    applyWhen odd $
      \continue n ->
         n : continue (3 * n + 1)

collatz :: Integral n => n -> [n]
collatz = runRound [halve, triplePlusOne]


-- Fizz-Buzz

newline = putStrLn ""

printNumber continue n =
  putStr (show n ++ ". ") >> continue n

printWhen predicate string =
  applyWhen predicate $
    \continue n ->
      do putStr string
         continue n

incrementUpTo max =
  stopWhen (>= max) (\_ -> newline) $
    \continue n ->
      newline >> continue (n + 1)

divisibleBy x y =
  y `mod` x == 0

fizzbuzz n =
  runRound [ printNumber
           , printWhen (divisibleBy 3) "fizz"
           , printWhen (divisibleBy 5) "buzz"
           , incrementUpTo n
           ] 0


-- SAT Solver in CPS (it's slow, but that's not the point)

data Formula = Var String
             | Lit Bool
             | Not Formula
             | Formula :&&: Formula
             | Formula :||: Formula
             | Formula :=>: Formula
             deriving Show

infixr 3 :&&:
infixr 2 :||:
infixr 0 :=>:

instance IsString Formula where
  fromString = Var

subst :: forall r. String -> Bool -> Formula -> (Formula -> r) -> r
subst v b formula k =
  case formula of
    Var v' | v' == v   -> k (Lit b)
           | otherwise -> k (Var v')
    Lit b     -> k (Lit b)
    Not f     -> unsubst Not f
    f :&&: g  -> binsubst (:&&:) f g
    f :||: g  -> binsubst (:||:) f g
    f :=>: g  -> binsubst (:=>:) f g
  where
    unsubst :: (Formula -> Formula) -> Formula -> r
    unsubst op f =
      subst v b f $ \f' ->
      k (op f')

    binsubst :: (Formula -> Formula -> Formula) -> Formula -> Formula -> r
    binsubst op f g =
      subst v b f $ \f' ->
      subst v b g $ \g' ->
      k (f' `op` g')

eval :: forall r. Formula -> (String -> r) -> (Bool -> r) -> r
eval formula onUnbound onResult =
  case formula of
    Var v     -> onUnbound v
    Lit b     -> onResult b
    Not f     -> unop not f
    f :&&: g  -> binop (&&) f g
    f :||: g  -> binop (||) f g
    f :=>: g  -> binop (\b c  -> not b || c) f g
  where
    unop :: (Bool -> Bool) -> Formula -> r
    unop op f =
      eval f onUnbound $ \b ->
      onResult (op b)

    binop :: (Bool -> Bool -> Bool) -> Formula -> Formula -> r
    binop op f g =
      eval f onUnbound $ \b1 ->
      eval g onUnbound $ \b2 ->
      onResult (b1 `op` b2)

solve :: forall r. Formula -> (Maybe [(String, Bool)] -> r) -> r
solve formula k =
  eval formula onUnbound onResult
  where
    onResult False = k Nothing
    onResult True  = k (Just [])

    onUnbound :: String -> r
    onUnbound v =
      subst v True formula $ \f' ->
      solve f' $ \case
        Just assigns -> k (Just ((v, True) : assigns))
        Nothing ->
          subst v False formula $ \f'' ->
          solve f'' $ \case
            Just assigns -> k (Just ((v, False) : assigns))
            Nothing      -> k Nothing

printSAT :: Maybe [(String, Bool)] -> IO ()
printSAT Nothing        = putStrLn "UNSAT."
printSAT (Just assigns) = do
  putStrLn "SAT:"
  case assigns of
    [] -> putStrLn "(no variables)"
    _ -> forM_ assigns $ \(var, val) ->
           putStrLn (var ++ " = " ++ show val)

-- Examples of use:
-- solve (Lit True :=>: "x" :&&: (Not "y")) printSAT
-- solve ("x" :=>: Lit False) printSAT
-- solve ("x" :&&: Not "x") printSAT

Thanks for sharing these fantastic examples of CPS!

Found a little counter-example of quicksort:
*Continuations> quicksort [1,2,0,7,4,5]
[0,4,1,5,2,7]

Suggested Fix: in line-46 and 47, replace "x" right after "partitionAt" with "pivot".

You're totally right, thanks!