diff --git a/lib/Differentiator/ReverseModeVisitor.cpp b/lib/Differentiator/ReverseModeVisitor.cpp index 3c35506c2..086a263d9 100644 --- a/lib/Differentiator/ReverseModeVisitor.cpp +++ b/lib/Differentiator/ReverseModeVisitor.cpp @@ -237,22 +237,78 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { return atomicAddCall; } + namespace { + // Recognises a C heap-memory builtin call and centralises the invariants + // reverse mode must preserve for it, so all memory-op reasoning goes through + // one place instead of ad-hoc name checks scattered across the visitor. + class AllocCallInfo { + public: + enum class Kind { None, Malloc, Calloc, Realloc, Free }; + + AllocCallInfo() = default; + + // Recognise E as a memory builtin; Kind::None if it is not one. Strips the + // C-style cast that wraps the call (e.g. `(double*)realloc(...)`), so + // IgnoreParenCasts, not IgnoreParenImpCasts, is required here. + [[nodiscard]] static AllocCallInfo recognize(clang::Expr* E) { + auto* CE = llvm::dyn_cast_or_null( + E ? E->IgnoreParenCasts() : nullptr); + if (!CE) + return {}; + const clang::FunctionDecl* FD = CE->getDirectCallee(); + // getName() asserts on non-identifier names (operators, constructors), + // which vector/STL code produces; the builtins are plain identifiers. + if (!FD || !FD->getDeclName().isIdentifier()) + return {}; + Kind k = llvm::StringSwitch(FD->getName()) + .Case("malloc", Kind::Malloc) + .Case("calloc", Kind::Calloc) + .Case("realloc", Kind::Realloc) + .Case("free", Kind::Free) + .Default(Kind::None); + return AllocCallInfo(k, CE); + } + + // The number-of-bytes operand that a following memset must zero: + // malloc(n) -> n, realloc(p, n) -> n. calloc self-zeroes and needs no + // memset, so it (and free/none) report null here. + [[nodiscard]] clang::Expr* memsetByteSize() const { + switch (m_Kind) { + case Kind::Malloc: + return m_Call->getArg(0); + case Kind::Realloc: + return m_Call->getArg(1); + default: + return nullptr; + } + } + + // True for an in-place `p = realloc(p, n)`: the LHS is realloc's own + // pointer argument. Only then may the reallocated pointer be kept across + // the call (realloc frees the old block, so a saved pointer would dangle). + [[nodiscard]] bool isInPlaceRealloc(const clang::Expr* LHS) const { + if (m_Kind != Kind::Realloc || m_Call->getNumArgs() == 0) + return false; + const auto* LDRE = + llvm::dyn_cast(LHS->IgnoreParenCasts()); + const auto* ArgDRE = llvm::dyn_cast( + m_Call->getArg(0)->IgnoreParenCasts()); + return LDRE && ArgDRE && LDRE->getDecl() == ArgDRE->getDecl(); + } + + private: + AllocCallInfo(Kind k, clang::CallExpr* c) : m_Kind(k), m_Call(c) {} + Kind m_Kind = Kind::None; + clang::CallExpr* m_Call = nullptr; + }; + } // namespace + // Both LHS (the memset destination) and the size expression are cloned here: // callers pass the derived pointer they also assign to and the malloc/realloc // size they also pass to that call, so cloning keeps the memset a distinct // subtree. Expr* ReverseModeVisitor::CheckAndBuildCallToMemset(Expr* LHS, Expr* RHS) { - Expr* size = nullptr; - if (auto* callExpr = dyn_cast_or_null(RHS)) - if (auto* declRef = - dyn_cast(callExpr->getCallee()->IgnoreImpCasts())) - if (auto* FD = dyn_cast(declRef->getDecl())) { - if (FD->getNameAsString() == "malloc") - size = callExpr->getArg(0); - else if (FD->getNameAsString() == "realloc") - size = callExpr->getArg(1); - } - + Expr* size = AllocCallInfo::recognize(RHS).memsetByteSize(); if (size) { llvm::SmallVector args = { CloneNode(LHS), getZeroInit(m_Context.IntTy), CloneNode(size)}; @@ -3174,18 +3230,9 @@ Expr* ReverseModeVisitor::getStdInitListSizeExpr(const Expr* E) { // in-place realloc is not handled -- its reverse sweep may read the // freed tail out of bounds, as it did before this change; no test // exercises it. - bool isReallocAssignment = false; - if (opCode == BO_Assign) - if (auto* RCall = dyn_cast(R->IgnoreParenCasts())) - if (const FunctionDecl* RFD = RCall->getDirectCallee()) - if (RFD->getNameAsString() == "realloc" && RCall->getNumArgs()) { - // Only in-place: the LHS must be realloc's own pointer argument. - const auto* LDRE = dyn_cast(L->IgnoreParenCasts()); - const auto* ArgDRE = - dyn_cast(RCall->getArg(0)->IgnoreParenCasts()); - isReallocAssignment = - LDRE && ArgDRE && LDRE->getDecl() == ArgDRE->getDecl(); - } + bool isReallocAssignment = + opCode == BO_Assign && + AllocCallInfo::recognize(R).isInPlaceRealloc(L); // Store the value of the LHS of the assignment in the forward pass // and restore it in the reverse pass