Merge pull request #41449 from lorentey/custom-traps

lorentey · web-flow · commit e24bb55abf40 · 2022-02-21T12:46:33.000-08:00
[IRGen] Allow configuring a panic function to call instead of trapping
diff --git a/include/swift/AST/IRGenOptions.h b/include/swift/AST/IRGenOptions.h
@@ -416,6 +416,10 @@ class IRGenOptions {
 
   JITDebugArtifact DumpJIT = JITDebugArtifact::None;
 
+  /// If not an empty string, trap intrinsics are lowered to calls to this
+  /// function instead of to trap instructions.
+  std::string TrapFuncName = "";
+
   IRGenOptions()
       : DWARFVersion(2),
         OutputKind(IRGenOutputKind::LLVMAssemblyAfterOptimization),
diff --git a/include/swift/Option/FrontendOptions.td b/include/swift/Option/FrontendOptions.td
@@ -594,6 +594,10 @@ def disable_preallocated_instantiation_caches : Flag<["-"], "disable-preallocate
   Flags<[FrontendOption, NoInteractiveOption, HelpHidden]>,
   HelpText<"Avoid preallocating metadata instantiation caches in globals">;
 
+def trap_function
+  : Separate<["-"], "trap-function">, MetaVarName<"<name>">,
+    HelpText<"Lower traps to calls to this function instead of trap instructions">;
+
 def disable_previous_implementation_calls_in_dynamic_replacements :
   Flag<["-"], "disable-previous-implementation-calls-in-dynamic-replacements">,
   Flags<[FrontendOption, NoInteractiveOption, HelpHidden]>,
diff --git a/lib/Frontend/CompilerInvocation.cpp b/lib/Frontend/CompilerInvocation.cpp
@@ -1929,6 +1929,9 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,
     Opts.StackPromotionSizeLimit = limit;
   }
 
+  if (Args.hasArg(OPT_trap_function))
+    Opts.TrapFuncName = Args.getLastArgValue(OPT_trap_function).str();
+
   Opts.FunctionSections = Args.hasArg(OPT_function_sections);
 
   if (Args.hasArg(OPT_autolink_force_load))
diff --git a/lib/IRGen/IRGenFunction.cpp b/lib/IRGen/IRGenFunction.cpp
@@ -496,6 +496,13 @@ llvm::CallInst *IRBuilder::CreateNonMergeableTrap(IRGenModule &IGM,
   }
   auto Call = IRBuilderBase::CreateCall(trapIntrinsic, {});
   setCallingConvUsingCallee(Call);
+
+  if (!IGM.IRGen.Opts.TrapFuncName.empty()) {
+    auto A = llvm::Attribute::get(getContext(), "trap-func-name",
+                                  IGM.IRGen.Opts.TrapFuncName);
+    Call->addFnAttr(A);
+  }
+
   return Call;
 }
 
diff --git a/lib/IRGen/IRGenModule.cpp b/lib/IRGen/IRGenModule.cpp
@@ -132,6 +132,9 @@ static clang::CodeGenerator *createClangCodeGenerator(ASTContext &Context,
     CGO.DwarfDebugFlags = Opts.getDebugFlags(PD);
     break;
   }
+  if (!Opts.TrapFuncName.empty()) {
+    CGO.TrapFuncName = Opts.TrapFuncName;
+  }
 
   auto &HSI = Importer->getClangPreprocessor()
                   .getHeaderSearchInfo()
diff --git a/test/IRGen/Inputs/usr/include/Gizmo.h b/test/IRGen/Inputs/usr/include/Gizmo.h
@@ -70,6 +70,8 @@ static inline int zeroRedeclared(void) { return innerZero(); }
 
 static int staticButNotInline(void) { return innerZero(); }
 
+static inline void ackbar(void) { __builtin_trap(); }
+
 @interface NSView : NSObject
 - (struct NSRect) convertRectFromBase: (struct NSRect) r;
 @end
diff --git a/test/IRGen/trap-function-clang.swift b/test/IRGen/trap-function-clang.swift
@@ -0,0 +1,23 @@
+// RUN: %empty-directory(%t)
+// RUN: %build-irgen-test-overlays
+// RUN: %target-swift-frontend -sdk %S/Inputs -primary-file %s -trap-function oopsie -enable-objc-interop -emit-ir -module-name trap_function -I %t | %FileCheck %s -check-prefix=TRAPFN
+// RUN: %target-swift-frontend -O -sdk %S/Inputs -primary-file %s -trap-function oopsie -enable-objc-interop -emit-ir -module-name trap_function -I %t | %FileCheck %s -check-prefix=TRAPFN_OPT
+// RUN: %target-swift-frontend -sdk %S/Inputs -primary-file %s -enable-objc-interop -emit-ir -module-name trap_function -I %t | %FileCheck %s -check-prefix=NOTRAPFN
+
+import gizmo
+
+// TRAPFN-LABEL: define hidden swiftcc void @"$s13trap_function18checkClangImporteryyF"
+// TRAPFN: call void @llvm.trap() [[ATTR0:#[0-9]+]]
+
+// TRAPFN_OPT-LABEL: define hidden swiftcc void @"$s13trap_function18checkClangImporteryyF"
+// TRAPFN_OPT: call void @llvm.trap() [[ATTR0:#[0-9]+]]
+
+// NOTRAPFN-LABEL: define hidden swiftcc void @"$s13trap_function18checkClangImporteryyF"
+// NOTRAPFN: call void @llvm.trap(){{$}}
+func checkClangImporter() {
+  ackbar()
+}
+
+// TRAPFN: attributes [[ATTR0]] = { {{.*}}"trap-func-name"="oopsie" }
+// TRAPFN_OPT: attributes [[ATTR0]] = { {{.*}}"trap-func-name"="oopsie" }
+// NOTRAPFN-NOT: attributes {{.*}} = { {{.*}}"trap-func-name"="oopsie" }
diff --git a/test/IRGen/trap-function.swift b/test/IRGen/trap-function.swift
@@ -0,0 +1,67 @@
+// RUN: %empty-directory(%t)
+// RUN: %build-irgen-test-overlays
+// RUN: %target-swift-frontend -primary-file %s -trap-function oopsie -emit-ir -module-name trap_function -I %t | %FileCheck %s -check-prefix=TRAPFN
+// RUN: %target-swift-frontend -O -primary-file %s -trap-function oopsie -emit-ir -module-name trap_function -I %t | %FileCheck %s -check-prefix=TRAPFN_OPT
+// RUN: %target-swift-frontend -primary-file %s -emit-ir -module-name trap_function -I %t | %FileCheck %s -check-prefix=NOTRAPFN
+
+// TRAPFN-LABEL: define hidden swiftcc {{.*}} @"$s13trap_function14checkOverflow1yS2i_SitF"
+// TRAPFN: call void @llvm.trap() [[ATTR0:#[0-9]+]]
+
+// TRAPFN_OPT-LABEL: define hidden swiftcc {{.*}} @"$s13trap_function14checkOverflow1yS2i_SitF"
+// TRAPFN_OPT: call void @llvm.trap() [[ATTR0:#[0-9]+]]
+
+// NOTRAPFN-LABEL: define hidden swiftcc {{.*}} @"$s13trap_function14checkOverflow1yS2i_SitF"
+// NOTRAPFN: call void @llvm.trap(){{$}}
+func checkOverflow1(_ a: Int, _ b: Int) -> Int {
+  a + b
+}
+
+// TRAPFN-LABEL: define hidden swiftcc {{.*}} @"$s13trap_function14checkOverflow2yySiz_SitF"
+// TRAPFN: call void @llvm.trap() [[ATTR0:#[0-9]+]]
+
+// TRAPFN_OPT-LABEL: define hidden swiftcc {{.*}} @"$s13trap_function14checkOverflow2yySiz_SitF"
+// TRAPFN_OPT: call void @llvm.trap() [[ATTR0:#[0-9]+]]
+
+// NOTRAPFN-LABEL: define hidden swiftcc {{.*}} @"$s13trap_function14checkOverflow2yySiz_SitF"
+// NOTRAPFN: call void @llvm.trap(){{$}}
+func checkOverflow2(_ a: inout Int, _ b: Int) {
+  a *= b
+}
+
+// TRAPFN-LABEL: define hidden swiftcc void @"$s13trap_function17checkPreconditionyySiF"
+// TRAPFN: call swiftcc void @"$ss17_assertionFailure__4file4line5flagss5NeverOs12StaticStringV_SSAHSus6UInt32VtF"
+// TRAPFN-NOT: call void @llvm.trap()
+
+// TRAPFN_OPT-LABEL: define hidden swiftcc void @"$s13trap_function17checkPreconditionyySiF"
+// TRAPFN_OPT: call void @llvm.trap() [[ATTR0]]
+
+// NOTRAPFN-LABEL: define hidden swiftcc void @"$s13trap_function17checkPreconditionyySiF"
+// NOTRAPFN: call swiftcc void @"$ss17_assertionFailure__4file4line5flagss5NeverOs12StaticStringV_SSAHSus6UInt32VtF"
+func checkPrecondition(_ a: Int) {
+  precondition(a != 23)
+}
+
+// TRAPFN-LABEL: define hidden swiftcc void @"$s13trap_function24checkPreconditionFailureyySiF"
+// TRAPFN: call swiftcc void @"$ss17_assertionFailure__4file4line5flagss5NeverOs12StaticStringV_SSAHSus6UInt32VtF"
+// TRAPFN-NOT: call void @llvm.trap()
+
+// TRAPFN_OPT-LABEL: define hidden swiftcc void @"$s13trap_function24checkPreconditionFailureyySiF"
+// TRAPFN_OPT: call void @llvm.trap() [[ATTR0]]
+
+// NOTRAPFN-LABEL: define hidden swiftcc void @"$s13trap_function24checkPreconditionFailureyySiF"
+// NOTRAPFN: call swiftcc void @"$ss17_assertionFailure__4file4line5flagss5NeverOs12StaticStringV_SSAHSus6UInt32VtF"
+func checkPreconditionFailure(_ a: Int) {
+  if a == 42 {
+    preconditionFailure()
+  }
+}
+
+// TRAPFN-LABEL: define hidden swiftcc void @"$s13trap_function10terminatoryyF"
+// TRAPFN_OPT-LABEL: define hidden swiftcc void @"$s13trap_function10terminatoryyF"
+// NOTRAPFN-LABEL: define hidden swiftcc void @"$s13trap_function10terminatoryyF"
+func terminator() {
+}
+
+// TRAPFN: attributes [[ATTR0]] = { {{.*}}"trap-func-name"="oopsie" }
+// TRAPFN_OPT: attributes [[ATTR0]] = { {{.*}}"trap-func-name"="oopsie" }
+// NOTRAPFN-NOT: attributes {{.*}} = { {{.*}}"trap-func-name"="oopsie" }

Original file line number	Diff line number	Diff line change
`@@ -1929,6 +1929,9 @@ static bool ParseIRGenArgs(IRGenOptions &Opts, ArgList &Args,`
`1929`	`1929`	`Opts.StackPromotionSizeLimit = limit;`
`1930`	`1930`	`}`
`1931`	`1931`
	`1932`	`+ if (Args.hasArg(OPT_trap_function))`
	`1933`	`+ Opts.TrapFuncName = Args.getLastArgValue(OPT_trap_function).str();`
	`1934`	`+`
`1932`	`1935`	`Opts.FunctionSections = Args.hasArg(OPT_function_sections);`
`1933`	`1936`
`1934`	`1937`	`if (Args.hasArg(OPT_autolink_force_load))`
Original file line number	Diff line number	Diff line change
`@@ -496,6 +496,13 @@ llvm::CallInst *IRBuilder::CreateNonMergeableTrap(IRGenModule &IGM,`
`496`	`496`	`}`
`497`	`497`	`auto Call = IRBuilderBase::CreateCall(trapIntrinsic, {});`
`498`	`498`	`setCallingConvUsingCallee(Call);`
	`499`	`+`
	`500`	`+ if (!IGM.IRGen.Opts.TrapFuncName.empty()) {`
	`501`	`+ auto A = llvm::Attribute::get(getContext(), "trap-func-name",`
	`502`	`+ IGM.IRGen.Opts.TrapFuncName);`
	`503`	`+ Call->addFnAttr(A);`
	`504`	`+ }`
	`505`	`+`
`499`	`506`	`return Call;`
`500`	`507`	`}`
`501`	`508`
Original file line number	Diff line number	Diff line change
`@@ -132,6 +132,9 @@ static clang::CodeGenerator *createClangCodeGenerator(ASTContext &Context,`
`132`	`132`	`CGO.DwarfDebugFlags = Opts.getDebugFlags(PD);`
`133`	`133`	`break;`
`134`	`134`	`}`
	`135`	`+ if (!Opts.TrapFuncName.empty()) {`
	`136`	`+ CGO.TrapFuncName = Opts.TrapFuncName;`
	`137`	`+ }`
`135`	`138`
`136`	`139`	`auto &HSI = Importer->getClangPreprocessor()`
`137`	`140`	`.getHeaderSearchInfo()`