Rewrite PRSelf when loading a dependency package (#3406)

When buidling simple project that has our favourite large project as a
dependency, this decreased
- total allocations from 63GB to 57GB
- run time from 34.0s to 31.5s

compiler/daml-lf-ast/src/DA/Daml/LF/Ast/World.hs

@@ -11,7 +11,6 @@ module DA.Daml.LF.Ast.World(
     initWorldSelf,
     extendWorldSelf,
     ExternalPackage(..),
-    rewriteSelfReferences,
     LookupError,
     lookupTemplate,
     lookupDataType,
@@ -31,7 +30,6 @@ import qualified Data.NameMap as NM
 import GHC.Generics
 
 import DA.Daml.LF.Ast.Base
-import DA.Daml.LF.Ast.Optics (moduleModuleRef)
 import DA.Daml.LF.Ast.Pretty ()
 import DA.Daml.LF.Ast.Version
 
@@ -63,14 +61,6 @@ data DalfPackage = DalfPackage
 
 instance NFData DalfPackage
 
--- | Rewrite all `PRSelf` references to `PRImport` references.
-rewriteSelfReferences :: PackageId -> Package -> ExternalPackage
-rewriteSelfReferences pkgId = ExternalPackage pkgId . rewrite
-    where
-        rewrite = over (_packageModules . NM.traverse . moduleModuleRef . _1) $ \case
-            PRSelf -> PRImport pkgId
-            ref@PRImport{} -> ref
-
 -- | Construct the 'World' from only the imported packages.
 initWorld :: [ExternalPackage] -> Version -> World
 initWorld importedPkgs version =

compiler/daml-lf-proto/src/DA/Daml/LF/Proto3/Archive.hs

@@ -10,6 +10,7 @@ module DA.Daml.LF.Proto3.Archive
   , encodeArchiveAndHash
   , encodePackageHash
   , ArchiveError(..)
+  , DecodingMode(..)
   ) where
 
 import           Control.Lens             (over, _Left)
@@ -36,9 +37,20 @@ data ArchiveError
     | HashMismatch !T.Text !T.Text
   deriving (Eq, Show)
 
+-- | Mode in which to decode the DALF. Currently, this only decides whether
+-- to rewrite occurrences of `PRSelf` with `PRImport packageId`.
+data DecodingMode
+    = DecodeAsMain
+      -- ^ Keep occurrences of `PRSelf` as is.
+    | DecodeAsDependency
+      -- ^ Replace `PRSelf` with `PRImport packageId`, where `packageId` is
+      -- the id of the package being decoded.
+    deriving (Eq, Show)
+
+
 -- | Decode a LF archive header, returing the hash and the payload
-decodeArchive :: BS.ByteString -> Either ArchiveError (LF.PackageId, LF.Package)
-decodeArchive bytes = do
+decodeArchive :: DecodingMode -> BS.ByteString -> Either ArchiveError (LF.PackageId, LF.Package)
+decodeArchive mode bytes = do
     archive <- over _Left (ProtobufError . show) $ Proto.fromByteString bytes
     let payloadBytes = ProtoLF.archivePayload archive
     let archiveHash = TL.toStrict (ProtoLF.archiveHash archive)
@@ -51,10 +63,14 @@ decodeArchive bytes = do
 
     when (computedHash /= archiveHash) $
       Left (HashMismatch archiveHash computedHash)
+    let packageId = LF.PackageId archiveHash
+    let selfPackageRef = case mode of
+            DecodeAsMain -> LF.PRSelf
+            DecodeAsDependency -> LF.PRImport packageId
 
     payload <- over _Left (ProtobufError . show) $ Proto.fromByteString payloadBytes
-    package <- over _Left (ProtobufError. show) $ Decode.decodePayload payload
-    return (LF.PackageId archiveHash, package)
+    package <- over _Left (ProtobufError. show) $ Decode.decodePayload selfPackageRef payload
+    return (packageId, package)
 
 
 -- | Encode a LFv1 package payload into a DAML-LF archive using the default

compiler/daml-lf-proto/src/DA/Daml/LF/Proto3/Decode.hs

@@ -7,14 +7,14 @@ module DA.Daml.LF.Proto3.Decode
   ) where
 
 import Com.Digitalasset.DamlLfDev.DamlLf (ArchivePayload(..), ArchivePayloadSum(..))
-import DA.Daml.LF.Ast (Package)
+import DA.Daml.LF.Ast (Package, PackageRef)
 import DA.Daml.LF.Proto3.Error
 import qualified DA.Daml.LF.Proto3.DecodeV1 as DecodeV1
 
-decodePayload :: ArchivePayload -> Either Error Package
-decodePayload payload = case archivePayloadSum payload of
+decodePayload :: PackageRef -> ArchivePayload -> Either Error Package
+decodePayload selfPackageRef payload = case archivePayloadSum payload of
     Just ArchivePayloadSumDamlLf0{} -> Left $ ParseError "Payload is DamlLf0"
-    Just (ArchivePayloadSumDamlLf1 package) -> DecodeV1.decodePackage minor package
+    Just (ArchivePayloadSumDamlLf1 package) -> DecodeV1.decodePackage minor selfPackageRef package
     Nothing -> Left $ ParseError "Empty payload"
     where
         minor = archivePayloadMinor payload

compiler/daml-lf-proto/src/DA/Daml/LF/Proto3/DecodeV1.hs

@@ -32,6 +32,7 @@ import qualified Proto3.Suite as Proto
 data DecodeEnv = DecodeEnv
     { internedStrings :: !(V.Vector T.Text)
     , internedDottedNames :: !(V.Vector [T.Text])
+    , selfPackageRef :: PackageRef
     }
 
 newtype Decode a = Decode{unDecode :: ReaderT DecodeEnv (Except Error) a}
@@ -134,7 +135,7 @@ decodeValName LF1.ValName{..} = do
 decodePackageRef :: LF1.PackageRef -> Decode PackageRef
 decodePackageRef (LF1.PackageRef pref) =
     mayDecode "packageRefSum" pref $ \case
-        LF1.PackageRefSumSelf _ -> pure PRSelf
+        LF1.PackageRefSumSelf _ -> asks selfPackageRef
         LF1.PackageRefSumPackageIdStr pkgId -> pure $ PRImport $ PackageId $ decodeString pkgId
         LF1.PackageRefSumPackageIdInternedStr strId -> PRImport . PackageId <$> lookupString strId
 
@@ -161,8 +162,8 @@ decodeInternedDottedName :: LF1.InternedDottedName -> Decode [T.Text]
 decodeInternedDottedName (LF1.InternedDottedName ids) =
     mapM lookupString $ V.toList ids
 
-decodePackage :: TL.Text -> LF1.Package -> Either Error Package
-decodePackage minorText (LF1.Package mods internedStringsV internedDottedNamesV) = do
+decodePackage :: TL.Text -> LF.PackageRef -> LF1.Package -> Either Error Package
+decodePackage minorText selfPackageRef (LF1.Package mods internedStringsV internedDottedNamesV) = do
   version <- decodeVersion (decodeString minorText)
   let internedStrings = V.map decodeString internedStringsV
   let internedDottedNames = V.empty
@@ -174,7 +175,7 @@ decodePackage minorText (LF1.Package mods internedStringsV internedDottedNamesV)
 
 decodeScenarioModule :: TL.Text -> LF1.Package -> Either Error Module
 decodeScenarioModule minorText protoPkg = do
-    Package _ modules <- decodePackage minorText protoPkg
+    Package _ modules <- decodePackage minorText PRSelf protoPkg
     pure $ head $ NM.toList modules
 
 decodeModule :: LF1.Module -> Decode Module

compiler/damlc/daml-ide-core/src/Development/IDE/Core/Rules/Daml.hs

@@ -409,9 +409,9 @@ generatePackageMap fps = do
         return $ do
           (pkgId, package) <-
             mapLeft (ideErrorPretty $ toNormalizedFilePath dalf) $
-            Archive.decodeArchive dalfBS
+            Archive.decodeArchive Archive.DecodeAsDependency dalfBS
           let unitId = stringToUnitId $ dropExtension $ takeFileName dalf
-          Right (unitId, LF.DalfPackage pkgId (LF.rewriteSelfReferences pkgId package) dalfBS)
+          Right (unitId, LF.DalfPackage pkgId (LF.ExternalPackage pkgId package) dalfBS)
   return (diags, Map.fromList pkgs)
 
 generatePackageMapRule :: Options -> Rules ()

compiler/damlc/daml-visual/src/DA/Daml/Visual.hs

@@ -189,12 +189,12 @@ moduleAndTemplates :: LF.World -> LF.Module -> [TemplateChoices]
 moduleAndTemplates world mod = map (\t -> TemplateChoices t (LF.moduleName mod) (templatePossibleUpdates world t)) $ NM.toList $ LF.moduleTemplates mod
 
 dalfBytesToPakage :: BSL.ByteString -> ExternalPackage
-dalfBytesToPakage bytes = case Archive.decodeArchive $ BSL.toStrict bytes of
-    Right (pkgId, pkg) -> rewriteSelfReferences pkgId pkg
+dalfBytesToPakage bytes = case Archive.decodeArchive Archive.DecodeAsDependency $ BSL.toStrict bytes of
+    Right (pkgId, pkg) -> ExternalPackage pkgId pkg
     Left err -> error (show err)
 
 darToWorld :: Dalfs -> LF.World
-darToWorld Dalfs{..} = case Archive.decodeArchive $ BSL.toStrict mainDalf of
+darToWorld Dalfs{..} = case Archive.decodeArchive Archive.DecodeAsMain $ BSL.toStrict mainDalf of
     Right (_, mainPkg) -> AST.initWorldSelf pkgs mainPkg
     Left err -> error (show err)
     where

compiler/damlc/lib/DA/Cli/Damlc.hs

@@ -554,7 +554,8 @@ createProjectPackageDb opts thisSdkVer deps0 dataDeps = do
         forM allDalfs $ \(name, dalf) -> do
             (pkgId, package) <-
                 either (fail . DA.Pretty.renderPretty) pure $
-                Archive.decodeArchive dalf
+                -- FIXME(MH): This keeps the old behaviour but seems wrong to me.
+                Archive.decodeArchive Archive.DecodeAsMain dalf
             pure (pkgId, package, dalf, stringToUnitId name)
     -- mapping from package id's to unit id's. if the same package is imported with
     -- different unit id's, we would loose a unit id here.
@@ -979,7 +980,7 @@ execInspect inFile outFile jsonOutput lvl =
          $ Proto.JSONPB.toAesonValue archive
       else do
         (pkgId, lfPkg) <- errorOnLeft "Cannot decode package" $
-                   Archive.decodeArchive bytes
+                   Archive.decodeArchive Archive.DecodeAsMain bytes
         writeOutput outFile $ render Plain $
           DA.Pretty.vsep
             [ DA.Pretty.keyword_ "package" DA.Pretty.<-> DA.Pretty.text (LF.unPackageId pkgId) DA.Pretty.<-> DA.Pretty.keyword_ "where"
@@ -1011,7 +1012,7 @@ execInspectDar inFile =
           (pkgId, _lfPkg) <-
               errorOnLeft
                   ("Cannot decode package " <> ZipArchive.eRelativePath dalfEntry)
-                  (Archive.decodeArchive dalf)
+                  (Archive.decodeArchive Archive.DecodeAsMain dalf)
           putStrLn $
               (dropExtension $ takeFileName $ ZipArchive.eRelativePath dalfEntry) <> " " <>
               show (LF.unPackageId pkgId)
@@ -1090,7 +1091,7 @@ execMigrate projectOpts inFile1_ inFile2_ mbDir =
     decode dalf =
         errorOnLeft
             "Cannot decode daml-lf archive"
-            (Archive.decodeArchive dalf)
+            (Archive.decodeArchive Archive.DecodeAsMain dalf)
     getModule modName pkg =
         maybe
             (fail $ T.unpack $ "Can't find module" <> LF.moduleNameString modName)
@@ -1139,7 +1140,7 @@ execGenerateSrc dalfFp = Command GenerateSrc effect
     unitId = stringToUnitId $ takeBaseName dalfFp
     effect = do
         bytes <- B.readFile dalfFp
-        case Archive.decodeArchive bytes of
+        case Archive.decodeArchive Archive.DecodeAsMain bytes of
             Left err -> fail $ DA.Pretty.renderPretty err
             Right (pkgId, pkg) -> do
                 let genSrcs =
@@ -1191,7 +1192,7 @@ execGenerateGenSrc darFp mbQual outDir = Command GenerateGenerics effect
             createDirectoryIfMissing True $ takeDirectory fp
             writeFileUTF8 fp src
 
-    decode = either (fail . DA.Pretty.renderPretty) pure . Archive.decodeArchive
+    decode = either (fail . DA.Pretty.renderPretty) pure . Archive.decodeArchive Archive.DecodeAsMain
 
 
 

language-support/hs/bindings/test/DA/Ledger/Tests.hs

@@ -10,7 +10,7 @@ import Control.Concurrent
 import Control.Monad
 import Control.Monad.IO.Class(liftIO)
 import DA.Bazel.Runfiles
-import DA.Daml.LF.Proto3.Archive (decodeArchive)
+import DA.Daml.LF.Proto3.Archive (DecodingMode(DecodeAsMain), decodeArchive)
 import DA.Daml.LF.Reader(Dalfs(..),readDalfs)
 import DA.Ledger.Sandbox (Sandbox,SandboxSpec(..),startSandbox,shutdownSandbox,withSandbox)
 import Data.List (elem,isPrefixOf,isInfixOf,(\\))
@@ -693,7 +693,7 @@ mainPackageId :: SandboxSpec -> IO PackageId
 mainPackageId (SandboxSpec dar) = do
     archive <- Zip.toArchive <$> BSL.readFile dar
     Dalfs { mainDalf } <- either fail pure $ readDalfs archive
-    case decodeArchive (BSL.toStrict mainDalf) of
+    case decodeArchive DecodeAsMain (BSL.toStrict mainDalf) of
         Left err -> fail $ show err
         Right (LF.PackageId pId, _) -> pure (PackageId $ Text.fromStrict pId)