A small fix and a note.
[ghc.git] / compiler / typecheck / TcTypeNats.hs
1 {-# LANGUAGE PatternGuards #-}
2 module TcTypeNats ( canonicalNum, NumericsResult(..) ) where
3
4 import TcSMonad ( TcS, Xi
5 , newWantedCoVar
6 , setWantedCoBind, setDictBind, newDictVar
7 , getWantedLoc
8 , tcsLookupTyCon, tcsLookupClass
9 , CanonicalCt (..), CanonicalCts
10 , mkFrozenError
11 , traceTcS
12 )
13 import TcRnTypes ( CtFlavor(..) )
14 import TcCanonical (mkCanonicals)
15 import HsBinds (EvTerm(..))
16 import Class ( Class, className, classTyCon )
17 import Type ( tcView, mkTyConApp, mkNumberTy, isNumberTy
18 , tcEqType, tcCmpType, pprType
19 )
20 import TypeRep (Type(..))
21 import TyCon (TyCon, tyConName)
22 import Var (EvVar)
23 import Coercion ( mkUnsafeCoercion )
24 import Outputable
25 import PrelNames ( lessThanEqualClassName
26 , addTyFamName, mulTyFamName, expTyFamName
27 )
28 import Bag (bagToList, emptyBag, unionManyBags, unionBags)
29 import Data.Maybe (fromMaybe, listToMaybe, maybeToList)
30 import Data.List (nub, partition, sort)
31 import Control.Monad (msum, mplus, zipWithM, (<=<), guard)
32
33 import Debug.Trace
34 import Unique(getUnique)
35 import Type(getTyVar)
36
37
38 data Term = Var Xi | Num Integer (Maybe Xi)
39 data Op = Add | Mul | Exp deriving Eq
40 data Prop = EqFun Op Term Term Term
41 | Leq Term Term
42 | Eq Term Term
43
44 commutative :: Op -> Bool
45 commutative op = op == Add || op == Mul
46
47 associative :: Op -> Bool
48 associative op = op == Add || op == Mul
49
50 num :: Integer -> Term
51 num n = Num n Nothing
52
53 opFun :: Op -> (Integer -> Integer -> Integer)
54 opFun x = case x of
55 Add -> (+)
56 Mul -> (*)
57 Exp -> (^)
58
59
60 instance Eq Term where
61 Var x == Var y = tcEqType x y
62 Num x _ == Num y _ = x == y
63 _ == _ = False
64
65 instance Ord Term where
66 compare (Num x _) (Num y _) = compare x y
67 compare (Num _ _) (Var _) = LT
68 compare (Var x) (Var y) = tcCmpType x y
69 compare (Var _) (Num _ _) = GT
70
71
72 --------------------------------------------------------------------------------
73 -- Interface with the type checker
74 --------------------------------------------------------------------------------
75
76 data NumericsResult = NumericsResult
77 { numNewWork :: CanonicalCts
78 , numInert :: Maybe CanonicalCts -- Nothing for "no change"
79 , numNext :: Maybe CanonicalCt
80 }
81
82 -- We keep the original type in numeric constants to preserve type synonyms.
83 toTerm :: Xi -> Term
84 toTerm xi = case mplus (isNumberTy xi) (isNumberTy =<< tcView xi) of
85 Just n -> Num n (Just xi)
86 _ -> Var xi
87
88 fromTerm :: Term -> Xi
89 fromTerm (Num n mb) = fromMaybe (mkNumberTy n) mb
90 fromTerm (Var xi) = xi
91
92 toProp :: CanonicalCt -> Prop
93 toProp (CDictCan { cc_class = c, cc_tyargs = [xi1,xi2] })
94 | className c == lessThanEqualClassName = Leq (toTerm xi1) (toTerm xi2)
95
96 toProp (CFunEqCan { cc_fun = tc, cc_tyargs = [xi11,xi12], cc_rhs = xi2 })
97 | tyConName tc == addTyFamName = EqFun Add t1 t2 t3
98 | tyConName tc == mulTyFamName = EqFun Mul t1 t2 t3
99 | tyConName tc == expTyFamName = EqFun Exp t1 t2 t3
100
101 where t1 = toTerm xi11
102 t2 = toTerm xi12
103 t3 = toTerm xi2
104
105 toProp p = panic $
106 "[TcTypeNats.toProp] Unexpected CanonicalCt: " ++ showSDoc (ppr p)
107
108
109 canonicalNum :: CanonicalCts -> CanonicalCts -> CanonicalCts -> CanonicalCt ->
110 TcS NumericsResult
111 canonicalNum given derived wanted prop =
112 case cc_flavor prop of
113 Wanted {} -> solveNumWanted given derived wanted prop
114 Derived {} -> addNumDerived given derived wanted prop
115 Given {} -> addNumGiven given derived wanted prop
116
117
118 solveNumWanted :: CanonicalCts -> CanonicalCts -> CanonicalCts -> CanonicalCt ->
119 TcS NumericsResult
120 solveNumWanted given derived wanted prop =
121 do let asmps = map toProp $ bagToList $ unionManyBags [given,derived,wanted]
122 goal = toProp prop
123
124 numTrace "solveNumWanted" (vmany asmps <+> text "|-" <+> ppr goal)
125
126 case solve asmps goal of
127
128 Simplified sgs ->
129 do numTrace "Simplified to" (vmany sgs)
130 defineDummy (cc_id prop) =<< fromProp goal
131 evs <- mapM (newSubGoal <=< fromProp) sgs
132 goals <- mkCanonicals (cc_flavor prop) evs
133 return NumericsResult
134 { numNext = Nothing, numInert = Nothing, numNewWork = goals }
135
136 -- XXX: The new wanted might imply some of the existing wanteds...
137 Improved is ->
138 do numTrace "Improved by" (vmany is)
139 evs <- mapM (newSubGoal <=< fromProp) is
140 goals <- mkCanonicals (cc_flavor prop) evs
141 return NumericsResult
142 { numNext = Just prop, numInert = Nothing, numNewWork = goals }
143
144 Impossible -> impossible prop
145
146
147 -- XXX: Need to understand derived work better.
148 addNumDerived :: CanonicalCts -> CanonicalCts -> CanonicalCts -> CanonicalCt ->
149 TcS NumericsResult
150 addNumDerived given derived wanted prop =
151 do let asmps = map toProp $ bagToList given
152 goal = toProp prop
153
154 numTrace "addNumDerived" (vmany asmps <+> text "|-" <+> ppr goal)
155
156 case solve asmps goal of
157
158 Simplified sgs ->
159 do numTrace "Simplified to" (vmany sgs)
160 defineDummy (cc_id prop) =<< fromProp goal
161 evs <- mapM (newSubGoal <=< fromProp) sgs
162 goals <- mkCanonicals (cc_flavor prop) evs
163 return NumericsResult
164 { numNext = Nothing, numInert = Nothing, numNewWork = goals }
165
166 -- XXX: watch out for cycles because of the wanteds being restarted:
167 -- W => D && D => W, we could solve W by W
168 -- if W <=> D, then we should simplify W to D, not make it derived.
169 Improved is ->
170 do numTrace "Improved by" (vmany is)
171 evs <- mapM (newSubGoal <=< fromProp) is
172 goals <- mkCanonicals (Derived (getWantedLoc prop)) evs
173 return NumericsResult
174 { numNext = Just prop, numInert = Just (unionBags given derived)
175 , numNewWork = unionBags wanted goals }
176
177 Impossible -> impossible prop
178
179
180 addNumGiven :: CanonicalCts -> CanonicalCts -> CanonicalCts -> CanonicalCt ->
181 TcS NumericsResult
182 addNumGiven given derived wanted prop =
183 do let asmps = map toProp (bagToList given)
184 goal = toProp prop
185
186 numTrace "addNumGiven" (vmany asmps <+> text " /\\ " <+> ppr goal)
187 case solve asmps goal of
188
189 Simplified sgs ->
190 do numTrace "Simplified to" (vmany sgs)
191 evs <- mapM (newFact <=< fromProp) sgs
192 facts <- mkCanonicals (cc_flavor prop) evs
193 return NumericsResult
194 { numNext = Nothing, numInert = Nothing, numNewWork = facts }
195
196 Improved is ->
197 do numTrace "Improved by" (vmany is)
198 evs <- mapM (newFact <=< fromProp) is
199 facts <- mkCanonicals (cc_flavor prop) evs
200 return NumericsResult
201 { numNext = Just prop, numInert = Just (unionBags given derived)
202 , numNewWork = unionBags wanted facts }
203
204 Impossible -> impossible prop
205
206 impossible :: CanonicalCt -> TcS NumericsResult
207 impossible c =
208 do numTrace "Impossible" empty
209 let err = mkFrozenError (cc_flavor c) (cc_id c)
210 return NumericsResult
211 { numNext = Just err, numInert = Nothing, numNewWork = emptyBag }
212
213
214
215 data CvtProp = CvtClass Class [Type]
216 | CvtCo Type Type
217
218 fromProp :: Prop -> TcS CvtProp
219 fromProp (Leq t1 t2) =
220 do cl <- tcsLookupClass lessThanEqualClassName
221 return (CvtClass cl [ fromTerm t1, fromTerm t2 ])
222
223 fromProp (Eq t1 t2) = return $ CvtCo (fromTerm t1) (fromTerm t2)
224
225 fromProp (EqFun op t1 t2 t3) =
226 do tc <- tcsLookupTyCon $ case op of
227 Add -> addTyFamName
228 Mul -> mulTyFamName
229 Exp -> expTyFamName
230 return $ CvtCo (mkTyConApp tc [fromTerm t1, fromTerm t2]) (fromTerm t3)
231
232
233 newSubGoal :: CvtProp -> TcS EvVar
234 newSubGoal (CvtClass c ts) = newDictVar c ts
235 newSubGoal (CvtCo t1 t2) = newWantedCoVar t1 t2
236
237 newFact :: CvtProp -> TcS EvVar
238 newFact prop =
239 do d <- newSubGoal prop
240 defineDummy d prop
241 return d
242
243
244 -- If we decided that we want to generate evidence terms,
245 -- here we would set the evidence properly. For now, we skip this
246 -- step because evidence terms are not used for anything, and they
247 -- get quite large, at least, if we start with a small set of axioms.
248 defineDummy :: EvVar -> CvtProp -> TcS ()
249 defineDummy d (CvtClass c ts) =
250 setDictBind d $ EvAxiom "<=" $ mkTyConApp (classTyCon c) ts
251
252 defineDummy c (CvtCo t1 t2) =
253 setWantedCoBind c $ mkUnsafeCoercion t1 t2
254
255
256
257
258 --------------------------------------------------------------------------------
259 -- The Solver
260 --------------------------------------------------------------------------------
261
262
263 data Result = Impossible -- We know that the goal cannot be solved.
264 | Simplified [Prop] -- We reformulated the goal.
265 | Improved [Prop] -- We learned some new facts.
266
267 solve :: [Prop] -> Prop -> Result
268 solve asmps (Leq a b) =
269 let ps = propsToOrd asmps
270 in case isLeq ps a b of
271 True -> Simplified []
272 False ->
273 case improveLeq ps a b of
274 Nothing -> Impossible
275 Just ps -> Improved ps
276
277 solve asmps prop =
278 case solveWanted1 prop of
279 Just sgs -> Simplified sgs
280 Nothing ->
281 case msum $ zipWith solveWanted2 asmps (repeat prop) of
282 Just sgs -> Simplified sgs
283 Nothing
284 | solveAC asmps prop -> Simplified []
285 | byDistr asmps prop -> Simplified []
286 | otherwise ->
287 case improve1 prop of
288 Nothing -> Impossible
289 Just eqs ->
290 case zipWithM improve2 asmps (repeat prop) of
291 Nothing -> Impossible
292 Just eqs1 -> Improved (concat (eqs : eqs1))
293
294
295
296
297 improve1 :: Prop -> Maybe [Prop]
298 improve1 prop =
299 case prop of
300
301 EqFun Add (Num m _) (Num n _) t -> Just [ Eq (num (m+n)) t ]
302 EqFun Add (Num 0 _) s t -> Just [ Eq s t ]
303 EqFun Add (Num m _) s (Num n _)
304 | m <= n -> Just [ Eq (num (n-m)) s ]
305 | otherwise -> Nothing
306 EqFun Add r s (Num 0 _) -> Just [ Eq (num 0) r, Eq (num 0) s ]
307 EqFun Add r s t
308 | r == t -> Just [ Eq (num 0) s ]
309 | s == t -> Just [ Eq (num 0) r ]
310
311 EqFun Mul (Num m _) (Num n _) t -> Just [ Eq (num (m*n)) t ]
312 EqFun Mul (Num 0 _) _ t -> Just [ Eq (num 0) t ]
313 EqFun Mul (Num 1 _) s t -> Just [ Eq s t ]
314 EqFun Mul (Num _ _) s t
315 | s == t -> Just [ Eq (num 0) s ]
316
317 EqFun Mul r s (Num 1 _) -> Just [ Eq (num 1) r, Eq (num 1) s ]
318 EqFun Mul (Num m _) s (Num n _)
319 | Just a <- divide n m -> Just [ Eq (num a) s ]
320 | otherwise -> Nothing
321
322
323 EqFun Exp (Num m _) (Num n _) t -> Just [ Eq (num (m ^ n)) t ]
324
325 EqFun Exp (Num 1 _) _ t -> Just [ Eq (num 1) t ]
326 EqFun Exp (Num _ _) s t
327 | s == t -> Nothing
328 EqFun Exp (Num m _) s (Num n _) -> do a <- descreteLog m n
329 return [ Eq (num a) s ]
330
331 EqFun Exp _ (Num 0 _) t -> Just [ Eq (num 1) t ]
332 EqFun Exp r (Num 1 _) t -> Just [ Eq r t ]
333 EqFun Exp r (Num m _) (Num n _) -> do a <- descreteRoot m n
334 return [ Eq (num a) r ]
335
336 _ -> Just []
337
338 improve2 :: Prop -> Prop -> Maybe [ Prop ]
339 improve2 asmp prop =
340 case asmp of
341
342 EqFun Add a1 b1 c1 ->
343 case prop of
344 EqFun Add a2 b2 c2
345 | a1 == a2 && b1 == b2 -> Just [ Eq c1 c2 ]
346 | a1 == b2 && b1 == a2 -> Just [ Eq c1 c2 ]
347 | c1 == c2 && a1 == a2 -> Just [ Eq b1 b2 ]
348 | c1 == c2 && a1 == b2 -> Just [ Eq b1 a2 ]
349 | c1 == c2 && b1 == b2 -> Just [ Eq a1 a2 ]
350 | c1 == c2 && b1 == a2 -> Just [ Eq a1 b2 ]
351 _ -> Just []
352
353
354 EqFun Mul a1 b1 c1 ->
355 case prop of
356 EqFun Mul a2 b2 c2
357 | a1 == a2 && b1 == b2 -> Just [ Eq c1 c2 ]
358 | a1 == b2 && b1 == a2 -> Just [ Eq c1 c2 ]
359 | c1 == c2 && b1 == b2, Num m _ <- a1, Num n _ <- a2, m /= n
360 -> Just [ Eq (num 0) b1, Eq (num 0) c1 ]
361 _ -> Just []
362
363
364 _ -> Just []
365
366
367 solveWanted1 :: Prop -> Maybe [ Prop ]
368 solveWanted1 prop =
369 case prop of
370
371 EqFun Add (Num m _) (Num n _) (Num mn _) | m + n == mn -> Just []
372 EqFun Add (Num 0 _) s t | s == t -> Just []
373 EqFun Add r s t | r == s ->
374 Just [ EqFun Mul (num 2) r t ]
375
376 EqFun Mul (Num m _) (Num n _) (Num mn _) | m * n == mn -> Just []
377 EqFun Mul (Num 0 _) _ (Num 0 _) -> Just []
378 EqFun Mul (Num 1 _) s t | s == t -> Just []
379 EqFun Mul r s t | r == s ->
380 Just [ EqFun Exp r (num 2) t ]
381
382 -- Simple normalization of commutative operators
383 EqFun op r@(Var _) s@(Num _ _) t | commutative op ->
384 Just [ EqFun op s r t ]
385
386 EqFun Exp (Num m _) (Num n _) (Num mn _) | m ^ n == mn -> Just []
387 EqFun Exp (Num 1 _) _ (Num 1 _) -> Just []
388 EqFun Exp _ (Num 0 _) (Num 1 _) -> Just []
389 EqFun Exp r (Num 1 _) t | r == t -> Just []
390 EqFun Exp r (Num _ _) t | r == t ->
391 Just [Leq r (num 1)]
392
393 _ -> Nothing
394
395
396 solveWanted2 :: Prop -> Prop -> Maybe [Prop]
397 solveWanted2 asmp prop =
398 case (asmp, prop) of
399
400 (EqFun op1 r1 s1 t1, EqFun op2 r2 s2 t2)
401 | op1 == op2 && t1 == t2 &&
402 ( r1 == r2 && s1 == s2
403 || commutative op1 && r1 == s2 && s1 == r2
404 ) -> Just []
405
406 (EqFun Add (Num m _) b c1, EqFun Add (Num n _) d c2)
407 | c1 == c2 -> if m >= n then Just [ EqFun Add (num (m - n)) b d ]
408 else Just [ EqFun Add (num (n - m)) d b ]
409
410 (EqFun Mul (Num m _) b c1, EqFun Mul (Num n _) d c2)
411 | c1 == c2, Just x <- divide m n -> Just [ EqFun Mul (num x) b d ]
412 | c1 == c2, Just x <- divide n m -> Just [ EqFun Mul (num x) d b ]
413
414 -- hm: m * b = c |- c + b = d <=> (m + 1) * b = d
415 (EqFun Mul (Num m _) s1 t1, EqFun Add r2 s2 t2)
416 | t1 == r2 && s1 == s2 -> Just [ EqFun Mul (num (m + 1)) s1 t2 ]
417 | t1 == s2 && s1 == r2 -> Just [ EqFun Mul (num (m + 1)) r2 t2 ]
418
419 _ -> Nothing
420
421
422 --------------------------------------------------------------------------------
423 -- Reasoning about ordering.
424 --------------------------------------------------------------------------------
425
426 -- This function assumes that the assumptions are acyclic.
427 isLeq :: [(Term, Term)] -> Term -> Term -> Bool
428 isLeq _ (Num 0 _) _ = True
429 isLeq _ (Num m _) (Num n _) = m <= n
430 isLeq _ a b | a == b = True
431 isLeq ps (Num m _) a = or [ isLeq ps b a | (Num x _, b) <- ps, m <= x ]
432 isLeq ps a (Num m _) = or [ isLeq ps a b | (b, Num x _) <- ps, x <= m ]
433 isLeq ps a b = or [ isLeq ps c b | (a',c) <- ps, a == a' ]
434
435 isGt :: [(Term, Term)] -> Term -> Term -> Bool
436 isGt _ (Num m _) (Num n _) = m > n
437 isGt ps (Num m _) a = (m > 0) && isLeq ps a (num (m - 1))
438 isGt ps a (Num m _) = isLeq ps (num (m + 1)) a
439 isGt _ _ _ = False
440
441 improveLeq :: [(Term,Term)] -> Term -> Term -> Maybe [Prop]
442 improveLeq ps a b | isLeq ps b a = Just [Eq a b]
443 | isGt ps a b = Nothing
444 | otherwise = Just []
445
446 -- Ordering constraints derived from numeric predicates.
447 -- We do not consider equlities because they should be substituted away.
448 propsToOrd :: [Prop] -> [(Term,Term)]
449 propsToOrd props = loop (step [] unconditional)
450 where
451 loop ps = let new = filter (`notElem` ps) (step ps conditional)
452 in if null new then ps else loop (new ++ ps)
453
454 step ps = nub . concatMap (toOrd ps)
455
456 isConditional (EqFun op _ _ _) = op == Mul || op == Exp
457 isConditional _ = False
458
459 (conditional,unconditional) = partition isConditional props
460
461 toOrd _ (Leq a b) = [(a,b)]
462 toOrd _ (Eq _ _) = [] -- Would lead to a cycle, should be subst. away
463 toOrd ps (EqFun op a b c) =
464 case op of
465 Add -> [(a,c),(b,c)]
466
467 Mul -> (guard (isLeq ps (num 1) a) >> return (b,c)) ++
468 (guard (isLeq ps (num 1) b) >> return (a,c))
469 Exp
470 | Num 0 _ <- a -> [(c, num 1)]
471 | a == c -> [(a, num 1)]
472 | otherwise -> (guard (isLeq ps (num 2) a) >> return (b,c)) ++
473 (guard (isLeq ps (num 1) b) >> return (a,c))
474
475 --------------------------------------------------------------------------------
476 -- Associative and Commutative Operators
477 --------------------------------------------------------------------------------
478
479 -- XXX: It'd be nice to go back to a "small-step" style solve.
480 -- If we can get it to work, it is likely to be simpler,
481 -- and it would be way easier to provide proofs.
482
483
484 -- XXX: recursion may non-terminate: x * y = x
485 -- XXX: does not do improvements
486
487
488
489 solveAC :: [Prop] -> Prop -> Bool
490 solveAC ps (EqFun op x y z)
491 | commutative op && associative op =
492 (xs_ys === z) || or [ add as xs_ys === r | (as,r) <- cancelCands z ]
493
494 where
495 xs_ys = add (sums x) (sums y)
496
497 candidates c = [ (a,b) | (a,b,c') <- asmps, c == c' ]
498
499 -- (xs,e) `elem` cancelCands g ==> xs + g = e
500 cancelCands :: Term -> [([[Term]],Term)]
501 cancelCands g = do (a,b,c) <- asmps
502 let us = filter (all mayCancel)
503 $ case () of
504 _ | a == g -> sums b
505 | b == g -> sums a
506 | otherwise -> [ ]
507 guard (not (null us))
508 (us,c) : [ (add us vs, e) | (vs,e) <- cancelCands c ]
509
510
511 mayCancel x = case op of
512 Add -> True
513 Mul -> isLeq ordProps (num 1) x
514 Exp -> False
515
516 -- xs `elem` sums a ==> sum xs = a
517 sums a = [a] : [ p | (b,c) <- candidates a, p <- add (sums b) (sums c)]
518
519 add :: [Terms] -> [Terms] -> [Terms]
520 add as bs = [ doOp2 (opFun op) u v | u <- as, v <- bs ]
521
522 (===) :: [[Term]] -> Term -> Bool
523 as === b = any (`elem` sums b) as
524
525 -- Facts in a more convenient from
526 asmps = [ (a,b,c) | EqFun op' a b c <- ps, op == op' ]
527 ordProps = propsToOrd ps
528
529
530 solveAC _ _ = False
531
532
533
534
535 -- Combining sequences of terms with the same operation.
536 -- We preserve a normal form with at most 1 constant and variables sorted.
537 -- Examples:
538 -- (1 + c + y) + (2 + b + z) = 3 + b + c + y + z
539 -- (2 * c * y) * (3 * b * z) = 6 * b * c * y * z
540
541 type Terms = [Term]
542
543 doOp2 :: (Integer -> Integer -> Integer) -> Terms -> Terms -> Terms
544 doOp2 op (Num m _ : as) (Num n _ : bs) = num (op m n) : merge as bs
545 doOp2 _ xs ys = merge xs ys
546
547 merge :: Ord a => [a] -> [a] -> [a]
548 merge xs@(a:as) ys@(b:bs)
549 | a <= b = a : merge as ys
550 | otherwise = b : merge xs bs
551 merge [] ys = ys
552 merge xs [] = xs
553
554
555 -- XXX: This duplicates a lot of work done by AC.
556 byDistr :: [Prop] -> Prop -> Bool
557 byDistr ps (EqFun op x y z)
558 -- (a * b + a * c) = x |- a * (b + c) = x
559 | op == Mul = let zs = sumsFor z
560 ps = sumProds x y
561 in -- pNumTrace "Distr *" (text "zs:" <+> vmany zs <> comma <+> text "ps:" <+> vmany ps) $
562 any (`elem` zs) ps
563
564 -- a * (b + c) = x |- a * b + a * c = x
565 | op == Add = let as = sumsFor x
566 bs = sumsFor y
567 cs = [ doOp2 (opFun op) a b | a <- as, b <- bs ]
568 ds = do (a,b,z') <- prods
569 guard (z == z')
570 sumProds a b
571 in -- pNumTrace "Distr +" (text "cs:" <+> vmany cs <> comma <+> text "ds:" <+> vmany ds) $
572 any (`elem` cs) ds
573
574 where
575 sums = [ (a,b,c) | EqFun Add a b c <- ps ]
576 prods = [ (a,b,c) | EqFun Mul a b c <- ps ]
577 sumsFor :: Term -> [[Term]]
578 sumsFor c = [c] : do (a,b,c') <- sums
579 guard (c == c')
580 u <- sumsFor a
581 v <- sumsFor b
582 return (doOp2 (opFun Mul) u v)
583
584 -- This is very ad-hock.
585 prod (Num 1 _, a) = Just a
586 prod (a, Num 1 _) = Just a
587 prod (a,b) = listToMaybe [ c | (a',b',c) <- prods, a == a', b == b' ]
588
589 sumProds u v = do s1 <- sumsFor u
590 s2 <- sumsFor v
591 maybeToList $
592 do s3 <- mapM prod [ (a,b) | a <- s1, b <- s2 ]
593 return (sort s3)
594 byDistr _ _ = False
595
596
597
598
599 --------------------------------------------------------------------------------
600 -- Descrete Math
601 --------------------------------------------------------------------------------
602
603 descreteRoot :: Integer -> Integer -> Maybe Integer
604 descreteRoot root num = search 0 num
605 where
606 search from to = let x = from + div (to - from) 2
607 a = x ^ root
608 in case compare a num of
609 EQ -> Just x
610 LT | x /= from -> search x to
611 GT | x /= to -> search from x
612 _ -> Nothing
613
614 descreteLog :: Integer -> Integer -> Maybe Integer
615 descreteLog _ 0 = Just 0
616 descreteLog base num | base == num = Just 1
617 descreteLog base num = case divMod num base of
618 (x,0) -> fmap (1+) (descreteLog base x)
619 _ -> Nothing
620
621 divide :: Integer -> Integer -> Maybe Integer
622 divide _ 0 = Nothing
623 divide x y = case divMod x y of
624 (a,0) -> Just a
625 _ -> Nothing
626
627
628 --------------------------------------------------------------------------------
629 -- Debugging
630 --------------------------------------------------------------------------------
631
632 numTrace :: String -> SDoc -> TcS ()
633 numTrace x y = traceTcS ("[numerics] " ++ x) y
634
635 pNumTrace :: String -> SDoc -> a -> a
636 pNumTrace x y = trace ("[numerics] " ++ x ++ " " ++ showSDoc y)
637
638 vmany :: Outputable a => [a] -> SDoc
639 vmany xs = braces $ hcat $ punctuate comma $ map ppr xs
640
641 instance Outputable Term where
642 ppr (Var xi) = pprType xi <> un
643 where un = brackets $ text $ show $ getUnique (getTyVar "numerics dbg" xi)
644 ppr (Num n _) = integer n
645
646 instance Outputable Prop where
647 ppr (EqFun op t1 t2 t3) = ppr t1 <+> ppr op <+> ppr t2 <+> char '~' <+> ppr t3
648 ppr (Leq t1 t2) = ppr t1 <+> text "<=" <+> ppr t2
649 ppr (Eq t1 t2) = ppr t1 <+> char '~' <+> ppr t2
650
651 instance Outputable Op where
652 ppr op = case op of
653 Add -> char '+'
654 Mul -> char '*'
655 Exp -> char '^'
656