//===-- W65816Layer2Gate.cpp - LTO x Layer 2 silent-miscompile gate -------===//
//
// 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
//
//===----------------------------------------------------------------------===//
//
// Phase 1.12 of GAP_CLOSURE_PLAN.md: prevent LTO from silently mixing TUs
// built with `-mllvm -w65816-dbr-safe-ptrs` (Layer 2) and TUs built without.
//
// The Layer 2 flag is a cl::opt -- it is per-clang-invocation, hence per-TU.
// When the linker merges bitcode from multiple TUs into one module for LTO,
// the cl::opt value seen by the final codegen pass reflects only the LTO
// driver's invocation; the per-TU value is gone.  A function compiled with
// Layer 2 on assumes the pointer's bank byte matches DBR; a function
// compiled with Layer 2 off does not.  Mix them in one module and you get
// silent wrong code in the worst possible place (struct-field deref hot
// paths in Lua / CoreMark).
//
// Strategy (two complementary passes):
//
//   1) W65816Layer2Stamp: stamp every Function in the TU with a
//      "w65816-layer2"="true"|"false" attribute reflecting the per-TU
//      cl::opt value.  Provided in three flavors so the attribute lands
//      in IR regardless of which pipeline clang is using:
//        - new-PM ModulePass W65816Layer2StampPass: registered via
//          PassBuilder's PipelineStartEPCallback so it runs at the very
//          start of clang's optimization pipeline.  This is the path
//          that matters for LTO: bitcode is written AFTER this point,
//          so the attribute lands in the .bc file.
//        - legacy FunctionPass W65816Layer2Stamp: registered into
//          TargetPassConfig::addIRPasses() for the codegen pipeline.
//          Belt-and-suspenders coverage for the non-LTO non-NPM path
//          (and any future pipeline that bypasses PassBuilder).
//        - the legacy and NPM versions share the same stamping
//          implementation (`stampFunction`) so behaviour is identical.
//
//   2) W65816Layer2Gate: walks every Function in the post-link module
//      and hard-fails if any two functions disagree on the attribute.
//      Available in both new-PM and legacy flavors:
//        - new-PM ModulePass W65816Layer2GatePass: registered via
//          PassBuilder's pipeline-parsing callback, can be invoked
//          via `opt -passes=w65816-layer2-gate <module.bc>`.
//        - legacy ModulePass W65816Layer2Gate: registered with the
//          legacy PassRegistry under the same name.
//      Phase 5.2's ltoLink.sh is the planned consumer; until that
//      script exists the gate is callable directly via opt.
//
//===----------------------------------------------------------------------===//

#include "W65816.h"
#include "llvm/ADT/Twine.h"
#include "llvm/IR/Function.h"
#include "llvm/IR/Module.h"
#include "llvm/IR/PassManager.h"
#include "llvm/Pass.h"
#include "llvm/Passes/PassBuilder.h"
#include "llvm/Support/CommandLine.h"
#include "llvm/Support/ErrorHandling.h"
#include "llvm/Support/raw_ostream.h"

using namespace llvm;

// W65816ISelLowering.cpp owns the cl::opt; we read it here.  The
// definition there lives at top-level (file scope, post-`using namespace
// llvm;`) -- NOT in namespace llvm -- so this extern matches.
extern cl::opt<bool> DbrSafePtrs;

#define STAMP_DEBUG_TYPE "w65816-layer2-stamp"
#define GATE_DEBUG_TYPE  "w65816-layer2-gate"

// The attribute name.  Single source of truth for both passes and any
// external consumer (ltoLink.sh helper, debugger, etc.).
static const char kLayer2Attr[] = "w65816-layer2";
static const char kLayer2True[] = "true";
static const char kLayer2False[] = "false";


// Core stamping operation.  Returns true if the function was modified.
// Declarations have no body, but we still stamp them so a Function that
// becomes a definition after llvm-link inherits a consistent attribute.
// Stamping a declaration costs nothing.
static bool stampFunction(Function &F) {
  StringRef Want = DbrSafePtrs ? kLayer2True : kLayer2False;
  Attribute Existing = F.getFnAttribute(kLayer2Attr);
  if (Existing.isStringAttribute() && Existing.getValueAsString() == Want) {
    return false;
  }
  F.removeFnAttr(kLayer2Attr);
  F.addFnAttr(kLayer2Attr, Want);
  return true;
}


// Format and raise a fatal-error report citing the two offending
// functions.  Used by both the new-PM and legacy gate passes.
[[noreturn]] static void reportGateError(const Module &M,
                                         const Function &FirstFn,
                                         const Function &OtherFn,
                                         const char *Detail) {
  Attribute FirstAttr = FirstFn.getFnAttribute(kLayer2Attr);
  Attribute OtherAttr = OtherFn.getFnAttribute(kLayer2Attr);

  std::string Msg;
  raw_string_ostream OS(Msg);
  OS << "W65816 Layer 2 LTO gate: " << Detail << "\n"
     << "  module:    " << M.getModuleIdentifier() << "\n"
     << "  function:  " << FirstFn.getName() << " -> 'w65816-layer2'="
     << (FirstAttr.isStringAttribute()
             ? FirstAttr.getValueAsString()
             : StringRef("<unset>"))
     << "\n"
     << "  function:  " << OtherFn.getName() << " -> 'w65816-layer2'="
     << (OtherAttr.isStringAttribute()
             ? OtherAttr.getValueAsString()
             : StringRef("<unset>"))
     << "\n"
     << "  See docs/GAP_CLOSURE_PLAN.md Phase 1.12.";
  report_fatal_error(Twine(Msg));
}


// Core gate operation, shared by new-PM and legacy variants.  Walks the
// module and hard-fails on any inconsistency.  Returns without effect
// when all functions agree (or only declarations lack the attribute).
//
// Two-pass algorithm so ordering of stamped vs unstamped functions in
// the module doesn't matter:
//   Pass 1: find the first stamped function (if any) and remember its
//           value.  Also validate the stamp value is recognized.
//   Pass 2: walk again.  Any definition lacking the attribute (when
//           Pass 1 found a stamped one) OR carrying a different value
//           is a fatal inconsistency.
//
// Declarations (no body) without the attribute are always skipped --
// they will be resolved at link time by whatever TU defines them, and
// that TU's stamp is what matters.
static void runGateOverModule(Module &M) {
  const Function *FirstStamped = nullptr;
  StringRef FirstValue;

  // Pass 1: establish the module-wide expected value.
  for (const Function &F : M) {
    Attribute Attr = F.getFnAttribute(kLayer2Attr);
    if (!Attr.isStringAttribute()) {
      continue;
    }
    StringRef Value = Attr.getValueAsString();
    if (Value != kLayer2True && Value != kLayer2False) {
      std::string Msg;
      raw_string_ostream OS(Msg);
      OS << "W65816 Layer 2 gate: function '" << F.getName()
         << "' has unrecognized 'w65816-layer2' value '" << Value
         << "' (expected 'true' or 'false').";
      report_fatal_error(Twine(Msg));
    }
    if (!FirstStamped) {
      FirstStamped = &F;
      FirstValue = Value;
    }
  }

  // No stamped function found at all -- module was assembled entirely
  // from non-W65816-codegen sources (asm-only objects, hand-written IR
  // for a test, etc.).  Nothing to enforce; let the build proceed.
  if (!FirstStamped) {
    return;
  }

  // Pass 2: enforce consistency.  Every function definition must either
  // be unstamped+declaration-only or stamped with FirstValue.
  for (const Function &F : M) {
    Attribute Attr = F.getFnAttribute(kLayer2Attr);

    if (!Attr.isStringAttribute()) {
      // Declarations without bodies are fine -- they get resolved at
      // link time by whichever TU defines them, and that TU's stamp
      // is what matters for the eventual codegen.
      if (F.isDeclaration()) {
        continue;
      }
      // Definition with no stamp + at least one stamped function in
      // the module = mixed provenance.  This usually means an older
      // clang (lacking the stamp pass) was used to compile one TU.
      // Cannot tell what Layer 2 mode the unstamped TU was built in,
      // so fail loudly.
      reportGateError(M, *FirstStamped, F,
                      "function lacks 'w65816-layer2' attribute "
                      "but another function in the module has it. "
                      "Mixed-provenance LTO is not safe -- rebuild "
                      "all TUs with the same compiler.");
    }

    StringRef Value = Attr.getValueAsString();
    if (Value != FirstValue) {
      reportGateError(M, *FirstStamped, F,
                      "Layer 2 mode disagreement between TUs.  "
                      "Rebuild all TUs in this LTO set with the same "
                      "'-mllvm -w65816-dbr-safe-ptrs' setting "
                      "(either all on or all off).");
    }
  }
}


namespace llvm {


// New-PM ModulePass that stamps every function in the module with the
// per-TU Layer 2 attribute.  Registered via PipelineStartEPCallback so
// it runs at the start of clang's optimization pipeline -- BEFORE
// bitcode is written for LTO -- ensuring the attribute survives into
// the .bc file.
class W65816Layer2StampPass : public PassInfoMixin<W65816Layer2StampPass> {
public:
  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM) {
    for (Function &F : M) {
      stampFunction(F);
    }
    // We only add a function attribute, no IR-level effects.  Preserve
    // everything; the inliner et al. will copy the attribute on inline.
    return PreservedAnalyses::all();
  }

  static bool isRequired() { return true; }
};


// New-PM ModulePass that hard-fails if any two functions in the module
// disagree on the "w65816-layer2" attribute.  Invocable via
// `opt -passes=w65816-layer2-gate`.
class W65816Layer2GatePass : public PassInfoMixin<W65816Layer2GatePass> {
public:
  PreservedAnalyses run(Module &M, ModuleAnalysisManager &AM) {
    runGateOverModule(M);
    return PreservedAnalyses::all();
  }

  static bool isRequired() { return true; }
};


// Registration entry point called from W65816TargetMachine::
// registerPassBuilderCallbacks.  Wires both the stamp and gate passes
// into the new pass manager.
void registerW65816Layer2GatePasses(PassBuilder &PB) {
  // Stamp every function at the very start of the optimization
  // pipeline.  PipelineStartEPCallback fires for the per-TU
  // optimization pipeline -- exactly where we need it for LTO bitcode
  // emission.  Per LLVM docs: "This does not apply to 'backend'
  // compiles (LTO and ThinLTO link-time pipelines)" -- which is what
  // we want, because the per-TU stamp must happen BEFORE LTO link,
  // not at LTO link time.
  PB.registerPipelineStartEPCallback(
      [](ModulePassManager &MPM, OptimizationLevel) {
        MPM.addPass(W65816Layer2StampPass());
      });

  // Make both passes invocable by name via `opt -passes=...`.  This is
  // how Phase 5.2's ltoLink.sh will call the gate on a freshly-linked
  // bitcode module before invoking llc.
  PB.registerPipelineParsingCallback(
      [](StringRef Name, ModulePassManager &PM,
         ArrayRef<PassBuilder::PipelineElement>) {
        if (Name == STAMP_DEBUG_TYPE) {
          PM.addPass(W65816Layer2StampPass());
          return true;
        }
        if (Name == GATE_DEBUG_TYPE) {
          PM.addPass(W65816Layer2GatePass());
          return true;
        }
        return false;
      });
}


} // end namespace llvm


// =====================================================================
// Legacy PM compatibility shims.  TargetPassConfig::addIRPasses still
// uses the legacy pass manager, so addIRPasses() needs a legacy
// FunctionPass.  Likewise the gate is exposed as a legacy ModulePass
// for callers that prefer the legacy `opt -passes=...` syntax.
// =====================================================================

namespace {


class W65816Layer2Stamp : public FunctionPass {
public:
  static char ID;

  W65816Layer2Stamp() : FunctionPass(ID) {}

  StringRef getPassName() const override {
    return "W65816 Layer 2 attribute stamp (legacy)";
  }

  bool runOnFunction(Function &F) override {
    return stampFunction(F);
  }
};


class W65816Layer2Gate : public ModulePass {
public:
  static char ID;

  W65816Layer2Gate() : ModulePass(ID) {}

  StringRef getPassName() const override {
    return "W65816 Layer 2 LTO consistency gate (legacy)";
  }

  bool runOnModule(Module &M) override {
    runGateOverModule(M);
    return false;
  }
};


} // end anonymous namespace


char W65816Layer2Stamp::ID = 0;
char W65816Layer2Gate::ID = 0;


INITIALIZE_PASS(W65816Layer2Stamp, STAMP_DEBUG_TYPE,
                "W65816 stamp Layer 2 attribute on functions",
                false, false)

INITIALIZE_PASS(W65816Layer2Gate, GATE_DEBUG_TYPE,
                "W65816 enforce Layer 2 attribute consistency under LTO",
                false, false)


namespace llvm {


FunctionPass *createW65816Layer2Stamp() {
  return new W65816Layer2Stamp();
}


ModulePass *createW65816Layer2Gate() {
  return new W65816Layer2Gate();
}


} // end namespace llvm
