Unity Build: Fix per-config sources in multi-config generators

Single-config generators already support unity builds with per-config
sources because they compute sources using `CMAKE_BUILD_TYPE` as the
configuration.  Each original source is either included in the unity
build source, or not.

Teach multi-config generators to compute the list of sources for
inclusion in unity builds using all configurations.  Previously they
only used the empty string as the configuration.  Each original source
may be included in some configurations, but not others.  Use
preprocessor conditions to guard their inclusion when necessary.

Fixes: #22892

7 files changed, 130 insertions, 31 deletions

diff --git a/Source/cmLocalGenerator.cxx b/Source/cmLocalGenerator.cxx
index 3ac77c1..c106137 100644
--- a/Source/cmLocalGenerator.cxx
+++ b/Source/cmLocalGenerator.cxx
@@ -2808,7 +2808,7 @@ inline void RegisterUnitySources(cmGeneratorTarget* target, cmSourceFile* sf,
 }
 
 cmLocalGenerator::UnitySource cmLocalGenerator::WriteUnitySource(
-  cmGeneratorTarget* target,
+  cmGeneratorTarget* target, std::vector<std::string> const& configs,
   cmRange<std::vector<UnityBatchedSource>::const_iterator> sources,
   cmValue beforeInclude, cmValue afterInclude, std::string filename) const
 {
@@ -2818,21 +2818,36 @@ cmLocalGenerator::UnitySource cmLocalGenerator::WriteUnitySource(
   file.SetCopyIfDifferent(true);
   file << "/* generated by CMake */\n\n";
 
+  bool perConfig = false;
   for (UnityBatchedSource const& ubs : sources) {
+    cm::optional<std::string> cond;
+    if (ubs.Configs.size() != configs.size()) {
+      perConfig = true;
+      cond = std::string();
+      cm::string_view sep;
+      for (size_t ci : ubs.Configs) {
+        cond = cmStrCat(*cond, sep, "defined(CMAKE_UNITY_CONFIG_",
+                        cmSystemTools::UpperCase(configs[ci]), ")");
+        sep = " || "_s;
+      }
+    }
     RegisterUnitySources(target, ubs.Source, filename);
-    WriteUnitySourceInclude(file, ubs.Source->ResolveFullPath(), beforeInclude,
-                            afterInclude, uniqueIdName);
+    WriteUnitySourceInclude(file, cond, ubs.Source->ResolveFullPath(),
+                            beforeInclude, afterInclude, uniqueIdName);
   }
 
-  return UnitySource(std::move(filename));
+  return UnitySource(std::move(filename), perConfig);
 }
 
-void cmLocalGenerator::WriteUnitySourceInclude(std::ostream& unity_file,
-                                               std::string const& sf_full_path,
-                                               cmValue beforeInclude,
-                                               cmValue afterInclude,
-                                               cmValue uniqueIdName) const
+void cmLocalGenerator::WriteUnitySourceInclude(
+  std::ostream& unity_file, cm::optional<std::string> const& cond,
+  std::string const& sf_full_path, cmValue beforeInclude, cmValue afterInclude,
+  cmValue uniqueIdName) const
 {
+  if (cond) {
+    unity_file << "#if " << *cond << "\n";
+  }
+
   if (cmNonempty(uniqueIdName)) {
     std::string pathToHash;
     auto PathEqOrSubDir = [](std::string const& a, std::string const& b) {
@@ -2867,12 +2882,16 @@ void cmLocalGenerator::WriteUnitySourceInclude(std::ostream& unity_file,
   if (afterInclude) {
     unity_file << *afterInclude << "\n";
   }
+  if (cond) {
+    unity_file << "#endif\n";
+  }
   unity_file << "\n";
 }
 
 std::vector<cmLocalGenerator::UnitySource>
 cmLocalGenerator::AddUnityFilesModeAuto(
   cmGeneratorTarget* target, std::string const& lang,
+  std::vector<std::string> const& configs,
   std::vector<UnityBatchedSource> const& filtered_sources,
   cmValue beforeInclude, cmValue afterInclude,
   std::string const& filename_base, size_t batchSize)
@@ -2891,9 +2910,9 @@ cmLocalGenerator::AddUnityFilesModeAuto(
                                     (lang == "C") ? "_c.c" : "_cxx.cxx");
     auto const begin = filtered_sources.begin() + batch * batchSize;
     auto const end = begin + chunk;
-    unity_files.emplace_back(
-      this->WriteUnitySource(target, cmMakeRange(begin, end), beforeInclude,
-                             afterInclude, std::move(filename)));
+    unity_files.emplace_back(this->WriteUnitySource(
+      target, configs, cmMakeRange(begin, end), beforeInclude, afterInclude,
+      std::move(filename)));
   }
   return unity_files;
 }
@@ -2901,6 +2920,7 @@ cmLocalGenerator::AddUnityFilesModeAuto(
 std::vector<cmLocalGenerator::UnitySource>
 cmLocalGenerator::AddUnityFilesModeGroup(
   cmGeneratorTarget* target, std::string const& lang,
+  std::vector<std::string> const& configs,
   std::vector<UnityBatchedSource> const& filtered_sources,
   cmValue beforeInclude, cmValue afterInclude,
   std::string const& filename_base)
@@ -2927,9 +2947,9 @@ cmLocalGenerator::AddUnityFilesModeGroup(
     auto const& name = item.first;
     std::string filename = cmStrCat(filename_base, "unity_", name,
                                     (lang == "C") ? "_c.c" : "_cxx.cxx");
-    unity_files.emplace_back(
-      this->WriteUnitySource(target, cmMakeRange(item.second), beforeInclude,
-                             afterInclude, std::move(filename)));
+    unity_files.emplace_back(this->WriteUnitySource(
+      target, configs, cmMakeRange(item.second), beforeInclude, afterInclude,
+      std::move(filename)));
   }
 
   return unity_files;
@@ -2943,19 +2963,24 @@ void cmLocalGenerator::AddUnityBuild(cmGeneratorTarget* target)
 
   std::vector<UnityBatchedSource> unitySources;
 
-  {
-    // FIXME: Handle all configurations in multi-config generators.
-    std::string config;
-    if (!this->GetGlobalGenerator()->IsMultiConfig()) {
-      config = this->Makefile->GetSafeDefinition("CMAKE_BUILD_TYPE");
-    }
+  std::vector<std::string> configs =
+    this->Makefile->GetGeneratorConfigs(cmMakefile::IncludeEmptyConfig);
+
+  std::map<cmSourceFile const*, size_t> index;
 
+  for (size_t ci = 0; ci < configs.size(); ++ci) {
     // FIXME: Refactor collection of sources to not evaluate object libraries.
     std::vector<cmSourceFile*> sources;
-    target->GetSourceFiles(sources, config);
-
+    target->GetSourceFiles(sources, configs[ci]);
     for (cmSourceFile* sf : sources) {
-      unitySources.emplace_back(sf);
+      auto mi = index.find(sf);
+      if (mi == index.end()) {
+        unitySources.emplace_back(sf);
+        std::map<cmSourceFile const*, size_t>::value_type entry(
+          sf, unitySources.size() - 1);
+        mi = index.insert(entry).first;
+      }
+      unitySources[mi->second].Configs.emplace_back(ci);
     }
   }
 
@@ -2990,13 +3015,13 @@ void cmLocalGenerator::AddUnityBuild(cmGeneratorTarget* target)
 
     std::vector<UnitySource> unity_files;
     if (!unityMode || *unityMode == "BATCH") {
-      unity_files =
-        AddUnityFilesModeAuto(target, lang, filtered_sources, beforeInclude,
-                              afterInclude, filename_base, unityBatchSize);
+      unity_files = AddUnityFilesModeAuto(
+        target, lang, configs, filtered_sources, beforeInclude, afterInclude,
+        filename_base, unityBatchSize);
     } else if (unityMode && *unityMode == "GROUP") {
       unity_files =
-        AddUnityFilesModeGroup(target, lang, filtered_sources, beforeInclude,
-                               afterInclude, filename_base);
+        AddUnityFilesModeGroup(target, lang, configs, filtered_sources,
+                               beforeInclude, afterInclude, filename_base);
     } else {
       // unity mode is set to an unsupported value
       std::string e("Invalid UNITY_BUILD_MODE value of " + *unityMode +
@@ -3010,6 +3035,10 @@ void cmLocalGenerator::AddUnityBuild(cmGeneratorTarget* target)
       target->AddSource(file.Path, true);
       unity->SetProperty("SKIP_UNITY_BUILD_INCLUSION", "ON");
       unity->SetProperty("UNITY_SOURCE_FILE", file.Path);
+      if (file.PerConfig) {
+        unity->SetProperty("COMPILE_DEFINITIONS",
+                           "CMAKE_UNITY_CONFIG_$<UPPER_CASE:$<CONFIG>>");
+      }
     }
   }
 }
diff --git a/Source/cmLocalGenerator.h b/Source/cmLocalGenerator.h
index 4c9714b..726817a 100644
--- a/Source/cmLocalGenerator.h
+++ b/Source/cmLocalGenerator.h
@@ -14,6 +14,8 @@
 #include <utility>
 #include <vector>
 
+#include <cm/optional>
+
 #include <cm3p/kwiml/int.h>
 
 #include "cmCustomCommandTypes.h"
@@ -664,6 +666,7 @@ private:
   struct UnityBatchedSource
   {
     cmSourceFile* Source = nullptr;
+    std::vector<size_t> Configs;
     UnityBatchedSource(cmSourceFile* sf)
       : Source(sf)
     {
@@ -672,27 +675,32 @@ private:
   struct UnitySource
   {
     std::string Path;
-    UnitySource(std::string path)
+    bool PerConfig = false;
+    UnitySource(std::string path, bool perConfig)
       : Path(std::move(path))
+      , PerConfig(perConfig)
     {
     }
   };
 
   UnitySource WriteUnitySource(
-    cmGeneratorTarget* target,
+    cmGeneratorTarget* target, std::vector<std::string> const& configs,
     cmRange<std::vector<UnityBatchedSource>::const_iterator> sources,
     cmValue beforeInclude, cmValue afterInclude, std::string filename) const;
   void WriteUnitySourceInclude(std::ostream& unity_file,
+                               cm::optional<std::string> const& cond,
                                std::string const& sf_full_path,
                                cmValue beforeInclude, cmValue afterInclude,
                                cmValue uniqueIdName) const;
   std::vector<UnitySource> AddUnityFilesModeAuto(
     cmGeneratorTarget* target, std::string const& lang,
+    std::vector<std::string> const& configs,
     std::vector<UnityBatchedSource> const& filtered_sources,
     cmValue beforeInclude, cmValue afterInclude,
     std::string const& filename_base, size_t batchSize);
   std::vector<UnitySource> AddUnityFilesModeGroup(
     cmGeneratorTarget* target, std::string const& lang,
+    std::vector<std::string> const& configs,
     std::vector<UnityBatchedSource> const& filtered_sources,
     cmValue beforeInclude, cmValue afterInclude,
     std::string const& filename_base);
diff --git a/Tests/RunCMake/UnityBuild/RunCMakeTest.cmake b/Tests/RunCMake/UnityBuild/RunCMakeTest.cmake
index 8019f09..e3643c0 100644
--- a/Tests/RunCMake/UnityBuild/RunCMakeTest.cmake
+++ b/Tests/RunCMake/UnityBuild/RunCMakeTest.cmake
@@ -26,6 +26,22 @@ run_build(unitybuild_anon_ns)
 run_build(unitybuild_anon_ns_no_unity_build)
 run_build(unitybuild_anon_ns_group_mode)
 
+function(run_per_config name)
+  set(RunCMake_TEST_BINARY_DIR ${RunCMake_BINARY_DIR}/${name}-build)
+  run_cmake(${name})
+  set(RunCMake_TEST_NO_CLEAN 1)
+  if(RunCMake_GENERATOR_IS_MULTI_CONFIG)
+    run_cmake_command(${name}-build-debug ${CMAKE_COMMAND} --build . --config Debug)
+    run_cmake_command(${name}-build-release ${CMAKE_COMMAND} --build . --config Release)
+  else()
+    run_cmake_command(${name}-build ${CMAKE_COMMAND} --build .)
+  endif()
+endfunction()
+
+if(NOT RunCMake_GENERATOR STREQUAL "Xcode")
+  run_per_config(per_config_c)
+endif()
+
 function(run_test name)
   set(RunCMake_TEST_BINARY_DIR ${RunCMake_BINARY_DIR}/${name}-build)
   run_cmake(${name})
diff --git a/Tests/RunCMake/UnityBuild/per_config_c.c b/Tests/RunCMake/UnityBuild/per_config_c.c
new file mode 100644
index 0000000..081c7d3
--- /dev/null
+++ b/Tests/RunCMake/UnityBuild/per_config_c.c
@@ -0,0 +1,16 @@
+#ifdef CFG_DEBUG
+extern void per_config_c_debug(void);
+#endif
+#ifdef CFG_OTHER
+extern void per_config_c_other(void);
+#endif
+int main(void)
+{
+#ifdef CFG_DEBUG
+  per_config_c_debug();
+#endif
+#ifdef CFG_OTHER
+  per_config_c_other();
+#endif
+  return 0;
+}
diff --git a/Tests/RunCMake/UnityBuild/per_config_c.cmake b/Tests/RunCMake/UnityBuild/per_config_c.cmake
new file mode 100644
index 0000000..9f2ee48
--- /dev/null
+++ b/Tests/RunCMake/UnityBuild/per_config_c.cmake
@@ -0,0 +1,12 @@
+enable_language(C)
+
+add_executable(per_config_c per_config_c.c
+  "$<$<CONFIG:Debug>:per_config_c_debug.c>"
+  "$<$<NOT:$<CONFIG:Debug>>:per_config_c_other.c>"
+  )
+
+set_target_properties(per_config_c PROPERTIES UNITY_BUILD ON)
+target_compile_definitions(per_config_c PRIVATE
+  "$<$<CONFIG:Debug>:CFG_DEBUG>"
+  "$<$<NOT:$<CONFIG:Debug>>:CFG_OTHER>"
+  )
diff --git a/Tests/RunCMake/UnityBuild/per_config_c_debug.c b/Tests/RunCMake/UnityBuild/per_config_c_debug.c
new file mode 100644
index 0000000..6d32ead
--- /dev/null
+++ b/Tests/RunCMake/UnityBuild/per_config_c_debug.c
@@ -0,0 +1,9 @@
+#ifndef CFG_DEBUG
+#  error "per_config_c_debug built without CFG_DEBUG"
+#endif
+#ifdef CFG_OTHER
+#  error "per_config_c_debug built with CFG_OTHER"
+#endif
+void per_config_c_debug(void)
+{
+}
diff --git a/Tests/RunCMake/UnityBuild/per_config_c_other.c b/Tests/RunCMake/UnityBuild/per_config_c_other.c
new file mode 100644
index 0000000..89c7a6b
--- /dev/null
+++ b/Tests/RunCMake/UnityBuild/per_config_c_other.c
@@ -0,0 +1,9 @@
+#ifdef CFG_DEBUG
+#  error "per_config_c_other built with CFG_DEBUG"
+#endif
+#ifndef CFG_OTHER
+#  error "per_config_c_other built without CFG_OTHER"
+#endif
+void per_config_c_other(void)
+{
+}