Add support for evaluation of type-level natural numbers.

This patch implements some simple evaluation of type-level expressions
featuring natural numbers.  We can evaluate *concrete* expressions that
use the built-in type families (+), (*), (^), and (<=?), declared in
GHC.TypeLits.   We can also do some type inference involving these
functions.  For example, if we encounter a constraint such as `(2 + x) ~ 5`
we can infer that `x` must be 3.  Note, however, this is used only to
resolve unification variables (i.e., as a form of a constraint improvement)
and not to generate new facts.  This is similar to how functional
dependencies work in GHC.

The patch adds a new form of coercion, `AxiomRuleCo`, which makes use
of a new form of axiom called `CoAxiomRule`.  This is the form of evidence
generate when we solve a constraint, such as `(1 + 2) ~ 3`.

The patch also adds support for built-in type-families, by adding a new
form of TyCon rhs: `BuiltInSynFamTyCon`.  such built-in type-family
constructors contain a record with functions that are used by the
constraint solver to simplify and improve constraints involving the
built-in function (see `TcInteract`).  The record in defined in `FamInst`.

The type constructors and rules for evaluating the type-level functions
are in a new module called `TcTypeNats`.

3 files changed, 84 insertions, 5 deletions

diff --git a/compiler/coreSyn/CoreLint.lhs b/compiler/coreSyn/CoreLint.lhs
index ffddd78516..9a9860f0fa 100644
--- a/compiler/coreSyn/CoreLint.lhs
+++ b/compiler/coreSyn/CoreLint.lhs
@@ -54,6 +54,7 @@ import OptCoercion ( checkAxInstCo )
 import Control.Monad
 import MonadUtils
 import Data.Maybe
+import Pair
 \end{code}
 
 Note [GHC Formalism]
@@ -987,6 +988,53 @@ lintCoercion co@(SubCo co')
   = do { (k,s,t,r) <- lintCoercion co'
        ; checkRole co Nominal r
        ; return (k,s,t,Representational) }
+
+
+lintCoercion this@(AxiomRuleCo co ts cs)
+  = do _ks <- mapM lintType ts
+       eqs <- mapM lintCoercion cs
+
+       let tyNum = length ts
+
+       case compare (coaxrTypeArity co) tyNum of
+         EQ -> return ()
+         LT -> err "Too many type arguments"
+                    [ txt "expected" <+> int (coaxrTypeArity co)
+                    , txt "provided" <+> int tyNum ]
+         GT -> err "Not enough type arguments"
+                    [ txt "expected" <+> int (coaxrTypeArity co)
+                          , txt "provided" <+> int tyNum ]
+       checkRoles 0 (coaxrAsmpRoles co) eqs
+
+       case coaxrProves co ts [ Pair l r | (_,l,r,_) <- eqs ] of
+         Nothing -> err "Malformed use of AxiomRuleCo" [ ppr this ]
+         Just (Pair l r) ->
+           do kL <- lintType l
+              kR <- lintType r
+              unless (eqKind kL kR)
+                $ err "Kind error in CoAxiomRule"
+                       [ppr kL <+> txt "/=" <+> ppr kR]
+              return (kL, l, r, coaxrRole co)
+  where
+  txt       = ptext . sLit
+  err m xs  = failWithL $
+                hang (txt m) 2 $ vcat (txt "Rule:" <+> ppr (coaxrName co) : xs)
+
+  checkRoles n (e : es) ((_,_,_,r) : rs)
+    | e == r    = checkRoles (n+1) es rs
+    | otherwise = err "Argument roles mismatch"
+                      [ txt "In argument:" <+> int (n+1)
+                      , txt "Expected:" <+> ppr e
+                      , txt "Found:" <+> ppr r ]
+  checkRoles _ [] []  = return ()
+  checkRoles n [] rs  = err "Too many coercion arguments"
+                          [ txt "Expected:" <+> int n
+                          , txt "Provided:" <+> int (n + length rs) ]
+
+  checkRoles n es []  = err "Not enough coercion arguments"
+                          [ txt "Expected:" <+> int (n + length es)
+                          , txt "Provided:" <+> int n ]
+
 \end{code}
 
 %************************************************************************
diff --git a/compiler/coreSyn/MkExternalCore.lhs b/compiler/coreSyn/MkExternalCore.lhs
index bdb54d80d6..6a6f0551ed 100644
--- a/compiler/coreSyn/MkExternalCore.lhs
+++ b/compiler/coreSyn/MkExternalCore.lhs
@@ -332,6 +332,8 @@ make_co dflags (NthCo d co)          = C.NthCoercion d (make_co dflags co)
 make_co dflags (LRCo lr co)          = C.LRCoercion (make_lr lr) (make_co dflags co)
 make_co dflags (InstCo co ty)        = C.InstCoercion (make_co dflags co) (make_ty dflags ty)
 make_co dflags (SubCo co)            = C.SubCoercion (make_co dflags co)
+make_co _ (AxiomRuleCo {})           = panic "make_co AxiomRuleCo: not yet implemented"
+
 
 make_lr :: LeftOrRight -> C.LeftOrRight
 make_lr CLeft  = C.CLeft
diff --git a/compiler/coreSyn/TrieMap.lhs b/compiler/coreSyn/TrieMap.lhs
index f8ad8da5f4..255ab899d7 100644
--- a/compiler/coreSyn/TrieMap.lhs
+++ b/compiler/coreSyn/TrieMap.lhs
@@ -18,7 +18,8 @@ module TrieMap(
    CoercionMap, 
    MaybeMap, 
    ListMap,
-   TrieMap(..)
+   TrieMap(..),
+   lookupTypeMapTyCon
  ) where
 
 import CoreSyn
@@ -27,10 +28,12 @@ import Literal
 import Name
 import Type
 import TypeRep
+import TyCon(TyCon)
 import Var
 import UniqFM
 import Unique( Unique )
 import FastString(FastString)
+import CoAxiom(CoAxiomRule(coaxrName))
 
 import qualified Data.Map    as Map
 import qualified Data.IntMap as IntMap
@@ -471,7 +474,10 @@ data CoercionMap a
        , km_left   :: CoercionMap a
        , km_right  :: CoercionMap a
        , km_inst   :: CoercionMap (TypeMap a) 
-       , km_sub    :: CoercionMap a }
+       , km_sub    :: CoercionMap a
+       , km_axiom_rule :: Map.Map FastString
+                                  (ListMap TypeMap (ListMap CoercionMap a))
+       }
 
 wrapEmptyKM :: CoercionMap a
 wrapEmptyKM = KM { km_refl = emptyTM, km_tc_app = emptyTM
@@ -479,7 +485,8 @@ wrapEmptyKM = KM { km_refl = emptyTM, km_tc_app = emptyTM
                  , km_var = emptyTM, km_axiom = emptyNameEnv
                  , km_univ = emptyTM, km_sym = emptyTM, km_trans = emptyTM
                  , km_nth = emptyTM, km_left = emptyTM, km_right = emptyTM
-                 , km_inst = emptyTM, km_sub = emptyTM }
+                 , km_inst = emptyTM, km_sub = emptyTM 
+                 , km_axiom_rule = emptyTM }
 
 instance TrieMap CoercionMap where
    type Key CoercionMap = Coercion
@@ -496,7 +503,8 @@ mapC f (KM { km_refl = krefl, km_tc_app = ktc
            , km_var = kvar, km_axiom = kax
            , km_univ   = kuniv  , km_sym = ksym, km_trans = ktrans
            , km_nth = knth, km_left = kml, km_right = kmr
-           , km_inst = kinst, km_sub = ksub })
+           , km_inst = kinst, km_sub = ksub
+           , km_axiom_rule = kaxr })
   = KM { km_refl   = mapTM (mapTM f) krefl
        , km_tc_app = mapTM (mapNameEnv (mapTM f)) ktc
        , km_app    = mapTM (mapTM f) kapp
@@ -510,7 +518,9 @@ mapC f (KM { km_refl = krefl, km_tc_app = ktc
        , km_left   = mapTM f kml
        , km_right  = mapTM f kmr
        , km_inst   = mapTM (mapTM f) kinst
-       , km_sub    = mapTM f ksub }
+       , km_sub    = mapTM f ksub
+       , km_axiom_rule = mapTM (mapTM (mapTM f)) kaxr
+       }
 
 lkC :: CmEnv -> Coercion -> CoercionMap a -> Maybe a
 lkC env co m 
@@ -531,6 +541,11 @@ lkC env co m
     go (LRCo CLeft  c)         = km_left   >.> lkC env c
     go (LRCo CRight c)         = km_right  >.> lkC env c
     go (SubCo c)               = km_sub    >.> lkC env c
+    go (AxiomRuleCo co ts cs)  = km_axiom_rule >.>
+                                    lookupTM (coaxrName co) >=>
+                                    lkList (lkT env) ts >=>
+                                    lkList (lkC env) cs
+
 
 xtC :: CmEnv -> Coercion -> XT a -> CoercionMap a -> CoercionMap a
 xtC env co f EmptyKM = xtC env co f wrapEmptyKM
@@ -549,6 +564,10 @@ xtC env (NthCo n c)             f m = m { km_nth    = km_nth m |> xtInt n |>> xt
 xtC env (LRCo CLeft  c)         f m = m { km_left   = km_left  m |> xtC env c f } 
 xtC env (LRCo CRight c)         f m = m { km_right  = km_right m |> xtC env c f }
 xtC env (SubCo c)               f m = m { km_sub    = km_sub m |> xtC env c f } 
+xtC env (AxiomRuleCo co ts cs)  f m = m { km_axiom_rule = km_axiom_rule m
+                                                        |>  alterTM (coaxrName co)
+                                                        |>> xtList (xtT env) ts
+                                                        |>> xtList (xtC env) cs f}
 
 fdC :: (a -> b -> b) -> CoercionMap a -> b -> b
 fdC _ EmptyKM = \z -> z
@@ -566,6 +585,7 @@ fdC k m = foldTM (foldTM k) (km_refl m)
         . foldTM k          (km_right m)
         . foldTM (foldTM k) (km_inst m)
         . foldTM k          (km_sub m)
+        . foldTM (foldTM (foldTM k)) (km_axiom_rule m)
 
 \end{code}
 
@@ -630,6 +650,15 @@ emptyTypeMap = EmptyTM
 lookupTypeMap :: TypeMap a -> Type -> Maybe a
 lookupTypeMap cm t = lkT emptyCME t cm
 
+-- Returns the type map entries that have keys starting with the given tycon.
+-- This only considers saturated applications (i.e. TyConApp ones).
+lookupTypeMapTyCon :: TypeMap a -> TyCon -> [a]
+lookupTypeMapTyCon EmptyTM _ = []
+lookupTypeMapTyCon TM { tm_tc_app = cs } tc =
+  case lookupUFM cs tc of
+    Nothing -> []
+    Just xs -> foldTM (:) xs []
+
 extendTypeMap :: TypeMap a -> Type -> a -> TypeMap a
 extendTypeMap m t v = xtT emptyCME t (\_ -> Just v) m