Move sortQuantVars to MkCore

3 files changed, 35 insertions, 19 deletions

diff --git a/compiler/coreSyn/MkCore.lhs b/compiler/coreSyn/MkCore.lhs
index 5d1c19bc5f..9e42290f7e 100644
--- a/compiler/coreSyn/MkCore.lhs
+++ b/compiler/coreSyn/MkCore.lhs
@@ -13,6 +13,7 @@ module MkCore (
         mkCoreApp, mkCoreApps, mkCoreConApps,
         mkCoreLams, mkWildCase, mkIfThenElse,
         mkWildValBinder, mkWildEvBinder,
+        sortQuantVars,
         
         -- * Constructing boxed literals
         mkWordExpr, mkWordExprWord,
@@ -84,7 +85,7 @@ import Outputable
 import FastString
 import UniqSupply
 import BasicTypes
-import Util             ( notNull, zipEqual )
+import Util             ( notNull, zipEqual, sortLe )
 import Pair
 import Constants
 
@@ -101,6 +102,23 @@ infixl 4 `mkCoreApp`, `mkCoreApps`
 %************************************************************************
 
 \begin{code}
+sortQuantVars :: [Var] -> [Var]
+-- Sort the variables (KindVars, TypeVars, and Ids) 
+-- into order: Kind, then Type, then Id
+sortQuantVars = sortLe le
+  where
+    v1 `le` v2 = case (is_tv v1, is_tv v2) of
+                   (True, False)  -> True
+                   (False, True)  -> False
+                   (True, True)   ->
+                     case (is_kv v1, is_kv v2) of
+                       (True, False) -> True
+                       (False, True) -> False
+                       _             -> v1 <= v2  -- Same family
+                   (False, False) -> v1 <= v2
+    is_tv v = isTyVar v
+    is_kv v = isKindVar v
+
 -- | Bind a binding group over an expression, using a @let@ or @case@ as
 -- appropriate (see "CoreSyn#let_app_invariant")
 mkCoreLet :: CoreBind -> CoreExpr -> CoreExpr
diff --git a/compiler/simplCore/SetLevels.lhs b/compiler/simplCore/SetLevels.lhs
index 6e0afb4415..394cd9801e 100644
--- a/compiler/simplCore/SetLevels.lhs
+++ b/compiler/simplCore/SetLevels.lhs
@@ -68,7 +68,9 @@ import CoreArity	( exprBotStrictness_maybe )
 import CoreFVs		-- all of it
 import Coercion         ( isCoVar )
 import CoreSubst	( Subst, emptySubst, extendInScope, substBndr, substRecBndrs,
-			  extendIdSubst, extendSubstWithVar, cloneBndr, cloneRecIdBndrs, substTy, substCo )
+			  extendIdSubst, extendSubstWithVar, cloneBndr, 
+                          cloneRecIdBndrs, substTy, substCo )
+import MkCore           ( sortQuantVars ) 
 import Id
 import IdInfo
 import Var
@@ -78,8 +80,7 @@ import Literal		( litIsTrivial )
 import Demand		( StrictSig, increaseStrictSigArity )
 import Name		( getOccName, mkSystemVarName )
 import OccName		( occNameString )
-import Type		( isUnLiftedType, Type, sortQuantVars, mkPiTypes )
-import Kind		( kiVarsOfKinds )
+import Type		( isUnLiftedType, Type, mkPiTypes )
 import BasicTypes	( Arity )
 import UniqSupply
 import Util
@@ -1000,9 +1001,9 @@ abstractVars :: Level -> LevelEnv -> VarSet -> [Var]
 	-- whose level is greater than the destination level
 	-- These are the ones we are going to abstract out
 abstractVars dest_lvl (LE { le_lvl_env = lvl_env, le_env = id_env }) fvs
-  = map zap $ uniq $ sortQuantVars  -- IA0_NOTE: centralizing sorting on variables
+  = map zap $ uniq $ sortQuantVars
 	[var | fv <- varSetElems fvs
-	     , var <- absVarsOf id_env fv
+	     , var <- varSetElems (absVarsOf id_env fv)
 	     , abstract_me var ]
 	-- NB: it's important to call abstract_me only on the OutIds the
 	-- come from absVarsOf (not on fv, which is an InId)
@@ -1025,7 +1026,7 @@ abstractVars dest_lvl (LE { le_lvl_env = lvl_env, le_env = id_env }) fvs
 		     setIdInfo v vanillaIdInfo
 	  | otherwise = v
 
-absVarsOf :: IdEnv ([Var], LevelledExpr) -> Var -> [Var]
+absVarsOf :: IdEnv ([Var], LevelledExpr) -> Var -> VarSet
 	-- If f is free in the expression, and f maps to poly_f a b c in the
 	-- current substitution, then we must report a b c as candidate type
 	-- variables
@@ -1033,20 +1034,16 @@ absVarsOf :: IdEnv ([Var], LevelledExpr) -> Var -> [Var]
 	-- Also, if x::a is an abstracted variable, then so is a; that is,
 	-- we must look in x's type. What's more, if a mentions kind variables,
 	-- we must also return those.
-	-- 
-	-- And similarly if x is a coercion variable.
 absVarsOf id_env v 
-  | isId v    = [av2 | av1 <- lookup_avs v
-		     , av2 <- add_tyvars av1]
-  | otherwise = ASSERT( isTyVar v ) [v]
+  | isId v, Just (abs_vars, _) <- lookupVarEnv id_env v
+  = foldr (unionVarSet . close) emptyVarSet abs_vars
+  | otherwise
+  = close v
   where
-    lookup_avs v = case lookupVarEnv id_env v of
-			Just (abs_vars, _) -> abs_vars
-			Nothing	           -> [v]
-
-    add_tyvars v = v : (varSetElems tyvars ++ varSetElems kivars)
-    tyvars = varTypeTyVars v
-    kivars = kiVarsOfKinds (map tyVarKind (varSetElems tyvars))
+    close :: Var -> VarSet  -- Result include the input variable itself
+    close v = foldVarSet (unionVarSet . close)
+                         (unitVarSet v)
+                         (varTypeTyVars v)
 \end{code}
 
 \begin{code}
diff --git a/compiler/simplCore/SimplUtils.lhs b/compiler/simplCore/SimplUtils.lhs
index ad6fe5488b..7da185a1ae 100644
--- a/compiler/simplCore/SimplUtils.lhs
+++ b/compiler/simplCore/SimplUtils.lhs
@@ -38,6 +38,7 @@ module SimplUtils (
 
 import SimplEnv
 import CoreMonad        ( SimplifierMode(..), Tick(..) )
+import MkCore           ( sortQuantVars )
 import DynFlags
 import StaticFlags
 import CoreSyn