Undo a shrinking in-place realloc in the reverse sweep.

include/clad/Differentiator/BuiltinDerivatives.h

@@ -12,6 +12,8 @@
 
 #include <algorithm>
 #include <cmath>
+#include <cstdlib>
+#include <cstring>
 #include <functional>
 
 #define elidable_reverse_forw __attribute__((annotate("elidable_reverse_forw")))
@@ -1600,6 +1602,30 @@ void constructor_pullback(ValueAndPushforward<T, U> rhs,
 }
 } // namespace class_functions
 } // namespace custom_derivatives
+
+// Reverse-mode helper for an in-place realloc: resize `ptr` back to the byte
+// size it had before the forward realloc. When that regrows a buffer the
+// forward pass shrank, zero the re-grown tail for adjoint buffers so fresh
+// derivatives start at 0; primal buffers pass zeroGrownTail=false, their tail
+// being overwritten by value restores. Returns the possibly-moved pointer.
+// NOLINTBEGIN(cppcoreguidelines-no-malloc)
+// NOLINTBEGIN(cppcoreguidelines-owning-memory)
+// NOLINTBEGIN(cppcoreguidelines-pro-bounds-pointer-arithmetic)
+inline void* reverse_realloc(void* ptr, size_t oldBytes, size_t newBytes,
+                             bool zeroGrownTail) {
+  void* resized = ::realloc(ptr, oldBytes);
+  if (!resized) // realloc failed: keep the original buffer instead of leaking
+                // it
+    return ptr;
+  ptr = resized;
+  if (zeroGrownTail && oldBytes > newBytes)
+    ::memset(static_cast<char*>(ptr) + newBytes, 0, oldBytes - newBytes);
+  return ptr;
+}
+// NOLINTEND(cppcoreguidelines-pro-bounds-pointer-arithmetic)
+// NOLINTEND(cppcoreguidelines-owning-memory)
+// NOLINTEND(cppcoreguidelines-no-malloc)
+
 } // namespace clad
 
   // FIXME: These math functions depend on promote_2 just like pow:

include/clad/Differentiator/ReverseModeVisitor.h

@@ -591,6 +591,18 @@ namespace clad {
     /// @returns The call to memset if the condition is met, otherwise nullptr.
     clang::Expr* CheckAndBuildCallToMemset(clang::Expr* LHS, clang::Expr* RHS);
 
+    /// The byte-size operand of a malloc/calloc/realloc call, in the primal's
+    /// terms: malloc(n) -> n, calloc(k, s) -> k*s, realloc(p, n) -> n. Returns
+    /// null if `allocExpr` is not one of those. Used to track allocation sizes
+    /// so a realloc can be undone in the reverse sweep.
+    clang::Expr* buildAllocByteSize(clang::Expr* allocExpr);
+
+    /// Pointer variables that are the target of an in-place `p = realloc(p, n)`
+    /// somewhere in the function -- only these need an allocation-size shadow;
+    /// one for any other allocated pointer would be dead. Filled by a pre-scan
+    /// of the body before differentiation, then read at each allocation.
+    llvm::SmallPtrSet<const clang::VarDecl*, 4> m_InPlaceReallocPtrs;
+
     static DeclDiff<clang::StaticAssertDecl>
     DifferentiateStaticAssertDecl(const clang::StaticAssertDecl* SAD);
 

include/clad/Differentiator/VisitorBase.h

@@ -242,6 +242,11 @@ namespace clad {
     struct AdjointInfo {
       clang::VarDecl* Decl = nullptr;
       enum WrapKind : std::uint8_t { Plain, Deref, ParenDeref } Wrap = Plain;
+      /// For a pointer allocated with malloc/calloc/realloc, the `size_t`
+      /// shadow holding its current allocation size in bytes; null otherwise.
+      /// Reverse mode uses it to undo an in-place realloc. See
+      /// ReverseModeVisitor's VisitBinaryOperator/DifferentiateVarDecl.
+      clang::VarDecl* AllocSize = nullptr;
     };
     /// Map used to keep track of variable declarations and match them
     /// with their derivatives.

lib/Differentiator/ReverseModeVisitor.cpp

@@ -269,6 +269,9 @@
       return AllocCallInfo(k, CE);
     }
 
+    [[nodiscard]] Kind getKind() const { return m_Kind; }
+    [[nodiscard]] clang::CallExpr* getCall() const { return m_Call; }
+
     // 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.
@@ -318,6 +321,21 @@
     return nullptr;
   }
 
+  Expr* ReverseModeVisitor::buildAllocByteSize(Expr* allocExpr) {
+    AllocCallInfo ac = AllocCallInfo::recognize(allocExpr);
+    CallExpr* call = ac.getCall();
+    switch (ac.getKind()) {
+    case AllocCallInfo::Kind::Malloc:
+      return call->getArg(0);
+    case AllocCallInfo::Kind::Realloc:
+      return call->getArg(1);
+    case AllocCallInfo::Kind::Calloc:
+      return BuildOp(BO_Mul, call->getArg(0), call->getArg(1));
+    default:
+      return nullptr;
+    }
+  }
+
   ReverseModeVisitor::ReverseModeVisitor(DerivativeBuilder& builder,
                                          const DiffRequest& request)
       : VisitorBase(builder, request) {}
@@ -552,6 +570,26 @@
       }
     }
 
+    // Pre-scan the body for in-place `p = realloc(p, n)`. Only those pointers
+    // get an allocation-size shadow at their malloc/calloc/realloc; a shadow
+    // for any other allocated pointer would be dead code.
+    for (llvm::SmallVector<const Stmt*, 32> worklist = {m_DiffReq->getBody()};
+         !worklist.empty();) {
+      const Stmt* S = worklist.pop_back_val();
+      if (!S)
+        continue;
+      if (const auto* BO = dyn_cast<BinaryOperator>(S))
+        if (BO->getOpcode() == BO_Assign) {
+          AllocCallInfo ac = AllocCallInfo::recognize(BO->getRHS());
+          if (ac.isInPlaceRealloc(BO->getLHS()))
+            m_InPlaceReallocPtrs.insert(cast<VarDecl>(
+                cast<DeclRefExpr>(BO->getLHS()->IgnoreParenCasts())
+                    ->getDecl()));
+        }
+      for (const Stmt* C : S->children())
+        worklist.push_back(C);
+    }
+
     // Start the visitation process which outputs the statements in the
     // current block.
     StmtDiff BodyDiff = Visit(m_DiffReq->getBody());
@@ -3225,15 +3263,60 @@
       // For an in-place `p = realloc(p, sz)`, realloc frees the old block, so
       // saving p to restore it in the reverse sweep would dangle (and be
       // double-freed at cleanup). Keep the reallocated pointer instead, for
-      // both p and its adjoint _d_p. This is sound for a growing or same-size
-      // realloc, where every reverse access stays in bounds. A shrinking
-      // 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.
+      // both p and its adjoint _d_p.
       bool isReallocAssignment =
           opCode == BO_Assign &&
           AllocCallInfo::recognize(R).isInPlaceRealloc(L);
 
+      // For an in-place realloc whose buffer size we tracked at allocation,
+      // undo the resize in the reverse sweep: restore p and its adjoint _d_p to
+      // their pre-realloc byte size so statements that ran before the realloc
+      // access valid indices again (clad::reverse_realloc also zeroes the
+      // re-grown adjoint tail). Without a tracked size we fall back to keeping
+      // the reallocated pointer, which is correct only for grow/same-size.
+      if (isReallocAssignment) {
+        const auto* pDRE = cast<DeclRefExpr>(L->IgnoreParenCasts());
+        // Resolve the pointer to its adjoint record via the primal clone (the
+        // key m_Variables uses), then read the allocation-size shadow off it.
+        auto replIt = m_DeclReplacements.find(cast<VarDecl>(pDRE->getDecl()));
+        auto varIt = replIt == m_DeclReplacements.end()
+                         ? m_Variables.end()
+                         : m_Variables.find(replIt->second);
+        if (varIt != m_Variables.end() && varIt->second.AllocSize) {
+          VarDecl* sizeVar = varIt->second.AllocSize;
+          Expr* newBytes = AllocCallInfo::recognize(R).getCall()->getArg(1);
+          // Capture both sizes in the forward pass. The reverse sweep must use
+          // the sizes as they were at the realloc, but the operands may refer
+          // to variables whose values differ by the time the sweep runs, so
+          // neither can be re-evaluated there. oldBytes reads the shadow before
+          // it advances; newBytes is stored once and reused for the update.
+          Expr* oldBytes =
+              StoreAndRef(BuildDeclRef(sizeVar), direction::forward, "_t",
+                          /*forceDeclCreation=*/true);
+          Expr* newBytesRef = StoreAndRef(Clone(newBytes), direction::forward,
+                                          "_t", /*forceDeclCreation=*/true);
+          addToCurrentBlock(
+              BuildOp(BO_Assign, BuildDeclRef(sizeVar), Clone(newBytesRef)),
+              direction::forward);
+          auto buildReverseRealloc = [&](Expr* ptr, bool zeroTail) {
+            Expr* zt = new (m_Context)
+                CXXBoolLiteralExpr(zeroTail, m_Context.BoolTy, noLoc);
+            llvm::SmallVector<Expr*, 4> args = {Clone(ptr), Clone(oldBytes),
+                                                Clone(newBytesRef), zt};
+            Expr* call = GetFunctionCall("reverse_realloc", "clad", args);
+            // reverse_realloc returns void*; cast back to the pointer's type.
+            TypeSourceInfo* TSI =
+                m_Context.getTrivialTypeSourceInfo(ptr->getType(), noLoc);
+            call = m_Sema.BuildCStyleCastExpr(noLoc, TSI, noLoc, call).get();
+            return BuildOp(BO_Assign, Clone(ptr), call);
+          };
+          addToCurrentBlock(buildReverseRealloc(L, /*zeroTail=*/false),
+                            direction::reverse);
+          addToCurrentBlock(buildReverseRealloc(ResultRef, /*zeroTail=*/true),
+                            direction::reverse);
+        }
+      }
+
       // Store the value of the LHS of the assignment in the forward pass
       // and restore it in the reverse pass
       if (m_DiffReq.shouldBeRecorded(L) && !isReallocAssignment) {
@@ -3830,6 +3913,18 @@
                     BuildDeclRef(VDDerived),
                     VDDerived->getInit()->IgnoreCasts()))
               memsetCalls.push_back(memsetCall);
+            // Track this pointer's allocation size in bytes so an in-place
+            // realloc of it can be undone in the reverse sweep. The shadow is
+            // recorded on the pointer's adjoint entry, keyed like every other
+            // m_Variables record by the primal clone.
+            if (m_InPlaceReallocPtrs.contains(VD))
+              if (Expr* bytes = buildAllocByteSize(VDDerived->getInit())) {
+                VarDecl* sizeVar = BuildVarDecl(
+                    m_Context.getSizeType(),
+                    "_" + VD->getNameAsString() + "_size", Clone(bytes));
+                memsetCalls.push_back(BuildDeclStmt(sizeVar));
+                m_Variables[VDDiff.getDecl()].AllocSize = sizeVar;
+              }
           }
         }
       } else if (auto* SAD = dyn_cast<StaticAssertDecl>(D)) {

test/Gradient/Pointers.C

@@ -331,31 +331,39 @@ double cStyleMemoryAlloc(double x, size_t n) {
 // CHECK-NEXT:     double _t0 = t->x;
 // CHECK-NEXT:     t->x = x;
 // CHECK-NEXT:     double *_d_p = (double *)calloc(1, sizeof(double));
+// CHECK-NEXT:     {{__size_t|size_t|unsigned long long|unsigned long|unsigned int}} _p_size0 = 1 * sizeof(double);
 // CHECK-NEXT:     double *p = (double *)calloc(1, sizeof(double));
 // CHECK-NEXT:     double _t1 = *p;
 // CHECK-NEXT:     *p = x;
 // CHECK-NEXT:     double _d_res = 0.;
 // CHECK-NEXT:     double res = t->x + *p;
+// CHECK-NEXT:     {{__size_t|size_t|unsigned long long|unsigned long|unsigned int}} _t2 = _p_size0;
+// CHECK-NEXT:     {{__size_t|size_t|unsigned long long|unsigned long|unsigned int}} _t3 = 2 * sizeof(double);
+// CHECK-NEXT:     _p_size0 = _t3;
 // CHECK-NEXT:     _d_p = (double *)realloc(_d_p, 2 * sizeof(double));
 // CHECK-NEXT:     memset(_d_p, 0, 2 * sizeof(double));
 // CHECK-NEXT:     p = (double *)realloc(p, 2 * sizeof(double));
-// CHECK-NEXT:     double _t2 = p[1];
+// CHECK-NEXT:     double _t4 = p[1];
 // CHECK-NEXT:     p[1] = 2 * x;
-// CHECK-NEXT:     double _t3 = res;
+// CHECK-NEXT:     double _t5 = res;
 // CHECK-NEXT:     res += p[1];
 // CHECK-NEXT:     _d_res += 1;
 // CHECK-NEXT:     {
-// CHECK-NEXT:         res = _t3;
+// CHECK-NEXT:         res = _t5;
 // CHECK-NEXT:         double _r_d3 = _d_res;
 // CHECK-NEXT:         _d_p[1] += _r_d3;
 // CHECK-NEXT:     }
 // CHECK-NEXT:     {
-// CHECK-NEXT:         p[1] = _t2;
+// CHECK-NEXT:         p[1] = _t4;
 // CHECK-NEXT:         double _r_d2 = _d_p[1];
 // CHECK-NEXT:         _d_p[1] = 0.;
 // CHECK-NEXT:         *_d_x += 2 * _r_d2;
 // CHECK-NEXT:     }
 // CHECK-NEXT:     {
+// CHECK-NEXT:         p = (double *)clad::reverse_realloc(p, _t2, _t3, false);
+// CHECK-NEXT:         _d_p = (double *)clad::reverse_realloc(_d_p, _t2, _t3, true);
+// CHECK-NEXT:     }
+// CHECK-NEXT:     {
 // CHECK-NEXT:         _d_t->x += _d_res;
 // CHECK-NEXT:         *_d_p += _d_res;
 // CHECK-NEXT:     }

test/Gradient/ReallocShrink.C

@@ -0,0 +1,29 @@
+// RUN: %cladclang %s -I%S/../../include -oReallocShrink.out 2>&1 | %filecheck %s
+// RUN: ./ReallocShrink.out | %filecheck_exec %s
+
+// Shrinking in-place realloc: the reverse sweep must resize the buffer (and its
+// adjoint) back to the pre-realloc size so the pre-realloc reads stay in bounds
+// and the adjoint accumulates correctly. Runs Memcheck-clean under the valgrind
+// CI row.
+
+#include "clad/Differentiator/Differentiator.h"
+#include <cstdlib>
+
+// res = (x + x*x) + x  ->  dres/dx = 2 + 2x ; at x = 3 -> 8
+double shrink_realloc(double x) {
+  double* p = (double*)malloc(2 * sizeof(double));
+  p[0] = x;
+  p[1] = x * x;
+  double res = p[0] + p[1];
+  p = (double*)realloc(p, 1 * sizeof(double)); // shrink 2 -> 1
+  res += p[0];
+  free(p);
+  return res;
+}
+
+int main() {
+  auto g = clad::gradient<clad::opts::disable_tbr>(shrink_realloc, "x");
+  double dx = 0;
+  g.execute(3.0, &dx);
+  printf("{%.2f}\n", dx); // CHECK-EXEC: {8.00}
+}