lambdabot-4.1: scripts/Djinn/Poly.hs
module Poly(Ninf, inf, ninfToN,
Poly, var, constp, coeffsOf, mustBeConst, substPolyVars,
EqnSystem, solveEqnSystem) where
import List(partition, sort)
import Util.Digraph(stronglyConnComp, SCC(..))
import Util(mapSnd, insert, makeSet)
---------------------------------------
----- The type naturals + infinity
data Ninf = I !Integer | Inf
inf :: Ninf
inf = Inf
ninfToN :: Ninf -> Maybe Integer
ninfToN Inf = Nothing
ninfToN (I i) = Just i
instance Show Ninf where
showsPrec p (I i) = showsPrec p i
showsPrec _ Inf = showString "inf"
instance Eq Ninf where
I i == I j = i == j
Inf == Inf = error "Ninf: Inf == Inf" --True -- XXX?
_ == _ = False
instance Ord Ninf where
I i <= I j = i <= j
I _ <= Inf = True
_ <= _ = False
instance Num Ninf where
I i + I j = I (i + j)
_ + _ = Inf
I i * I j = I (i * j)
I 0 * Inf = I 0
Inf * I 0 = I 0
_ * _ = Inf
negate _ = error "Ninf: negate"
abs _ = error "Ninf: abs"
signum _ = error "Ninf: signum"
fromInteger i | i < 0 = error "Ninf: fromInteger"
| otherwise = I i
---------------------------------------
----- Polynomials over Ninf
-- Parametrized over the variable type
data Poly a = P ![Summand a] -- summands, sorted
deriving (Eq, Ord)
data Summand a = S ![(a, Integer)] !Ninf -- variables with their exponents and scalar factor, sorted
deriving (Eq, Ord)
instance (Show a) => Show (Poly a) where
showsPrec _ (P []) = showString "0"
showsPrec _ (P ss) = mix (map showS ss) " + "
where mix [] _ = error "Poly.mix"
mix [x] _ = x
mix (x:xs) m = x . showString m . mix xs m
showS (S [] i) = showsPrec 0 i
showS (S vs 1) = mix (map showV vs) "*"
showS (S vs n) = mix (showsPrec 0 n : map showV vs) "*"
showV (v, 1) = showsPrec 0 v
showV (v, n) = showsPrec 0 v . showString "^" . showsPrec 0 n
mkS :: [(a, Integer)] -> Ninf -> Summand a
mkS ves Inf = S [(v,1) | (v,_) <- ves] Inf
mkS ves k = S ves k
-- Make a polynomial that is just a variable
var :: a -> Poly a
var x = P [S [(x, 1)] 1]
constp :: Ninf -> Poly a
constp 0 = P []
constp i = P [S [] i]
isConstPoly :: Poly a -> Bool
isConstPoly (P []) = True
isConstPoly (P [S [] _]) = True
isConstPoly _ = False
instance (Show a, Ord a) => Num (Poly a) where
P ss1 + P ss2 = P $ add ss1 ss2
P ss1 * P ss2 = sum [ P [muls s1 s2] | s1 <- ss1, s2 <- ss2 ]
negate _ = error "Poly: negate"
abs _ = error "Poly: abs"
signum _ = error "Poly: signum"
fromInteger i | i < 0 = error "Poly: fromInteger"
| i == 0 = P []
| otherwise = P [S [] (fromInteger i)]
add :: (Ord a) => [Summand a] -> [Summand a] -> [Summand a]
add [] ss = ss
add ss [] = ss
add (s1@(S vs1 k1) : ss1) (s2@(S vs2 k2) : ss2) =
case compare vs1 vs2 of
LT -> s1 : add ss1 (s2:ss2)
EQ -> mkS vs1 (k1 + k2) : add ss1 ss2
GT -> s2 : add (s1:ss1) ss2
muls :: (Ord a) => Summand a -> Summand a -> Summand a
muls (S vs1 k1) (S vs2 k2) = mkS (merge vs1 vs2) (k1 * k2)
where merge [] ys = ys
merge xs [] = xs
merge (x@(v1, e1):xs) (y@(v2, e2):ys) =
case compare v1 v2 of
LT -> x : merge xs (y:ys)
EQ -> (v1, e1+e2) : merge xs ys
GT -> y : merge (x:xs) ys
-----
-- Return all the coefficints for a certain variable of
-- a polynomial. Returns (exponent, coefficient) list, sorted by exponent.
coeffsOf :: (Show a, Ord a) => a -> Poly a -> [(Integer, Poly a)]
coeffsOf x (P ss) = sort $ foldr coeff [] ss
where coeff s@(S vs k) cs =
case partition (\ (v, _) -> v == x) vs of
([(_, a)], vs') -> addCoeff a (P [S vs' k]) cs
_ -> addCoeff 0 (P [s]) cs
addCoeff a c [] = [(a, c)]
addCoeff a c ((a', c'):cs) | a == a' = (a, c + c') : cs
addCoeff a c (k:cs) = k : addCoeff a c cs
mustBeConst :: (Show a) => Poly a -> Ninf
mustBeConst (P []) = 0
mustBeConst (P [S [] k]) = k
mustBeConst p = error $ "Poly.mustBeConst: " ++ show p
getPolyVars :: (Ord a) => Poly a -> [a]
getPolyVars (P ss) = makeSet [ x | S vs _ <- ss, (x, _) <- vs ]
substPolyVars :: (Eq a) => [(a, b)] -> Poly a -> Poly b
substPolyVars ons (P ss) = P [ S [ (new v, e) | (v, e) <- ves ] k | S ves k <- ss ]
where new v = maybe (error "substPolyVars") id $ lookup v ons
---------------------------------------
----- Find the least fixed point solution of an equation system
-- There is always a solution because +, * are monotonic and
-- Ninf is a complete lattice.
-- Uses Zaionc's algorithm from
-- "Fixpoint Technique for Counting Terms in Typed $lambda$-calculus",
-- by Marek Zaionc.
-- The equations are of the form
-- x_1 = poly_1(x_1, ..., x_n)
-- ...
-- x_n = poly_n(x_1, ..., x_n)
type EqnSystem var = [(var, Poly var)]
solveEqnSystem :: (Ord var, Show var) => EqnSystem var -> [(var, Ninf)]
solveEqnSystem eqns =
-- The definition could be mapSnd mustBeConst $ solve eqns
-- but this can be very slow.
let groups = map flatten $ stronglyConnComp [ (eqn, v, getPolyVars rhs) | eqn@(v, rhs) <- eqns ]
where flatten (CyclicSCC es) = es
flatten (AcyclicSCC e) = [e]
solveGroup :: (Ord var, Show var) => [(var, Ninf)] -> EqnSystem var -> [(var, Ninf)]
solveGroup sols es =
let newsols = mapSnd mustBeConst $ solve (mapSnd (evalPolyConst sols) es)
in foldr insert sols newsols
in foldl solveGroup [] groups
solve :: (Ord var, Show var) => EqnSystem var -> [(var, Poly var)]
solve [] = []
solve ((v,f):xfs) =
let cs = coeffsOf v f
a0 = getCoeff 0 cs
a = sum [ ai | (i, ai) <- cs, i /= 0 ]
g = a0 + a0 * a * constp inf
xfs' = [ (y, evalPoly [(v, g)] p) | (y, p) <- xfs ]
sol = solve xfs'
x = evalPoly sol g
in (v, x) : sol
evalPoly :: (Ord var, Show var) => [(var, Poly var)] -> Poly var -> Poly var
evalPoly _ poly | isConstPoly poly = poly -- just a speedup
evalPoly vps poly = foldr eval1 poly vps
where eval1 (v, p) q = sum [ c * p^e | (e, c) <- coeffsOf v q ]
-- evalPolyConst is just a specialized version of evalConst when
-- the values are constants. It's also optimized a lot.
evalPolyConst :: (Ord var, Show var) => [(var, Ninf)] -> Poly var -> Poly var
evalPolyConst _ poly | isConstPoly poly = poly -- just a speedup
-- Assume vps is sorted in the same order as vs.
evalPolyConst vcs (P ss) = P (concatMap evalS ss)
where evalS (S ves k) = loop vcs ves [] k
loop [] ves rves k = [mkS (reverse rves ++ ves) k]
loop _ [] rves k = [mkS (reverse rves) k]
loop ovps@((v,c) : vps) oves@(ve@(v',e) : ves) rves k =
case v `compare` v' of
LT -> loop vps oves rves k
EQ -> if c == 0 then [] else loop vps ves rves (k * c^e)
GT -> loop ovps ves (ve:rves) k
getCoeff :: (Ord a, Show a) => Integer -> [(Integer, Poly a)] -> Poly a
getCoeff i ips = maybe 0 id $ lookup i ips