{-# LANGUAGE InstanceSigs, PartialTypeSignatures #-}
module Nondeterminism
( PrologStream
, newThread
, takeNextSolution, takeAllRemainingSolutions
, reversibly, liftPrologIO
) where
import Primitives
import Classes
import Control.Applicative 
import Control.Monad
import Data.IORef

data Suspend a = Yield a | Suspend
data PrologStream a
  = PMaybe (Suspend a) (PrologIO (PrologStream a))
  | PNo

reversibly :: PrologIO a -> PrologIO () -> PrologStream a
reversibly forward backward = 
  PMaybe Suspend $ forward >>= \fwdValue -> return $
    PMaybe (Yield fwdValue) $ backward >>= \_ -> return $
      PNo

liftPrologIO :: PrologIO a -> PrologStream a
liftPrologIO forward = reversibly forward (return ())

instance LiftIO PrologStream where
  liftIO :: IO a -> PrologStream a
  liftIO x = liftPrologIO (liftIO x)

instance Functor Suspend where
  fmap f (Yield a) = Yield (f a)
  fmap _ Suspend = Suspend

instance Functor PrologStream where
  fmap :: (a -> b) -> PrologStream a -> PrologStream b
  fmap f (PMaybe a as) 
    = PMaybe (fmap f a) (fmap (fmap f) as)
  fmap _ PNo = PNo

instance Monad PrologStream where
  return :: a -> PrologStream a
  return = pure

  (>>=) :: PrologStream a -> (a -> PrologStream b) -> PrologStream b
  PNo >>= _ = PNo
  (>>=) (PMaybe (Yield x) xs) f = f x <|> PMaybe Suspend (fmap (>>= f) xs)
  (>>=) (PMaybe Suspend xs) f = PMaybe Suspend (fmap (>>= f) xs)

instance Applicative PrologStream where
  pure :: a -> PrologStream a
  pure x = PMaybe (Yield x) (pure PNo)

  (<*>) :: PrologStream (a -> b) -> PrologStream a -> PrologStream b
  fs <*> xs = do { f <- fs; x <- xs; return (f x) }

instance Alternative PrologStream where
  empty :: PrologStream a
  empty = PNo

  (<|>) :: PrologStream a -> PrologStream a -> PrologStream a
  (PMaybe x xsFuture) <|> next = PMaybe x merged
    where
    merged = do { xs <- xsFuture; return (xs <|> next) }
  PNo <|> next = next

instance MonadPlus PrologStream where
  mzero :: PrologStream a
  mzero = empty

  mplus :: PrologStream a -> PrologStream a -> PrologStream a
  mplus = (<|>)

squeeze :: PrologStream a -> PrologIO (Maybe a, PrologStream a)
squeeze PNo = return (Nothing, PNo)
squeeze (PMaybe Suspend rest) = do { v <- rest; squeeze v }
squeeze (PMaybe (Yield x) rest) 
  = return (Just x, PMaybe Suspend rest)

data Thread a = Thread
  { threadGlobals :: Globals
  , threadStream :: IORef (PrologStream a)
  }

newThread :: PrologStream a -> IO (Thread a)
newThread stream = do
  globals <- newGlobals
  streamRef <- newIORef stream
  return $ Thread
    { threadGlobals = globals
    , threadStream = streamRef
    }

takeNextSolution :: Thread a -> IO (Maybe a)
takeNextSolution thread = do
  stream <- readIORef (threadStream thread)
  (value, newStream) <- runPrologIO (squeeze stream) (threadGlobals thread)
  writeIORef (threadStream thread) newStream
  return value

takeAllRemainingSolutions :: Thread a -> IO [a]
takeAllRemainingSolutions thread = do
  value <- takeNextSolution thread
  case value of
    Nothing -> return []
    Just x -> fmap (x:) (takeAllRemainingSolutions thread)