Merge pull request #9851 from eeckstein/fix-kp-demangling

diff --git a/include/swift/Driver/Compilation.h b/include/swift/Driver/Compilation.h
index cc3104c..bada9fe 100644
--- a/include/swift/Driver/Compilation.h
+++ b/include/swift/Driver/Compilation.h
@@ -57,6 +57,13 @@
   Parseable,
 };
 
+/// Indicates whether a temporary file should always be preserved if a part of
+/// the compilation crashes.
+enum class PreserveOnSignal : bool {
+  No,
+  Yes
+};
+
 class Compilation {
   friend class PerformJobsState;
 private:
@@ -87,8 +94,8 @@
 
   /// Temporary files that should be cleaned up after the compilation finishes.
   ///
-  /// These apply whether the compilation succeeds or fails.
-  std::vector<std::string> TempFilePaths;
+  /// These apply whether the compilation succeeds or fails. If the
+  llvm::StringMap<PreserveOnSignal> TempFilePaths;
 
   /// Write information about this compilation to this file.
   ///
@@ -174,14 +181,13 @@
   }
   Job *addJob(std::unique_ptr<Job> J);
 
-  void addTemporaryFile(StringRef file) {
-    TempFilePaths.push_back(file.str());
+  void addTemporaryFile(StringRef file,
+                        PreserveOnSignal preserve = PreserveOnSignal::No) {
+    TempFilePaths[file] = preserve;
   }
 
   bool isTemporaryFile(StringRef file) {
-    // TODO: Use a set instead of a linear search.
-    return std::find(TempFilePaths.begin(), TempFilePaths.end(), file) !=
-             TempFilePaths.end();
+    return TempFilePaths.count(file);
   }
 
   const llvm::opt::DerivedArgList &getArgs() const { return *TranslatedArgs; }
@@ -255,9 +261,12 @@
 private:
   /// \brief Perform all jobs.
   ///
-  /// \returns exit code of the first failed Job, or 0 on success. A return
-  /// value of -2 indicates that a Job crashed during execution.
-  int performJobsImpl();
+  /// \param[out] abnormalExit Set to true if any job exits abnormally (i.e.
+  /// crashes).
+  ///
+  /// \returns exit code of the first failed Job, or 0 on success. If a Job
+  /// crashes during execution, a negative value will be returned.
+  int performJobsImpl(bool &abnormalExit);
 
   /// \brief Performs a single Job by executing in place, if possible.
   ///
diff --git a/lib/Driver/Compilation.cpp b/lib/Driver/Compilation.cpp
index 6dece64..80b0f3b 100644
--- a/lib/Driver/Compilation.cpp
+++ b/lib/Driver/Compilation.cpp
@@ -165,6 +165,9 @@
     /// Cumulative result of PerformJobs(), accumulated from subprocesses.
     int Result = EXIT_SUCCESS;
 
+    /// True if any Job crashed.
+    bool AnyAbnormalExit = false;
+
     /// Timers for monitoring execution time of subprocesses.
     llvm::TimerGroup DriverTimerGroup {"driver", "Driver Compilation Time"};
     llvm::SmallDenseMap<const Job *, std::unique_ptr<llvm::Timer>, 16>
@@ -470,6 +473,7 @@
 
       // Since the task signalled, unconditionally set result to -2.
       Result = -2;
+      AnyAbnormalExit = true;
 
       return TaskFinishedResponse::StopExecution;
     }
@@ -689,6 +693,10 @@
         Result = Comp.Diags.hadAnyError();
       return Result;
     }
+
+    bool hadAnyAbnormalExit() {
+      return AnyAbnormalExit;
+    }
   };
 } // namespace driver
 } // namespace swift
@@ -816,8 +824,7 @@
   return true;
 }
 
-int Compilation::performJobsImpl() {
-
+int Compilation::performJobsImpl(bool &abnormalExit) {
   PerformJobsState State(*this);
 
   State.scheduleInitialJobs();
@@ -833,6 +840,7 @@
                            InputInfo);
   }
 
+  abnormalExit = State.hadAnyAbnormalExit();
   return State.getResult();
 }
 
@@ -917,15 +925,14 @@
   if (!TaskQueue::supportsParallelExecution() && NumberOfParallelCommands > 1) {
     Diags.diagnose(SourceLoc(), diag::warning_parallel_execution_not_supported);
   }
-  
-  int result = performJobsImpl();
+
+  bool abnormalExit;
+  int result = performJobsImpl(abnormalExit);
 
   if (!SaveTemps) {
-    // FIXME: Do we want to be deleting temporaries even when a child process
-    // crashes?
-    for (auto &path : TempFilePaths) {
-      // Ignore the error code for removing temporary files.
-      (void)llvm::sys::fs::remove(path);
+    for (const auto &pathPair : TempFilePaths) {
+      if (!abnormalExit || pathPair.getValue() == PreserveOnSignal::No)
+        (void)llvm::sys::fs::remove(pathPair.getKey());
     }
   }
 
@@ -945,7 +952,7 @@
       llvm::report_fatal_error("unable to create list of input sources");
     }
     auto *mutableThis = const_cast<Compilation *>(this);
-    mutableThis->addTemporaryFile(Buffer.str());
+    mutableThis->addTemporaryFile(Buffer.str(), PreserveOnSignal::Yes);
     mutableThis->AllSourceFilesPath = getArgs().MakeArgString(Buffer);
   }
   return AllSourceFilesPath;
diff --git a/lib/Driver/ToolChain.cpp b/lib/Driver/ToolChain.cpp
index 011390c..2e1a004 100644
--- a/lib/Driver/ToolChain.cpp
+++ b/lib/Driver/ToolChain.cpp
@@ -57,7 +57,7 @@
     llvm::report_fatal_error("unable to create temporary file for filelist");
   }
 
-  C.addTemporaryFile(buffer.str());
+  C.addTemporaryFile(buffer.str(), PreserveOnSignal::Yes);
   // We can't just reference the data in the TemporaryFiles vector because
   // that could theoretically get copied to a new address.
   return C.getArgs().MakeArgString(buffer.str());
diff --git a/test/Driver/Dependencies/Inputs/update-dependencies-bad.py b/test/Driver/Dependencies/Inputs/update-dependencies-bad.py
index b2b72f8..97585bd 100755
--- a/test/Driver/Dependencies/Inputs/update-dependencies-bad.py
+++ b/test/Driver/Dependencies/Inputs/update-dependencies-bad.py
@@ -12,8 +12,8 @@
 # ----------------------------------------------------------------------------
 #
 # Fails if the input file is named "bad.swift" or "crash.swift"; otherwise
-# dispatches to update-dependencies.py. "crash.swift" gives an exit code
-# other than 1.
+# dispatches to update-dependencies.py. "crash.swift" results in an
+# exit-by-SIGKILL
 #
 # ----------------------------------------------------------------------------
 
@@ -21,6 +21,7 @@
 
 import os
 import shutil
+import signal
 import sys
 
 assert sys.argv[1] == '-frontend'
@@ -42,7 +43,8 @@
     if os.path.basename(primaryFile) == 'bad.swift':
         sys.exit(1)
     else:
-        sys.exit(129)
+        sys.stdout.flush()
+        os.kill(os.getpid(), signal.SIGKILL)
 
 execDir = os.path.dirname(os.path.abspath(__file__))
 exec(open(os.path.join(execDir, "update-dependencies.py")).read())
diff --git a/test/Driver/Dependencies/whole-module-build-record.swift b/test/Driver/Dependencies/whole-module-build-record.swift
index 57866c1..4c53260 100644
--- a/test/Driver/Dependencies/whole-module-build-record.swift
+++ b/test/Driver/Dependencies/whole-module-build-record.swift
@@ -14,6 +14,6 @@
 
 
 // RUN: touch -t 201401240006 %t/other.swift
-// RUN: cd %t && not %swiftc_driver -c -driver-use-frontend-path %S/Inputs/fail.py -output-file-map %t/output.json -whole-module-optimization ./main.swift ./other.swift -module-name main -j1 -v 2>&1
+// RUN: cd %t && not %swiftc_driver -c -driver-use-frontend-path %S/../Inputs/fail.py -output-file-map %t/output.json -whole-module-optimization ./main.swift ./other.swift -module-name main -j1 -v 2>&1
 
 // Just don't crash.
diff --git a/test/Driver/Inputs/crash.py b/test/Driver/Inputs/crash.py
new file mode 100755
index 0000000..6dd356a
--- /dev/null
+++ b/test/Driver/Inputs/crash.py
@@ -0,0 +1,14 @@
+#!/usr/bin/env python
+# crash.py - Sends SIGKILL to self. -*- python -*-
+#
+# This source file is part of the Swift.org open source project
+#
+# Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
+# Licensed under Apache License v2.0 with Runtime Library Exception
+#
+# See https://swift.org/LICENSE.txt for license information
+# See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
+
+import os
+import signal
+os.kill(os.getpid(), signal.SIGKILL)
diff --git a/test/Driver/Dependencies/Inputs/fail.py b/test/Driver/Inputs/fail.py
similarity index 100%
rename from test/Driver/Dependencies/Inputs/fail.py
rename to test/Driver/Inputs/fail.py
diff --git a/test/Driver/filelists.swift b/test/Driver/filelists.swift
index 90b3e79..390256a 100644
--- a/test/Driver/filelists.swift
+++ b/test/Driver/filelists.swift
@@ -32,6 +32,15 @@
 // CHECK-WMO-THREADED-NEXT: ...with output!
 // CHECK-WMO-THREADED-NOT: Handled
 
+// RUN: mkdir -p %t/tmp-fail/
+// RUN: (cd %t && not env TMPDIR="%t/tmp-fail/" %swiftc_driver_plain -driver-use-frontend-path %S/Inputs/fail.py -c ./a.swift ./b.swift ./c.swift -module-name main -target x86_64-apple-macosx10.9 -driver-use-filelists -output-file-map=%S/Inputs/filelists/output.json -force-single-frontend-invocation -num-threads 1)
+// RUN: not ls %t/tmp-fail/sources-*
+// RUN: not ls %t/tmp-fail/outputs-*
+
+// RUN: mkdir -p %t/tmp-crash/
+// RUN: (cd %t && not env TMPDIR="%t/tmp-crash/" %swiftc_driver_plain -driver-use-frontend-path %S/Inputs/crash.py -c ./a.swift ./b.swift ./c.swift -module-name main -target x86_64-apple-macosx10.9 -driver-use-filelists -output-file-map=%S/Inputs/filelists/output.json -force-single-frontend-invocation -num-threads 1)
+// RUN: ls %t/tmp-crash/sources-* %t/tmp-crash/outputs-*
+
 
 // RUN: (cd %t && env PATH=%t/bin/:$PATH %swiftc_driver_plain -driver-use-frontend-path %S/Inputs/filelists/check-filelist-abc.py -emit-library ./a.swift ./b.swift ./c.swift -module-name main -target x86_64-apple-macosx10.9 -driver-use-filelists -output-file-map=%S/Inputs/filelists/output.json 2>&1 | %FileCheck -check-prefix=CHECK-LINK %s)
 // RUN: (cd %t && env PATH=%t/bin/:$PATH %swiftc_driver_plain -driver-use-frontend-path %S/Inputs/filelists/check-filelist-abc.py -emit-library ./a.swift ./b.swift ./c.swift -module-name main -target x86_64-apple-macosx10.9 -driver-use-filelists -output-file-map=%S/Inputs/filelists/output.json -force-single-frontend-invocation -num-threads 1 2>&1 | %FileCheck -check-prefix=CHECK-LINK %s)