{-# LANGUAGE InstanceSigs #-}
module Nondeterminism
( PrologStream
, newThread
, takeNextSolution, takeAllRemainingSolutions
, reversibly
, liftPrologIO
) where
import Primitives (Globals, PrologIO(..), newGlobals)
import Control.Applicative
import Control.Monad
import Classes (LiftIO (..))
import GHC.IORef

data Suspend a = Yield a | Suspend

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

data PrologStream a 
  = PYes (Suspend a) (PrologIO (PrologStream a))
  | PNo

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

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

  (<*>) :: PrologStream (a -> b) -> PrologStream a -> PrologStream b
  f <*> x = f >>= (\f' -> x >>= (\x' -> return (f' x')))

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

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

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

  (<|>) :: PrologStream a -> PrologStream a -> PrologStream a
  (PYes x xs) <|> next = PYes x (liftM2 (<|>) xs (return next))
  PNo <|> next = next

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

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

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 { threadGlobals = globals, threadStream = streamRef }) = do
  stream <- readIORef streamRef
  (value, newStream) <- runPrologIO (squeeze stream) globals
  writeIORef streamRef 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)

squeeze :: PrologStream a -> PrologIO (Maybe a, PrologStream a)
squeeze PNo = return (Nothing, PNo)
squeeze (PYes Suspend rest) = rest >>= squeeze
squeeze (PYes (Yield x) rest) = fmap ((,) (Just x)) rest

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

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

instance LiftIO PrologStream where
  liftIO x = liftPrologIO (liftIO x)