build-and-test: Make it possible to customize stage 2

Same possibilities as stage 1: subprojects, CMake definitions and build
flags.

Change-Id: Ib10be5d2f7687040979c79113e7eeb0af5903df1
diff --git a/scripts/llvm.py b/scripts/llvm.py
index 53ffd89..bdeae85 100644
--- a/scripts/llvm.py
+++ b/scripts/llvm.py
@@ -200,6 +200,9 @@
     stage1BuildFlags = args.stage1BuildFlags
 
     stage2BuildDir = args.stage2
+    stage2Subprojs = args.stage2Subprojs
+    stage2Defs = cmake_flags_from_args(args.stage2Defs)
+    stage2BuildFlags = args.stage2BuildFlags
 
     enableTestSuite = args.enableTestSuite
     sandboxDir = args.sandbox
@@ -229,6 +232,13 @@
         testedBuildDir = stage1BuildDir
 
         if stage2BuildDir is not None:
+            if stage2Subprojs:
+                toAdd = list(set(stage2Subprojs) - set(stage1Subprojs))
+                toRemove = list(set(stage1Subprojs) - set(stage2Subprojs))
+                sourceConfig.update(
+                    subproj_to_repo_map(toAdd, proj, llvmRepos, args.dry),
+                    toRemove)
+
             if not dryRun and not os.path.exists(stage2BuildDir):
                 os.makedirs(stage2BuildDir)
 
@@ -238,11 +248,11 @@
             # TODO: Make sure clang is actually built in this config (preferably
             # before reaching this point)
             buildConfig2.cmake(
-                [
+                stage2Defs + [
                     "-DCMAKE_C_COMPILER={}/bin/clang".format(stage1BuildDir),
                     "-DCMAKE_CXX_COMPILER={}/bin/clang++".format(stage1BuildDir)],
                 "Ninja")
-            buildConfig2.build()
+            buildConfig2.build(stage2BuildFlags)
             testedBuildDir = stage2BuildDir
 
         if enableTestSuite:
@@ -518,6 +528,36 @@
     dest='stage2',
     help="Path to the build directory for stage 2.")
 buildAndTest.add_argument(
+    '--stage2-subproject',
+    dest='stage2Subprojs',
+    metavar='SUBPROJ',
+    choices=valid_subprojects,
+    default=[],
+    action='append',
+    help="Subprojects to enable for stage 2 of the build. Can be passed "
+         "multiple times. Valid values for the subproject are: {}. "
+         "If this is a 2-stage build, the same subprojects will be used for "
+         "both stages unless other subprojects are explicitly requested for "
+         "stage 2.".format(" ".join(valid_subprojects)))
+buildAndTest.add_argument(
+    '--stage2-cmake-def',
+    dest='stage2Defs',
+    metavar='VAR=VALUE',
+    default=[],
+    action='append',
+    help="Additional CMake definitions for stage 2, e.g. "
+         "CMAKE_BUILD_TYPE=Release. Can be passed multiple times. "
+         "The -D is added automatically.")
+buildAndTest.add_argument(
+    '--stage2-build-flag',
+    dest='stage2BuildFlags',
+    metavar='FLAG',
+    default=[],
+    action='append',
+    help="Additional flags for the stage 2 build command (e.g. targets to "
+         "build). Can be passed multiple times. If your flag starts with "
+         "a '-', use '--stage2-build-flag=-FLAG' to pass it.")
+buildAndTest.add_argument(
     "--enable-test-suite",
     dest='enableTestSuite',
     action='store_true',
diff --git a/tests/cli/testbuildandtest.py b/tests/cli/testbuildandtest.py
index 1475958..f7996ad 100644
--- a/tests/cli/testbuildandtest.py
+++ b/tests/cli/testbuildandtest.py
@@ -219,6 +219,69 @@
             commands[3],
             "{stage2}\$ ninja".format(stage2=buildDir2))
 
+    def test_custom_both_stages(self):
+        """
+        Test that we get the correct commands when trying to customize both
+        stage 1 and stage 2.
+        """
+        reposDir = "path-to-repos"
+        sourceDir = "path-to-sources"
+        buildDir1 = "path-to-stage1"
+        buildDir2 = "path-to-stage2"
+
+        output = self.run_with_output(
+            self.llvm_build_and_test(
+                "--dry-run",
+                "--repos-dir", reposDir,
+                "--source-dir", sourceDir,
+                "--stage1-build-dir", buildDir1,
+                "--stage1-subproject", "clang",
+                "--stage1-subproject", "compiler-rt",
+                "--stage1-cmake-def", "CMAKE_CXX_FLAGS=-marm",
+                "--stage1-cmake-def", "LLVM_ENABLE_ASSERTIONS=True",
+                "--stage1-build-flag=-j8",
+                "--stage2-build-dir", buildDir2,
+                "--stage2-subproject", "clang",
+                "--stage2-subproject", "libcxx",
+                "--stage2-cmake-def", "CMAKE_BUILD_TYPE=MinSizeRel",
+                "--stage2-build-flag", "check-all"))
+
+        commands = output.splitlines()
+
+        self.assertRegex(commands[0],
+                         "{build}\$ cmake -G Ninja .* {sources}".format(
+                             build=buildDir1, sources=sourceDir))
+
+        self.assertIn("-DLLVM_TOOL_CLANG_BUILD=ON", commands[0])
+        self.assertIn("-DLLVM_TOOL_COMPILER_RT_BUILD=ON", commands[0])
+        self.assertIn("-DLLVM_TOOL_LIBCXX_BUILD=OFF", commands[0])
+        self.assertIn("-DLLVM_TOOL_LLDB_BUILD=OFF", commands[0])
+
+        self.assertIn("-DCMAKE_CXX_FLAGS=-marm", commands[0])
+        self.assertIn("-DLLVM_ENABLE_ASSERTIONS=True", commands[0])
+        self.assertNotIn("-DCMAKE_BUILD_TYPE=MinSizeRel", commands[0])
+
+        self.assertRegex(commands[1],
+                         "{build}\$ ninja -j8".format(build=buildDir1))
+
+        self.assertRegex(
+            commands[2], "{stage2}\$ cmake -G Ninja .* "
+            "-DCMAKE_C_COMPILER={stage1}/bin/clang "
+            "-DCMAKE_CXX_COMPILER={stage1}/bin/clang\+\+ {sources}".format(
+                stage1=buildDir1, stage2=buildDir2, sources=sourceDir))
+
+        self.assertIn("-DLLVM_TOOL_CLANG_BUILD=ON", commands[2])
+        self.assertIn("-DLLVM_TOOL_COMPILER_RT_BUILD=OFF", commands[2])
+        self.assertIn("-DLLVM_TOOL_LIBCXX_BUILD=ON", commands[2])
+        self.assertIn("-DLLVM_TOOL_LLDB_BUILD=OFF", commands[2])
+
+        self.assertNotIn("-DCMAKE_CXX_FLAGS=-marm", commands[2])
+        self.assertNotIn("-DLLVM_ENABLE_ASSERTIONS=True", commands[2])
+        self.assertIn("-DCMAKE_BUILD_TYPE=MinSizeRel", commands[2])
+
+        self.assertRegex(
+            commands[3],
+            "{stage2}\$ ninja check-all".format(stage2=buildDir2))
 
     def test_stage2_and_testsuite(self):
         """