[Localization] Delay loading localization

- Add `LocalizationProducerState` to manage the states of LocalizationProducers.
- Add `initializeImpl` and `initializeIfNeeded` to manage lazily initialization of `LocalizationProducer`s.
- Move constructing a localization producer from DiagEngine to `LocalizationProrducer` itself.

Author: HassanElDesouky

include/swift/AST/DiagnosticEngine.h

@@ -802,29 +802,11 @@ namespace swift {
 
     bool isPrettyPrintingDecl() const { return IsPrettyPrintingDecl; }
 
-    void setLocalization(std::string locale, std::string path) {
+    void setLocalization(StringRef locale, StringRef path) {
       assert(!locale.empty());
       assert(!path.empty());
-      llvm::SmallString<128> filePath(path);
-      llvm::sys::path::append(filePath, locale);
-      llvm::sys::path::replace_extension(filePath, ".db");
-
-      // If the serialized diagnostics file not available,
-      // fallback to the `YAML` file.
-      if (llvm::sys::fs::exists(filePath)) {
-        if (auto file = llvm::MemoryBuffer::getFile(filePath)) {
-          localization = std::make_unique<diag::SerializedLocalizationProducer>(
-              std::move(file.get()), getPrintDiagnosticNames());
-        }
-      } else {
-        llvm::sys::path::replace_extension(filePath, ".yaml");
-        // In case of missing localization files, we should fallback to messages
-        // from `.def` files.
-        if (llvm::sys::fs::exists(filePath)) {
-          localization = std::make_unique<diag::YAMLLocalizationProducer>(
-              filePath.str(), getPrintDiagnosticNames());
-        }
-      }
+      localization = diag::LocalizationProducer::producerFor(
+          locale, path, getPrintDiagnosticNames());
     }
 
     void ignoreDiagnostic(DiagID id) {

include/swift/Localization/LocalizationFormat.h

@@ -44,6 +44,12 @@ namespace diag {
 
 using namespace llvm::support;
 
+enum LocalizationProducerState : uint8_t {
+  NotInitialized,
+  Initialized,
+  FailedInitialization
+};
+
 class DefToYAMLConverter {
   llvm::ArrayRef<const char *> IDs;
   llvm::ArrayRef<const char *> Messages;
@@ -162,6 +168,7 @@ class LocalizationProducer {
   llvm::BumpPtrAllocator localizationAllocator;
   llvm::StringSaver localizationSaver;
   bool printDiagnosticNames;
+  LocalizationProducerState state = NotInitialized;
 
 public:
   LocalizationProducer(bool printDiagnosticNames = false)
@@ -173,16 +180,33 @@ class LocalizationProducer {
   virtual llvm::StringRef getMessageOr(swift::DiagID id,
                                        llvm::StringRef defaultMessage);
 
+  /// \returns a `SerializedLocalizationProducer` pointer if the serialized
+  /// diagnostics file available, otherwise returns a `YAMLLocalizationProducer`
+  /// if the `YAML` file is available. If both files aren't available returns a
+  /// `nullptr`.
+  static std::unique_ptr<LocalizationProducer>
+  producerFor(llvm::StringRef locale, llvm::StringRef path,
+              bool printDiagnosticNames);
+
   virtual ~LocalizationProducer() {}
 
 protected:
+  LocalizationProducerState getState() const;
+
+  /// Used to lazily initialize `LocalizationProducer`s.
+  /// \returns true if the producer is successfully initialized, false
+  /// otherwise.
+  virtual bool initializeImpl() = 0;
+  virtual void initializeIfNeeded() final;
+
   /// Retrieve a message for the given diagnostic id.
   /// \returns empty string if message couldn't be found.
   virtual llvm::StringRef getMessage(swift::DiagID id) const = 0;
 };
 
 class YAMLLocalizationProducer final : public LocalizationProducer {
   std::vector<std::string> diagnostics;
+  std::string filePath;
 
 public:
   /// The diagnostics IDs that are no longer available in `.def`
@@ -194,9 +218,10 @@ class YAMLLocalizationProducer final : public LocalizationProducer {
   /// maintained by this producer, callback gets each translation
   /// with its unique identifier.
   void forEachAvailable(
-      llvm::function_ref<void(swift::DiagID, llvm::StringRef)> callback) const;
+      llvm::function_ref<void(swift::DiagID, llvm::StringRef)> callback);
 
 protected:
+  bool initializeImpl() override;
   llvm::StringRef getMessage(swift::DiagID id) const override;
 };
 
@@ -213,6 +238,7 @@ class SerializedLocalizationProducer final : public LocalizationProducer {
       bool printDiagnosticNames = false);
 
 protected:
+  bool initializeImpl() override;
   llvm::StringRef getMessage(swift::DiagID id) const override;
 };
 

lib/AST/DiagnosticEngine.cpp

@@ -1160,9 +1160,9 @@ DiagnosticEngine::diagnosticStringFor(const DiagID id,
   auto defaultMessage = printDiagnosticNames
                             ? debugDiagnosticStrings[(unsigned)id]
                             : diagnosticStrings[(unsigned)id];
-  if (localization) {
-    auto localizedMessage =
-        localization.get()->getMessageOr(id, defaultMessage);
+
+  if (auto producer = localization.get()) {
+    auto localizedMessage = producer->getMessageOr(id, defaultMessage);
     return localizedMessage;
   }
   return defaultMessage;

lib/Localization/LocalizationFormat.cpp

@@ -93,9 +93,24 @@ bool SerializedLocalizationWriter::emit(llvm::StringRef filePath) {
   return OS.has_error();
 }
 
+void LocalizationProducer::initializeIfNeeded() {
+  if (state != NotInitialized)
+    return;
+
+  if (initializeImpl())
+    state = Initialized;
+  else
+    state = FailedInitialization;
+}
+
 llvm::StringRef
 LocalizationProducer::getMessageOr(swift::DiagID id,
                                    llvm::StringRef defaultMessage) {
+  initializeIfNeeded();
+  if (getState() == FailedInitialization) {
+    return defaultMessage;
+  }
+
   auto localizedMessage = getMessage(id);
   if (localizedMessage.empty())
     return defaultMessage;
@@ -108,14 +123,22 @@ LocalizationProducer::getMessageOr(swift::DiagID id,
   return localizedMessage;
 }
 
+LocalizationProducerState LocalizationProducer::getState() const {
+  return state;
+}
+
 SerializedLocalizationProducer::SerializedLocalizationProducer(
     std::unique_ptr<llvm::MemoryBuffer> buffer, bool printDiagnosticNames)
     : LocalizationProducer(printDiagnosticNames), Buffer(std::move(buffer)) {
+}
+
+bool SerializedLocalizationProducer::initializeImpl() {
   auto base =
       reinterpret_cast<const unsigned char *>(Buffer.get()->getBufferStart());
   auto tableOffset = endian::read<offset_type>(base, little);
   SerializedTable.reset(SerializedLocalizationTable::Create(
       base + tableOffset, base + sizeof(offset_type), base));
+  return true;
 }
 
 llvm::StringRef
@@ -128,27 +151,64 @@ SerializedLocalizationProducer::getMessage(swift::DiagID id) const {
 
 YAMLLocalizationProducer::YAMLLocalizationProducer(llvm::StringRef filePath,
                                                    bool printDiagnosticNames)
-    : LocalizationProducer(printDiagnosticNames) {
+    : LocalizationProducer(printDiagnosticNames), filePath(filePath) {
+}
+
+bool YAMLLocalizationProducer::initializeImpl() {
   auto FileBufOrErr = llvm::MemoryBuffer::getFileOrSTDIN(filePath);
   llvm::MemoryBuffer *document = FileBufOrErr->get();
   diag::LocalizationInput yin(document->getBuffer());
   yin >> diagnostics;
   unknownIDs = std::move(yin.unknownIDs);
+  return true;
 }
 
 llvm::StringRef YAMLLocalizationProducer::getMessage(swift::DiagID id) const {
   return diagnostics[(unsigned)id];
 }
 
 void YAMLLocalizationProducer::forEachAvailable(
-    llvm::function_ref<void(swift::DiagID, llvm::StringRef)> callback) const {
+    llvm::function_ref<void(swift::DiagID, llvm::StringRef)> callback) {
+  initializeIfNeeded();
+  if (getState() == FailedInitialization) {
+    return;
+  }
+
   for (uint32_t i = 0, n = diagnostics.size(); i != n; ++i) {
     auto translation = diagnostics[i];
     if (!translation.empty())
       callback(static_cast<swift::DiagID>(i), translation);
   }
 }
 
+std::unique_ptr<LocalizationProducer>
+LocalizationProducer::producerFor(llvm::StringRef locale, llvm::StringRef path,
+                                  bool printDiagnosticNames) {
+  std::unique_ptr<LocalizationProducer> producer;
+  llvm::SmallString<128> filePath(path);
+  llvm::sys::path::append(filePath, locale);
+  llvm::sys::path::replace_extension(filePath, ".db");
+
+  // If the serialized diagnostics file not available,
+  // fallback to the `YAML` file.
+  if (llvm::sys::fs::exists(filePath)) {
+    if (auto file = llvm::MemoryBuffer::getFile(filePath)) {
+      producer = std::make_unique<diag::SerializedLocalizationProducer>(
+          std::move(file.get()), printDiagnosticNames);
+    }
+  } else {
+    llvm::sys::path::replace_extension(filePath, ".yaml");
+    // In case of missing localization files, we should fallback to messages
+    // from `.def` files.
+    if (llvm::sys::fs::exists(filePath)) {
+      producer = std::make_unique<diag::YAMLLocalizationProducer>(
+          filePath.str(), printDiagnosticNames);
+    }
+  }
+
+  return producer;
+}
+
 llvm::Optional<uint32_t> LocalizationInput::readID(llvm::yaml::IO &io) {
   LocalDiagID diagID;
   io.mapRequired("id", diagID);