Put one-Shot info in the interface

Differential Revision: https://phabricator.haskell.org/D391

5 files changed, 69 insertions, 19 deletions

diff --git a/compiler/coreSyn/CoreTidy.lhs b/compiler/coreSyn/CoreTidy.lhs
index 56da4944e3..810a71ca6c 100644
--- a/compiler/coreSyn/CoreTidy.lhs
+++ b/compiler/coreSyn/CoreTidy.lhs
@@ -153,6 +153,8 @@ tidyIdBndr env@(tidy_env, var_env) id
         -- Note [Tidy IdInfo]
         new_info = vanillaIdInfo `setOccInfo` occInfo old_info
                                  `setUnfoldingInfo` new_unf
+                                  -- see Note [Preserve OneShotInfo]
+                                 `setOneShotInfo` oneShotInfo old_info
         old_info = idInfo id
         old_unf  = unfoldingInfo old_info
         new_unf | isEvaldUnfolding old_unf = evaldUnfolding
@@ -256,6 +258,17 @@ preserve the evaluated-ness on 'y' in tidyBndr.
 (Another alternative would be to tidy unboxed lets into cases,
 but that seems more indirect and surprising.)
 
+Note [Preserve OneShotInfo]
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+We keep the OneShotInfo because we want it to propagate into the interface.
+Not all OneShotInfo is determined by a compiler analysis; some is added by a
+call of GHC.Exts.oneShot, which is then discarded before the end of of the
+optimisation pipeline, leaving only the OneShotInfo on the lambda. Hence we
+must preserve this info in inlinings.
+
+This applies to lambda binders only, hence it is stored in IfaceLamBndr.
+
 
 \begin{code}
 (=:) :: a -> (a -> b) -> b
diff --git a/compiler/iface/IfaceSyn.lhs b/compiler/iface/IfaceSyn.lhs
index 7e2d6f2172..5cfe773dc8 100644
--- a/compiler/iface/IfaceSyn.lhs
+++ b/compiler/iface/IfaceSyn.lhs
@@ -482,7 +482,7 @@ data IfaceExpr
   | IfaceType   IfaceType
   | IfaceCo     IfaceCoercion
   | IfaceTuple  TupleSort [IfaceExpr]   -- Saturated; type arguments omitted
-  | IfaceLam    IfaceBndr IfaceExpr
+  | IfaceLam    IfaceLamBndr IfaceExpr
   | IfaceApp    IfaceExpr IfaceExpr
   | IfaceCase   IfaceExpr IfLclName [IfaceAlt]
   | IfaceECase  IfaceExpr IfaceType     -- See Note [Empty case alternatives]
@@ -981,7 +981,7 @@ pprIfaceExpr add_par app@(IfaceApp _ _) = add_par (pprIfaceApp app [])
 pprIfaceExpr _       (IfaceTuple c as)  = tupleParens c (interpp'SP as)
 
 pprIfaceExpr add_par i@(IfaceLam _ _)
-  = add_par (sep [char '\\' <+> sep (map ppr bndrs) <+> arrow,
+  = add_par (sep [char '\\' <+> sep (map pprIfaceLamBndr bndrs) <+> arrow,
                   pprIfaceExpr noParens body])
   where
     (bndrs,body) = collect [] i
@@ -1273,16 +1273,16 @@ freeNamesIfUnfold (IfInlineRule _ _ _ e) = freeNamesIfExpr e
 freeNamesIfUnfold (IfDFunUnfold bs es)   = fnList freeNamesIfBndr bs &&& fnList freeNamesIfExpr es
 
 freeNamesIfExpr :: IfaceExpr -> NameSet
-freeNamesIfExpr (IfaceExt v)      = unitNameSet v
-freeNamesIfExpr (IfaceFCall _ ty) = freeNamesIfType ty
-freeNamesIfExpr (IfaceType ty)    = freeNamesIfType ty
-freeNamesIfExpr (IfaceCo co)      = freeNamesIfCoercion co
-freeNamesIfExpr (IfaceTuple _ as) = fnList freeNamesIfExpr as
-freeNamesIfExpr (IfaceLam b body) = freeNamesIfBndr b &&& freeNamesIfExpr body
-freeNamesIfExpr (IfaceApp f a)    = freeNamesIfExpr f &&& freeNamesIfExpr a
-freeNamesIfExpr (IfaceCast e co)  = freeNamesIfExpr e &&& freeNamesIfCoercion co
-freeNamesIfExpr (IfaceTick _ e)   = freeNamesIfExpr e
-freeNamesIfExpr (IfaceECase e ty) = freeNamesIfExpr e &&& freeNamesIfType ty
+freeNamesIfExpr (IfaceExt v)          = unitNameSet v
+freeNamesIfExpr (IfaceFCall _ ty)     = freeNamesIfType ty
+freeNamesIfExpr (IfaceType ty)        = freeNamesIfType ty
+freeNamesIfExpr (IfaceCo co)          = freeNamesIfCoercion co
+freeNamesIfExpr (IfaceTuple _ as)     = fnList freeNamesIfExpr as
+freeNamesIfExpr (IfaceLam (b,_) body) = freeNamesIfBndr b &&& freeNamesIfExpr body
+freeNamesIfExpr (IfaceApp f a)        = freeNamesIfExpr f &&& freeNamesIfExpr a
+freeNamesIfExpr (IfaceCast e co)      = freeNamesIfExpr e &&& freeNamesIfCoercion co
+freeNamesIfExpr (IfaceTick _ e)       = freeNamesIfExpr e
+freeNamesIfExpr (IfaceECase e ty)     = freeNamesIfExpr e &&& freeNamesIfType ty
 freeNamesIfExpr (IfaceCase s _ alts)
   = freeNamesIfExpr s &&& fnList fn_alt alts &&& fn_cons alts
   where
@@ -1741,9 +1741,10 @@ instance Binary IfaceExpr where
         putByte bh 3
         put_ bh ac
         put_ bh ad
-    put_ bh (IfaceLam ae af) = do
+    put_ bh (IfaceLam (ae, os) af) = do
         putByte bh 4
         put_ bh ae
+        put_ bh os
         put_ bh af
     put_ bh (IfaceApp ag ah) = do
         putByte bh 5
@@ -1793,8 +1794,9 @@ instance Binary IfaceExpr where
                     ad <- get bh
                     return (IfaceTuple ac ad)
             4 -> do ae <- get bh
+                    os <- get bh
                     af <- get bh
-                    return (IfaceLam ae af)
+                    return (IfaceLam (ae, os) af)
             5 -> do ag <- get bh
                     ah <- get bh
                     return (IfaceApp ag ah)
diff --git a/compiler/iface/IfaceType.lhs b/compiler/iface/IfaceType.lhs
index c55edc6185..aae61c47ed 100644
--- a/compiler/iface/IfaceType.lhs
+++ b/compiler/iface/IfaceType.lhs
@@ -12,7 +12,7 @@ module IfaceType (
 
         IfaceType(..), IfacePredType, IfaceKind, IfaceTyCon(..), IfaceCoercion(..),
         IfaceTyLit(..), IfaceTcArgs(..),
-        IfaceContext, IfaceBndr(..), IfaceTvBndr, IfaceIdBndr,
+        IfaceContext, IfaceBndr(..), IfaceOneShot(..), IfaceLamBndr, IfaceTvBndr, IfaceIdBndr,
 
         -- Conversion from Type -> IfaceType
         toIfaceType, toIfaceTypes, toIfaceKind, toIfaceTyVar,
@@ -28,7 +28,7 @@ module IfaceType (
 
         -- Printing
         pprIfaceType, pprParendIfaceType, pprIfaceContext, pprIfaceContextArr,
-        pprIfaceIdBndr, pprIfaceTvBndr, pprIfaceTvBndrs,
+        pprIfaceIdBndr, pprIfaceLamBndr, pprIfaceTvBndr, pprIfaceTvBndrs,
         pprIfaceBndrs, pprIfaceTcArgs, pprParendIfaceTcArgs,
         pprIfaceForAllPart, pprIfaceForAll, pprIfaceSigmaType,
         pprIfaceCoercion, pprParendIfaceCoercion,
@@ -84,6 +84,14 @@ data IfaceBndr          -- Local (non-top-level) binders
 type IfaceIdBndr  = (IfLclName, IfaceType)
 type IfaceTvBndr  = (IfLclName, IfaceKind)
 
+
+data IfaceOneShot    -- see Note [Preserve OneShotInfo]
+  = IfaceNoOneShot
+  | IfaceOneShot
+
+type IfaceLamBndr
+  = (IfaceBndr, IfaceOneShot)
+
 -------------------------------
 type IfaceKind     = IfaceType
 
@@ -139,6 +147,8 @@ data IfaceCoercion
   | IfaceInstCo      IfaceCoercion IfaceType
   | IfaceSubCo       IfaceCoercion
   | IfaceAxiomRuleCo IfLclName [IfaceType] [IfaceCoercion]
+
+
 \end{code}
 
 %************************************************************************
@@ -335,6 +345,10 @@ instance Outputable IfaceBndr where
 pprIfaceBndrs :: [IfaceBndr] -> SDoc
 pprIfaceBndrs bs = sep (map ppr bs)
 
+pprIfaceLamBndr :: IfaceLamBndr -> SDoc
+pprIfaceLamBndr (b, IfaceNoOneShot) = ppr b
+pprIfaceLamBndr (b, IfaceOneShot)   = ppr b <> text "[OneShot]"
+
 pprIfaceIdBndr :: (IfLclName, IfaceType) -> SDoc
 pprIfaceIdBndr (name, ty) = hsep [ppr name, dcolon, ppr ty]
 
@@ -360,6 +374,17 @@ instance Binary IfaceBndr where
                       return (IfaceIdBndr aa)
               _ -> do ab <- get bh
                       return (IfaceTvBndr ab)
+
+instance Binary IfaceOneShot where
+    put_ bh IfaceNoOneShot = do
+            putByte bh 0
+    put_ bh IfaceOneShot = do
+            putByte bh 1
+    get bh = do
+            h <- getByte bh
+            case h of
+              0 -> do return IfaceNoOneShot
+              _ -> do return IfaceOneShot
 \end{code}
 
 ----------------------------- Printing IfaceType ------------------------------------
diff --git a/compiler/iface/MkIface.lhs b/compiler/iface/MkIface.lhs
index 88d4e0aa9f..78111b299e 100644
--- a/compiler/iface/MkIface.lhs
+++ b/compiler/iface/MkIface.lhs
@@ -1978,7 +1978,7 @@ toIfaceExpr (Var v)         = toIfaceVar v
 toIfaceExpr (Lit l)         = IfaceLit l
 toIfaceExpr (Type ty)       = IfaceType (toIfaceType ty)
 toIfaceExpr (Coercion co)   = IfaceCo   (toIfaceCoercion co)
-toIfaceExpr (Lam x b)       = IfaceLam (toIfaceBndr x) (toIfaceExpr b)
+toIfaceExpr (Lam x b)       = IfaceLam (toIfaceBndr x, toIfaceOneShot x) (toIfaceExpr b)
 toIfaceExpr (App f a)       = toIfaceApp f [a]
 toIfaceExpr (Case s x ty as)
   | null as                 = IfaceECase (toIfaceExpr s) (toIfaceType ty)
@@ -1989,6 +1989,13 @@ toIfaceExpr (Tick t e)
   | Just t' <- toIfaceTickish t = IfaceTick t' (toIfaceExpr e) 
   | otherwise                   = toIfaceExpr e
 
+toIfaceOneShot :: Id -> IfaceOneShot
+toIfaceOneShot id | isId id
+                  , OneShotLam <- oneShotInfo (idInfo id)
+                  = IfaceOneShot
+                  | otherwise
+                  = IfaceNoOneShot
+
 ---------------------
 toIfaceTickish :: Tickish Id -> Maybe IfaceTickish
 toIfaceTickish (ProfNote cc tick push) = Just (IfaceSCC cc tick push)
diff --git a/compiler/iface/TcIface.lhs b/compiler/iface/TcIface.lhs
index 424a46c232..4e2cfd5a76 100644
--- a/compiler/iface/TcIface.lhs
+++ b/compiler/iface/TcIface.lhs
@@ -1093,9 +1093,12 @@ tcIfaceExpr (IfaceTuple boxity args)  = do
     con_id = dataConWorkId (tupleCon boxity arity)
 
 
-tcIfaceExpr (IfaceLam bndr body)
+tcIfaceExpr (IfaceLam (bndr, os) body)
   = bindIfaceBndr bndr $ \bndr' ->
-    Lam bndr' <$> tcIfaceExpr body
+    Lam (tcIfaceOneShot os bndr') <$> tcIfaceExpr body
+  where
+    tcIfaceOneShot IfaceOneShot b = setOneShotLambda b
+    tcIfaceOneShot _            b = b
 
 tcIfaceExpr (IfaceApp fun arg)
   = tcIfaceApps fun arg