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