Added normalisation of closed type families

compiler/main/GHC.hs

@@ -157,7 +157,7 @@ module GHC (
         TyCon, 
         tyConTyVars, tyConDataCons, tyConArity,
         isClassTyCon, isSynTyCon, isNewTyCon, isPrimTyCon, isFunTyCon,
-        isFamilyTyCon, tyConClass_maybe,
+        isFamilyTyCon, isOpenFamilyTyCon, tyConClass_maybe,
         synTyConRhs_maybe, synTyConDefn_maybe, synTyConResKind,
 
         -- ** Type variables

compiler/typecheck/FamInst.lhs

@@ -211,7 +211,7 @@ which implies that :R42T was declared as 'data instance T [a]'.
 \begin{code}
 tcLookupFamInst :: TyCon -> [Type] -> TcM (Maybe FamInstMatch)
 tcLookupFamInst tycon tys
-  | not (isFamilyTyCon tycon)
+  | not (isOpenFamilyTyCon tycon)
   = return Nothing
   | otherwise
   = do { instEnv <- tcGetFamInstEnvs
@@ -229,7 +229,7 @@ tcLookupDataFamInst :: TyCon -> [Type] -> TcM (TyCon, [Type])
 -- Find the instance of a data family
 -- Note [Looking up family instances for deriving]
 tcLookupDataFamInst tycon tys
-  | not (isFamilyTyCon tycon)
+  | not (isOpenFamilyTyCon tycon)
   = return (tycon, tys)
   | otherwise
   = ASSERT( isAlgTyCon tycon )

compiler/typecheck/TcInteract.lhs

@@ -20,7 +20,6 @@ import VarSet
 import Type
 import Unify
 import FamInstEnv
-import Coercion( mkUnbranchedAxInstRHS )
 
 import Var
 import TcType
@@ -1482,16 +1481,13 @@ doTopReactFunEq _ct fl fun_tc args xi loc
     do { match_res <- matchFam fun_tc args   -- See Note [MATCHING-SYNONYMS]
        ; case match_res of {
            Nothing -> return NoTopInt ;
-           Just (FamInstMatch { fim_instance = famInst
-                              , fim_tys      = rep_tys }) -> 
+           Just (co, ty) ->
 
     -- Found a top-level instance
     do {    -- Add it to the solved goals
          unless (isDerived fl) (addSolvedFunEq fam_ty fl xi)
 
-       ; let coe_ax = famInstAxiom famInst 
-       ; succeed_with "Fun/Top" (mkTcUnbranchedAxInstCo coe_ax rep_tys)
-                      (mkUnbranchedAxInstRHS coe_ax rep_tys) } } } } }
+       ; succeed_with "Fun/Top" co ty } } } } }
   where
     fam_ty = mkTyConApp fun_tc args
 
@@ -1759,7 +1755,7 @@ matchClassInst _ clas [ k, ty ] _
     case unwrapNewTyCon_maybe (classTyCon clas) of
       Just (_,dictRep, axDict)
         | Just tcSing <- tyConAppTyCon_maybe dictRep ->
-           do mbInst <- matchFam tcSing [k,ty]
+           do mbInst <- matchOpenFam tcSing [k,ty]
               case mbInst of
                 Just FamInstMatch
                   { fim_instance = FamInst { fi_axiom  = axDataFam

compiler/typecheck/TcRnDriver.lhs

@@ -1765,7 +1765,7 @@ lookupInsts (ATyCon tc)
   = do  { inst_envs <- tcGetInstEnvs
         ; return (classInstances inst_envs cls, []) }
 
-  | isFamilyTyCon tc || isTyConAssoc tc
+  | isOpenFamilyTyCon tc || isTyConAssoc tc
   = do  { inst_envs <- tcGetFamInstEnvs
         ; return ([], familyInstances inst_envs tc) }
 

compiler/typecheck/TcSMonad.lhs

@@ -88,7 +88,7 @@ module TcSMonad (
 
     getDefaultInfo, getDynFlags,
 
-    matchClass, matchFam, MatchInstResult (..), 
+    matchClass, matchFam, matchOpenFam, MatchInstResult (..), 
     checkWellStagedDFun, 
     pprEq                                    -- Smaller utils, re-exported from TcM
                                              -- TODO (DV): these are only really used in the 
@@ -135,6 +135,7 @@ import TcRnTypes
 import Unique 
 import UniqFM
 import Maybes ( orElse, catMaybes, firstJust )
+import Pair ( pSnd )
 
 import Control.Monad( unless, when, zipWithM )
 import Data.IORef
@@ -1690,8 +1691,30 @@ matchClass clas tys
 	}
         }
 
-matchFam :: TyCon -> [Type] -> TcS (Maybe FamInstMatch)
-matchFam tycon args = wrapTcS $ tcLookupFamInst tycon args
+matchOpenFam :: TyCon -> [Type] -> TcS (Maybe FamInstMatch)
+matchOpenFam tycon args = wrapTcS $ tcLookupFamInst tycon args
+
+matchFam :: TyCon -> [Type] -> TcS (Maybe (TcCoercion, TcType))
+matchFam tycon args
+  | isOpenSynFamilyTyCon tycon
+  = do { maybe_match <- matchOpenFam tycon args
+       ; case maybe_match of
+           Nothing -> return Nothing
+           Just (FamInstMatch { fim_instance = famInst
+                              , fim_tys      = inst_tys })
+             -> let co = mkTcUnbranchedAxInstCo (famInstAxiom famInst) inst_tys
+                    ty = pSnd $ tcCoercionKind co
+                in return $ Just (co, ty) }
+
+  | Just ax <- isClosedSynFamilyTyCon_maybe tycon
+  , Just (ind, inst_tys) <- chooseBranch ax args
+  = let co = mkTcAxInstCo ax ind inst_tys
+        ty = pSnd (tcCoercionKind co)
+    in return $ Just (co, ty)
+
+  | otherwise
+  = return Nothing
+       
 \end{code}
 
 \begin{code}

compiler/typecheck/TcSplice.lhs

@@ -1027,7 +1027,7 @@ reifyInstances th_nm th_tys
                -> do { inst_envs <- tcGetInstEnvs
                      ; let (matches, unifies, _) = lookupInstEnv inst_envs cls tys
                      ; mapM reifyClassInstance (map fst matches ++ unifies) }
-               | isFamilyTyCon tc
+               | isOpenFamilyTyCon tc
                -> do { inst_envs <- tcGetFamInstEnvs
                      ; let matches = lookupFamInstEnv inst_envs tc tys
                      ; mapM (reifyFamilyInstance . fim_instance) matches }

compiler/types/CoAxiom.lhs

@@ -20,7 +20,7 @@ module CoAxiom (
 
        toBranchedAxiom, toUnbranchedAxiom,
        coAxiomName, coAxiomArity, coAxiomBranches,
-       coAxiomTyCon, isImplicitCoAxiom,
+       coAxiomTyCon, isImplicitCoAxiom, coAxiomNumPats,
        coAxiomNthBranch, coAxiomSingleBranch_maybe,
        coAxiomSingleBranch, coAxBranchTyVars, coAxBranchLHS,
        coAxBranchRHS, coAxBranchSpan
@@ -178,12 +178,12 @@ brListMapM f (NextBranch h t) = do { fh <- f h
 
 brListFoldlM_ :: forall a b m br. Monad m
               => (a -> b -> m a) -> a -> BranchList b br -> m ()
-brListFoldlM_ f z brs = do { _ <- go brs
+brListFoldlM_ f z brs = do { _ <- go z brs
                            ; return () }
-  where go :: forall br'. Monad m => BranchList b br' -> m a
-        go (FirstBranch b)  = f z b
-        go (NextBranch h t) = do { t' <- go t
-                                 ; f t' h }
+  where go :: forall br'. Monad m => a -> BranchList b br' -> m a
+        go acc (FirstBranch b)  = f acc b
+        go acc (NextBranch h t) = do { fh <- f acc h
+                                     ; go fh t }
 
 -- zipWith
 brListZipWith :: (a -> b -> c) -> BranchList a br1 -> BranchList b br2 -> [c]
@@ -241,6 +241,9 @@ toUnbranchedAxiom :: CoAxiom br -> CoAxiom Unbranched
 toUnbranchedAxiom (CoAxiom unique name tc branches implicit)
   = CoAxiom unique name tc (toUnbranchedList branches) implicit
 
+coAxiomNumPats :: CoAxiom br -> Int
+coAxiomNumPats = length . coAxBranchLHS . (flip coAxiomNthBranch 0)
+
 coAxiomNthBranch :: CoAxiom br -> BranchIndex -> CoAxBranch
 coAxiomNthBranch (CoAxiom { co_ax_branches = bs }) index
   = brListNth bs index

compiler/types/Coercion.lhs

@@ -682,7 +682,7 @@ mkCoVarCo cv
 mkReflCo :: Type -> Coercion
 mkReflCo = Refl
 
-mkAxInstCo :: CoAxiom br -> Int -> [Type] -> Coercion
+mkAxInstCo :: CoAxiom br -> BranchIndex -> [Type] -> Coercion
 -- mkAxInstCo can legitimately be called over-staturated; 
 -- i.e. with more type arguments than the coercion requires
 mkAxInstCo ax index tys

compiler/types/FamInstEnv.lhs

@@ -22,7 +22,7 @@ module FamInstEnv (
         isDominatedBy,
 	
 	-- Normalisation
-	topNormaliseType, normaliseType, normaliseTcApp
+	chooseBranch, topNormaliseType, normaliseType, normaliseTcApp
     ) where
 
 #include "HsVersions.h"
@@ -42,6 +42,7 @@ import UniqFM
 import Outputable
 import Maybes
 import Util
+import Pair
 import NameSet
 import FastString
 \end{code}
@@ -472,7 +473,7 @@ lookup_fam_inst_env' 	      -- The worker, local to this module
     -> TyCon -> [Type]		-- What we are looking for
     -> [FamInstMatch] 	        -- Successful matches
 lookup_fam_inst_env' match_fun one_sided ie fam tys
-  | not (isFamilyTyCon fam) 
+  | not (isOpenFamilyTyCon fam) 
   = []
   | otherwise
   = ASSERT2( n_tys >= arity, ppr fam <+> ppr tys )	-- Family type applications must be saturated
@@ -571,6 +572,73 @@ isDominatedBy branch branches
         = isJust $ tcMatchTys (mkVarSet tvs) tys lhs
 \end{code}
 
+%************************************************************************
+%*									*
+                Choosing an axiom application
+%*									*
+%************************************************************************
+
+The lookupFamInstEnv function does a nice job for *open* type families,
+but we also need to handle closed ones when normalising a type:
+
+\begin{code}
+
+-- The TyCon can be oversaturated. This works on both open and closed families
+chooseAxiom :: FamInstEnvs -> TyCon -> [Type] -> Maybe (Coercion, Type)
+chooseAxiom envs tc tys
+  | isOpenFamilyTyCon tc
+  , [FamInstMatch { fim_instance = fam_inst
+                  , fim_tys =      inst_tys }] <- lookupFamInstEnv envs tc tys
+  = let co = mkUnbranchedAxInstCo (famInstAxiom fam_inst) inst_tys
+        ty = pSnd (coercionKind co)
+    in Just (co, ty)
+
+  | Just ax <- isClosedSynFamilyTyCon_maybe tc
+  , Just (ind, inst_tys) <- chooseBranch ax tys
+  = let co = mkAxInstCo ax ind inst_tys
+        ty = pSnd (coercionKind co)
+    in Just (co, ty)
+
+  | otherwise
+  = Nothing
+
+-- The axiom can be oversaturated. (Closed families only.)
+chooseBranch :: CoAxiom Branched -> [Type] -> Maybe (BranchIndex, [Type])
+chooseBranch axiom tys
+  = do { let num_pats = coAxiomNumPats axiom
+             (target_tys, extra_tys) = splitAt num_pats tys
+             branches = coAxiomBranches axiom
+       ; (ind, inst_tys) <- findBranch [] (fromBranchList branches) 0 target_tys
+       ; return (ind, inst_tys ++ extra_tys) }
+
+-- The axiom must *not* be oversaturated
+findBranch :: [CoAxBranch]             -- branches seen so far
+           -> [CoAxBranch]             -- branches to check
+           -> BranchIndex              -- index of current branch
+           -> [Type]                   -- target types
+           -> Maybe (BranchIndex, [Type])
+findBranch prev_branches
+           (cur@CoAxBranch { cab_tvs = tpl_tvs, cab_lhs = tpl_lhs }
+              : rest) ind target_tys
+  = case tcMatchTys (mkVarSet tpl_tvs) tpl_lhs target_tys of
+      Just subst -- matching worked. now, check for apartness.
+        |  all (isSurelyApart . tcApartTys instanceBindFun target_tys . coAxBranchLHS)
+             prev_branches
+        -> -- matching worked & we're apart from all incompatible branches. success
+           Just (ind, substTyVars subst tpl_tvs)
+
+      -- failure. keep looking
+      _ -> findBranch (cur : prev_branches) rest (ind+1) target_tys
+
+  where isSurelyApart SurelyApart = True
+        isSurelyApart _           = False
+
+-- fail if no branches left
+findBranch _ [] _ _ = Nothing
+
+\end{code}
+
+
 %************************************************************************
 %*									*
 		Looking up a family instance
@@ -607,7 +675,7 @@ topNormaliseType env ty
           else let nt_co = mkUnbranchedAxInstCo (newTyConCo tc) tys
                in add_co nt_co rec_nts' nt_rhs
 
-	| isFamilyTyCon tc		-- Expand open tycons
+	| isFamilyTyCon tc		-- Expand family tycons
 	, (co, ty) <- normaliseTcApp env tc tys
 		-- Note that normaliseType fully normalises 'tys', 
 		-- It has do to so to be sure that nested calls like
@@ -633,12 +701,8 @@ normaliseTcApp :: FamInstEnvs -> TyCon -> [Type] -> (Coercion, Type)
 normaliseTcApp env tc tys
   | isFamilyTyCon tc
   , tyConArity tc <= length tys	   -- Unsaturated data families are possible
-  , [FamInstMatch { fim_instance = fam_inst
-                  , fim_tys =      inst_tys }] <- lookupFamInstEnv env tc ntys 
-  = let    -- A matching family instance exists
-        ax              = famInstAxiom fam_inst
-        co              = mkUnbranchedAxInstCo  ax inst_tys
-        rhs             = mkUnbranchedAxInstRHS ax inst_tys
+  , Just (co, rhs) <- chooseAxiom env tc ntys
+  = let    -- A reduction is possible
 	first_coi       = mkTransCo tycon_coi co
 	(rest_coi,nty)  = normaliseType env rhs
 	fix_coi         = mkTransCo first_coi rest_coi

compiler/types/TyCon.lhs

@@ -45,7 +45,8 @@ module TyCon(
         isDataTyCon, isProductTyCon, isDataProductTyCon_maybe,
         isEnumerationTyCon,
         isNewTyCon, isAbstractTyCon,
-        isFamilyTyCon, isSynFamilyTyCon, isDataFamilyTyCon,
+        isFamilyTyCon, isOpenFamilyTyCon,
+        isSynFamilyTyCon, isDataFamilyTyCon,
         isOpenSynFamilyTyCon, isClosedSynFamilyTyCon_maybe,
         isUnLiftedTyCon,
         isGadtSyntaxTyCon, isDistinctTyCon, isDistinctAlgRhs,
@@ -1161,6 +1162,13 @@ isFamilyTyCon (SynTyCon {synTcRhs = ClosedSynFamilyTyCon {} }) = True
 isFamilyTyCon (AlgTyCon {algTcRhs = DataFamilyTyCon {}})       = True
 isFamilyTyCon _ = False
 
+-- | Is this a 'TyCon', synonym or otherwise, that defines an family with
+-- instances?
+isOpenFamilyTyCon :: TyCon -> Bool
+isOpenFamilyTyCon (SynTyCon {synTcRhs = OpenSynFamilyTyCon }) = True
+isOpenFamilyTyCon (AlgTyCon {algTcRhs = DataFamilyTyCon })    = True
+isOpenFamilyTyCon _ = False
+
 -- | Is this a synonym 'TyCon' that can have may have further instances appear?
 isSynFamilyTyCon :: TyCon -> Bool
 isSynFamilyTyCon (SynTyCon {synTcRhs = OpenSynFamilyTyCon {}})   = True

compiler/vectorise/Vectorise/Monad/InstEnv.hs

@@ -67,7 +67,7 @@ lookupInst cls tys
 --
 lookupFamInst :: TyCon -> [Type] -> VM FamInstMatch
 lookupFamInst tycon tys
-  = ASSERT( isFamilyTyCon tycon )
+  = ASSERT( isOpenFamilyTyCon tycon )
     do { instEnv <- readGEnv global_fam_inst_env
        ; case lookupFamInstEnv instEnv tycon tys of
            [match] -> return match