Merge pull request #2340 from ucsd-progsys/fd/noinline4

Add NOINLINE pragmas ahead of desugaring

liquidhaskell-boot/liquidhaskell-boot.cabal

@@ -82,7 +82,6 @@ library
                       Language.Haskell.Liquid.Termination.Structural
                       Language.Haskell.Liquid.Transforms.ANF
                       Language.Haskell.Liquid.Transforms.CoreToLogic
-                      Language.Haskell.Liquid.Transforms.Rec
                       Language.Haskell.Liquid.Transforms.RefSplit
                       Language.Haskell.Liquid.Transforms.Rewrite
                       Language.Haskell.Liquid.Transforms.Simplify

liquidhaskell-boot/src-ghc/Liquid/GHC/API.hs

@@ -523,6 +523,7 @@ import GHC.Tc.Utils.Monad             as Ghc
     , pushTcLevelM
     , reportDiagnostic
     , reportDiagnostics
+    , updEnv
     )
 import GHC.Tc.Utils.TcType            as Ghc (tcSplitDFunTy, tcSplitMethodTy)
 import GHC.Tc.Zonk.Type               as Ghc

liquidhaskell-boot/src-ghc/Liquid/GHC/API/Extra.hs

@@ -6,6 +6,7 @@
 module Liquid.GHC.API.Extra (
     module StableModule
   , ApiComment(..)
+  , addNoInlinePragmasToLocalBinds
   , apiComments
   , apiCommentsParsedSource
   , dataConSig
@@ -34,8 +35,8 @@ module Liquid.GHC.API.Extra (
 import Control.Monad.IO.Class
 import           Liquid.GHC.API.StableModule      as StableModule
 import GHC
-import Data.Data (Data, gmapQr)
-import Data.Generics (extQ)
+import Data.Data (Data, gmapQr, gmapT)
+import Data.Generics (extQ, extT)
 import Data.Foldable                  (asum)
 import Data.List                      (sortOn)
 import qualified Data.Map as Map
@@ -46,6 +47,7 @@ import GHC.Core.Coercion              as Ghc
 import GHC.Core.DataCon               as Ghc
 import GHC.Core.Make                  (pAT_ERROR_ID)
 import GHC.Core.Type                  as Ghc hiding (typeKind , isPredTy, extendCvSubst, linear)
+import GHC.Data.Bag                   (bagToList)
 import GHC.Data.FastString            as Ghc
 import qualified GHC.Data.EnumSet as EnumSet
 import GHC.Data.Maybe
@@ -54,8 +56,10 @@ import GHC.Driver.Env
 import GHC.Driver.Main
 import GHC.Driver.Session             as Ghc
 import GHC.Tc.Types
+import GHC.Types.Basic
 import GHC.Types.Name                 (isSystemName, nameModule_maybe, occNameFS)
-import GHC.Types.Name.Reader          (nameRdrName)
+import GHC.Types.Name.Reader          (nameRdrName, rdrNameOcc)
+import GHC.Types.SourceText           (SourceText(..))
 import GHC.Types.SrcLoc               as Ghc
 import GHC.Types.TypeEnv
 import GHC.Types.Unique               (getUnique, hasKey)
@@ -217,6 +221,47 @@ apiCommentsParsedSource ps =
     spanToLineColumn =
       fmap (\s -> (srcSpanStartLine s, srcSpanStartCol s)) . srcSpanToRealSrcSpan
 
+-- | Adds NOINLINE pragmas to all local bindings in the module.
+--
+-- This prevents the simple optimizer from inlining such bindings which might
+-- have specs that would otherwise be left dangling.
+--
+-- https://gitlab.haskell.org/ghc/ghc/-/issues/24386
+--
+addNoInlinePragmasToLocalBinds :: ParsedModule -> ParsedModule
+addNoInlinePragmasToLocalBinds ps =
+    ps { pm_parsed_source = go (pm_parsed_source ps) }
+  where
+    go :: forall a. Data a => a -> a
+    go = gmapT (go `extT` addNoInlinePragmas)
+
+    addNoInlinePragmas :: HsValBinds GhcPs -> HsValBinds GhcPs
+    addNoInlinePragmas = \case
+      ValBinds x binds sigs ->
+          ValBinds x binds (newSigs ++ dropInlinePragmas sigs)
+        where
+          dropInlinePragmas = filter (not . isInlinePragma . unLoc)
+
+          isInlinePragma InlineSig{} = True
+          isInlinePragma _ = False
+
+          newSigs = concatMap (traverse noInlinePragmaForBind) $ bagToList binds
+
+          noInlinePragmaForBind :: HsBindLR GhcPs GhcPs -> [Sig GhcPs]
+          noInlinePragmaForBind bndr =
+              [ InlineSig
+                  []
+                  (noLocA b)
+                  defaultInlinePragma
+                    { inl_inline = NoInline $ SourceText $ occNameFS $ rdrNameOcc b
+                    , inl_act = NeverActive
+                    }
+              | b <- collectHsBindBinders CollNoDictBinders bndr
+              ]
+
+      bnds -> bnds
+
+
 lookupModSummary :: HscEnv -> ModuleName -> Maybe ModSummary
 lookupModSummary hscEnv mdl = do
    let mg = hsc_mod_graph hscEnv

liquidhaskell-boot/src/Language/Haskell/Liquid/GHC/Plugin.hs

@@ -34,7 +34,6 @@ import           Language.Haskell.Liquid.GHC.Plugin.SpecFinder
                                                          as SpecFinder
 
 import           Language.Haskell.Liquid.GHC.Types       (MGIModGuts(..), miModGuts)
-import           Language.Haskell.Liquid.Transforms.Rec (transformRecExpr)
 import           Language.Haskell.Liquid.Transforms.InlineAux (inlineAux)
 import           Language.Haskell.Liquid.Transforms.Rewrite (rewriteBinds)
 import           GHC.LanguageExtensions
@@ -166,16 +165,22 @@ customDynFlags df =
        `gopt_set` Opt_PIC
        `gopt_set` Opt_DeferTypedHoles
        `gopt_set` Opt_KeepRawTokenStream
-       -- Opt_InsertBreakpoints is used during desugaring to prevent the
-       -- simple optimizer from inlining local bindings to which we might want
-       -- to attach specifications.
-       --
-       -- https://gitlab.haskell.org/ghc/ghc/-/issues/24386
-       `gopt_set` Opt_InsertBreakpoints
        `xopt_set` MagicHash
        `xopt_set` DeriveGeneric
        `xopt_set` StandaloneDeriving
 
+maybeInsertBreakPoints :: Config -> DynFlags -> DynFlags
+maybeInsertBreakPoints cfg dflags =
+    if insertCoreBreakPoints cfg then
+      -- Opt_InsertBreakpoints is used during desugaring to prevent the
+      -- simple optimizer from inlining local bindings to which we might want
+      -- to attach specifications.
+      --
+      -- https://gitlab.haskell.org/ghc/ghc/-/issues/24386
+      dflags `gopt_set` Opt_InsertBreakpoints
+    else
+      dflags
+
 --------------------------------------------------------------------------------
 -- | \"Unoptimising\" things ----------------------------------------------------
 --------------------------------------------------------------------------------
@@ -211,16 +216,15 @@ unoptimiseDynFlags df = updOptLevel 0 df
 typecheckHook :: Config -> ModSummary -> TcGblEnv -> TcM (Either LiquidCheckException TcGblEnv)
 typecheckHook cfg0 modSummary0 tcGblEnv = do
   debugLog $ "We are in module: " <> show (toStableModule thisModule)
-  let modSummary = modSummary0
-        { ms_hspp_opts =
-            customDynFlags $ unoptimiseDynFlags (ms_hspp_opts modSummary0)
-        }
+  let modSummary =
+        updateModSummaryDynFlags (customDynFlags . unoptimiseDynFlags) modSummary0
       thisFile = LH.modSummaryHsFile modSummary
 
   env0 <- env_top <$> getEnv
   let env = env0 { hsc_dflags = ms_hspp_opts modSummary }
-  parsed           <- liftIO $ parseModuleIO env (LH.keepRawTokenStream modSummary)
-  let specComments = map mkSpecComment $ LH.extractSpecComments parsed
+  parsed0 <- liftIO $ parseModuleIO env (LH.keepRawTokenStream modSummary)
+  let specComments = map mkSpecComment $ LH.extractSpecComments parsed0
+      parsed = addNoInlinePragmasToLocalBinds parsed0
 
   case parseSpecComments (coerce specComments) of
     Left errors ->
@@ -229,20 +233,25 @@ typecheckHook cfg0 modSummary0 tcGblEnv = do
 
       withPragmas cfg0 thisFile [s | Pragma s <- specs] $ \cfg -> do
 
+        let modSummary2 = updateModSummaryDynFlags (maybeInsertBreakPoints cfg) modSummary
+            parsed2 = parsed { pm_mod_summary = modSummary2 }
+            env2 = env { hsc_dflags = ms_hspp_opts modSummary2 }
+
         -- The LH plugin itself calls the type checker (see following line). This
         -- would lead to a loop if we didn't remove the plugin when calling the type
         -- checker.
-        typechecked     <- liftIO $ typecheckModuleIO (dropPlugins env) (LH.ignoreInline parsed)
-        resolvedNames   <- liftIO $ LH.lookupTyThings env tcGblEnv
-        availTyCons     <- liftIO $ LH.availableTyCons env tcGblEnv (tcg_exports tcGblEnv)
-        availVars       <- liftIO $ LH.availableVars env tcGblEnv (tcg_exports tcGblEnv)
+        typechecked     <- liftIO $ typecheckModuleIO (dropPlugins env2) (LH.ignoreInline parsed2)
+        resolvedNames   <- liftIO $ LH.lookupTyThings env2 tcGblEnv
+        availTyCons     <- liftIO $ LH.availableTyCons env2 tcGblEnv (tcg_exports tcGblEnv)
+        availVars       <- liftIO $ LH.availableVars env2 tcGblEnv (tcg_exports tcGblEnv)
 
-        unoptimisedGuts <- liftIO $ desugarModuleIO env modSummary typechecked
+        unoptimisedGuts <- liftIO $ desugarModuleIO env2 modSummary2 typechecked
 
         let tcData = mkTcData (tcg_rn_imports tcGblEnv) resolvedNames availTyCons availVars
         let pipelineData = PipelineData unoptimisedGuts tcData specs
 
-        liquidHaskellCheckWithConfig cfg pipelineData modSummary tcGblEnv
+        updEnv (\e -> e {env_top = env2}) $
+          liquidHaskellCheckWithConfig cfg pipelineData modSummary2 tcGblEnv
 
   where
     thisModule :: Module
@@ -252,6 +261,8 @@ typecheckHook cfg0 modSummary0 tcGblEnv = do
 
     continue = pure $ Left (ErrorsOccurred [])
 
+    updateModSummaryDynFlags f ms = ms { ms_hspp_opts = f (ms_hspp_opts ms) }
+
 serialiseSpec :: Module -> TcGblEnv -> LiquidLib -> TcM TcGblEnv
 serialiseSpec thisModule tcGblEnv liquidLib = do
   -- ---
@@ -593,9 +604,8 @@ makeTargetSrc cfg dynFlags file tcData modGuts hscEnv = do
         deriv   = Just $ instEnvElts $ mg_inst_env modGuts
 
     preNormalize :: Config -> ModGuts -> [CoreBind]
-    preNormalize cfg modGuts = rewriteBinds cfg orig_cbs
+    preNormalize cfg modGuts = rewriteBinds cfg inl_cbs
       where
-        orig_cbs = transformRecExpr inl_cbs
         inl_cbs  = inlineAux cfg (mg_module modGuts) (mg_binds modGuts)
 
 getFamInstances :: ModGuts -> [FamInst]

liquidhaskell-boot/src/Language/Haskell/Liquid/Transforms/Rewrite.hs

@@ -32,6 +32,7 @@ import           Language.Fixpoint.Misc       ({- mapFst, -}  mapSnd)
 import           Language.Haskell.Liquid.Misc (Nat)
 import           Language.Haskell.Liquid.GHC.Play (sub, substExpr)
 import           Language.Haskell.Liquid.GHC.Misc (unTickExpr, isTupleId, mkAlive)
+import           Language.Haskell.Liquid.Types.Errors (impossible)
 import           Language.Haskell.Liquid.UX.Config  (Config, noSimplifyCore)
 import qualified Data.List as L
 import qualified Data.HashMap.Strict as M
@@ -47,7 +48,8 @@ rewriteBinds cfg
        . tidyTuples
        . rewriteBindWith inlineLoopBreakerTx
        . inlineLoopBreaker
-       . rewriteBindWith strictifyLazyLets)
+       . rewriteBindWith strictifyLazyLets
+       . inlineFailCases)
   | otherwise
   = id
 
@@ -374,3 +376,43 @@ collectNonRecLets :: Expr t -> ([Bind t], Expr t)
 collectNonRecLets = go []
   where go bs (Let b@(NonRec _ _) e') = go (b:bs) e'
         go bs e'                      = (reverse bs, e')
+
+-- | Inlines bindings of the form
+--
+-- > let v = \x -> e0
+-- >  in e1
+--
+-- whenever all of the following hold:
+--  * "fail" is a prefix of variable @v@,
+--  * @x@ is not free in @e0@, and
+--  * v is applied to some value in @e1@.
+--
+-- In addition to inlining, this function also beta reduces
+-- the resulting expressions @(\x -> e0) a@ by replacing them
+-- with @e0@.
+--
+inlineFailCases :: CoreBind -> CoreBind
+inlineFailCases = go []
+  where
+    go su (Rec xes)    = Rec (mapSnd (go' su) <$> xes)
+    go su (NonRec x e) = NonRec x (go' su e)
+
+    go' su (App (Var x) _)       | isFailId x, Just e <- getFailExpr x su = e
+    go' su (Let (NonRec x ex) e) | isFailId x   = go' (addFailExpr x (go' su ex) su) e
+
+    go' su (App e1 e2)      = App (go' su e1) (go' su e2)
+    go' su (Lam x e)        = Lam x (go' su e)
+    go' su (Let xs e)       = Let (go su xs) (go' su e)
+    go' su (Case e x t alt) = Case (go' su e) x t (goalt su <$> alt)
+    go' su (Cast e c)       = Cast (go' su e) c
+    go' su (Tick t e)       = Tick t (go' su e)
+    go' _  e                = e
+
+    goalt su (Alt c xs e)   = Alt c xs (go' su e)
+
+    isFailId x  = isLocalId x && isSystemName (varName x) && L.isPrefixOf "fail" (getOccString x)
+    getFailExpr = L.lookup
+
+    addFailExpr x (Lam v e) su
+      | not (elemVarSet v $ exprFreeVars e)  = (x, e):su
+    addFailExpr _ _         _  = impossible Nothing "internal error" -- this cannot happen

liquidhaskell-boot/src/Language/Haskell/Liquid/UX/CmdLine.hs

@@ -475,6 +475,10 @@ config = cmdArgsMode $ Config {
     = def &= help "Dump pre-normalized core (before a-normalization)"
           &= name "dump-pre-normalized-core"
           &= explicit
+  , insertCoreBreakPoints
+    = def &= help "Insert break point when desugaring Haskell (enabled by default)"
+          &= name "insert-core-break-points"
+          &= explicit
   } &= program "liquid"
     &= help    "Refinement Types for Haskell"
     &= summary copyright
@@ -740,6 +744,7 @@ defConfig = Config
   , excludeAutomaticAssumptionsFor = []
   , dumpOpaqueReflections    = False
   , dumpPreNormalizedCore    = False
+  , insertCoreBreakPoints    = True
   }
 
 -- | Write the annotations (i.e. the files in the \".liquid\" hidden folder) and

liquidhaskell-boot/src/Language/Haskell/Liquid/UX/Config.hs

@@ -110,6 +110,7 @@ data Config = Config
   , excludeAutomaticAssumptionsFor :: [String]
   , dumpOpaqueReflections    :: Bool       -- Dumps all opaque reflections to the stdout
   , dumpPreNormalizedCore    :: Bool       -- Dumps the prenormalized core (before a-normalization)
+  , insertCoreBreakPoints    :: Bool       -- Insert break points in core
   } deriving (Generic, Data, Typeable, Show, Eq)
 
 allowPLE :: Config -> Bool