Amb

module Amb (AmbT, Amb, amb, cut, runAmbT, runAmb) where
 
import Control.Monad.Cont
import Control.Monad.State
import Control.Monad.Identity
 
type Point r s m = () -> AmbT r s m s
newtype AmbT r s m a = AmbT { unAmbT :: StateT [Point r s m] (ContT r m) a }
type Amb r s = AmbT r s Identity
 
instance MonadTrans (AmbT r s) where
    lift = AmbT . lift . lift
 
instance (Monad m) => Monad (AmbT r s m) where
    AmbT a >>= b = AmbT $ a >>= unAmbT . b
    return = AmbT . return
 
backtrack :: (Monad m) => AmbT r s m a
backtrack = do xss <- AmbT get
               case xss of
                 [] -> fail "amb tree exhausted"
                 (f:xs) -> do AmbT $ put xs
                              f ()
                              return undefined
 
addPoint :: (Monad m) => Point r s m -> AmbT r s m ()
addPoint x = AmbT $ modify (x:)
 
amb :: (Monad m) => [a] -> AmbT r s m a
amb []     = backtrack
amb (x:xs) = ambCC $ \exit -> do
               ambCC $ \k -> addPoint k >> exit x
               amb xs
    where ambCC f = AmbT $ callCC $ \k -> unAmbT $ f $ AmbT . k
 
cut :: (Monad m) => AmbT r s m ()
cut = AmbT $ put []
 
runAmbT :: (Monad m) => AmbT r s m r -> m r
runAmbT (AmbT a) = runContT (evalStateT a []) return
 
runAmb :: Amb r s r -> r
runAmb = runIdentity . runAmbT

example :: Amb r Integer (Integer,Integer)
example = do x <- amb [1,2,3]
             y <- amb [4,5,6]
             if x*y == 8
               then return (x,y)
               else amb []
 
factor :: Integer -> Amb r Integer (Integer,Integer)
factor a = do x <- amb [2..]
              y <- amb [2..x]
              if x*y == a
                then return (x,y)
                else amb []
 
factorIO :: Integer -> AmbT r Integer IO (Integer,Integer)
factorIO a = do lift $ putStrLn $ "Factoring " ++ show a
                x <- amb [2..]
                y <- amb [2..x]
                lift $ putStrLn $ "Trying " ++ show x ++ " and " ++ show y
                if x*y == a
                  then do lift $ putStrLn "Found it!"
                          return (x,y)
                  else do lift $ putStrLn $ "Nope (" ++ show (x*y) ++ ")"
                          amb []