feat: LLVM backend: implement the equivalent of -fstack-clash-protection

src/Lean/Compiler/IR/EmitLLVM.lean

@@ -27,6 +27,27 @@ namespace LLVM
 -- TODO(bollu): instantiate target triple and find out what size_t is.
 def size_tType (llvmctx : LLVM.Context) : IO (LLVM.LLVMType llvmctx) :=
   LLVM.i64Type llvmctx
+
+-- Helper to add a function if it does not exist, and to return the function handle if it does.
+def getOrAddFunction (m : LLVM.Module ctx) (name : String) (type : LLVM.LLVMType ctx) : BaseIO (LLVM.Value ctx) :=  do
+  match (← LLVM.getNamedFunction m name) with
+  | some fn => return fn
+  | none =>
+    /-
+    By the evidence shown in: https://github.com/leanprover/lean4/issues/2373#issuecomment-1658743284
+    this is how clang implements `-fstack-clash-protection` in the LLVM IR, we want this feature
+    for robust stack overflow detection.
+    -/
+    let fn ← LLVM.addFunction m name type
+    let attr ← LLVM.createStringAttribute "probe-stack" "inline-asm"
+    LLVM.addAttributeAtIndex fn LLVM.AttributeIndex.AttributeFunctionIndex attr
+    return fn
+
+def getOrAddGlobal (m : LLVM.Module ctx) (name : String) (type : LLVM.LLVMType ctx) : BaseIO (LLVM.Value ctx) :=  do
+  match (← LLVM.getNamedGlobal m name) with
+  | .some fn => return fn
+  | .none => LLVM.addGlobal m name type
+
 end LLVM
 
 namespace EmitLLVM

src/Lean/Compiler/IR/LLVMBindings.lean

@@ -29,6 +29,14 @@ def IntPredicate.EQ : IntPredicate := { val := 32 }
 def IntPredicate.NE : IntPredicate := { val := IntPredicate.EQ.val + 1 }
 def IntPredicate.UGT : IntPredicate := { val := IntPredicate.NE.val + 1 }
 
+-- https://github.com/llvm/llvm-project/blob/c3e073bcbdc523b0f758d44a89a6333e38bff863/llvm/include/llvm-c/Core.h#L457
+structure AttributeIndex where
+  private mk :: val : UInt64
+
+def AttributeIndex.AttributeReturnIndex : AttributeIndex := { val := 0 }
+-- This value is ~0 for 64 bit
+def AttributeIndex.AttributeFunctionIndex : AttributeIndex := { val := 18446744073709551615 }
+
 structure BasicBlock (ctx : Context)  where
   private mk :: ptr : USize
 instance : Nonempty (BasicBlock ctx) := ⟨{ ptr := default }⟩
@@ -79,6 +87,10 @@ def Value.isNull (v : Value ctx) : Bool := v.ptr == 0
 @[extern "lean_llvm_get_value_name2"]
 opaque Value.getName {ctx : Context} (value : Value ctx) : BaseIO String
 
+structure Attribute (ctx : Context) where
+  private mk :: ptr : USize
+instance : Nonempty (Attribute ctx) := ⟨{ ptr := default }⟩
+
 @[extern "lean_llvm_initialize_target_info"]
 opaque llvmInitializeTargetInfo : BaseIO (Unit)
 
@@ -314,6 +326,12 @@ opaque disposeTargetMachine (tm : TargetMachine ctx) : BaseIO Unit
 @[extern "lean_llvm_dispose_module"]
 opaque disposeModule (m : Module ctx) : BaseIO Unit
 
+@[extern "lean_llvm_create_string_attribute"]
+opaque createStringAttribute (key : String) (value : String) : BaseIO (Attribute ctx)
+
+@[extern "lean_llvm_add_attribute_at_index"]
+opaque addAttributeAtIndex (fn : Value ctx) (idx: AttributeIndex) (attr: Attribute ctx) : BaseIO Unit
+
 
 -- https://github.com/llvm/llvm-project/blob/c3e073bcbdc523b0f758d44a89a6333e38bff863/llvm/include/llvm-c/Core.h#L198
 structure Visibility where
@@ -379,18 +397,6 @@ def Linkage.linkerPrivateWeak : Linkage := { val := 16 }
 opaque setLinkage {ctx : Context} (value : Value ctx) (linkage : Linkage) : BaseIO Unit
 
 
--- Helper to add a function if it does not exist, and to return the function handle if it does.
-def getOrAddFunction(m : Module ctx) (name : String) (type : LLVMType ctx) : BaseIO (Value ctx) :=  do
-  match (← getNamedFunction m name) with
-  | .some fn => return fn
-  | .none => addFunction m name type
-
-def getOrAddGlobal(m : Module ctx) (name : String) (type : LLVMType ctx) : BaseIO (Value ctx) :=  do
-  match (← getNamedGlobal m name) with
-  | .some fn => return fn
-  | .none => addGlobal m name type
-
-
 def i1Type (ctx : LLVM.Context) : BaseIO (LLVM.LLVMType ctx) :=
   LLVM.intTypeInContext ctx 1
 

src/library/compiler/llvm.cpp

@@ -151,6 +151,15 @@ static inline size_t PassManagerBuilder_to_lean(LLVMPassManagerBuilderRef s) {
 static inline LLVMPassManagerBuilderRef lean_to_PassManagerBuilder(size_t s) {
     return reinterpret_cast<LLVMPassManagerBuilderRef>(s);
 }
+
+// == LLVM <-> Lean: AttributeRef ==
+static inline size_t Attribute_to_lean(LLVMAttributeRef s) {
+    return reinterpret_cast<size_t>(s);
+}
+
+static inline LLVMAttributeRef lean_to_Attribute(size_t s) {
+    return reinterpret_cast<LLVMAttributeRef>(s);
+}
 #else
 typedef int LLVMBasicBlockRef;
 typedef int LLVMContextRef;
@@ -1309,6 +1318,29 @@ extern "C" LEAN_EXPORT lean_object *lean_llvm_set_dll_storage_class(size_t ctx,
 #endif  // LEAN_LLVM
 }
 
+extern "C" LEAN_EXPORT lean_object *lean_llvm_create_string_attribute(size_t ctx, lean_object* key, lean_object* value,
+    lean_object * /* w */) {
+#ifndef LEAN_LLVM
+    lean_always_assert(
+        false && ("Please build a version of Lean4 with -DLLVM=ON to invoke "
+                  "the LLVM backend function."));
+#else
+    LLVMAttributeRef attr = LLVMCreateStringAttribute(lean_to_Context(ctx), lean_string_cstr(key), lean_string_len(key), lean_string_cstr(value), lean_string_len(value));
+    return lean_io_result_mk_ok(lean_box_usize(Attribute_to_lean(attr)));
+#endif  // LEAN_LLVM
+}
+
+extern "C" LEAN_EXPORT lean_object *lean_llvm_add_attribute_at_index(size_t ctx, size_t fn, uint64_t idx, size_t attr,
+    lean_object * /* w */) {
+#ifndef LEAN_LLVM
+    lean_always_assert(
+        false && ("Please build a version of Lean4 with -DLLVM=ON to invoke "
+                  "the LLVM backend function."));
+#else
+    LLVMAddAttributeAtIndex(lean_to_Value(fn), idx, lean_to_Attribute(attr));
+    return lean_io_result_mk_ok(lean_box(0));
+#endif  // LEAN_LLVM
+}
 
 extern "C" LEAN_EXPORT lean_object *lean_llvm_get_first_global(size_t ctx, size_t mod,
     lean_object * /* w */) {