{-# LANGUAGE BangPatterns #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE PatternGuards #-}
{-# LANGUAGE Safe #-}
{- |
This module implements a decision procedure for quantifier-free linear
arithmetic. The algorithm is based on the following paper:
An Online Proof-Producing Decision Procedure for
Mixed-Integer Linear Arithmetic
by
Sergey Berezin, Vijay Ganesh, and David L. Dill
-}
module Data.Integer.SAT
( PropSet,
noProps,
checkSat,
assert,
Prop (..),
Expr (..),
BoundType (..),
getExprBound,
getExprRange,
Name,
toName,
fromName,
-- * Iterators
allSolutions,
slnCurrent,
slnNextVal,
slnNextVar,
slnEnumerate,
-- * Debug
dotPropSet,
sizePropSet,
allInerts,
ppInerts,
-- * For QuickCheck
iPickBounded,
Bound (..),
tConst,
)
where
import Control.Applicative (Alternative (..), Applicative (..), (<$>))
import Control.Monad (MonadPlus (..), ap, guard, liftM)
import Data.List (partition)
import Data.Map (Map)
import qualified Data.Map as Map
import Data.Maybe (fromMaybe, mapMaybe, maybeToList)
import Text.PrettyPrint
import Prelude hiding ((<>))
infixr 2 :||
infixr 3 :&&
infix 4 :==, :/=, :<, :<=, :>, :>=
infixl 6 :+, :-
infixl 7 :*
--------------------------------------------------------------------------------
-- Solver interface
-- | A collection of propositions.
newtype PropSet = State (Answer RW)
deriving (Show)
dotPropSet :: PropSet -> Doc
dotPropSet (State a) = dotAnswer (ppInerts . inerts) a
sizePropSet :: PropSet -> (Integer, Integer, Integer)
sizePropSet (State a) = answerSize a
-- | An empty collection of propositions.
noProps :: PropSet
noProps = State $ return initRW
-- | Add a new proposition to an existing collection.
assert :: Prop -> PropSet -> PropSet
assert p (State rws) = State $ fmap snd $ m =<< rws
where
S m = prop p
{- | Extract a model from a consistent set of propositions.
Returns 'Nothing' if the assertions have no model.
If a variable does not appear in the assignment, then it is 0 (?).
-}
checkSat :: PropSet -> Maybe [(Int, Integer)]
checkSat (State m) = go m
where
go None = mzero
go (One rw) = return [(x, v) | (UserName x, v) <- iModel (inerts rw)]
go (Choice m1 m2) = mplus (go m1) (go m2)
allInerts :: PropSet -> [Inerts]
allInerts (State m) = map inerts (toList m)
allSolutions :: PropSet -> [Solutions]
allSolutions = map startIter . allInerts
{- | Computes bounds on the expression that are compatible with the model.
Returns `Nothing` if the bound is not known.
-}
getExprBound :: BoundType -> Expr -> PropSet -> Maybe Integer
getExprBound bt e (State s) =
do
let S m = expr e
check (t, s1) = iTermBound bt t (inerts s1)
bs <- mapM check $ toList $ s >>= m
case bs of
[] -> Nothing
_ -> Just (maximum bs)
{- | Compute the range of possible values for an expression.
Returns `Nothing` if the bound is not known.
-}
getExprRange :: Expr -> PropSet -> Maybe [Integer]
getExprRange e (State s) =
do
let S m = expr e
check (t, s1) = do
l <- iTermBound Lower t (inerts s1)
u <- iTermBound Upper t (inerts s1)
return (l, u)
bs <- mapM check $ toList $ s >>= m
case bs of
[] -> Nothing
_ ->
let (ls, us) = unzip bs
in Just [x | x <- [minimum ls .. maximum us]]
-- | The type of proposition.
data Prop
= PTrue
| PFalse
| Prop :|| Prop
| Prop :&& Prop
| Not Prop
| Expr :== Expr
| Expr :/= Expr
| Expr :< Expr
| Expr :> Expr
| Expr :<= Expr
| Expr :>= Expr
deriving (Read, Show, Ord, Eq)
{- | The type of integer expressions.
Variable names must be non-negative.
-}
data Expr
= -- | Addition
Expr :+ Expr
| -- | Subtraction
Expr :- Expr
| -- | Multiplication by a constant
Integer :* Expr
| -- | Negation
Negate Expr
| -- | Variable
Var Name
| -- | Constant
K Integer
| -- | A conditional expression
If Prop Expr Expr
| -- | Division, rounds down
Div Expr Integer
| -- | Non-negative remainder
Mod Expr Integer
| -- | Minimum of two arguments
Min Expr Expr
| -- | Maximum of two arguments
Max Expr Expr
deriving (Read, Show, Ord, Eq)
prop :: Prop -> S ()
prop PTrue = return ()
prop PFalse = mzero
prop (p1 :|| p2) = prop p1 `mplus` prop p2
prop (p1 :&& p2) = prop p1 >> prop p2
prop (Not p) = prop (neg p)
where
neg PTrue = PFalse
neg PFalse = PTrue
neg (p1 :&& p2) = neg p1 :|| neg p2
neg (p1 :|| p2) = neg p1 :&& neg p2
neg (Not q) = q
neg (e1 :== e2) = e1 :/= e2
neg (e1 :/= e2) = e1 :== e2
neg (e1 :< e2) = e1 :>= e2
neg (e1 :<= e2) = e1 :> e2
neg (e1 :> e2) = e1 :<= e2
neg (e1 :>= e2) = e1 :< e2
prop (e1 :== e2) = do
t1 <- expr e1
t2 <- expr e2
solveIs0 (t1 |-| t2)
prop (e1 :/= e2) = do
t1 <- expr e1
t2 <- expr e2
let t = t1 |-| t2
solveIsNeg t `orElse` solveIsNeg (tNeg t)
prop (e1 :< e2) = do
t1 <- expr e1
t2 <- expr e2
solveIsNeg (t1 |-| t2)
prop (e1 :<= e2) = do
t1 <- expr e1
t2 <- expr e2
let t = t1 |-| t2 |-| tConst 1
solveIsNeg t
prop (e1 :> e2) = prop (e2 :< e1)
prop (e1 :>= e2) = prop (e2 :<= e1)
expr :: Expr -> S Term
expr (e1 :+ e2) = (|+|) <$> expr e1 <*> expr e2
expr (e1 :- e2) = (|-|) <$> expr e1 <*> expr e2
expr (k :* e2) = (k |*|) <$> expr e2
expr (Negate e) = tNeg <$> expr e
expr (Var x) = pure (tVar x)
expr (K x) = pure (tConst x)
expr (Min a b) = expr (If (a :<= b) a b)
expr (Max a b) = expr (If (a :<= b) b a)
expr (If p e1 e2) = do
x <- newVar
prop (p :&& Var x :== e1 :|| Not p :&& Var x :== e2)
return (tVar x)
expr (Div e k) = fmap fst $ exprDivMod e k
expr (Mod e k) = fmap snd $ exprDivMod e k
exprDivMod :: Expr -> Integer -> S (Term, Term)
exprDivMod e k =
do
guard (k /= 0) -- Always unsat
q <- newVar
r <- newVar
let er = Var r
prop (k :* Var q :+ er :== e :&& er :< K k :&& K 0 :<= er)
return (tVar q, tVar r)
--------------------------------------------------------------------------------
data RW = RW
{ nameSource :: !Int
, inerts :: Inerts
}
deriving (Show)
initRW :: RW
initRW = RW {nameSource = 0, inerts = iNone}
--------------------------------------------------------------------------------
-- Constraints and Bound on Variables
ctLt :: Term -> Term -> Term
ctLt t1 t2 = t1 |-| t2
ctEq :: Term -> Term -> Term
ctEq t1 t2 = t1 |-| t2
data Bound
= -- | The integer is strictly positive
Bound Integer Term
deriving (Show)
data BoundType = Lower | Upper
deriving (Show)
toCt :: BoundType -> Name -> Bound -> Term
toCt Lower x (Bound c t) = ctLt t (c |*| tVar x)
toCt Upper x (Bound c t) = ctLt (c |*| tVar x) t
--------------------------------------------------------------------------------
-- Inert set
-- | The inert contains the solver state on one possible path.
data Inerts = Inerts
{ -- | Known lower and upper bounds for variables.
-- Each bound @(c,t)@ in the first list asserts that @t < c * x@
-- Each bound @(c,t)@ in the second list asserts that @c * x < t@
bounds :: NameMap ([Bound], [Bound])
, -- | Definitions for resolved variables.
-- These form an idempotent substitution.
solved :: NameMap Term
}
deriving (Show)
ppInerts :: Inerts -> Doc
ppInerts is =
vcat $
[ppLower x b | (x, (ls, _)) <- bnds, b <- ls]
++ [ppUpper x b | (x, (_, us)) <- bnds, b <- us]
++ [ppEq e | e <- Map.toList (solved is)]
where
bnds = Map.toList (bounds is)
ppT c x = ppTerm (c |*| tVar x)
ppLower x (Bound c t) = ppTerm t <+> text "<" <+> ppT c x
ppUpper x (Bound c t) = ppT c x <+> text "<" <+> ppTerm t
ppEq (x, t) = ppName x <+> text "=" <+> ppTerm t
-- | An empty inert set.
iNone :: Inerts
iNone =
Inerts
{ bounds = Map.empty
, solved = Map.empty
}
-- | Rewrite a term using the definitions from an inert set.
iApSubst :: Inerts -> Term -> Term
iApSubst i t = foldr apS t $ Map.toList $ solved i
where
apS (x, t1) t2 = tLet x t1 t2
{- | Add a definition. Upper and lower bound constraints that mention
the variable are "kicked-out" so that they can be reinserted in the
context of the new knowledge.
* Assumes substitution has already been applied.
* The kicked-out constraints are NOT rewritten, this happens
when they get inserted in the work queue.
-}
iSolved :: Name -> Term -> Inerts -> ([Term], Inerts)
iSolved x t i =
( kickedOut
, Inerts
{ bounds = otherBounds
, solved = Map.insert x t $ Map.map (tLet x t) $ solved i
}
)
where
(kickedOut, otherBounds) =
-- First, we eliminate all entries for `x`
let (mb, mp1) = Map.updateLookupWithKey (\_ _ -> Nothing) x (bounds i)
-- Next, we elminate all constraints that mentiond `x` in bounds
mp2 = Map.mapWithKey extractBounds mp1
in ( [ ct | (lbs, ubs) <- maybeToList mb, ct <- map (toCt Lower x) lbs ++ map (toCt Upper x) ubs
]
++ [ct | (_, cts) <- Map.elems mp2, ct <- cts]
, fmap fst mp2
)
extractBounds y (lbs, ubs) =
let (lbsStay, lbsKick) = partition stay lbs
(ubsStay, ubsKick) = partition stay ubs
in ( (lbsStay, ubsStay)
, map (toCt Lower y) lbsKick
++ map (toCt Upper y) ubsKick
)
stay (Bound _ bnd) = not (tHasVar x bnd)
{- | Given some lower and upper bounds, find the interval the satisfies them.
Note the upper and lower bounds are strict (i.e., < and >)
-}
boundInterval :: [Bound] -> [Bound] -> Maybe (Maybe Integer, Maybe Integer)
boundInterval lbs ubs =
do
ls <- mapM (normBound Lower) lbs
us <- mapM (normBound Upper) ubs
let lb = case ls of
[] -> Nothing
_ -> Just (maximum ls + 1)
ub = case us of
[] -> Nothing
_ -> Just (minimum us - 1)
case (lb, ub) of
(Just l, Just u) -> guard (l <= u)
_ -> return ()
return (lb, ub)
where
normBound Lower (Bound c t) = do
k <- isConst t
return (div (k + c - 1) c)
normBound Upper (Bound c t) = do
k <- isConst t
return (div k c)
data Solutions
= Done
| TopVar Name Integer (Maybe Integer) (Maybe Integer) Inerts
| FixedVar Name Integer Solutions
deriving (Show)
slnCurrent :: Solutions -> [(Int, Integer)]
slnCurrent s = [(x, v) | (UserName x, v) <- go s]
where
go Done = []
go (TopVar x v _ _ is) = (x, v) : iModel (iLet x v is)
go (FixedVar x v i) = (x, v) : go i
{- | Replace occurances of a variable with an integer.
WARNING: The integer should be a valid value for the variable.
-}
iLet :: Name -> Integer -> Inerts -> Inerts
iLet x v is =
Inerts
{ bounds = fmap updBs (bounds is)
, solved = fmap (tLetNum x v) (solved is)
}
where
updB (Bound c t) = Bound c (tLetNum x v t)
updBs (ls, us) = (map updB ls, map updB us)
startIter :: Inerts -> Solutions
startIter is =
case Map.maxViewWithKey (bounds is) of
Nothing ->
case Map.maxViewWithKey (solved is) of
Nothing -> Done
Just ((x, t), mp1) ->
case [y | y <- tVarList t] of
y : _ -> TopVar y 0 Nothing Nothing is
[] ->
let v = tConstPart t
in TopVar x v (Just v) (Just v) $ is {solved = mp1}
Just ((x, (lbs, ubs)), mp1) ->
case [y | Bound _ t <- lbs ++ ubs, y <- tVarList t] of
y : _ -> TopVar y 0 Nothing Nothing is
[] -> case boundInterval lbs ubs of
Nothing -> error "bug: cannot compute interval?"
Just (lb, ub) ->
let v = fromMaybe 0 (mplus lb ub)
in TopVar x v lb ub $ is {bounds = mp1}
slnEnumerate :: Solutions -> [Solutions]
slnEnumerate s0 = go s0 []
where
go s k = case slnNextVar s of
Nothing -> hor s k
Just s1 -> go s1 $ case slnNextVal s of
Nothing -> k
Just s2 -> go s2 k
hor s k =
s :
case slnNextVal s of
Nothing -> k
Just s1 -> hor s1 k
slnNextVal :: Solutions -> Maybe Solutions
slnNextVal Done = Nothing
slnNextVal (FixedVar x v i) = FixedVar x v `fmap` slnNextVal i
slnNextVal it@(TopVar _ _ lb _ _) =
case lb of
Just _ -> slnNextValWith (+ 1) it
Nothing -> slnNextValWith (subtract 1) it
slnNextValWith :: (Integer -> Integer) -> Solutions -> Maybe Solutions
slnNextValWith _ Done = Nothing
slnNextValWith f (FixedVar x v i) = FixedVar x v `fmap` slnNextValWith f i
slnNextValWith f (TopVar x v lb ub is) =
do
let v1 = f v
case lb of
Just l -> guard (l <= v1)
Nothing -> return ()
case ub of
Just u -> guard (v1 <= u)
Nothing -> return ()
return $ TopVar x v1 lb ub is
slnNextVar :: Solutions -> Maybe Solutions
slnNextVar Done = Nothing
slnNextVar (TopVar x v _ _ is) = Just $ FixedVar x v $ startIter $ iLet x v is
slnNextVar (FixedVar x v i) = FixedVar x v `fmap` slnNextVar i
-- Given a list of lower (resp. upper) bounds, compute the least (resp. largest)
-- value that satisfies them all.
iPickBounded :: BoundType -> [Bound] -> Maybe Integer
iPickBounded _ [] = Nothing
iPickBounded bt bs =
do
xs <- mapM (normBound bt) bs
return $ case bt of
Lower -> maximum xs
Upper -> minimum xs
where
-- t < c*x
-- <=> t+1 <= c*x
-- <=> (t+1)/c <= x
-- <=> ceil((t+1)/c) <= x
-- <=> t `div` c + 1 <= x
normBound Lower (Bound c t) = do
k <- isConst t
return (k `div` c + 1)
-- c*x < t
-- <=> c*x <= t-1
-- <=> x <= (t-1)/c
-- <=> x <= floor((t-1)/c)
-- <=> x <= (t-1) `div` c
normBound Upper (Bound c t) = do
k <- isConst t
return (div (k -1) c)
{- | The largest (resp. least) upper (resp. lower) bound on a term
that will satisfy the model
-}
iTermBound :: BoundType -> Term -> Inerts -> Maybe Integer
iTermBound bt (T k xs) is = do
ks <- mapM summand (Map.toList xs)
return $ sum $ k : ks
where
summand (x, c) = fmap (c *) (iVarBound (newBt c) x is)
newBt c =
if c > 0
then bt
else case bt of
Lower -> Upper
Upper -> Lower
{- | The largest (resp. least) upper (resp. lower) bound on a variable
that will satisfy the model.
-}
iVarBound :: BoundType -> Name -> Inerts -> Maybe Integer
iVarBound bt x is
| Just t <- Map.lookup x (solved is) = iTermBound bt t is
iVarBound bt x is =
do
both <- Map.lookup x (bounds is)
case mapMaybe fromBound (chooseBounds both) of
[] -> Nothing
bs -> return (combineBounds bs)
where
fromBound (Bound c t) = fmap (scaleBound c) (iTermBound bt t is)
combineBounds = case bt of
Upper -> minimum
Lower -> maximum
chooseBounds = case bt of
Upper -> snd
Lower -> fst
scaleBound c b = case bt of
Upper -> div (b -1) c
Lower -> div b c + 1
iModel :: Inerts -> [(Name, Integer)]
iModel i = goBounds [] (bounds i)
where
goBounds su mp =
case Map.maxViewWithKey mp of
Nothing -> goEqs su $ Map.toList $ solved i
Just ((x, (lbs0, ubs0)), mp1) ->
let lbs = [Bound c (tLetNums su t) | Bound c t <- lbs0]
ubs = [Bound c (tLetNums su t) | Bound c t <- ubs0]
sln =
fromMaybe 0 $
mplus (iPickBounded Lower lbs) (iPickBounded Upper ubs)
in goBounds ((x, sln) : su) mp1
goEqs su [] = su
goEqs su ((x, t) : more) =
let t1 = tLetNums su t
vs = tVarList t1
su1 = [(v, 0) | v <- vs] ++ (x, tConstPart t1) : su
in goEqs su1 more
--------------------------------------------------------------------------------
-- Solving constraints
solveIs0 :: Term -> S ()
solveIs0 t = solveIs0' =<< apSubst t
{- | Solve a constraint if the form @t = 0@.
Assumes substitution has already been applied.
-}
solveIs0' :: Term -> S ()
solveIs0' t
-- A == 0
| Just a <- isConst t = guard (a == 0)
-- A + B * x = 0
| Just (a, b, x) <- tIsOneVar t =
case divMod (- a) b of
(q, 0) -> addDef x (tConst q)
_ -> mzero
-- x + S = 0
-- -x + S = 0
| Just (xc, x, s) <- tGetSimpleCoeff t =
addDef x (if xc > 0 then tNeg s else s)
-- A * S = 0
| Just (_, s) <- tFactor t = solveIs0 s
-- See Section 3.1 of paper for details.
-- We obtain an equivalent formulation but with smaller coefficients.
| Just (ak, xk, s) <- tLeastAbsCoeff t =
do
let m = abs ak + 1
v <- newVar
let sgn = signum ak
soln =
(negate sgn * m) |*| tVar v
|+| tMapCoeff (\c -> sgn * modulus c m) s
addDef xk soln
let upd i = div (2 * i + m) (2 * m) + modulus i m
solveIs0 (negate (abs ak) |*| tVar v |+| tMapCoeff upd s)
| otherwise = error "solveIs0: unreachable"
modulus :: Integer -> Integer -> Integer
modulus a m = a - m * div (2 * a + m) (2 * m)
solveIsNeg :: Term -> S ()
solveIsNeg t = solveIsNeg' =<< apSubst t
{- | Solve a constraint of the form @t < 0@.
Assumes that substitution has been applied
-}
solveIsNeg' :: Term -> S ()
solveIsNeg' t
-- A < 0
| Just a <- isConst t = guard (a < 0)
-- A * S < 0
| Just (_, s) <- tFactor t = solveIsNeg s
-- See Section 5.1 of the paper
| Just (xc, x, s) <- tLeastVar t =
do
ctrs <-
if xc < 0
then -- -XC*x + S < 0
-- S < XC*x
do
ubs <- getBounds Upper x
let b = negate xc
beta = s
addBound Lower x (Bound b beta)
return [(a, alpha, b, beta) | Bound a alpha <- ubs]
else -- XC*x + S < 0
-- XC*x < -S
do
lbs <- getBounds Lower x
let a = xc
alpha = tNeg s
addBound Upper x (Bound a alpha)
return [(a, alpha, b, beta) | Bound b beta <- lbs]
-- See Note [Shadows]
mapM_
( \(a, alpha, b, beta) ->
do
let real = ctLt (a |*| beta) (b |*| alpha)
dark = ctLt (tConst (a * b)) (b |*| alpha |-| a |*| beta)
gray =
[ ctEq (b |*| tVar x) (tConst i |+| beta)
| i <- [1 .. b - 1]
]
solveIsNeg real
foldl orElse (solveIsNeg dark) (map solveIs0 gray)
)
ctrs
| otherwise = error "solveIsNeg: unreachable"
orElse :: S () -> S () -> S ()
orElse x y = mplus x y
{- Note [Shadows]
P: beta < b * x
Q: a * x < alpha
real: a * beta < b * alpha
beta < b * x -- from P
a * beta < a * b * x -- (a *)
a * beta < b * alpha -- comm. and Q
dark: b * alpha - a * beta > a * b
gray: b * x = beta + 1 \/
b * x = beta + 2 \/
...
b * x = beta + (b-1)
We stop at @b - 1@ because if:
> b * x >= beta + b
> a * b * x >= a * (beta + b) -- (a *)
> a * b * x >= a * beta + a * b -- distrib.
> b * alpha > a * beta + a * b -- comm. and Q
> b * alpha - a * beta > a * b -- subtract (a * beta)
which is covered by the dark shadow.
-}
--------------------------------------------------------------------------------
-- Monads
data Answer a = None | One a | Choice (Answer a) (Answer a)
deriving (Show)
answerSize :: Answer a -> (Integer, Integer, Integer)
answerSize = go 0 0 0
where
go !n !o !c ans =
case ans of
None -> (n + 1, o, c)
One _ -> (n, o + 1, c)
Choice x y ->
case go n o (c + 1) x of
(n', o', c') -> go n' o' c' y
dotAnswer :: (a -> Doc) -> Answer a -> Doc
dotAnswer pp g0 = vcat [text "digraph {", nest 2 (fst $ go 0 g0), text "}"]
where
node x d =
integer x <+> brackets (text "label=" <> text (show d))
<> semi
edge x y = integer x <+> text "->" <+> integer y
go x None =
let x' = x + 1
in seq x' (node x "", x')
go x (One a) =
let x' = x + 1
in seq x' (node x (show (pp a)), x')
go x (Choice c1 c2) =
let x' = x + 1
(ls1, x1) = go x' c1
(ls2, x2) = go x1 c2
in seq
x'
( vcat
[ node x "|"
, edge x x'
, edge x x1
, ls1
, ls2
]
, x2
)
toList :: Answer a -> [a]
toList a = go a []
where
go (Choice xs ys) zs = go xs (go ys zs)
go (One x) xs = x : xs
go None xs = xs
instance Monad Answer where
#if !MIN_VERSION_ghc(8,8,1)
fail _ = None
#endif
None >>= _ = None
One a >>= k = k a
Choice m1 m2 >>= k = mplus (m1 >>= k) (m2 >>= k)
instance Alternative Answer where
empty = mzero
(<|>) = mplus
instance MonadPlus Answer where
mzero = None
mplus None x = x
-- mplus (Choice x y) z = mplus x (mplus y z)
mplus x y = Choice x y
instance Functor Answer where
fmap _ None = None
fmap f (One x) = One (f x)
fmap f (Choice x1 x2) = Choice (fmap f x1) (fmap f x2)
instance Applicative Answer where
pure = One
(<*>) = ap
newtype S a = S (RW -> Answer (a, RW))
instance Monad S where
S m >>= k = S $ \s -> do
(a, s1) <- m s
let S m1 = k a
m1 s1
instance Alternative S where
empty = mzero
(<|>) = mplus
instance MonadPlus S where
mzero = S $ \_ -> mzero
mplus (S m1) (S m2) = S $ \s -> mplus (m1 s) (m2 s)
instance Functor S where
fmap = liftM
instance Applicative S where
pure a = S $ \s -> pure (a, s)
(<*>) = ap
updS :: (RW -> (a, RW)) -> S a
updS f = S $ \s -> return (f s)
updS_ :: (RW -> RW) -> S ()
updS_ f = updS $ \rw -> ((), f rw)
get :: (RW -> a) -> S a
get f = updS $ \rw -> (f rw, rw)
newVar :: S Name
newVar = updS $ \rw ->
( SysName (nameSource rw)
, rw {nameSource = nameSource rw + 1}
)
-- | Get lower ('fst'), or upper ('snd') bounds for a variable.
getBounds :: BoundType -> Name -> S [Bound]
getBounds f x = get $ \rw -> case Map.lookup x $ bounds $ inerts rw of
Nothing -> []
Just bs -> case f of
Lower -> fst bs
Upper -> snd bs
addBound :: BoundType -> Name -> Bound -> S ()
addBound bt x b = updS_ $ \rw ->
let i = inerts rw
entry = case bt of
Lower -> ([b], [])
Upper -> ([], [b])
jn (newL, newU) (oldL, oldU) = (newL ++ oldL, newU ++ oldU)
in rw {inerts = i {bounds = Map.insertWith jn x entry (bounds i)}}
{- | Add a new definition.
Assumes substitution has already been applied
-}
addDef :: Name -> Term -> S ()
addDef x t =
do
newWork <- updS $ \rw ->
let (newWork, newInerts) = iSolved x t (inerts rw)
in (newWork, rw {inerts = newInerts})
mapM_ solveIsNeg newWork
apSubst :: Term -> S Term
apSubst t =
do
i <- get inerts
return (iApSubst i t)
--------------------------------------------------------------------------------
data Name = UserName !Int | SysName !Int
deriving (Read, Show, Eq, Ord)
ppName :: Name -> Doc
ppName (UserName x) = text "u" <> int x
ppName (SysName x) = text "s" <> int x
toName :: Int -> Name
toName = UserName
fromName :: Name -> Maybe Int
fromName (UserName x) = Just x
fromName (SysName _) = Nothing
type NameMap = Map Name
{- | The type of terms. The integer is the constant part of the term,
and the `Map` maps variables (represented by @Int@ to their coefficients).
The term is a sum of its parts.
INVARIANT: the `Map` does not map anything to 0.
-}
data Term = T !Integer (NameMap Integer)
deriving (Eq, Ord)
infixl 6 |+|, |-|
infixr 7 |*|
-- | A constant term.
tConst :: Integer -> Term
tConst k = T k Map.empty
-- | Construct a term with a single variable.
tVar :: Name -> Term
tVar x = T 0 (Map.singleton x 1)
(|+|) :: Term -> Term -> Term
T n1 m1 |+| T n2 m2 =
T (n1 + n2) $
if Map.null m1
then m2
else
if Map.null m2
then m1
else Map.filter (/= 0) $ Map.unionWith (+) m1 m2
(|*|) :: Integer -> Term -> Term
0 |*| _ = tConst 0
1 |*| t = t
k |*| T n m = T (k * n) (fmap (k *) m)
tNeg :: Term -> Term
tNeg t = (-1) |*| t
(|-|) :: Term -> Term -> Term
t1 |-| t2 = t1 |+| tNeg t2
-- | Replace a variable with a term.
tLet :: Name -> Term -> Term -> Term
tLet x t1 t2 =
let (a, t) = tSplitVar x t2
in a |*| t1 |+| t
-- | Replace a variable with a constant.
tLetNum :: Name -> Integer -> Term -> Term
tLetNum x k t =
let (c, T n m) = tSplitVar x t
in T (c * k + n) m
-- | Replace the given variables with constants.
tLetNums :: [(Name, Integer)] -> Term -> Term
tLetNums xs t = foldr (\(x, i) t1 -> tLetNum x i t1) t xs
instance Show Term where
showsPrec c t = showsPrec c (show (ppTerm t))
ppTerm :: Term -> Doc
ppTerm (T k m) =
case Map.toList m of
[] -> integer k
xs | k /= 0 -> hsep (integer k : map ppProd xs)
x : xs -> hsep (ppFst x : map ppProd xs)
where
ppFst (x, 1) = ppName x
ppFst (x, -1) = text "-" <> ppName x
ppFst (x, n) = ppMul n x
ppProd (x, 1) = text "+" <+> ppName x
ppProd (x, -1) = text "-" <+> ppName x
ppProd (x, n)
| n > 0 = text "+" <+> ppMul n x
| otherwise = text "-" <+> ppMul (abs n) x
ppMul n x = integer n <+> text "*" <+> ppName x
{- | Remove a variable from the term, and return its coefficient.
If the variable is not present in the term, the coefficient is 0.
-}
tSplitVar :: Name -> Term -> (Integer, Term)
tSplitVar x t@(T n m) =
case Map.updateLookupWithKey (\_ _ -> Nothing) x m of
(Nothing, _) -> (0, t)
(Just k, m1) -> (k, T n m1)
-- | Does the term contain this varibale?
tHasVar :: Name -> Term -> Bool
tHasVar x (T _ m) = Map.member x m
-- | Is this terms just an integer.
isConst :: Term -> Maybe Integer
isConst (T n m)
| Map.null m = Just n
| otherwise = Nothing
tConstPart :: Term -> Integer
tConstPart (T n _) = n
-- | Returns: @Just (a, b, x)@ if the term is the form: @a + b * x@
tIsOneVar :: Term -> Maybe (Integer, Integer, Name)
tIsOneVar (T a m) = case Map.toList m of
[(x, b)] -> Just (a, b, x)
_ -> Nothing
{- | Spots terms that contain variables with unit coefficients
(i.e., of the form @x + t@ or @t - x@).
Returns (coeff, var, rest of term)
-}
tGetSimpleCoeff :: Term -> Maybe (Integer, Name, Term)
tGetSimpleCoeff (T a m) =
do
let (m1, m2) = Map.partition (\x -> x == 1 || x == -1) m
((x, xc), m3) <- Map.minViewWithKey m1
return (xc, x, T a (Map.union m3 m2))
tVarList :: Term -> [Name]
tVarList (T _ m) = Map.keys m
{- | Try to factor-out a common consant (> 1) from a term.
For example, @2 + 4x@ becomes @2 * (1 + 2x)@.
-}
tFactor :: Term -> Maybe (Integer, Term)
tFactor (T c m) =
do
d <- common (c : Map.elems m)
return (d, T (div c d) (fmap (`div` d) m))
where
common :: [Integer] -> Maybe Integer
common [] = Nothing
common [x] = Just x
common (x : y : zs) =
case gcd x y of
1 -> Nothing
n -> common (n : zs)
-- | Extract a variable with a coefficient whose absolute value is minimal.
tLeastAbsCoeff :: Term -> Maybe (Integer, Name, Term)
tLeastAbsCoeff (T c m) = do
(xc, x, m1) <- Map.foldrWithKey step Nothing m
return (xc, x, T c m1)
where
step x xc Nothing = Just (xc, x, Map.delete x m)
step x xc (Just (yc, _, _))
| abs xc < abs yc = Just (xc, x, Map.delete x m)
step _ _ it = it
-- | Extract the least variable from a term
tLeastVar :: Term -> Maybe (Integer, Name, Term)
tLeastVar (T c m) =
do
((x, xc), m1) <- Map.minViewWithKey m
return (xc, x, T c m1)
-- | Apply a function to all coefficients, including the constnat
tMapCoeff :: (Integer -> Integer) -> Term -> Term
tMapCoeff f (T c m) = T (f c) (fmap f m)