Minizinc.hs

{-# LANGUAGE DeriveFunctor #-}
{-# LANGUAGE LambdaCase #-}
{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE ScopedTypeVariables #-}

{- | A set of types and functions to help calling Minizinc as an external binary.

Current strategy is to use JSON encode/decoding for passing in inputs and
reading outputs.
At this time, only a primitive output parser is supported.
-}
module Process.Minizinc (
    MiniZinc (..),
    simpleMiniZinc,
    withArgs,
    Solver (..),
    SolverName,
    MilliSeconds,
    SearchState (..),
    result,
    ResultHandler (..),
    runMinizincJSON,
    collectResults,
    keepLast,
    runLastMinizincJSON,
    cleanTmpFile,
)
where

import Control.Applicative ((<|>))
import Control.Exception (bracket)
import Data.Aeson (FromJSON, Result (..), ToJSON, Value, encode, fromJSON)
import Data.Aeson.Parser.Internal (json')
import Data.Attoparsec.ByteString (IResult (..), Parser, parse)
import Data.Attoparsec.Combinator (try)
import Data.ByteString (ByteString)
import qualified Data.ByteString as ByteString
import qualified Data.ByteString.Char8 as C8
import qualified Data.ByteString.Lazy as LByteString
import Data.Hashable (Hashable, hash)
import GHC.IO.Handle (Handle, hIsEOF)
import System.Directory (removeFile)
import System.Process (StdStream (CreatePipe), cleanupProcess, createProcess, proc, std_out)

-- | Type alias asking for milliseconds.
type MilliSeconds a = Int

-- | Name of a solver to be passed to the minizinc binary.
type SolverName = String

-- | Supported solvers or 'Other'.
data Solver = Chuffed | COIN_BC | CPLEX | Gecode | Gurobi | SCIP | Xpress | Other SolverName
    deriving (Show)

-- | An object helping to run MiniZinc.
data MiniZinc input answer
    = MiniZinc
    { model :: FilePath
    -- ^ a file path to a model
    , mkTmpDataPath :: input -> FilePath
    -- ^ a file path to hold, must be writable and readable
    , mkTimeLimit :: input -> MilliSeconds Int
    -- ^ a timelimit in seconds (an Int)
    , mkSolver :: input -> Solver
    -- ^ the solver to use (see `minizinc --solvers`)
    , mkExtraArgs :: input -> [String]
    -- ^ other arguments that get appended before the path to the model and data paths
    }

-- | Removes the temporary data file created as input before running minizinc.
cleanTmpFile :: MiniZinc input a -> input -> IO ()
cleanTmpFile mzn = removeFile . mkTmpDataPath mzn

-- | A constructor for MiniZinc object for simple situations.
simpleMiniZinc ::
    (Hashable input) =>
    FilePath ->
    MilliSeconds Int ->
    Solver ->
    MiniZinc input answer
simpleMiniZinc path timeout solver =
    MiniZinc
        path
        (\obj -> "minizinc-process-" ++ show (hash obj) ++ ".json")
        (const timeout)
        (const solver)
        (const [])

-- | Helper to set arguments.
withArgs :: [String] -> MiniZinc input answer -> MiniZinc input answer
withArgs args mzn = mzn{mkExtraArgs = const args}

-- | Runs MiniZinc on the input and parses output for the last answer.
runLastMinizincJSON ::
    (ToJSON input, FromJSON answer) =>
    MiniZinc input answer ->
    input ->
    IO (Maybe answer)
runLastMinizincJSON minizinc obj = do
    fmap adapt $ runMinizincJSON minizinc obj Nothing keepLast
  where
    adapt :: Maybe (SearchState a) -> Maybe a
    adapt x = x >>= result

data SearchState a
    = Exhausted a
    | Incomplete a
    | Unsatisfiable
    | InternalError String
    deriving (Show, Eq, Ord, Functor)

result :: SearchState a -> Maybe a
result (Incomplete a) = Just a
result (Exhausted a) = Just a
result _ = Nothing

reduce :: (FromJSON a) => SearchState Value -> SearchState a
reduce Unsatisfiable = Unsatisfiable
reduce (InternalError s) = InternalError s
reduce (Exhausted val) = case fromJSON val of
    Success obj -> Exhausted obj
    Error err -> InternalError err
reduce (Incomplete val) = case fromJSON val of
    Success obj -> Incomplete obj
    Error err -> InternalError err

data ResultHandler obj b
    = ResultHandler
    { handleNext :: b -> SearchState obj -> IO (b, Maybe (ResultHandler obj b))
    }

{- | Collect all results in memory.
The resulting list is in reverse order (best are first elements in case of optimizations).
-}
collectResults :: ResultHandler obj [SearchState obj]
collectResults = ResultHandler go
  where
    go xs x = seq x $ pure (x : xs, Just $ collectResults)

-- | Keep only the latest result in memory.
keepLast :: ResultHandler obj (Maybe (SearchState obj))
keepLast = ResultHandler go
  where
    go _ x = seq x $ pure (Just x, Just $ keepLast)

runMinizincJSON ::
    forall input answer b.
    (ToJSON input, FromJSON answer) =>
    MiniZinc input answer ->
    input ->
    b ->
    ResultHandler answer b ->
    IO b
runMinizincJSON minizinc obj v0 resultHandler = do
    LByteString.writeFile fullPath $ encode obj
    let start = createProcess (proc "minizinc" args){std_out = CreatePipe}
    bracket
        start
        cleanupProcess
        (\(_, Just out, _, _) -> go out (parse oneResult) "" v0 resultHandler)
  where
    go ::
        Handle ->
        (ByteString -> IResult ByteString (SearchState Value)) ->
        ByteString ->
        b ->
        ResultHandler answer b ->
        IO b
    go out parsebuf buf v1 handler = do
        case parsebuf buf of
            Done remainderBuf stateVal -> do
                (v2, nextHandler) <- (handleNext handler) v1 (reduce stateVal)
                case nextHandler of
                    Nothing -> userFinished v2
                    Just resultHandler -> go out (parse oneResult) remainderBuf v2 resultHandler
            Fail _ _ err -> do
                (v2, _) <- (handleNext handler) v1 (InternalError err)
                finalizeFailure v2
            Partial f -> do
                -- NOTE: on End-Of-File we attempt to feed an empty byte to
                -- conclude the parse.
                -- Thus the real "final marked" is when the Handle is at EOF and
                -- the input chunk is empty.
                -- A more explicit final-result marker could benefit the
                -- implementation of ResultHandler.
                eof <- hIsEOF out
                case (eof, ByteString.null buf) of
                    (True, True) -> inputFinished v1
                    (True, False) -> go out f "" v1 handler
                    _ -> do
                        dat <- C8.hGetLine out
                        go out f dat v1 handler

    inputFinished = pure
    userFinished = pure
    finalizeFailure = pure

    fullPath :: FilePath
    fullPath = mkTmpDataPath minizinc obj
    args :: [String]
    args =
        [ "--time-limit"
        , show (mkTimeLimit minizinc obj)
        , "--solver"
        , showSolver (mkSolver minizinc obj)
        , "--output-mode"
        , "json"
        ]
            ++ (mkExtraArgs minizinc obj)
            ++ [ model minizinc
               , fullPath
               ]

showSolver :: Solver -> String
showSolver = \case
    Chuffed -> "Chuffed"
    COIN_BC -> "COIN-BC"
    CPLEX -> "CPLEX"
    Gecode -> "Gecode"
    Gurobi -> "Gurobi"
    SCIP -> "SCIP"
    Xpress -> "Xpress"
    Other n -> n

{- | NOTE: the parser is fed with hGetLine (stripping EOL markers)
this assumption simplifies the grammar below
-}
oneResult :: Parser (SearchState Value)
oneResult =
    try unsat
        <|> sat
  where
    unsat = unsatMark *> pure Unsatisfiable

    sat = reverseAp <$> json' <*> searchstate
    reverseAp val constructor = constructor val

    searchstate =
        try (resultMark *> exhaustiveMark *> pure Exhausted)
            <|> try (resultMark *> pure Incomplete)

    resultMark = "----------"
    exhaustiveMark = "=========="
    unsatMark = "=====UNSATISFIABLE====="