Cache exposed package modules in Pantry DB

.hlint.yaml

@@ -40,11 +40,6 @@
   ]
 }
 
-# Pretty-printing
-- error: {lhs: "fromString . toFilePath", rhs: "display"}
-- ignore: {name: "Use display", within: "warnMultiple"}
-- ignore: {name: "Use display", within: "Text.PrettyPrint.Leijen.Extended"}
-
 - error: {lhs: "Network.HTTP.Client.MultipartFormData.formDataBody", rhs: "Network.HTTP.StackClient.formDataBody"}
 - error: {lhs: "Network.HTTP.Client.MultipartFormData.partBS", rhs: "Network.HTTP.StackClient.partBS"}
 - error: {lhs: "Network.HTTP.Client.MultipartFormData.partFileRequestBody", rhs: "Network.HTTP.StackClient.partFileRequestBody"}

src/Stack/Build/Source.hs

@@ -19,14 +19,14 @@ module Stack.Build.Source
 import              Stack.Prelude
 import qualified    Pantry.SHA256 as SHA256
 import qualified    Data.ByteString as S
+import              Data.ByteString.Builder (toLazyByteString)
 import              Conduit (ZipSink (..), withSourceFile)
 import qualified    Data.Conduit.List as CL
 import qualified    Distribution.PackageDescription as C
 import              Data.List
 import qualified    Data.Map as Map
 import qualified    Data.Map.Strict as M
 import qualified    Data.Set as Set
-import qualified    Data.Text as T
 import              Foreign.C.Types (CTime)
 import              Stack.Build.Cache
 import              Stack.Build.Haddock (shouldHaddockDeps)
@@ -42,9 +42,6 @@ import              Stack.Types.SourceMap
 import              System.FilePath (takeFileName)
 import              System.IO.Error (isDoesNotExistError)
 import              System.PosixCompat.Files (modificationTime, getFileStatus)
-import qualified    RIO.ByteString as B
-import qualified    RIO.ByteString.Lazy as BL
-import              RIO.Process (proc, readProcess_)
 
 -- | loads and returns project packages
 projectLocalPackages :: HasEnvConfig env
@@ -146,39 +143,33 @@ hashSourceMapData
     -> Map PackageName DepPackage
     -> RIO env SourceMapHash
 hashSourceMapData bc boptsCli wc smDeps = do
-    compilerPath <- encodeUtf8 . T.pack . toFilePath <$> getCompilerPath wc
-    let compilerExe =
-            case wc of
-                Ghc -> "ghc"
-                Ghcjs -> "ghcjs"
-    compilerInfo <- BL.toStrict . fst <$> proc compilerExe ["--info"] readProcess_
+    compilerPath <- getUtf8Builder . fromString . toFilePath <$> getCompilerPath wc
+    compilerInfo <- getCompilerInfo wc
     immDeps <- forM (Map.elems smDeps) depPackageHashableContent
     let -- extra bytestring specifying GHC options supposed to be applied to
         -- GHC boot packages so we'll have differrent hashes when bare
         -- resolver 'ghc-X.Y.Z' is used, no extra-deps and e.g. user wants builds
         -- with profiling or without
-        bootGhcOpts = B.concat $ map encodeUtf8 (generalGhcOptions bc boptsCli False False)
-        hashedContent = compilerPath:compilerInfo:bootGhcOpts:immDeps
-    return $ SourceMapHash (SHA256.hashLazyBytes $ BL.fromChunks hashedContent)
+        bootGhcOpts = map display (generalGhcOptions bc boptsCli False False)
+        hashedContent = toLazyByteString $ compilerPath <> compilerInfo <>
+            getUtf8Builder (mconcat bootGhcOpts) <> mconcat immDeps
+    return $ SourceMapHash (SHA256.hashLazyBytes hashedContent)
 
-depPackageHashableContent :: (HasConfig env) => DepPackage -> RIO env ByteString
+depPackageHashableContent :: (HasConfig env) => DepPackage -> RIO env Builder
 depPackageHashableContent DepPackage {..} = do
     case dpLocation of
         PLMutable _ -> return ""
         PLImmutable pli -> do
             let flagToBs (f, enabled) =
                     if enabled
                         then ""
-                        else "-" <> encodeUtf8 (T.pack $ C.unFlagName f)
+                        else "-" <> fromString (C.unFlagName f)
                 flags = map flagToBs $ Map.toList (cpFlags dpCommon)
-                locationTreeKey (PLIHackage _ _ tk) = tk
-                locationTreeKey (PLIArchive _ pm) = pmTreeKey pm
-                locationTreeKey (PLIRepo _ pm) = pmTreeKey pm
-                treeKeyToBs (TreeKey (BlobKey sha _)) = SHA256.toHexBytes sha
-                ghcOptions = map encodeUtf8 (cpGhcOptions dpCommon)
+                ghcOptions = map display (cpGhcOptions dpCommon)
                 haddocks = if cpHaddocks dpCommon then "haddocks" else ""
-                hash = treeKeyToBs $ locationTreeKey pli
-            return $ B.concat ([hash, haddocks] ++ flags ++ ghcOptions)
+                hash = immutableLocSha pli
+            return $ hash <> haddocks <> getUtf8Builder (mconcat flags) <>
+                getUtf8Builder (mconcat ghcOptions)
 
 -- | All flags for a local package.
 getLocalFlags

src/Stack/Script.hs

@@ -7,13 +7,19 @@ module Stack.Script
     ) where
 
 import           Stack.Prelude
+import           Data.ByteString.Builder    (toLazyByteString)
 import qualified Data.ByteString.Char8      as S8
 import qualified Data.Conduit.List          as CL
 import           Data.List.Split            (splitWhen)
 import qualified Data.Map.Strict            as Map
 import qualified Data.Set                   as Set
+import           Distribution.Compiler      (CompilerFlavor (..))
 import qualified Distribution.PackageDescription as PD
+import qualified Distribution.Types.CondTree as C
 import           Distribution.Types.PackageName (mkPackageName)
+import           Distribution.Types.VersionRange (withinRange)
+import           Distribution.System        (Platform (..))
+import qualified Pantry.SHA256 as SHA256
 import           Path
 import           Path.IO
 import qualified Stack.Build
@@ -24,6 +30,7 @@ import           Stack.PackageDump
 import           Stack.Options.ScriptParser
 import           Stack.Runners
 import           Stack.Setup                (withNewLocalBuildTargets)
+import           Stack.SourceMap            (getCompilerInfo, immutableLocSha)
 import           Stack.Types.BuildPlan
 import           Stack.Types.Compiler
 import           Stack.Types.Config
@@ -33,6 +40,21 @@ import qualified RIO.Directory as Dir
 import           RIO.Process
 import qualified RIO.Text as T
 
+data StackScriptException
+    = MutableDependenciesForScript [PackageName]
+    | AmbiguousModuleName ModuleName [PackageName]
+  deriving Typeable
+
+instance Exception StackScriptException
+
+instance Show StackScriptException where
+    show (MutableDependenciesForScript names) = unlines
+        $ "No mutable packages are allowed in the `script` command. Mutable packages found:"
+        : map (\name -> "- " ++ packageNameString name) names
+    show (AmbiguousModuleName mname pkgs) = unlines
+        $ ("Module " ++ moduleNameString mname ++ " appears in multiple packages: ")
+        : [unwords $ map packageNameString pkgs ]
+
 -- | Run a Stack Script
 scriptCmd :: ScriptOpts -> RIO Runner ()
 scriptCmd opts = do
@@ -87,8 +109,7 @@ scriptCmd opts = do
             case soPackages opts of
                 [] -> do
                     -- Using the import parser
-                    moduleInfo <- getModuleInfo
-                    getPackagesFromModuleInfo moduleInfo (soFile opts)
+                    getPackagesFromImports (soFile opts)
                 packages -> do
                     let targets = concatMap wordsComma packages
                     targets' <- mapM parsePackageNameThrowing targets
@@ -154,31 +175,95 @@ scriptCmd opts = do
       then replaceExtension fp "exe"
       else dropExtension fp
 
-getPackagesFromModuleInfo
-  :: ModuleInfo
-  -> FilePath -- ^ script filename
+getPackagesFromImports
+  :: FilePath -- ^ script filename
   -> RIO EnvConfig (Set PackageName)
-getPackagesFromModuleInfo mi scriptFP = do
-    (pns1, mns) <- liftIO $ parseImports <$> S8.readFile scriptFP
-    pns2 <-
-        if Set.null mns
-            then return Set.empty
-            else do
-                pns <- forM (Set.toList mns) $ \mn ->
-                    case Map.lookup mn $ miModules mi of
-                        Just pns ->
-                            case Set.toList pns of
-                                [] -> assert False $ return Set.empty
-                                [pn] -> return $ Set.singleton pn
-                                pns' -> throwString $ concat
-                                    [ "Module "
-                                    , moduleNameString mn
-                                    , " appears in multiple packages: "
-                                    , unwords $ map packageNameString pns'
-                                    ]
-                        Nothing -> return Set.empty
-                return $ Set.unions pns `Set.difference` blacklist
-    return $ Set.union pns1 pns2
+getPackagesFromImports scriptFP = do
+    (pns, mns) <- liftIO $ parseImports <$> S8.readFile scriptFP
+    if Set.null mns
+        then return pns
+        else Set.union pns <$> getPackagesFromModuleNames mns
+
+getPackagesFromModuleNames
+  :: Set ModuleName
+  -> RIO EnvConfig (Set PackageName)
+getPackagesFromModuleNames mns = do
+    hash <- hashSnapshot
+    withSnapshotCache hash mapSnapshotPackageModules $ \getModulePackages -> do
+        pns <- forM (Set.toList mns) $ \mn -> do
+            pkgs <- getModulePackages mn
+            case pkgs of
+                [] -> return Set.empty
+                [pn] -> return $ Set.singleton pn
+                _ -> throwM $ AmbiguousModuleName mn pkgs
+        return $ Set.unions pns `Set.difference` blacklist
+
+hashSnapshot :: RIO EnvConfig SnapshotCacheHash
+hashSnapshot = do
+    sourceMap <- view $ envConfigL . to envConfigSourceMap
+    let wc = whichCompiler $ smCompiler sourceMap
+    compilerInfo <- getCompilerInfo wc
+    let eitherPliHash (pn, dep) | PLImmutable pli <- dpLocation dep =
+                                    Right $ immutableLocSha pli
+                                | otherwise =
+                                    Left pn
+        deps = Map.toList (smDeps sourceMap)
+    case partitionEithers (map eitherPliHash deps) of
+        ([], pliHashes) -> do
+            let hashedContent = mconcat $ compilerInfo : pliHashes
+            pure $ SnapshotCacheHash (SHA256.hashLazyBytes $ toLazyByteString hashedContent)
+        (mutables, _) ->
+            throwM $ MutableDependenciesForScript mutables
+
+mapSnapshotPackageModules :: RIO EnvConfig (Map PackageName (Set ModuleName))
+mapSnapshotPackageModules = do
+    sourceMap <- view $ envConfigL . to envConfigSourceMap
+    installMap <- toInstallMap sourceMap
+    (_installedMap, globalDumpPkgs, snapshotDumpPkgs, _localDumpPkgs) <-
+        getInstalled installMap
+    let globals = dumpedPackageModules (smGlobal sourceMap) globalDumpPkgs
+        notHidden = Map.filter (not . dpHidden)
+        notHiddenDeps = notHidden $ smDeps sourceMap
+        installedDeps = dumpedPackageModules notHiddenDeps snapshotDumpPkgs
+        dumpPkgs = Set.fromList $ map (pkgName . dpPackageIdent) snapshotDumpPkgs
+        notInstalledDeps = Map.withoutKeys notHiddenDeps dumpPkgs
+    otherDeps <- for notInstalledDeps $ \dep -> do
+        gpd <- liftIO $ cpGPD (dpCommon dep)
+        Set.fromList <$> allExposedModules gpd
+    -- source map construction process should guarantee unique package names
+    -- in these maps
+    return $ globals <> installedDeps <> otherDeps
+
+dumpedPackageModules :: Map PackageName a
+                     -> [DumpPackage]
+                     -> Map PackageName (Set ModuleName)
+dumpedPackageModules pkgs dumpPkgs =
+    let pnames = Map.keysSet pkgs `Set.difference` blacklist
+    in Map.fromList
+           [ (pn, dpExposedModules)
+           | DumpPackage {..} <- dumpPkgs
+           , let PackageIdentifier pn _ = dpPackageIdent
+           , pn `Set.member` pnames
+           ]
+
+allExposedModules :: PD.GenericPackageDescription -> RIO EnvConfig [ModuleName]
+allExposedModules gpd = do
+  Platform curArch curOs <- view platformL
+  curCompiler <- view actualCompilerVersionL
+  let checkCond (PD.OS os) = pure $ os == curOs
+      checkCond (PD.Arch arch) = pure $ arch == curArch
+      checkCond (PD.Impl compiler range) = case curCompiler of
+        ACGhc version ->
+          pure $ compiler == GHC && version `withinRange` range
+        ACGhcjs version _ghcVersion ->
+          pure $ compiler == GHCJS && version `withinRange` range
+      -- currently we don't do flag checking here
+      checkCond other = Left other
+      mlibrary = snd . C.simplifyCondTree checkCond <$> PD.condLibrary gpd
+  pure $ case mlibrary  of
+    Just lib -> PD.exposedModules lib ++
+                map PD.moduleReexportName (PD.reexportedModules lib)
+    Nothing  -> mempty
 
 -- | The Stackage project introduced the concept of hidden packages,
 -- to deal with conflicting module names. However, this is a
@@ -232,40 +317,6 @@ blacklist = Set.fromList
     , mkPackageName "cryptohash-sha256"
     ]
 
-getModuleInfo :: HasEnvConfig env => RIO env ModuleInfo
-getModuleInfo = do
-    sourceMap <- view $ envConfigL . to envConfigSourceMap
-    installMap <- toInstallMap sourceMap
-    (_installedMap, globalDumpPkgs, snapshotDumpPkgs, _localDumpPkgs) <-
-        getInstalled installMap
-    let globals = toModuleInfo (smGlobal sourceMap) globalDumpPkgs
-        notHiddenDeps = notHidden $ smDeps sourceMap
-        installedDeps = toModuleInfo notHiddenDeps snapshotDumpPkgs
-        dumpPkgs = Set.fromList $ map (pkgName . dpPackageIdent) snapshotDumpPkgs
-        notInstalledDeps = Map.withoutKeys notHiddenDeps dumpPkgs
-    otherDeps <- liftIO $
-                 fmap (Map.fromListWith mappend . concat) $
-                 forM (Map.toList notInstalledDeps) $ \(pname, dep) -> do
-        gpd <- cpGPD (dpCommon dep)
-        let modules = maybe [] PD.exposedModules $
-              maybe (PD.library $ PD.packageDescription gpd) (Just . PD.condTreeData) $
-              PD.condLibrary gpd
-        return [ (m, Set.singleton pname) | m <- modules ]
-    return $ globals <> installedDeps <> ModuleInfo otherDeps
-  where
-    notHidden = Map.filter (not . dpHidden)
-    toModuleInfo pkgs dumpPkgs =
-        let pnames = Map.keysSet pkgs `Set.difference` blacklist
-            modules =
-                Map.fromListWith mappend
-                    [ (m, Set.singleton pn)
-                    | DumpPackage {..} <- dumpPkgs
-                    , let PackageIdentifier pn _ = dpPackageIdent
-                    , pn `Set.member` pnames
-                    , m <- Set.toList dpExposedModules
-                    ]
-        in ModuleInfo modules
-
 parseImports :: ByteString -> (Set PackageName, Set ModuleName)
 parseImports =
     fold . mapMaybe (parseLine . stripCR') . S8.lines

src/Stack/SourceMap.hs

@@ -15,12 +15,16 @@ module Stack.SourceMap
     , globalCondCheck
     , pruneGlobals
     , globalsFromHints
+    , getCompilerInfo
+    , immutableLocSha
     ) where
 
+import Data.ByteString.Builder (byteString, lazyByteString)
 import qualified Data.Conduit.List as CL
 import qualified Distribution.PackageDescription as PD
 import Distribution.System (Platform(..))
 import Pantry
+import qualified Pantry.SHA256 as SHA256
 import qualified RIO
 import qualified RIO.Map as Map
 import qualified RIO.Set as Set
@@ -231,3 +235,19 @@ pruneGlobals globals deps =
             dpGhcPkgId dpDepends deps
   in Map.map (GlobalPackage . pkgVersion . dpPackageIdent) keptGlobals <>
      Map.map ReplacedGlobalPackage prunedGlobals
+
+getCompilerInfo :: (HasConfig env) => WhichCompiler -> RIO env Builder
+getCompilerInfo wc = do
+    let compilerExe =
+            case wc of
+                Ghc -> "ghc"
+                Ghcjs -> "ghcjs"
+    lazyByteString . fst <$> proc compilerExe ["--info"] readProcess_
+
+immutableLocSha :: PackageLocationImmutable -> Builder
+immutableLocSha = byteString . treeKeyToBs . locationTreeKey
+  where
+    locationTreeKey (PLIHackage _ _ tk) = tk
+    locationTreeKey (PLIArchive _ pm) = pmTreeKey pm
+    locationTreeKey (PLIRepo _ pm) = pmTreeKey pm
+    treeKeyToBs (TreeKey (BlobKey sha _)) = SHA256.toHexBytes sha

subs/pantry/src/Pantry.hs

@@ -160,6 +160,8 @@ module Pantry
   , getHackageTypoCorrections
   , loadGlobalHints
   , partitionReplacedDependencies
+  , SnapshotCacheHash (..)
+  , withSnapshotCache
   ) where
 
 import RIO
@@ -1501,3 +1503,21 @@ prunePackageWithDeps pkgs getName getDeps (pname, a)  = do
       else do
         modify' $ first (Map.insert pname prunedDeps)
       return $ not (null prunedDeps)
+
+withSnapshotCache
+  :: (HasPantryConfig env, HasLogFunc env)
+  => SnapshotCacheHash
+  -> RIO env (Map PackageName (Set ModuleName))
+  -> ((ModuleName -> RIO env [PackageName]) -> RIO env a)
+  -> RIO env a
+withSnapshotCache hash getModuleMapping f = do
+  mres <- withStorage $ getSnapshotCacheByHash hash
+  cacheId <- case mres of
+    Nothing -> do
+      scId <- withStorage $ getSnapshotCacheId hash
+      packageModules <- getModuleMapping
+      logWarn "Populating snapshot module name cache"
+      withStorage $ storeSnapshotModuleCache scId packageModules
+      return scId
+    Just scId -> pure scId
+  f $ withStorage . loadExposedModulePackages cacheId