Cache the analysis of the data con typewip/T12618

for faster compression/decompression.

5 files changed, 52 insertions, 30 deletions

diff --git a/compiler/basicTypes/DataCon.hs b/compiler/basicTypes/DataCon.hs
index 47b05c90aa..14795e8c1a 100644
--- a/compiler/basicTypes/DataCon.hs
+++ b/compiler/basicTypes/DataCon.hs
@@ -28,6 +28,7 @@ module DataCon (
 
         -- ** Type deconstruction
         dataConRepType, dataConSig, dataConInstSig, dataConFullSig,
+        dataConCompressScheme,
         dataConName, dataConIdentity, dataConTag, dataConTyCon,
         dataConOrigTyCon, dataConUserType,
         dataConUnivTyVars, dataConUnivTyVarBinders,
@@ -66,6 +67,7 @@ import ForeignCall ( CType )
 import Coercion
 import Unify
 import TyCon
+import CompressArgs
 import FieldLabel
 import Class
 import Name
@@ -407,6 +409,8 @@ data DataCon
         -- and use that to check the pattern.  Mind you, this is really only
         -- used in CoreLint.
 
+        dcCompressScheme  :: CompressScheme,
+
 
         dcInfix :: Bool,        -- True <=> declared infix
                                 -- Used for Template Haskell and 'deriving' only
@@ -797,6 +801,7 @@ mkDataCon name declared_infix prom_info
                   dcRepTyCon = rep_tycon,
                   dcSrcBangs = arg_stricts,
                   dcFields = fields, dcTag = tag, dcRepType = rep_ty,
+                  dcCompressScheme = genCompressScheme rep_ty,
                   dcWorkId = work_id,
                   dcRep = rep,
                   dcSourceArity = length orig_arg_tys,
@@ -882,6 +887,9 @@ dataConOrigTyCon dc
 dataConRepType :: DataCon -> Type
 dataConRepType = dcRepType
 
+dataConCompressScheme :: DataCon -> CompressScheme
+dataConCompressScheme = dcCompressScheme
+
 -- | Should the 'DataCon' be presented infix?
 dataConIsInfix :: DataCon -> Bool
 dataConIsInfix = dcInfix
diff --git a/compiler/coreSyn/CoreFVs.hs b/compiler/coreSyn/CoreFVs.hs
index 12544b89d2..5f1fad2f2e 100644
--- a/compiler/coreSyn/CoreFVs.hs
+++ b/compiler/coreSyn/CoreFVs.hs
@@ -74,7 +74,7 @@ import Type
 import TyCoRep
 import TyCon
 import CompressArgs
-import DataCon ( dataConRepType, dataConWorkId )
+import DataCon ( dataConRepType, dataConCompressScheme, dataConWorkId )
 import CoAxiom
 import FamInstEnv
 import TysPrim( funTyConName )
@@ -752,7 +752,7 @@ freeVars = go
         , AnnConApp dc cargs' )
       where
         cargs'   = map go cargs
-        args     = uncompressArgs exprTypeFV (go . Type) dc_ty cargs'
+        args     = uncompressArgs exprTypeFV (go . Type) (dataConCompressScheme dc) cargs'
         dc_ty    = dataConRepType dc
         res_ty   = foldl applyTypeToArg dc_ty (map deAnnotate args)
         -- Why does this not work? Isn't piResultTys just iterated application
diff --git a/compiler/coreSyn/CoreSyn.hs b/compiler/coreSyn/CoreSyn.hs
index 54a62ef984..b47b21c5a4 100644
--- a/compiler/coreSyn/CoreSyn.hs
+++ b/compiler/coreSyn/CoreSyn.hs
@@ -1501,7 +1501,7 @@ mkCoApps  f args = foldl (\ e a -> App e (Coercion a)) f args
 mkVarApps f vars = foldl (\ e a -> App e (varToCoreExpr a)) f vars
 mkConApp dc args =
     ASSERT2 ( dataConRepFullArity dc == length args, text "mkConApp: artiy mismatch" $$ ppr dc )
-    ConApp dc (compressArgs (dataConRepType dc) args)
+    ConApp dc (compressArgs (dataConCompressScheme dc) args)
 
 mkTyApps  f args = foldl (\ e a -> App e (typeOrCoercion a)) f args
   where
diff --git a/compiler/coreSyn/CoreUtils.hs b/compiler/coreSyn/CoreUtils.hs
index e71055b26d..89499e3aad 100644
--- a/compiler/coreSyn/CoreUtils.hs
+++ b/compiler/coreSyn/CoreUtils.hs
@@ -212,7 +212,7 @@ applyTypeToArgs e op_ty args
 -}
 
 collectConArgs :: CoreExpr -> [CoreArg]
-collectConArgs (ConApp dc cargs) = uncompressArgs exprTypeOrKind Type (dataConRepType dc) cargs
+collectConArgs (ConApp dc cargs) = uncompressArgs exprTypeOrKind Type (dataConCompressScheme dc) cargs
 collectConArgs _ = panic "conAppArgs"
 
 
diff --git a/compiler/types/CompressArgs.hs b/compiler/types/CompressArgs.hs
index 099ce20474..ccbc3578ce 100644
--- a/compiler/types/CompressArgs.hs
+++ b/compiler/types/CompressArgs.hs
@@ -1,44 +1,58 @@
-module CompressArgs (compressArgs, uncompressArgs) where
+module CompressArgs (
+    CompressScheme, -- abstract
+    genCompressScheme,
+    compressArgs,
+    uncompressArgs
+    ) where
 
 import Type
 import TyCoRep
 import Panic
 
-import Data.List ( findIndex )
+import Data.List ( findIndex, dropWhileEnd )
+import Data.Maybe ( isNothing )
 
-compressArgs    ::                               Type -> [a] -> [a]
-uncompressArgs  :: (a -> Type) -> (Type -> a) -> Type -> [a] -> [a]
+-- We want to analyze the data con type only once. The resulting information
+-- is given by a list of offsets.
+-- The list may be shorted.
+-- Abstract by design.
+newtype CompressScheme = CS ([Maybe Int])
 
-compressArgs funTy args = go pis args
+genCompressScheme :: Type -> CompressScheme
+genCompressScheme funTy = CS $ shorten $ go pis
  where
     (pis,_) = splitPiTys funTy
 
-    -- Remove redundant type type arguments
-    go (Named tyBndr : pis) (_ : args)
-      | any (isRedundandTyVar (binderVar tyBndr)) pis
-      = go pis args
+    shorten = dropWhileEnd isNothing
 
-    go (_ : pis) (a : args) = a : go pis args
-    go [] [] = []
-    -- Error conditions below
-    go [] _ = panic "compressArgs: not enough arrows in type"
-    go _ [] = panic "compressArgs: not enough args"
+    go (Named tyBndr : pis)
+      | Just i <- findIndex (isRedundandTyVar (binderVar tyBndr)) pis
+      = Just i : go pis
+    go (_ : pis)
+      = Nothing : go pis
+    go []
+      = []
 
-uncompressArgs typeOf mkType funTy args = go pis args
- where
-    (pis,_) = splitPiTys funTy
 
-    go (Named tyBndr : pis) args
-      | Just i <- findIndex (isRedundandTyVar (binderVar tyBndr)) pis
-      -- This is a type argument we have to recover
-      = let args' = go pis args
-        in mkType (typeOf (args' !! i)) : args'
+compressArgs    ::                               CompressScheme -> [a] -> [a]
+uncompressArgs  :: (a -> Type) -> (Type -> a) -> CompressScheme -> [a] -> [a]
+
+compressArgs (CS cs) args = go cs args
+ where
+    go (Just _  : pis) (_ : args) =     go pis args
+    go (Nothing : pis) (a : args) = a : go pis args
+    go []              args       = args
+    go _               []         = panic "compressArgs: not enough args"
 
-    go (_ : pis) (a : args) = a : go pis args
-    go [] [] = []
+uncompressArgs typeOf mkType (CS cs) args = go cs args
+ where
+    go (Just i  : pis) args       = mkType (typeOf (args' !! i)) : args'
+      where args' = go pis args
+    go (Nothing : pis) (a : args) = a : args'
+      where args' = go pis args
+    go []              args       = args
     -- Error conditions below
-    go [] _ = panic "uncompressArgs: not enough arrows in type"
-    go _ [] = panic "uncompressArgs: not enough args"
+    go _               []         = panic "uncompressArgs: not enough args"
 
 isRedundandTyVar :: TyVar -> TyBinder -> Bool
 isRedundandTyVar v (Anon t) | Just v' <- getTyVar_maybe t, v == v' = True