{-# OPTIONS_GHC -Wall #-}
-----------------------------------------------------------------------------
-- |
-- Module      :  ToySolver.Arith.Simplex.Simple
-- Copyright   :  (c) Masahiro Sakai 2016
-- License     :  BSD-style
--
-- Maintainer  :  masahiro.sakai@gmail.com
-- Stability   :  provisional
-- Portability :  portable
--
-----------------------------------------------------------------------------
module ToySolver.Arith.Simplex.Simple
  ( Model
  , OptDir (..)
  , OptResult (..)
  , check
  , optimize
  ) where

import Control.Monad
import Control.Monad.ST
import Data.Default.Class
import qualified Data.IntMap.Strict as IntMap
import qualified Data.IntSet as IntSet
import qualified ToySolver.Data.LA as LA
import ToySolver.Data.IntVar hiding (Model)
import qualified ToySolver.Arith.Simplex as Simplex
import ToySolver.Arith.Simplex hiding (check, optimize)

check :: VarSet -> [LA.Atom Rational] -> Maybe Model
check :: VarSet -> [Atom Rational] -> Maybe Model
check VarSet
vs [Atom Rational]
as = (forall s. ST s (Maybe Model)) -> Maybe Model
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (Maybe Model)) -> Maybe Model)
-> (forall s. ST s (Maybe Model)) -> Maybe Model
forall a b. (a -> b) -> a -> b
$ do
  GenericSolverM (ST s) (Delta Rational)
solver <- ST s (GenericSolverM (ST s) (Delta Rational))
forall (m :: * -> *) v.
(PrimMonad m, SolverValue v) =>
m (GenericSolverM m v)
Simplex.newSolver
  IntMap Var
s <- ([(Var, Var)] -> IntMap Var)
-> ST s [(Var, Var)] -> ST s (IntMap Var)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM [(Var, Var)] -> IntMap Var
forall a. [(Var, a)] -> IntMap a
IntMap.fromAscList (ST s [(Var, Var)] -> ST s (IntMap Var))
-> ST s [(Var, Var)] -> ST s (IntMap Var)
forall a b. (a -> b) -> a -> b
$ [Var] -> (Var -> ST s (Var, Var)) -> ST s [(Var, Var)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (VarSet -> [Var]
IntSet.toAscList VarSet
vs) ((Var -> ST s (Var, Var)) -> ST s [(Var, Var)])
-> (Var -> ST s (Var, Var)) -> ST s [(Var, Var)]
forall a b. (a -> b) -> a -> b
$ \Var
v -> do
    Var
v2 <- GenericSolverM (ST s) (Delta Rational) -> ST s Var
forall (m :: * -> *) v.
(PrimMonad m, SolverValue v) =>
GenericSolverM m v -> m Var
Simplex.newVar GenericSolverM (ST s) (Delta Rational)
solver
    (Var, Var) -> ST s (Var, Var)
forall (m :: * -> *) a. Monad m => a -> m a
return (Var
v, Var
v2)
  let s' :: IntMap (Expr Rational)
s' = (Var -> Expr Rational) -> IntMap Var -> IntMap (Expr Rational)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Var -> Expr Rational
forall r. Num r => Var -> Expr r
LA.var IntMap Var
s
      mtrans :: IntMap b -> IntMap b
mtrans IntMap b
m = (Var -> b) -> IntMap Var -> IntMap b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (IntMap b
m IntMap b -> Var -> b
forall a. IntMap a -> Var -> a
IntMap.!) IntMap Var
s
  [Atom Rational] -> (Atom Rational -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Atom Rational]
as ((Atom Rational -> ST s ()) -> ST s ())
-> (Atom Rational -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Atom Rational
a -> do
    GenericSolverM (ST s) (Delta Rational) -> Atom Rational -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
GenericSolverM m (Delta Rational) -> Atom Rational -> m ()
Simplex.assertAtomEx GenericSolverM (ST s) (Delta Rational)
solver ((Expr Rational -> Expr Rational) -> Atom Rational -> Atom Rational
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (IntMap (Expr Rational) -> Expr Rational -> Expr Rational
forall r. (Num r, Eq r) => VarMap (Expr r) -> Expr r -> Expr r
LA.applySubst IntMap (Expr Rational)
s') Atom Rational
a)
  Bool
ret <- GenericSolverM (ST s) (Delta Rational) -> ST s Bool
forall (m :: * -> *) v.
(PrimMonad m, SolverValue v) =>
GenericSolverM m v -> m Bool
Simplex.check GenericSolverM (ST s) (Delta Rational)
solver
  if Bool
ret then do
    Model
m <- GenericSolverM (ST s) (Delta Rational) -> ST s Model
forall v (m :: * -> *).
(SolverValue v, PrimMonad m) =>
GenericSolverM m v -> m Model
Simplex.getModel GenericSolverM (ST s) (Delta Rational)
solver
    Maybe Model -> ST s (Maybe Model)
forall (m :: * -> *) a. Monad m => a -> m a
return (Maybe Model -> ST s (Maybe Model))
-> Maybe Model -> ST s (Maybe Model)
forall a b. (a -> b) -> a -> b
$ Model -> Maybe Model
forall a. a -> Maybe a
Just (Model -> Maybe Model) -> Model -> Maybe Model
forall a b. (a -> b) -> a -> b
$ Model -> Model
forall b. IntMap b -> IntMap b
mtrans Model
m
  else
    Maybe Model -> ST s (Maybe Model)
forall (m :: * -> *) a. Monad m => a -> m a
return Maybe Model
forall a. Maybe a
Nothing

optimize :: VarSet -> OptDir -> LA.Expr Rational -> [LA.Atom Rational] -> (OptResult, Maybe Model)
optimize :: VarSet
-> OptDir
-> Expr Rational
-> [Atom Rational]
-> (OptResult, Maybe Model)
optimize VarSet
vs OptDir
dir Expr Rational
obj [Atom Rational]
as = (forall s. ST s (OptResult, Maybe Model))
-> (OptResult, Maybe Model)
forall a. (forall s. ST s a) -> a
runST ((forall s. ST s (OptResult, Maybe Model))
 -> (OptResult, Maybe Model))
-> (forall s. ST s (OptResult, Maybe Model))
-> (OptResult, Maybe Model)
forall a b. (a -> b) -> a -> b
$ do
  GenericSolverM (ST s) Rational
solver <- ST s (GenericSolverM (ST s) Rational)
forall (m :: * -> *) v.
(PrimMonad m, SolverValue v) =>
m (GenericSolverM m v)
Simplex.newSolver
  IntMap Var
s <- ([(Var, Var)] -> IntMap Var)
-> ST s [(Var, Var)] -> ST s (IntMap Var)
forall (m :: * -> *) a1 r. Monad m => (a1 -> r) -> m a1 -> m r
liftM [(Var, Var)] -> IntMap Var
forall a. [(Var, a)] -> IntMap a
IntMap.fromAscList (ST s [(Var, Var)] -> ST s (IntMap Var))
-> ST s [(Var, Var)] -> ST s (IntMap Var)
forall a b. (a -> b) -> a -> b
$ [Var] -> (Var -> ST s (Var, Var)) -> ST s [(Var, Var)]
forall (t :: * -> *) (m :: * -> *) a b.
(Traversable t, Monad m) =>
t a -> (a -> m b) -> m (t b)
forM (VarSet -> [Var]
IntSet.toAscList VarSet
vs) ((Var -> ST s (Var, Var)) -> ST s [(Var, Var)])
-> (Var -> ST s (Var, Var)) -> ST s [(Var, Var)]
forall a b. (a -> b) -> a -> b
$ \Var
v -> do
    Var
v2 <- GenericSolverM (ST s) Rational -> ST s Var
forall (m :: * -> *) v.
(PrimMonad m, SolverValue v) =>
GenericSolverM m v -> m Var
Simplex.newVar GenericSolverM (ST s) Rational
solver
    (Var, Var) -> ST s (Var, Var)
forall (m :: * -> *) a. Monad m => a -> m a
return (Var
v, Var
v2)
  let s' :: IntMap (Expr Rational)
s' = (Var -> Expr Rational) -> IntMap Var -> IntMap (Expr Rational)
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap Var -> Expr Rational
forall r. Num r => Var -> Expr r
LA.var IntMap Var
s
      mtrans :: IntMap b -> IntMap b
mtrans IntMap b
m = (Var -> b) -> IntMap Var -> IntMap b
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (IntMap b
m IntMap b -> Var -> b
forall a. IntMap a -> Var -> a
IntMap.!) IntMap Var
s
  [Atom Rational] -> (Atom Rational -> ST s ()) -> ST s ()
forall (t :: * -> *) (m :: * -> *) a b.
(Foldable t, Monad m) =>
t a -> (a -> m b) -> m ()
forM_ [Atom Rational]
as ((Atom Rational -> ST s ()) -> ST s ())
-> (Atom Rational -> ST s ()) -> ST s ()
forall a b. (a -> b) -> a -> b
$ \Atom Rational
a -> do
    GenericSolverM (ST s) Rational -> Atom Rational -> ST s ()
forall (m :: * -> *).
PrimMonad m =>
GenericSolverM m Rational -> Atom Rational -> m ()
assertAtom GenericSolverM (ST s) Rational
solver ((Expr Rational -> Expr Rational) -> Atom Rational -> Atom Rational
forall (f :: * -> *) a b. Functor f => (a -> b) -> f a -> f b
fmap (IntMap (Expr Rational) -> Expr Rational -> Expr Rational
forall r. (Num r, Eq r) => VarMap (Expr r) -> Expr r -> Expr r
LA.applySubst IntMap (Expr Rational)
s') Atom Rational
a)
  GenericSolverM (ST s) Rational -> OptDir -> ST s ()
forall (m :: * -> *) v.
PrimMonad m =>
GenericSolverM m v -> OptDir -> m ()
Simplex.setOptDir GenericSolverM (ST s) Rational
solver OptDir
dir
  GenericSolverM (ST s) Rational -> Expr Rational -> ST s ()
forall (m :: * -> *) v.
(PrimMonad m, SolverValue v) =>
GenericSolverM m v -> Expr Rational -> m ()
Simplex.setObj GenericSolverM (ST s) Rational
solver Expr Rational
obj
  OptResult
ret <- GenericSolverM (ST s) Rational -> Options -> ST s OptResult
forall (m :: * -> *).
PrimMonad m =>
GenericSolverM m Rational -> Options -> m OptResult
Simplex.optimize GenericSolverM (ST s) Rational
solver Options
forall a. Default a => a
def
  case OptResult
ret of
    OptResult
Optimum -> do
      Model
m <- GenericSolverM (ST s) Rational -> ST s Model
forall v (m :: * -> *).
(SolverValue v, PrimMonad m) =>
GenericSolverM m v -> m Model
Simplex.getModel GenericSolverM (ST s) Rational
solver
      (OptResult, Maybe Model) -> ST s (OptResult, Maybe Model)
forall (m :: * -> *) a. Monad m => a -> m a
return ((OptResult, Maybe Model) -> ST s (OptResult, Maybe Model))
-> (OptResult, Maybe Model) -> ST s (OptResult, Maybe Model)
forall a b. (a -> b) -> a -> b
$ (OptResult
ret, Model -> Maybe Model
forall a. a -> Maybe a
Just (Model -> Model
forall b. IntMap b -> IntMap b
mtrans Model
m))
    OptResult
Unsat -> do
      (OptResult, Maybe Model) -> ST s (OptResult, Maybe Model)
forall (m :: * -> *) a. Monad m => a -> m a
return ((OptResult, Maybe Model) -> ST s (OptResult, Maybe Model))
-> (OptResult, Maybe Model) -> ST s (OptResult, Maybe Model)
forall a b. (a -> b) -> a -> b
$ (OptResult
ret, Maybe Model
forall a. Maybe a
Nothing)
    OptResult
Unbounded -> do
      Model
m <- GenericSolverM (ST s) Rational -> ST s Model
forall v (m :: * -> *).
(SolverValue v, PrimMonad m) =>
GenericSolverM m v -> m Model
Simplex.getModel GenericSolverM (ST s) Rational
solver
      (OptResult, Maybe Model) -> ST s (OptResult, Maybe Model)
forall (m :: * -> *) a. Monad m => a -> m a
return ((OptResult, Maybe Model) -> ST s (OptResult, Maybe Model))
-> (OptResult, Maybe Model) -> ST s (OptResult, Maybe Model)
forall a b. (a -> b) -> a -> b
$ (OptResult
ret, Model -> Maybe Model
forall a. a -> Maybe a
Just (Model -> Model
forall b. IntMap b -> IntMap b
mtrans Model
m))
    OptResult
ObjLimit -> do
      [Char] -> ST s (OptResult, Maybe Model)
forall a. HasCallStack => [Char] -> a
error [Char]
"should not happen"