Add support for binary/integer contraints;

Change to hash maps to speed up variable LUT for large number of
variables.

1 files changed, 14 insertions, 3 deletions

diff --git a/Math/LinProg/Types.hs b/Math/LinProg/Types.hs
index 2a81918..4819dd3 100644
--- a/Math/LinProg/Types.hs
+++ b/Math/LinProg/Types.hs
@@ -26,14 +26,17 @@ module Math.LinProg.Types (
   ,(<:)
   ,(=:)
   ,(>:)
+  ,bin
+  ,int
 ) where
 
 import Data.Functor.Foldable
 import Control.Monad.Free
-import qualified Data.Map.Strict as M
+import qualified Data.HashMap.Strict as M
 import Test.QuickCheck
 import Control.Applicative
 import Data.List
+import Data.Hashable
 
 -- | Base AST for expressions.  Expressions have factors or type t and
 -- variables referenced by ids of type v.
@@ -117,7 +120,7 @@ rewrite = cata rewrite' where
   rewrite' a = Fix a
 
 -- | Reduces an expression to the variable terms
-varTerms :: (Num t, Eq t, Ord v) => LinExpr t v -> [(v, t)]
+varTerms :: (Num t, Eq t, Hashable v, Eq v) => LinExpr t v -> [(v, t)]
 varTerms = M.toList . cata go . rewrite where
   go (Wvar w a) = M.fromList [(a, w)]
   go (Add a b) = M.unionWith (+) a b
@@ -141,6 +144,8 @@ prettyPrint = cata prettyPrint' where
 -- in the data type).
 data LinProg' t v a =
   Objective !(LinExpr t v) !a
+  | Integer !v !a
+  | Binary !v !a
   | EqConstraint !(LinExpr t v) !(LinExpr t v) !a
   | LeqConstraint !(LinExpr t v) !(LinExpr t v) !a
   deriving (Show, Eq, Functor)
@@ -155,10 +160,16 @@ a =: b = liftF (EqConstraint a b ())
 
 -- | Define an inequality (less than equal) contraint
 a <: b = liftF (LeqConstraint a b ())
---
+
 -- | Define an inequality (greater than equal) contraint
 b >: a = liftF (LeqConstraint a b ())
 
+-- | Declare a variable to be binary
+bin (Fix (Var v)) = liftF (Binary v ())
+
+-- | Declare a variable to be integral
+int (Fix (Var v)) = liftF (Integer v ())
+
 infix 4 =:
 infix 4 <:
 infix 4 >: