Refactor the topNormaliseNewType story, fixing Trac #8467

A bit of a mess had accumulated, with unclear invariants.

* Remove splitNewTypeRepCo_maybe, in favour of topNormaliseNewType_maybe
  (which had the same signature but behaved subtly differently).

* Make topNormaliseNewType_maybe guaranteed to return a non-newtype
  if it says (Just ty).  This is what was causing the loop in #8467

* Apply similar tidying up to FamInstEnv.topNormaliseType

compiler/coreSyn/CoreArity.lhs

@@ -902,7 +902,7 @@ mkEtaWW orig_n orig_expr in_scope orig_ty
            -- Avoid free vars of the original expression
        = go (n-1) subst' res_ty (EtaVar eta_id' : eis)
        				   
-       | Just(ty',co) <- splitNewTypeRepCo_maybe ty
+       | Just (co, ty') <- topNormaliseNewType_maybe ty
        = 	-- Given this:
        		-- 	newtype T = MkT ([T] -> Int)
        		-- Consider eta-expanding this

compiler/deSugar/DsCCall.lhs

@@ -143,7 +143,7 @@ unboxArg arg
   = return (arg, \body -> body)
 
   -- Recursive newtypes
-  | Just(_rep_ty, co) <- splitNewTypeRepCo_maybe arg_ty
+  | Just(co, _rep_ty) <- topNormaliseNewType_maybe arg_ty
   = unboxArg (mkCast arg co)
       
   -- Booleans
@@ -344,8 +344,8 @@ resultWrapper result_ty
                                    [(DEFAULT                    ,[],Var trueDataConId ),
                                     (LitAlt (mkMachInt dflags 0),[],Var falseDataConId)])
 
-  -- Recursive newtypes
-  | Just (rep_ty, co) <- splitNewTypeRepCo_maybe result_ty
+  -- Newtypes
+  | Just (co, rep_ty) <- topNormaliseNewType_maybe result_ty
   = do (maybe_ty, wrapper) <- resultWrapper rep_ty
        return (maybe_ty, \e -> mkCast (wrapper e) (mkSymCo co))
 

compiler/stranal/WwLib.lhs

@@ -112,7 +112,7 @@ the unusable strictness-info into the interfaces.
 
 \begin{code}
 mkWwBodies :: DynFlags
-       -> Type				-- Type of original function
+           -> Type				-- Type of original function
 	   -> [Demand]				-- Strictness of original function
 	   -> DmdResult				-- Info about function result
 	   -> [Bool]				-- One-shot-ness of the function
@@ -285,42 +285,9 @@ mkWWargs :: TvSubst		-- Freshening substitution to apply to the type
 		     Type)			-- Type of wrapper body
 
 mkWWargs subst fun_ty arg_info
-  | Just (rep_ty, co) <- splitNewTypeRepCo_maybe fun_ty
-   	-- The newtype case is for when the function has
-	-- a recursive newtype after the arrow (rare)
-	-- We check for arity >= 0 to avoid looping in the case
-	-- of a function whose type is, in effect, infinite
-	-- [Arity is driven by looking at the term, not just the type.]
-	--
-	-- It's also important when we have a function returning (say) a pair
-	-- wrapped in a recursive newtype, at least if CPR analysis can look 
-	-- through such newtypes, which it probably can since they are 
-	-- simply coerces.
-	--
-	-- Note (Sept 08): This case applies even if demands is empty.
-	--		   I'm not quite sure why; perhaps it makes it
-	--		   easier for CPR
-  = do { (wrap_args, wrap_fn_args, work_fn_args, res_ty)
-	    <-  mkWWargs subst rep_ty arg_info
- 	; return (wrap_args,
-	     	  \e -> Cast (wrap_fn_args e) (mkSymCo co),
-     		  \e -> work_fn_args (Cast e co),
-     		  res_ty) } 
-
   | null arg_info
   = return ([], id, id, substTy subst fun_ty)
 
-  | Just (tv, fun_ty') <- splitForAllTy_maybe fun_ty
-  = do 	{ let (subst', tv') = substTyVarBndr subst tv
-		-- This substTyVarBndr clones the type variable when necy
-		-- See Note [Freshen type variables]
-  	; (wrap_args, wrap_fn_args, work_fn_args, res_ty)
-	     <- mkWWargs subst' fun_ty' arg_info
-	; return (tv' : wrap_args,
-        	  Lam tv' . wrap_fn_args,
-        	  work_fn_args . (`App` Type (mkTyVarTy tv')),
-        	  res_ty) }
-
   | ((dmd,one_shot):arg_info') <- arg_info
   , Just (arg_ty, fun_ty') <- splitFunTy_maybe fun_ty
   = do	{ uniq <- getUniqueM
@@ -333,6 +300,33 @@ mkWWargs subst fun_ty arg_info
         	  work_fn_args . (`App` varToCoreExpr id),
         	  res_ty) }
 
+  | Just (tv, fun_ty') <- splitForAllTy_maybe fun_ty
+  = do 	{ let (subst', tv') = substTyVarBndr subst tv
+		-- This substTyVarBndr clones the type variable when necy
+		-- See Note [Freshen type variables]
+  	; (wrap_args, wrap_fn_args, work_fn_args, res_ty)
+	     <- mkWWargs subst' fun_ty' arg_info
+	; return (tv' : wrap_args,
+        	  Lam tv' . wrap_fn_args,
+        	  work_fn_args . (`App` Type (mkTyVarTy tv')),
+        	  res_ty) }
+
+  | Just (co, rep_ty) <- topNormaliseNewType_maybe fun_ty
+   	-- The newtype case is for when the function has
+	-- a newtype after the arrow (rare)
+	--
+	-- It's also important when we have a function returning (say) a pair
+	-- wrapped in a  newtype, at least if CPR analysis can look
+	-- through such newtypes, which it probably can since they are
+	-- simply coerces.
+
+  = do { (wrap_args, wrap_fn_args, work_fn_args, res_ty)
+	    <-  mkWWargs subst rep_ty arg_info
+ 	; return (wrap_args,
+	     	  \e -> Cast (wrap_fn_args e) (mkSymCo co),
+     		  \e -> work_fn_args (Cast e co),
+     		  res_ty) } 
+
   | otherwise
   = WARN( True, ppr fun_ty )			-- Should not happen: if there is a demand
     return ([], id, id, substTy subst fun_ty) 	-- then there should be a function arrow
@@ -455,14 +449,14 @@ mkWWstr_one dflags arg
   | isStrictDmd dmd
   , Just cs <- splitProdDmd_maybe dmd
       -- See Note [Unpacking arguments with product and polymorphic demands]
-  , Just (data_con, inst_tys, inst_con_arg_tys, co) 
+  , Just (data_con, inst_tys, inst_con_arg_tys, co)
              <- deepSplitProductType_maybe (idType arg)
   =  do { (uniq1:uniqs) <- getUniquesM
 	; let   unpk_args      = zipWith mk_ww_local uniqs inst_con_arg_tys
 	        unpk_args_w_ds = zipWithEqual "mkWWstr" set_worker_arg_info unpk_args cs
-	        unbox_fn       = mkUnpackCase (Var arg `mkCast` co) uniq1
+	        unbox_fn       = mkUnpackCase (Var arg) co uniq1
                                               data_con unpk_args
-	        rebox_fn       = Let (NonRec arg con_app) 
+	        rebox_fn       = Let (NonRec arg con_app)
 	        con_app        = mkConApp2 data_con inst_tys unpk_args `mkCast` mkSymCo co
 	 ; (worker_args, wrap_fn, work_fn) <- mkWWstr dflags unpk_args_w_ds
 	 ; return (worker_args, unbox_fn . wrap_fn, work_fn . rebox_fn) }
@@ -489,17 +483,21 @@ nop_fn body = body
 \begin{code}
 deepSplitProductType_maybe :: Type -> Maybe (DataCon, [Type], [Type], Coercion)
 -- If    deepSplitProductType_maybe ty = Just (dc, tys, arg_tys, co)
--- then  dc @ tys (args::arg_tys)  |> co :: ty
+-- then  dc @ tys (args::arg_tys) :: rep_ty
+--       co :: ty ~ rep_ty
 deepSplitProductType_maybe ty
-  | let (ty1, co) = topNormaliseNewType ty `orElse` (ty, mkReflCo Representational ty)
+  | let (co, ty1) = topNormaliseNewType_maybe ty `orElse` (mkReflCo Representational ty, ty)
   , Just (tc, tc_args) <- splitTyConApp_maybe ty1
   , Just con <- isDataProductTyCon_maybe tc
   = Just (con, tc_args, dataConInstArgTys con tc_args, co)
 deepSplitProductType_maybe _ = Nothing
 
 deepSplitCprType_maybe :: ConTag -> Type -> Maybe (DataCon, [Type], [Type], Coercion)
+-- If    deepSplitCprType_maybe n ty = Just (dc, tys, arg_tys, co)
+-- then  dc @ tys (args::arg_tys) :: rep_ty
+--       co :: ty ~ rep_ty
 deepSplitCprType_maybe con_tag ty
-  | let (ty1, co) = topNormaliseNewType ty `orElse` (ty, mkReflCo Representational ty)
+  | let (co, ty1) = topNormaliseNewType_maybe ty `orElse` (mkReflCo Representational ty, ty)
   , Just (tc, tc_args) <- splitTyConApp_maybe ty1
   , isDataTyCon tc
   , let cons = tyConDataCons tc
@@ -540,7 +538,7 @@ mkWWcpr body_ty res
                     |  otherwise
                     -> WARN( True, text "mkWWcpr: non-algebraic or open body type" <+> ppr body_ty )
                        return (id, id, body_ty)
-          
+
 mkWWcpr_help :: (DataCon, [Type], [Type], Coercion)
              -> UniqSM (CoreExpr -> CoreExpr, CoreExpr -> CoreExpr, Type)
 
@@ -553,10 +551,10 @@ mkWWcpr_help (data_con, inst_tys, arg_tys, co)
 	-- Worker:	case (   ..body..    ) of C x -> x
   = do { (work_uniq : arg_uniq : _) <- getUniquesM
        ; let arg       = mk_ww_local arg_uniq  arg_ty1
-	     con_app   = mkConApp2 data_con inst_tys [arg] `mkCast` co
+	     con_app   = mkConApp2 data_con inst_tys [arg] `mkCast` mkSymCo co
 
        ; return ( \ wkr_call -> Case wkr_call arg (exprType con_app) [(DEFAULT, [], con_app)]
-                , \ body     -> mkUnpackCase body work_uniq data_con [arg] (Var arg)
+                , \ body     -> mkUnpackCase body co work_uniq data_con [arg] (Var arg)
                 , arg_ty1 ) }
 
   | otherwise 	-- The general case
@@ -567,25 +565,25 @@ mkWWcpr_help (data_con, inst_tys, arg_tys, co)
 	     ubx_tup_con  = tupleCon UnboxedTuple (length arg_tys)
 	     ubx_tup_ty	  = exprType ubx_tup_app
 	     ubx_tup_app  = mkConApp2 ubx_tup_con arg_tys args
-             con_app	  = mkConApp2 data_con inst_tys args `mkCast` co
+             con_app	  = mkConApp2 data_con inst_tys args `mkCast` mkSymCo co
 
        ; return ( \ wkr_call -> Case wkr_call wrap_wild (exprType con_app)  [(DataAlt ubx_tup_con, args, con_app)]
-		, \ body     -> mkUnpackCase body work_uniq data_con args ubx_tup_app
+		, \ body     -> mkUnpackCase body co work_uniq data_con args ubx_tup_app
 		, ubx_tup_ty ) }
 
-mkUnpackCase ::  CoreExpr -> Unique -> DataCon -> [Id] -> CoreExpr -> CoreExpr
--- (mkUnpackCase e bndr Con args body)
+mkUnpackCase ::  CoreExpr -> Coercion -> Unique -> DataCon -> [Id] -> CoreExpr -> CoreExpr
+-- (mkUnpackCase e co uniq Con args body)
 --      returns
--- case e of bndr { Con args -> body }
--- 
--- the type of the bndr passed in is irrelevent
-
-mkUnpackCase (Tick tickish e) uniq con args body   -- See Note [Profiling and unpacking]
-  = Tick tickish (mkUnpackCase e uniq con args body)
-mkUnpackCase scrut uniq boxing_con unpk_args body
-  = Case scrut 
-         (mk_ww_local uniq (exprType scrut)) (exprType body) 
+-- case e |> co of bndr { Con args -> body }
+
+mkUnpackCase (Tick tickish e) co uniq con args body   -- See Note [Profiling and unpacking]
+  = Tick tickish (mkUnpackCase e co uniq con args body)
+mkUnpackCase scrut co uniq boxing_con unpk_args body
+  = Case casted_scrut bndr (exprType body)
          [(DataAlt boxing_con, unpk_args, body)]
+  where
+    casted_scrut = scrut `mkCast` co
+    bndr = mk_ww_local uniq (exprType casted_scrut)
 \end{code}
 
 Note [Profiling and unpacking]

compiler/types/Coercion.lhs

@@ -38,8 +38,8 @@ module Coercion (
         mkAxiomRuleCo,
 
         -- ** Decomposition
-        splitNewTypeRepCo_maybe, instNewTyCon_maybe, 
-        topNormaliseNewType, topNormaliseNewTypeX,
+        instNewTyCon_maybe, 
+        topNormaliseNewType_maybe, 
 
         decomposeCo, getCoVar_maybe,
         splitAppCo_maybe,
@@ -1194,42 +1194,39 @@ instNewTyCon_maybe tc tys
   | otherwise
   = Nothing
 
-splitNewTypeRepCo_maybe :: Type -> Maybe (Type, Coercion)  
+topNormaliseNewType_maybe :: Type -> Maybe (Coercion, Type)
 -- ^ Sometimes we want to look through a @newtype@ and get its associated coercion.
--- This function only strips *one layer* of @newtype@ off, so the caller will usually call
--- itself recursively. If
+-- This function strips off @newtype@ layers enough to reveal something that isn't
+-- a @newtype@.  Specifically, here's the invariant:
 --
--- > splitNewTypeRepCo_maybe ty = Just (ty', co)
+-- > topNormaliseNewType_maybe rec_nts ty = Just (co, ty')
 --
--- then  @co : ty ~ ty'@.  The function returns @Nothing@ for non-@newtypes@, 
+-- then (a)  @co : ty0 ~ ty'@.
+--      (b)  ty' is not a newtype.
+--
+-- The function returns @Nothing@ for non-@newtypes@,
 -- or unsaturated applications
-splitNewTypeRepCo_maybe ty 
-  | Just ty' <- coreView ty 
-  = splitNewTypeRepCo_maybe ty'
-splitNewTypeRepCo_maybe (TyConApp tc tys)
-  = instNewTyCon_maybe tc tys
-splitNewTypeRepCo_maybe _
-  = Nothing
-
-topNormaliseNewType :: Type -> Maybe (Type, Coercion)
-topNormaliseNewType ty
-  = case topNormaliseNewTypeX initRecTc ty of
-      Just (_, co, ty) -> Just (ty, co)
-      Nothing          -> Nothing
-
-topNormaliseNewTypeX :: RecTcChecker -> Type -> Maybe (RecTcChecker, Coercion, Type)
-topNormaliseNewTypeX rec_nts ty
-  | Just ty' <- coreView ty         -- Expand predicates and synonyms
-  = topNormaliseNewTypeX rec_nts ty'
-
-topNormaliseNewTypeX rec_nts (TyConApp tc tys)
-  | Just rec_nts'     <- checkRecTc rec_nts tc  -- See Note [Expanding newtypes] in TyCon
-  , Just (rep_ty, co) <- instNewTyCon_maybe tc tys
-  = case topNormaliseNewTypeX rec_nts' rep_ty of
-       Nothing                       -> Just (rec_nts', co,                 rep_ty)
-       Just (rec_nts', co', rep_ty') -> Just (rec_nts', co `mkTransCo` co', rep_ty')
-
-topNormaliseNewTypeX _ _ = Nothing
+topNormaliseNewType_maybe ty
+  = go initRecTc Nothing ty
+  where
+    go rec_nts mb_co1 ty
+       | Just (tc, tys) <- splitTyConApp_maybe ty
+       , Just (ty', co2) <- instNewTyCon_maybe tc tys
+       , let co' = case mb_co1 of
+                      Nothing  -> co2
+                      Just co1 -> mkTransCo co1 co2
+       = case checkRecTc rec_nts tc of
+           Just rec_nts' -> go rec_nts' (Just co') ty'
+           Nothing       -> Nothing
+                  -- Return Nothing overall if we get stuck
+                  -- so that the return invariant is satisfied
+                  -- See Note [Expanding newtypes] in TyCon
+
+       | Just co1 <- mb_co1     -- Progress, but stopped on a non-newtype
+       = Just (co1, ty)
+
+       | otherwise              -- No progress
+       = Nothing
 \end{code}
 
 

compiler/types/FamInstEnv.lhs

@@ -771,27 +771,38 @@ 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}
+reduceTyFamApp_maybe :: FamInstEnvs -> Role -> TyCon -> [Type] -> Maybe (Coercion, Type)
+-- Attempt to do a *one-step* reduction of a type-family application
+-- It first normalises the type arguments, wrt functions but *not* newtypes,
+-- to be sure that nested calls like
+--    F (G Int)
+-- are correctly reduced
+--
+-- The TyCon can be oversaturated.
+-- Works on both open and closed families
 
--- The TyCon can be oversaturated. This works on both open and closed families
-chooseAxiom :: FamInstEnvs -> Role -> TyCon -> [Type] -> Maybe (Coercion, Type)
-chooseAxiom envs role tc tys
+reduceTyFamApp_maybe envs role tc tys
   | isOpenFamilyTyCon tc
   , [FamInstMatch { fim_instance = fam_inst
-                  , fim_tys =      inst_tys }] <- lookupFamInstEnv envs tc tys
+                  , fim_tys =      inst_tys }] <- lookupFamInstEnv envs tc ntys
   = let ax     = famInstAxiom fam_inst
         co     = mkUnbranchedAxInstCo role ax inst_tys
         ty     = pSnd (coercionKind co)
-    in Just (co, ty)
+    in Just (args_co `mkTransCo` co, ty)
 
   | Just ax <- isClosedSynFamilyTyCon_maybe tc
-  , Just (ind, inst_tys) <- chooseBranch ax tys
+  , Just (ind, inst_tys) <- chooseBranch ax ntys
   = let co     = mkAxInstCo role ax ind inst_tys
         ty     = pSnd (coercionKind co)
-    in Just (co, ty)
+    in Just (args_co `mkTransCo` co, ty)
 
   | otherwise
   = Nothing
 
+  where
+    (args_co, ntys) = normaliseTcArgs envs role tc tys
+
+
 -- The axiom can be oversaturated. (Closed families only.)
 chooseBranch :: CoAxiom Branched -> [Type] -> Maybe (BranchIndex, [Type])
 chooseBranch axiom tys
@@ -848,69 +859,75 @@ topNormaliseType_maybe :: FamInstEnvs
                        -> Type
                        -> Maybe (Coercion, Type)
 
--- Get rid of *outermost* (or toplevel) 
---      * type functions 
+-- Get rid of *outermost* (or toplevel)
+--      * type function redex
 --      * newtypes
--- using appropriate coercions.
--- By "outer" we mean that toplevelNormaliseType guarantees to return
--- a type that does not have a reducible redex (F ty1 .. tyn) as its
--- outermost form.  It *can* return something like (Maybe (F ty)), where
+-- using appropriate coercions.  Specifically, if
+--   topNormaliseType_maybe env ty = Maybe (co, ty')
+-- then
+--   (a) co :: ty ~ ty'
+--   (b) ty' is not a newtype, and is not a type-family redex
+--
+-- However, ty' can be something like (Maybe (F ty)), where
 -- (F ty) is a redex.
 
 -- Its a bit like Type.repType, but handles type families too
 -- The coercion returned is always an R coercion
 
 topNormaliseType_maybe env ty
-  = go initRecTc ty
+  = go initRecTc Nothing ty
   where
-    go :: RecTcChecker -> Type -> Maybe (Coercion, Type)
-    go rec_nts ty 
-        | Just ty' <- coreView ty     -- Expand synonyms
-        = go rec_nts ty'
-
-        | Just (rec_nts', nt_co, nt_rhs) <- topNormaliseNewTypeX rec_nts ty
-        = add_co nt_co rec_nts' nt_rhs
-
-    go rec_nts (TyConApp tc tys) 
-        | isFamilyTyCon tc              -- Expand family tycons
-        , (co, ty) <- normaliseTcApp env Representational tc tys
-                -- Note that normaliseType fully normalises 'tys',
-                -- wrt type functions but *not* newtypes
-                -- It has do to so to be sure that nested calls like
-                --    F (G Int)
-                -- are correctly top-normalised
-        , not (isReflCo co)
-        = add_co co rec_nts ty
-
-    go _ _ = Nothing
-
-    add_co co rec_nts ty 
-        = case go rec_nts ty of
-                Nothing         -> Just (co, ty)
-                Just (co', ty') -> Just (mkTransCo co co', ty')
-         
+    go :: RecTcChecker -> Maybe Coercion -> Type -> Maybe (Coercion, Type)
+    go rec_nts mb_co1 ty
+       = case splitTyConApp_maybe ty of
+           Just (tc, tys) -> go_tc tc tys
+           Nothing        -> all_done mb_co1
+       where
+         all_done Nothing   = Nothing
+         all_done (Just co) = Just (co, ty)
+
+         go_tc tc tys
+            | Just (ty', co2) <- instNewTyCon_maybe tc tys
+            = case checkRecTc rec_nts tc of
+                Just rec_nts' -> go rec_nts' (mb_co1 `trans` co2) ty'
+                Nothing       -> Nothing  -- No progress at all!
+                                          -- cf topNormaliseNewType_maybe
+
+            | Just (co2, rhs) <- reduceTyFamApp_maybe env Representational tc tys
+            = go rec_nts (mb_co1 `trans` co2) rhs
+
+            | otherwise
+            = all_done mb_co1
+
+    Nothing    `trans` co2 = Just co2
+    (Just co1) `trans` co2 = Just (co1 `mkTransCo` co2)
+
 
 ---------------
 normaliseTcApp :: FamInstEnvs -> Role -> TyCon -> [Type] -> (Coercion, Type)
 normaliseTcApp env role tc tys
-  | isFamilyTyCon tc
-  , Just (co, rhs) <- chooseAxiom env role tc ntys
+  | Just (first_co, ty') <- reduceTyFamApp_maybe env role tc tys
   = let    -- A reduction is possible
-        first_coi       = mkTransCo tycon_coi co
-        (rest_coi,nty)  = normaliseType env role rhs
-        fix_coi         = mkTransCo first_coi rest_coi
-    in 
-    (fix_coi, nty)
+        (rest_co,nty) = normaliseType env role ty'
+    in
+    (first_co `mkTransCo` rest_co, nty)
 
   | otherwise   -- No unique matching family instance exists;
                 -- we do not do anything
-  = (tycon_coi, TyConApp tc ntys)
+  = (Refl role ty, ty)
+  where
+    ty = mkTyConApp tc tys
 
+---------------
+normaliseTcArgs :: FamInstEnvs            -- environment with family instances
+                -> Role                   -- desired role of output coercion
+                -> TyCon -> [Type]        -- tc tys
+                -> (Coercion, [Type])      -- (co, new_tys), where
+                                          -- co :: tc tys ~ tc new_tys
+normaliseTcArgs env role tc tys
+  = (mkTyConAppCo role tc cois, ntys)
   where
-        -- Normalise the arg types so that they'll match 
-        -- when we lookup in in the instance envt
     (cois, ntys) = zipWithAndUnzip (normaliseType env) (tyConRolesX role tc) tys
-    tycon_coi    = mkTyConAppCo role tc cois
 
 ---------------
 normaliseType :: FamInstEnvs            -- environment with family instances