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