https://github.com/zyedidia updated 
https://github.com/llvm/llvm-project/pull/172906

>From 7899f4d20c91e0810da1c53b73d039804264cf29 Mon Sep 17 00:00:00 2001
From: Zachary Yedidia <[email protected]>
Date: Wed, 21 Jan 2026 17:35:59 -0500
Subject: [PATCH 1/4] [LFI] Add MCLFIRewriter infrastructure

---
 clang/tools/driver/cc1as_main.cpp             |  3 +
 llvm/include/llvm/MC/MCLFI.h                  | 23 ++++++
 llvm/include/llvm/MC/MCLFIRewriter.h          | 70 +++++++++++++++++
 .../llvm/MC/MCParser/MCAsmParserExtension.h   | 16 ++--
 llvm/include/llvm/MC/MCStreamer.h             |  7 ++
 llvm/include/llvm/MC/TargetRegistry.h         | 21 ++++++
 llvm/lib/MC/CMakeLists.txt                    |  2 +
 llvm/lib/MC/MCAsmStreamer.cpp                 |  5 ++
 llvm/lib/MC/MCLFI.cpp                         | 75 +++++++++++++++++++
 llvm/lib/MC/MCLFIRewriter.cpp                 | 61 +++++++++++++++
 llvm/lib/MC/MCObjectStreamer.cpp              |  5 ++
 llvm/lib/MC/MCParser/AsmParser.cpp            |  9 ++-
 llvm/lib/MC/MCParser/CMakeLists.txt           |  1 +
 llvm/lib/MC/MCParser/LFIAsmParser.cpp         | 69 +++++++++++++++++
 llvm/lib/MC/MCStreamer.cpp                    |  3 +
 llvm/tools/llvm-mc/llvm-mc.cpp                | 14 ++++
 16 files changed, 375 insertions(+), 9 deletions(-)
 create mode 100644 llvm/include/llvm/MC/MCLFI.h
 create mode 100644 llvm/include/llvm/MC/MCLFIRewriter.h
 create mode 100644 llvm/lib/MC/MCLFI.cpp
 create mode 100644 llvm/lib/MC/MCLFIRewriter.cpp
 create mode 100644 llvm/lib/MC/MCParser/LFIAsmParser.cpp

diff --git a/clang/tools/driver/cc1as_main.cpp 
b/clang/tools/driver/cc1as_main.cpp
index ccc48a77891d7..6cbccbc3d7408 100644
--- a/clang/tools/driver/cc1as_main.cpp
+++ b/clang/tools/driver/cc1as_main.cpp
@@ -28,6 +28,7 @@
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCInstPrinter.h"
 #include "llvm/MC/MCInstrInfo.h"
+#include "llvm/MC/MCLFI.h"
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCObjectWriter.h"
 #include "llvm/MC/MCParser/MCAsmParser.h"
@@ -593,6 +594,8 @@ static bool ExecuteAssemblerImpl(AssemblerInvocation &Opts,
     Triple T(Opts.Triple);
     Str.reset(TheTarget->createMCObjectStreamer(
         T, Ctx, std::move(MAB), std::move(OW), std::move(CE), *STI));
+    if (T.isLFI())
+      initializeLFIMCStreamer(*Str.get(), Ctx, T);
     if (T.isOSBinFormatMachO() && T.isOSDarwin()) {
       Triple *TVT = Opts.DarwinTargetVariantTriple
                         ? &*Opts.DarwinTargetVariantTriple
diff --git a/llvm/include/llvm/MC/MCLFI.h b/llvm/include/llvm/MC/MCLFI.h
new file mode 100644
index 0000000000000..6598453952642
--- /dev/null
+++ b/llvm/include/llvm/MC/MCLFI.h
@@ -0,0 +1,23 @@
+//===- MCLFI.h - LFI-specific code for MC -----------------------*- C++ 
-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM 
Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// This file was written by the LFI and Native Client authors.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/Support/CommandLine.h"
+#include "llvm/Support/Compiler.h"
+
+namespace llvm {
+
+class MCContext;
+class MCStreamer;
+class Triple;
+
+LLVM_ABI void initializeLFIMCStreamer(MCStreamer &Streamer, MCContext &Ctx,
+                                      const Triple &TheTriple);
+
+} // namespace llvm
diff --git a/llvm/include/llvm/MC/MCLFIRewriter.h 
b/llvm/include/llvm/MC/MCLFIRewriter.h
new file mode 100644
index 0000000000000..06e1d03b21efe
--- /dev/null
+++ b/llvm/include/llvm/MC/MCLFIRewriter.h
@@ -0,0 +1,70 @@
+//===- llvm/MC/MCLFIRewriter.h ----------------------------------*- C++ 
-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM 
Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// This file was written by the LFI and Native Client authors.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file declares the MCLFIRewriter class. This is an abstract
+// class that encapsulates the rewriting logic for MCInsts.
+//
+//===----------------------------------------------------------------------===//
+
+#ifndef LLVM_MC_MCLFIREWRITER_H
+#define LLVM_MC_MCLFIREWRITER_H
+
+#include "llvm/MC/MCContext.h"
+#include "llvm/MC/MCInstrInfo.h"
+#include "llvm/MC/MCRegisterInfo.h"
+#include "llvm/Support/Compiler.h"
+
+namespace llvm {
+class MCInst;
+class MCSubtargetInfo;
+class MCStreamer;
+class MCSymbol;
+
+class MCLFIRewriter {
+private:
+  MCContext &Ctx;
+  bool Enabled = true;
+
+protected:
+  std::unique_ptr<MCInstrInfo> InstInfo;
+  std::unique_ptr<MCRegisterInfo> RegInfo;
+
+public:
+  MCLFIRewriter(MCContext &Ctx, std::unique_ptr<MCRegisterInfo> &&RI,
+                std::unique_ptr<MCInstrInfo> &&II)
+      : Ctx(Ctx), InstInfo(std::move(II)), RegInfo(std::move(RI)) {}
+
+  LLVM_ABI void error(const MCInst &Inst, const char Msg[]);
+
+  LLVM_ABI void disable();
+  LLVM_ABI void enable();
+  LLVM_ABI bool isEnabled();
+
+  LLVM_ABI bool isCall(const MCInst &Inst) const;
+  LLVM_ABI bool isBranch(const MCInst &Inst) const;
+  LLVM_ABI bool isIndirectBranch(const MCInst &Inst) const;
+  LLVM_ABI bool isReturn(const MCInst &Inst) const;
+
+  LLVM_ABI bool mayLoad(const MCInst &Inst) const;
+  LLVM_ABI bool mayStore(const MCInst &Inst) const;
+
+  LLVM_ABI bool mayModifyRegister(const MCInst &Inst, MCRegister Reg) const;
+
+  virtual ~MCLFIRewriter() = default;
+  virtual bool rewriteInst(const MCInst &Inst, MCStreamer &Out,
+                           const MCSubtargetInfo &STI) = 0;
+
+  // Called when a label is emitted. Used for optimizations that require
+  // information about jump targets, such as guard elimination.
+  virtual void onLabel(const MCSymbol *Symbol) {}
+};
+
+} // namespace llvm
+#endif
diff --git a/llvm/include/llvm/MC/MCParser/MCAsmParserExtension.h 
b/llvm/include/llvm/MC/MCParser/MCAsmParserExtension.h
index 66fd28827065b..eb9a995c0a2be 100644
--- a/llvm/include/llvm/MC/MCParser/MCAsmParserExtension.h
+++ b/llvm/include/llvm/MC/MCParser/MCAsmParserExtension.h
@@ -11,6 +11,7 @@
 
 #include "llvm/ADT/STLFunctionalExtras.h"
 #include "llvm/ADT/StringRef.h"
+#include "llvm/MC/MCLFIRewriter.h"
 #include "llvm/MC/MCParser/MCAsmParser.h"
 #include "llvm/Support/Compiler.h"
 #include "llvm/Support/SMLoc.h"
@@ -120,13 +121,14 @@ class LLVM_ABI MCAsmParserExtension {
   /// @}
 };
 
-MCAsmParserExtension *createDarwinAsmParser();
-MCAsmParserExtension *createELFAsmParser();
-MCAsmParserExtension *createCOFFAsmParser();
-MCAsmParserExtension *createCOFFMasmParser();
-MCAsmParserExtension *createGOFFAsmParser();
-MCAsmParserExtension *createXCOFFAsmParser();
-MCAsmParserExtension *createWasmAsmParser();
+LLVM_ABI MCAsmParserExtension *createDarwinAsmParser();
+LLVM_ABI MCAsmParserExtension *createELFAsmParser();
+LLVM_ABI MCAsmParserExtension *createCOFFAsmParser();
+LLVM_ABI MCAsmParserExtension *createCOFFMasmParser();
+LLVM_ABI MCAsmParserExtension *createGOFFAsmParser();
+LLVM_ABI MCAsmParserExtension *createXCOFFAsmParser();
+LLVM_ABI MCAsmParserExtension *createWasmAsmParser();
+LLVM_ABI MCAsmParserExtension *createLFIAsmParser(MCLFIRewriter *Exp);
 
 } // end namespace llvm
 
diff --git a/llvm/include/llvm/MC/MCStreamer.h 
b/llvm/include/llvm/MC/MCStreamer.h
index 7dcc915d7b1cf..530d2f34469be 100644
--- a/llvm/include/llvm/MC/MCStreamer.h
+++ b/llvm/include/llvm/MC/MCStreamer.h
@@ -19,6 +19,7 @@
 #include "llvm/ADT/StringRef.h"
 #include "llvm/MC/MCDirectives.h"
 #include "llvm/MC/MCDwarf.h"
+#include "llvm/MC/MCLFIRewriter.h"
 #include "llvm/MC/MCLinkerOptimizationHint.h"
 #include "llvm/MC/MCPseudoProbe.h"
 #include "llvm/MC/MCSection.h"
@@ -290,6 +291,8 @@ class LLVM_ABI MCStreamer {
   /// Returns true if the .cv_loc directive is in the right section.
   bool checkCVLocSection(unsigned FuncId, unsigned FileNo, SMLoc Loc);
 
+  std::unique_ptr<MCLFIRewriter> LFIRewriter;
+
 public:
   MCStreamer(const MCStreamer &) = delete;
   MCStreamer &operator=(const MCStreamer &) = delete;
@@ -307,6 +310,10 @@ class LLVM_ABI MCStreamer {
     return StartTokLocPtr ? *StartTokLocPtr : SMLoc();
   }
 
+  void setLFIRewriter(MCLFIRewriter *Exp) { LFIRewriter.reset(Exp); }
+
+  MCLFIRewriter *getLFIRewriter() { return LFIRewriter.get(); }
+
   /// State management
   ///
   virtual void reset();
diff --git a/llvm/include/llvm/MC/TargetRegistry.h 
b/llvm/include/llvm/MC/TargetRegistry.h
index 234c587c8e951..88ec00e74af54 100644
--- a/llvm/include/llvm/MC/TargetRegistry.h
+++ b/llvm/include/llvm/MC/TargetRegistry.h
@@ -46,6 +46,7 @@ class MCDisassembler;
 class MCInstPrinter;
 class MCInstrAnalysis;
 class MCInstrInfo;
+class MCLFIRewriter;
 class MCObjectWriter;
 class MCRegisterInfo;
 class MCRelocationInfo;
@@ -237,6 +238,11 @@ class Target {
       mca::InstrumentManager *(*)(const MCSubtargetInfo &STI,
                                   const MCInstrInfo &MCII);
 
+  using MCLFIRewriterCtorTy =
+      MCLFIRewriter *(*)(MCStreamer & S,
+                         std::unique_ptr<MCRegisterInfo> &&RegInfo,
+                         std::unique_ptr<MCInstrInfo> &&InstInfo);
+
 private:
   /// Next - The next registered target in the linked list, maintained by the
   /// TargetRegistry.
@@ -351,6 +357,10 @@ class Target {
   /// InstrumentManager, if registered (default = nullptr).
   InstrumentManagerCtorTy InstrumentManagerCtorFn = nullptr;
 
+  // MCLFIRewriterCtorFn - Construction function for this target's
+  // MCLFIRewriter, if registered (default = nullptr).
+  MCLFIRewriterCtorTy MCLFIRewriterCtorFn = nullptr;
+
 public:
   Target() = default;
 
@@ -592,6 +602,13 @@ class Target {
     return nullptr;
   }
 
+  void createMCLFIRewriter(MCStreamer &S,
+                           std::unique_ptr<MCRegisterInfo> &&RegInfo,
+                           std::unique_ptr<MCInstrInfo> &&InstInfo) const {
+    if (MCLFIRewriterCtorFn)
+      MCLFIRewriterCtorFn(S, std::move(RegInfo), std::move(InstInfo));
+  }
+
   // TODO(boomanaiden154): Remove this function after LLVM 22 branches.
   [[deprecated("Use overload accepting Triple instead")]]
   MCRelocationInfo *createMCRelocationInfo(StringRef TT, MCContext &Ctx) const 
{
@@ -1064,6 +1081,10 @@ struct TargetRegistry {
     T.InstrumentManagerCtorFn = Fn;
   }
 
+  static void RegisterMCLFIRewriter(Target &T, Target::MCLFIRewriterCtorTy Fn) 
{
+    T.MCLFIRewriterCtorFn = Fn;
+  }
+
   /// @}
 };
 
diff --git a/llvm/lib/MC/CMakeLists.txt b/llvm/lib/MC/CMakeLists.txt
index 1388f130bb806..7a9e26af415c6 100644
--- a/llvm/lib/MC/CMakeLists.txt
+++ b/llvm/lib/MC/CMakeLists.txt
@@ -31,6 +31,8 @@ add_llvm_component_library(LLVMMC
   MCInstrAnalysis.cpp
   MCInstrDesc.cpp
   MCInstrInfo.cpp
+  MCLFI.cpp
+  MCLFIRewriter.cpp
   MCLabel.cpp
   MCLinkerOptimizationHint.cpp
   MCMachOStreamer.cpp
diff --git a/llvm/lib/MC/MCAsmStreamer.cpp b/llvm/lib/MC/MCAsmStreamer.cpp
index bc20833977707..adc5f5a03722c 100644
--- a/llvm/lib/MC/MCAsmStreamer.cpp
+++ b/llvm/lib/MC/MCAsmStreamer.cpp
@@ -19,6 +19,7 @@
 #include "llvm/MC/MCExpr.h"
 #include "llvm/MC/MCInst.h"
 #include "llvm/MC/MCInstPrinter.h"
+#include "llvm/MC/MCLFIRewriter.h"
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCObjectWriter.h"
 #include "llvm/MC/MCPseudoProbe.h"
@@ -2438,6 +2439,10 @@ void MCAsmStreamer::AddEncodingComment(const MCInst 
&Inst,
 
 void MCAsmStreamer::emitInstruction(const MCInst &Inst,
                                     const MCSubtargetInfo &STI) {
+  if (LFIRewriter && LFIRewriter->isEnabled() &&
+      LFIRewriter->rewriteInst(Inst, *this, STI))
+    return;
+
   if (CurFrag) {
     MCSection *Sec = getCurrentSectionOnly();
     Sec->setHasInstructions(true);
diff --git a/llvm/lib/MC/MCLFI.cpp b/llvm/lib/MC/MCLFI.cpp
new file mode 100644
index 0000000000000..6082edac1a781
--- /dev/null
+++ b/llvm/lib/MC/MCLFI.cpp
@@ -0,0 +1,75 @@
+//===- lib/MC/MCLFI.cpp - LFI-specific MC implementation 
------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM 
Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// This file was written by the LFI and Native Client authors.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/MC/MCLFI.h"
+#include "llvm/BinaryFormat/ELF.h"
+#include "llvm/MC/MCContext.h"
+#include "llvm/MC/MCLFIRewriter.h"
+#include "llvm/MC/MCSectionELF.h"
+#include "llvm/MC/MCStreamer.h"
+#include "llvm/MC/TargetRegistry.h"
+#include "llvm/Support/Alignment.h"
+#include "llvm/TargetParser/Triple.h"
+
+static const char NoteNamespace[] = "LFI";
+
+namespace llvm {
+
+cl::opt<bool> FlagEnableRewriting("lfi-enable-rewriter",
+                                  cl::desc("Don't enable rewriting for LFI."),
+                                  cl::init(true));
+
+void initializeLFIMCStreamer(MCStreamer &Streamer, MCContext &Ctx,
+                             const Triple &TheTriple) {
+  assert(TheTriple.isLFI());
+  const char *NoteName;
+  const char *NoteArch;
+  switch (TheTriple.getArch()) {
+  case Triple::aarch64:
+    NoteName = ".note.LFI.ABI.aarch64";
+    NoteArch = "aarch64";
+    break;
+  default:
+    report_fatal_error("Unsupported architecture for LFI");
+  }
+
+  std::string Error; // empty
+  const Target *TheTarget = TargetRegistry::lookupTarget(TheTriple, Error);
+
+  // Create the Target specific MCLFIRewriter.
+  assert(TheTarget != nullptr);
+  if (FlagEnableRewriting) {
+    TheTarget->createMCLFIRewriter(
+        Streamer,
+        std::unique_ptr<MCRegisterInfo>(TheTarget->createMCRegInfo(TheTriple)),
+        std::unique_ptr<MCInstrInfo>(TheTarget->createMCInstrInfo()));
+  }
+
+  // Emit an ELF Note section in its own COMDAT group which identifies LFI
+  // object files.
+  MCSectionELF *Note = Ctx.getELFSection(NoteName, ELF::SHT_NOTE,
+                                         ELF::SHF_ALLOC | ELF::SHF_GROUP, 0,
+                                         NoteName, /*IsComdat=*/true);
+
+  Streamer.pushSection();
+  Streamer.switchSection(Note);
+  Streamer.emitIntValue(strlen(NoteNamespace) + 1, 4);
+  Streamer.emitIntValue(strlen(NoteArch) + 1, 4);
+  Streamer.emitIntValue(ELF::NT_VERSION, 4);
+  Streamer.emitBytes(NoteNamespace);
+  Streamer.emitIntValue(0, 1); // NUL terminator
+  Streamer.emitValueToAlignment(Align(4));
+  Streamer.emitBytes(NoteArch);
+  Streamer.emitIntValue(0, 1); // NUL terminator
+  Streamer.emitValueToAlignment(Align(4));
+  Streamer.popSection();
+}
+
+} // namespace llvm
diff --git a/llvm/lib/MC/MCLFIRewriter.cpp b/llvm/lib/MC/MCLFIRewriter.cpp
new file mode 100644
index 0000000000000..5295a5c3e3771
--- /dev/null
+++ b/llvm/lib/MC/MCLFIRewriter.cpp
@@ -0,0 +1,61 @@
+//===- MCLFIRewriter.cpp ----------------------------------------*- C++ 
-*-===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM 
Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+// This file was written by the LFI and Native Client authors.
+//
+//===----------------------------------------------------------------------===//
+//
+// This file implements the MCLFIRewriter class. This is a base
+// class that encapsulates the rewriting logic for MCInsts.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/MC/MCLFIRewriter.h"
+#include "llvm/MC/MCInst.h"
+#include "llvm/MC/MCInstrInfo.h"
+#include "llvm/MC/MCRegisterInfo.h"
+
+namespace llvm {
+
+void MCLFIRewriter::error(const MCInst &Inst, const char Msg[]) {
+  Ctx.reportError(Inst.getLoc(), Msg);
+}
+
+void MCLFIRewriter::disable() { Enabled = false; }
+
+void MCLFIRewriter::enable() { Enabled = true; }
+
+bool MCLFIRewriter::isEnabled() { return Enabled; }
+
+bool MCLFIRewriter::isCall(const MCInst &Inst) const {
+  return InstInfo->get(Inst.getOpcode()).isCall();
+}
+
+bool MCLFIRewriter::isBranch(const MCInst &Inst) const {
+  return InstInfo->get(Inst.getOpcode()).isBranch();
+}
+
+bool MCLFIRewriter::isIndirectBranch(const MCInst &Inst) const {
+  return InstInfo->get(Inst.getOpcode()).isIndirectBranch();
+}
+
+bool MCLFIRewriter::isReturn(const MCInst &Inst) const {
+  return InstInfo->get(Inst.getOpcode()).isReturn();
+}
+
+bool MCLFIRewriter::mayLoad(const MCInst &Inst) const {
+  return InstInfo->get(Inst.getOpcode()).mayLoad();
+}
+
+bool MCLFIRewriter::mayStore(const MCInst &Inst) const {
+  return InstInfo->get(Inst.getOpcode()).mayStore();
+}
+
+bool MCLFIRewriter::mayModifyRegister(const MCInst &Inst,
+                                      MCRegister Reg) const {
+  return InstInfo->get(Inst.getOpcode()).hasDefOfPhysReg(Inst, Reg, *RegInfo);
+}
+} // namespace llvm
diff --git a/llvm/lib/MC/MCObjectStreamer.cpp b/llvm/lib/MC/MCObjectStreamer.cpp
index 261e9a37ecb55..3e249b2b2c396 100644
--- a/llvm/lib/MC/MCObjectStreamer.cpp
+++ b/llvm/lib/MC/MCObjectStreamer.cpp
@@ -15,6 +15,7 @@
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCDwarf.h"
 #include "llvm/MC/MCExpr.h"
+#include "llvm/MC/MCLFIRewriter.h"
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCObjectWriter.h"
 #include "llvm/MC/MCSFrame.h"
@@ -390,6 +391,10 @@ bool MCObjectStreamer::mayHaveInstructions(MCSection &Sec) 
const {
 
 void MCObjectStreamer::emitInstruction(const MCInst &Inst,
                                        const MCSubtargetInfo &STI) {
+  if (LFIRewriter && LFIRewriter->isEnabled() &&
+      LFIRewriter->rewriteInst(Inst, *this, STI))
+    return;
+
   MCStreamer::emitInstruction(Inst, STI);
 
   MCSection *Sec = getCurrentSectionOnly();
diff --git a/llvm/lib/MC/MCParser/AsmParser.cpp 
b/llvm/lib/MC/MCParser/AsmParser.cpp
index 3277994701b78..fdce7999f492e 100644
--- a/llvm/lib/MC/MCParser/AsmParser.cpp
+++ b/llvm/lib/MC/MCParser/AsmParser.cpp
@@ -120,6 +120,7 @@ class AsmParser : public MCAsmParser {
   SourceMgr::DiagHandlerTy SavedDiagHandler;
   void *SavedDiagContext;
   std::unique_ptr<MCAsmParserExtension> PlatformParser;
+  std::unique_ptr<MCAsmParserExtension> LFIParser;
   SMLoc StartTokLoc;
   std::optional<SMLoc> CFIStartProcLoc;
 
@@ -745,8 +746,8 @@ extern cl::opt<unsigned> AsmMacroMaxNestingDepth;
 
 AsmParser::AsmParser(SourceMgr &SM, MCContext &Ctx, MCStreamer &Out,
                      const MCAsmInfo &MAI, unsigned CB = 0)
-    : MCAsmParser(Ctx, Out, SM, MAI), CurBuffer(CB ? CB : SM.getMainFileID()),
-      MacrosEnabledFlag(true) {
+    : MCAsmParser(Ctx, Out, SM, MAI), LFIParser(nullptr),
+      CurBuffer(CB ? CB : SM.getMainFileID()), MacrosEnabledFlag(true) {
   HadError = false;
   // Save the old handler.
   SavedDiagHandler = SrcMgr.getDiagHandler();
@@ -788,6 +789,10 @@ AsmParser::AsmParser(SourceMgr &SM, MCContext &Ctx, 
MCStreamer &Out,
   }
 
   PlatformParser->Initialize(*this);
+  if (Out.getLFIRewriter()) {
+    LFIParser.reset(createLFIAsmParser(Out.getLFIRewriter()));
+    LFIParser->Initialize(*this);
+  }
   initializeDirectiveKindMap();
   initializeCVDefRangeTypeMap();
 }
diff --git a/llvm/lib/MC/MCParser/CMakeLists.txt 
b/llvm/lib/MC/MCParser/CMakeLists.txt
index 008a50e9da660..c911874fc6540 100644
--- a/llvm/lib/MC/MCParser/CMakeLists.txt
+++ b/llvm/lib/MC/MCParser/CMakeLists.txt
@@ -6,6 +6,7 @@ add_llvm_component_library(LLVMMCParser
   GOFFAsmParser.cpp
   DarwinAsmParser.cpp
   ELFAsmParser.cpp
+  LFIAsmParser.cpp
   MCAsmParser.cpp
   MCAsmParserExtension.cpp
   MCTargetAsmParser.cpp
diff --git a/llvm/lib/MC/MCParser/LFIAsmParser.cpp 
b/llvm/lib/MC/MCParser/LFIAsmParser.cpp
new file mode 100644
index 0000000000000..4601052e0d65f
--- /dev/null
+++ b/llvm/lib/MC/MCParser/LFIAsmParser.cpp
@@ -0,0 +1,69 @@
+//===- LFIAsmParser.cpp - LFI Assembly Parser 
-----------------------------===//
+//
+// Part of the LLVM Project, under the Apache License v2.0 with LLVM 
Exceptions.
+// See https://llvm.org/LICENSE.txt for license information.
+// SPDX-License-Identifier: Apache-2.0 WITH LLVM-exception
+//
+//===----------------------------------------------------------------------===//
+//
+// This file was written by the LFI and Native Client authors.
+//
+//===----------------------------------------------------------------------===//
+
+#include "llvm/MC/MCLFIRewriter.h"
+#include "llvm/MC/MCParser/MCAsmParserExtension.h"
+#include "llvm/MC/MCStreamer.h"
+
+using namespace llvm;
+
+class LFIAsmParser : public MCAsmParserExtension {
+  MCLFIRewriter *Rewriter;
+  template <bool (LFIAsmParser::*HandlerMethod)(StringRef, SMLoc)>
+  void addDirectiveHandler(StringRef Directive) {
+    MCAsmParser::ExtensionDirectiveHandler Handler =
+        std::make_pair(this, HandleDirective<LFIAsmParser, HandlerMethod>);
+
+    getParser().addDirectiveHandler(Directive, Handler);
+  }
+
+public:
+  LFIAsmParser(MCLFIRewriter *Exp) : Rewriter(Exp) {}
+  void Initialize(MCAsmParser &Parser) override {
+    // Call the base implementation.
+    MCAsmParserExtension::Initialize(Parser);
+    addDirectiveHandler<&LFIAsmParser::parseRewriteDisable>(
+        ".lfi_rewrite_disable");
+    addDirectiveHandler<&LFIAsmParser::parseRewriteEnable>(
+        ".lfi_rewrite_enable");
+  }
+
+  /// ::= {.lfi_rewrite_disable}
+  bool parseRewriteDisable(StringRef Directive, SMLoc Loc) {
+    getParser().checkForValidSection();
+    if (getLexer().isNot(AsmToken::EndOfStatement))
+      return TokError("unexpected token");
+    Lex();
+
+    Rewriter->disable();
+
+    return false;
+  }
+
+  /// ::= {.lfi_rewrite_enable}
+  bool parseRewriteEnable(StringRef Directive, SMLoc Loc) {
+    getParser().checkForValidSection();
+    if (getLexer().isNot(AsmToken::EndOfStatement))
+      return TokError("unexpected token");
+    Lex();
+
+    Rewriter->enable();
+
+    return false;
+  }
+};
+
+namespace llvm {
+MCAsmParserExtension *createLFIAsmParser(MCLFIRewriter *Exp) {
+  return new LFIAsmParser(Exp);
+}
+} // namespace llvm
diff --git a/llvm/lib/MC/MCStreamer.cpp b/llvm/lib/MC/MCStreamer.cpp
index 4196d4af7d09e..beaf030ac4e5a 100644
--- a/llvm/lib/MC/MCStreamer.cpp
+++ b/llvm/lib/MC/MCStreamer.cpp
@@ -406,6 +406,9 @@ void MCStreamer::emitLabel(MCSymbol *Symbol, SMLoc Loc) {
 
   Symbol->setFragment(&getCurrentSectionOnly()->getDummyFragment());
 
+  if (LFIRewriter)
+    LFIRewriter->onLabel(Symbol);
+
   MCTargetStreamer *TS = getTargetStreamer();
   if (TS)
     TS->emitLabel(Symbol);
diff --git a/llvm/tools/llvm-mc/llvm-mc.cpp b/llvm/tools/llvm-mc/llvm-mc.cpp
index baf3b5536138b..b1f73433b1bb3 100644
--- a/llvm/tools/llvm-mc/llvm-mc.cpp
+++ b/llvm/tools/llvm-mc/llvm-mc.cpp
@@ -21,6 +21,7 @@
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCInstPrinter.h"
 #include "llvm/MC/MCInstrInfo.h"
+#include "llvm/MC/MCLFI.h"
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCObjectWriter.h"
 #include "llvm/MC/MCParser/AsmLexer.h"
@@ -626,6 +627,12 @@ int main(int argc, char **argv) {
     Str.reset(TheTarget->createAsmStreamer(Ctx, std::move(FOut), std::move(IP),
                                            std::move(CE), std::move(MAB)));
 
+    Triple T(TripleName);
+    if (T.isLFI()) {
+      Str->initSections(NoExecStack, *STI);
+      initializeLFIMCStreamer(*Str.get(), Ctx, T);
+    }
+
   } else if (FileType == OFT_Null) {
     Str.reset(TheTarget->createNullStreamer(Ctx));
   } else {
@@ -646,6 +653,13 @@ int main(int argc, char **argv) {
     if (NoExecStack)
       Str->switchSection(
           Ctx.getAsmInfo()->getStackSection(Ctx, /*Exec=*/false));
+
+    Triple T(TripleName);
+    if (T.isLFI()) {
+      Str->initSections(NoExecStack, *STI);
+      initializeLFIMCStreamer(*Str.get(), Ctx, T);
+    }
+
     Str->emitVersionForTarget(TheTriple, VersionTuple(), nullptr,
                               VersionTuple());
   }

>From 2fd9b26f53b6a796971a4ace79defc3f01135edd Mon Sep 17 00:00:00 2001
From: Zachary Yedidia <[email protected]>
Date: Mon, 26 Jan 2026 17:18:48 -0500
Subject: [PATCH 2/4] [LFI] Small fixes based on code review

- Use `unique_ptr` for `setLFIRewriter`.
- Fix `lfi-enable-rewriter` flag description.
- No need to explicitly initialize LFIRewriter to null.
---
 llvm/include/llvm/MC/MCStreamer.h  | 4 +++-
 llvm/lib/MC/MCLFI.cpp              | 2 +-
 llvm/lib/MC/MCParser/AsmParser.cpp | 4 ++--
 llvm/tools/llvm-mc/llvm-mc.cpp     | 1 -
 4 files changed, 6 insertions(+), 5 deletions(-)

diff --git a/llvm/include/llvm/MC/MCStreamer.h 
b/llvm/include/llvm/MC/MCStreamer.h
index 530d2f34469be..9f7591d32f0d8 100644
--- a/llvm/include/llvm/MC/MCStreamer.h
+++ b/llvm/include/llvm/MC/MCStreamer.h
@@ -310,7 +310,9 @@ class LLVM_ABI MCStreamer {
     return StartTokLocPtr ? *StartTokLocPtr : SMLoc();
   }
 
-  void setLFIRewriter(MCLFIRewriter *Exp) { LFIRewriter.reset(Exp); }
+  void setLFIRewriter(std::unique_ptr<MCLFIRewriter> Rewriter) {
+    LFIRewriter = std::move(Rewriter);
+  }
 
   MCLFIRewriter *getLFIRewriter() { return LFIRewriter.get(); }
 
diff --git a/llvm/lib/MC/MCLFI.cpp b/llvm/lib/MC/MCLFI.cpp
index 6082edac1a781..af89d9bf3dc85 100644
--- a/llvm/lib/MC/MCLFI.cpp
+++ b/llvm/lib/MC/MCLFI.cpp
@@ -23,7 +23,7 @@ static const char NoteNamespace[] = "LFI";
 namespace llvm {
 
 cl::opt<bool> FlagEnableRewriting("lfi-enable-rewriter",
-                                  cl::desc("Don't enable rewriting for LFI."),
+                                  cl::desc("Enable rewriting for LFI."),
                                   cl::init(true));
 
 void initializeLFIMCStreamer(MCStreamer &Streamer, MCContext &Ctx,
diff --git a/llvm/lib/MC/MCParser/AsmParser.cpp 
b/llvm/lib/MC/MCParser/AsmParser.cpp
index fdce7999f492e..3452708bcec8a 100644
--- a/llvm/lib/MC/MCParser/AsmParser.cpp
+++ b/llvm/lib/MC/MCParser/AsmParser.cpp
@@ -746,8 +746,8 @@ extern cl::opt<unsigned> AsmMacroMaxNestingDepth;
 
 AsmParser::AsmParser(SourceMgr &SM, MCContext &Ctx, MCStreamer &Out,
                      const MCAsmInfo &MAI, unsigned CB = 0)
-    : MCAsmParser(Ctx, Out, SM, MAI), LFIParser(nullptr),
-      CurBuffer(CB ? CB : SM.getMainFileID()), MacrosEnabledFlag(true) {
+    : MCAsmParser(Ctx, Out, SM, MAI), CurBuffer(CB ? CB : SM.getMainFileID()),
+      MacrosEnabledFlag(true) {
   HadError = false;
   // Save the old handler.
   SavedDiagHandler = SrcMgr.getDiagHandler();
diff --git a/llvm/tools/llvm-mc/llvm-mc.cpp b/llvm/tools/llvm-mc/llvm-mc.cpp
index b1f73433b1bb3..6542612b5768c 100644
--- a/llvm/tools/llvm-mc/llvm-mc.cpp
+++ b/llvm/tools/llvm-mc/llvm-mc.cpp
@@ -632,7 +632,6 @@ int main(int argc, char **argv) {
       Str->initSections(NoExecStack, *STI);
       initializeLFIMCStreamer(*Str.get(), Ctx, T);
     }
-
   } else if (FileType == OFT_Null) {
     Str.reset(TheTarget->createNullStreamer(Ctx));
   } else {

>From 7cd80bdb05b8e3e3c6fdaa5012d4041f34727193 Mon Sep 17 00:00:00 2001
From: Zachary Yedidia <[email protected]>
Date: Mon, 26 Jan 2026 17:28:22 -0500
Subject: [PATCH 3/4] [LFI] initializeMCLFIStreamer in createMCObjectStreamer

- Still needs manual initialization for use with createAsmStreamer.
---
 clang/tools/driver/cc1as_main.cpp | 3 ---
 llvm/lib/MC/TargetRegistry.cpp    | 3 +++
 llvm/tools/llvm-mc/llvm-mc.cpp    | 6 ------
 3 files changed, 3 insertions(+), 9 deletions(-)

diff --git a/clang/tools/driver/cc1as_main.cpp 
b/clang/tools/driver/cc1as_main.cpp
index 6cbccbc3d7408..ccc48a77891d7 100644
--- a/clang/tools/driver/cc1as_main.cpp
+++ b/clang/tools/driver/cc1as_main.cpp
@@ -28,7 +28,6 @@
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCInstPrinter.h"
 #include "llvm/MC/MCInstrInfo.h"
-#include "llvm/MC/MCLFI.h"
 #include "llvm/MC/MCObjectFileInfo.h"
 #include "llvm/MC/MCObjectWriter.h"
 #include "llvm/MC/MCParser/MCAsmParser.h"
@@ -594,8 +593,6 @@ static bool ExecuteAssemblerImpl(AssemblerInvocation &Opts,
     Triple T(Opts.Triple);
     Str.reset(TheTarget->createMCObjectStreamer(
         T, Ctx, std::move(MAB), std::move(OW), std::move(CE), *STI));
-    if (T.isLFI())
-      initializeLFIMCStreamer(*Str.get(), Ctx, T);
     if (T.isOSBinFormatMachO() && T.isOSDarwin()) {
       Triple *TVT = Opts.DarwinTargetVariantTriple
                         ? &*Opts.DarwinTargetVariantTriple
diff --git a/llvm/lib/MC/TargetRegistry.cpp b/llvm/lib/MC/TargetRegistry.cpp
index 9263dda65a8b0..db743ef9bdcb6 100644
--- a/llvm/lib/MC/TargetRegistry.cpp
+++ b/llvm/lib/MC/TargetRegistry.cpp
@@ -13,6 +13,7 @@
 #include "llvm/MC/MCCodeEmitter.h"
 #include "llvm/MC/MCContext.h"
 #include "llvm/MC/MCInstPrinter.h"
+#include "llvm/MC/MCLFI.h"
 #include "llvm/MC/MCObjectStreamer.h"
 #include "llvm/MC/MCObjectWriter.h"
 #include "llvm/Support/raw_ostream.h"
@@ -76,6 +77,8 @@ MCStreamer *Target::createMCObjectStreamer(
   }
   if (ObjectTargetStreamerCtorFn)
     ObjectTargetStreamerCtorFn(*S, STI);
+  if (T.isLFI())
+    initializeLFIMCStreamer(*S, Ctx, T);
   return S;
 }
 
diff --git a/llvm/tools/llvm-mc/llvm-mc.cpp b/llvm/tools/llvm-mc/llvm-mc.cpp
index 6542612b5768c..ea2c8c3b6427c 100644
--- a/llvm/tools/llvm-mc/llvm-mc.cpp
+++ b/llvm/tools/llvm-mc/llvm-mc.cpp
@@ -653,12 +653,6 @@ int main(int argc, char **argv) {
       Str->switchSection(
           Ctx.getAsmInfo()->getStackSection(Ctx, /*Exec=*/false));
 
-    Triple T(TripleName);
-    if (T.isLFI()) {
-      Str->initSections(NoExecStack, *STI);
-      initializeLFIMCStreamer(*Str.get(), Ctx, T);
-    }
-
     Str->emitVersionForTarget(TheTriple, VersionTuple(), nullptr,
                               VersionTuple());
   }

>From 9b5fbe593628be8b7b8ef5f704fc79315ebdeeed Mon Sep 17 00:00:00 2001
From: Zachary Yedidia <[email protected]>
Date: Tue, 27 Jan 2026 13:19:48 -0500
Subject: [PATCH 4/4] [LFI] Rely on rewriteInst to track Enabled

---
 llvm/include/llvm/MC/MCLFIRewriter.h | 3 +--
 llvm/lib/MC/MCAsmStreamer.cpp        | 3 +--
 llvm/lib/MC/MCLFIRewriter.cpp        | 2 --
 llvm/lib/MC/MCObjectStreamer.cpp     | 3 +--
 4 files changed, 3 insertions(+), 8 deletions(-)

diff --git a/llvm/include/llvm/MC/MCLFIRewriter.h 
b/llvm/include/llvm/MC/MCLFIRewriter.h
index 06e1d03b21efe..3736eeac2c10d 100644
--- a/llvm/include/llvm/MC/MCLFIRewriter.h
+++ b/llvm/include/llvm/MC/MCLFIRewriter.h
@@ -30,9 +30,9 @@ class MCSymbol;
 class MCLFIRewriter {
 private:
   MCContext &Ctx;
-  bool Enabled = true;
 
 protected:
+  bool Enabled = true;
   std::unique_ptr<MCInstrInfo> InstInfo;
   std::unique_ptr<MCRegisterInfo> RegInfo;
 
@@ -45,7 +45,6 @@ class MCLFIRewriter {
 
   LLVM_ABI void disable();
   LLVM_ABI void enable();
-  LLVM_ABI bool isEnabled();
 
   LLVM_ABI bool isCall(const MCInst &Inst) const;
   LLVM_ABI bool isBranch(const MCInst &Inst) const;
diff --git a/llvm/lib/MC/MCAsmStreamer.cpp b/llvm/lib/MC/MCAsmStreamer.cpp
index adc5f5a03722c..390fcbc14c065 100644
--- a/llvm/lib/MC/MCAsmStreamer.cpp
+++ b/llvm/lib/MC/MCAsmStreamer.cpp
@@ -2439,8 +2439,7 @@ void MCAsmStreamer::AddEncodingComment(const MCInst &Inst,
 
 void MCAsmStreamer::emitInstruction(const MCInst &Inst,
                                     const MCSubtargetInfo &STI) {
-  if (LFIRewriter && LFIRewriter->isEnabled() &&
-      LFIRewriter->rewriteInst(Inst, *this, STI))
+  if (LFIRewriter && LFIRewriter->rewriteInst(Inst, *this, STI))
     return;
 
   if (CurFrag) {
diff --git a/llvm/lib/MC/MCLFIRewriter.cpp b/llvm/lib/MC/MCLFIRewriter.cpp
index 5295a5c3e3771..261afe1db23ba 100644
--- a/llvm/lib/MC/MCLFIRewriter.cpp
+++ b/llvm/lib/MC/MCLFIRewriter.cpp
@@ -28,8 +28,6 @@ void MCLFIRewriter::disable() { Enabled = false; }
 
 void MCLFIRewriter::enable() { Enabled = true; }
 
-bool MCLFIRewriter::isEnabled() { return Enabled; }
-
 bool MCLFIRewriter::isCall(const MCInst &Inst) const {
   return InstInfo->get(Inst.getOpcode()).isCall();
 }
diff --git a/llvm/lib/MC/MCObjectStreamer.cpp b/llvm/lib/MC/MCObjectStreamer.cpp
index 3e249b2b2c396..6d0d4ef94db43 100644
--- a/llvm/lib/MC/MCObjectStreamer.cpp
+++ b/llvm/lib/MC/MCObjectStreamer.cpp
@@ -391,8 +391,7 @@ bool MCObjectStreamer::mayHaveInstructions(MCSection &Sec) 
const {
 
 void MCObjectStreamer::emitInstruction(const MCInst &Inst,
                                        const MCSubtargetInfo &STI) {
-  if (LFIRewriter && LFIRewriter->isEnabled() &&
-      LFIRewriter->rewriteInst(Inst, *this, STI))
+  if (LFIRewriter && LFIRewriter->rewriteInst(Inst, *this, STI))
     return;
 
   MCStreamer::emitInstruction(Inst, STI);

_______________________________________________
cfe-commits mailing list
[email protected]
https://lists.llvm.org/cgi-bin/mailman/listinfo/cfe-commits

Reply via email to