feat: release CKKS scheme caches and trim allocator

src/core/include/utils/memory.h

@@ -63,6 +63,27 @@ void MoveAppend(std::vector<X>& dst, std::vector<X>& src) {
  */
 void secure_memset(volatile void* mem, uint8_t c, size_t len);
 
+/**
+ * @brief Ask the C allocator to return as many free pages to the OS as it can.
+ *
+ * Calling the cache-clear APIs (ClearBootstrapPrecom, ClearSchemeSwitchPrecom,
+ * ClearAllCKKSCaches, ReleaseAllContexts) frees memory back to the allocator,
+ * but whether that translates into RSS shrinking is allocator-dependent:
+ * glibc malloc keeps freed pages in its arena, mimalloc and jemalloc have
+ * their own purge policies. This helper calls the allocator's trim/purge
+ * primitive where available so the drop is visible at the OS level.
+ *
+ * Implementations:
+ *   - glibc:    malloc_trim(0)
+ *   - Apple:    malloc_zone_pressure_relief on every default zone
+ *   - MSVC:     _heapmin()
+ *   - others:   no-op
+ *
+ * Safe to call at any time. Returns true if the platform has a trim primitive,
+ * false if the call was a no-op on this build (issue #533).
+ */
+bool TrimAllocator();
+
 }  // namespace lbcrypto
 
 #endif  // LBCRYPTO_UTILS_MEMORY_H

src/core/lib/utils/memory.cpp

@@ -30,6 +30,15 @@
 //==================================================================================
 #include "utils/memory.h"
 
+#if defined(__APPLE__)
+    #include <mach/mach.h>
+    #include <malloc/malloc.h>
+#elif defined(__GLIBC__)
+    #include <malloc.h>
+#elif defined(_MSC_VER)
+    #include <malloc.h>
+#endif
+
 namespace lbcrypto {
 
 void secure_memset(volatile void* mem, uint8_t c, size_t len) {
@@ -38,4 +47,30 @@ void secure_memset(volatile void* mem, uint8_t c, size_t len) {
         *(ptr + i) = c;
 }
 
+bool TrimAllocator() {
+#if defined(__GLIBC__)
+    // Returns non-zero if any memory was released; absence of trim is not an error.
+    malloc_trim(0);
+    return true;
+#elif defined(__APPLE__)
+    // Walk every registered zone and pressure-relieve it. goal=0 asks libmalloc
+    // to return as much as it can without capping.
+    vm_address_t* zones = nullptr;
+    unsigned count      = 0;
+    if (malloc_get_all_zones(mach_task_self(), nullptr, &zones, &count) == KERN_SUCCESS && zones) {
+        for (unsigned i = 0; i < count; ++i) {
+            auto* z = reinterpret_cast<malloc_zone_t*>(zones[i]);
+            if (z)
+                malloc_zone_pressure_relief(z, 0);
+        }
+    }
+    return true;
+#elif defined(_MSC_VER)
+    _heapmin();
+    return true;
+#else
+    return false;
+#endif
+}
+
 }  // namespace lbcrypto

src/pke/examples/ckks-release-memory.cpp

@@ -0,0 +1,121 @@
+//==================================================================================
+// BSD 2-Clause License
+//
+// Copyright (c) 2014-2026, NJIT, Duality Technologies Inc. and other contributors
+//
+// All rights reserved.
+//
+// Author TPOC: [email protected]
+//
+// Redistribution and use in source and binary forms, with or without
+// modification, are permitted provided that the following conditions are met:
+//
+// 1. Redistributions of source code must retain the above copyright notice, this
+//    list of conditions and the following disclaimer.
+//
+// 2. Redistributions in binary form must reproduce the above copyright notice,
+//    this list of conditions and the following disclaimer in the documentation
+//    and/or other materials provided with the distribution.
+//
+// THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+// AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+// IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+// DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+// FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+// DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+// SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+// CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+// OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+// OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+//==================================================================================
+
+/*
+  Releasing CKKS scheme-level memory without destroying the CryptoContext
+  (issue #533).
+
+  CKKS bootstrap caches (FHECKKSRNS::m_bootPrecomMap) and scheme-switch caches
+  (SWITCHCKKSRNS: ccLWE, ccKS, switching keys, decoding-matrix plaintexts) are
+  owned by the scheme, which lives as long as the CryptoContext. Before these
+  APIs, the only way to free them was to drop every CryptoContext handle plus
+  call CryptoContextFactory::ReleaseAllContexts() — if any caller kept a
+  context ref alive, the caches survived.
+
+  This example demonstrates the three common patterns:
+
+    1. Clear one slot's bootstrap precom and reuse the context
+    2. Clear all bootstrap precom except a hot slot count
+    3. Full release + allocator trim at shutdown
+*/
+
+#include "openfhe.h"
+#include "utils/memory.h"
+
+#include <iostream>
+
+using namespace lbcrypto;
+
+static CryptoContext<DCRTPoly> BuildBootstrapContext() {
+    CCParams<CryptoContextCKKSRNS> parameters;
+    SecretKeyDist skDist = UNIFORM_TERNARY;
+    parameters.SetSecretKeyDist(skDist);
+    parameters.SetSecurityLevel(HEStd_NotSet);
+    parameters.SetRingDim(1 << 12);
+
+#if NATIVEINT == 128
+    parameters.SetScalingModSize(78);
+    parameters.SetFirstModSize(89);
+    parameters.SetScalingTechnique(FIXEDAUTO);
+#else
+    parameters.SetScalingModSize(59);
+    parameters.SetFirstModSize(60);
+    parameters.SetScalingTechnique(FLEXIBLEAUTO);
+#endif
+
+    std::vector<uint32_t> levelBudget = {4, 4};
+    uint32_t depth = 10 + FHECKKSRNS::GetBootstrapDepth(levelBudget, skDist);
+    parameters.SetMultiplicativeDepth(depth);
+
+    auto cc = GenCryptoContext(parameters);
+    cc->Enable(PKE);
+    cc->Enable(KEYSWITCH);
+    cc->Enable(LEVELEDSHE);
+    cc->Enable(ADVANCEDSHE);
+    cc->Enable(FHE);
+    return cc;
+}
+
+int main() {
+    auto cc           = BuildBootstrapContext();
+    uint32_t ringDim  = cc->GetRingDimension();
+    uint32_t hotSlots = ringDim / 2;
+    uint32_t tmpSlots = ringDim / 4;
+
+    std::cout << "ring dim = " << ringDim << "\n";
+
+    // Pattern 1: transient slot count, then clear just that entry.
+    cc->EvalBootstrapSetup({4, 4}, {0, 0}, tmpSlots);
+    std::cout << "after transient setup: "
+              << "bootstrap cache present? " << std::boolalpha << cc->HasBootstrapPrecom(tmpSlots) << "\n";
+    cc->ClearBootstrapPrecom(tmpSlots);
+    std::cout << "after per-slot clear:   " << cc->HasBootstrapPrecom(tmpSlots) << "\n";
+
+    // Pattern 2: hot-slot workload with occasional transient bootstraps.
+    cc->EvalBootstrapSetup({4, 4}, {0, 0}, hotSlots);
+    cc->EvalBootstrapSetup({4, 4}, {0, 0}, tmpSlots);
+    std::cout << "with two entries, keep-except hot slots:\n";
+    cc->ClearBootstrapPrecomExcept(hotSlots);
+    std::cout << "  hot still present? " << cc->HasBootstrapPrecom(hotSlots) << "\n";
+    std::cout << "  tmp gone?          " << !cc->HasBootstrapPrecom(tmpSlots) << "\n";
+
+    // Pattern 3: shutdown — release everything and ask the allocator to
+    // return pages to the OS. Useful right before a memory-sensitive phase
+    // of the surrounding application (e.g. forking a child process).
+    CryptoContextFactory<DCRTPoly>::ReleaseAllContextsAndTrim();
+    // After this call, static eval-key maps are empty, every context's
+    // scheme-level caches are released, and the allocator has been asked to
+    // trim. `cc` is still a valid handle but its caches are gone.
+    std::cout << "after ReleaseAllContextsAndTrim:\n";
+    std::cout << "  cc bootstrap caches? " << cc->HasBootstrapPrecom() << "\n";
+
+    return 0;
+}

src/pke/include/cryptocontext.h

@@ -620,6 +620,82 @@ class CryptoContextImpl : public Serializable {
     */
     static void ClearStaticMapsAndVectors();
 
+    /**
+    * @brief Release the scheme-level bootstrap precomputations held by this
+    * CryptoContext (e.g. FHECKKSRNS::m_bootPrecomMap). Lets callers free
+    * this memory without destroying the context itself, since the
+    * existing ClearEval*Keys / ReleaseAllContexts APIs leave scheme-level
+    * caches alive while any user reference keeps the context alive
+    * (issue #533). No-op for schemes without bootstrap caches.
+    */
+    void ClearBootstrapPrecom() noexcept {
+        if (m_scheme)
+            m_scheme->ClearBootstrapPrecom();
+    }
+
+    /**
+    * @brief Release the bootstrap precomputation entry associated with a
+    * specific number of slots. Useful for callers who re-use a CryptoContext
+    * with multiple slot counts and want to free only the stale ones
+    * (issue #533). No-op if no entry exists for the given slot count or the
+    * scheme does not cache per-slot bootstrap data.
+    */
+    void ClearBootstrapPrecom(uint32_t slots) noexcept {
+        if (m_scheme)
+            m_scheme->ClearBootstrapPrecom(slots);
+    }
+
+    /**
+    * @brief Release every bootstrap precomputation entry except the one keyed
+    * by @p keepSlots. Handy for workloads that repeatedly bootstrap a fixed
+    * slot count but occasionally bootstrap transient slot counts — reclaims
+    * the transient caches without invalidating the hot path (issue #533).
+    */
+    void ClearBootstrapPrecomExcept(uint32_t keepSlots) noexcept {
+        if (m_scheme)
+            m_scheme->ClearBootstrapPrecomExcept(keepSlots);
+    }
+
+    /**
+    * @brief Release scheme-switch precomputations (intermediate CKKS/FHEW
+    * contexts, switching keys, and linear-transform plaintexts) held by this
+    * CryptoContext (issue #533). No-op for schemes that do not support
+    * CKKS<->FHEW switching.
+    */
+    void ClearSchemeSwitchPrecom() noexcept {
+        if (m_scheme)
+            m_scheme->ClearSchemeSwitchPrecom();
+    }
+
+    /**
+    * @brief Convenience: release all scheme-level caches held by this
+    * CryptoContext (bootstrap + scheme-switch). This does NOT touch the
+    * static eval-key maps — call ClearEvalMultKeys / ClearEvalAutomorphismKeys
+    * for those (issue #533).
+    */
+    void ClearAllCKKSCaches() noexcept {
+        ClearBootstrapPrecom();
+        ClearSchemeSwitchPrecom();
+    }
+
+    /**
+    * @brief Inspectors for scheme-level caches. Return false for schemes that
+    * do not maintain the corresponding cache, and whenever the cache is empty
+    * (issue #533). Intended for callers that need to reason about memory
+    * without reaching into internal types.
+    */
+    bool HasBootstrapPrecom() const noexcept {
+        return m_scheme && m_scheme->HasBootstrapPrecom();
+    }
+
+    bool HasBootstrapPrecom(uint32_t slots) const noexcept {
+        return m_scheme && m_scheme->HasBootstrapPrecom(slots);
+    }
+
+    bool HasSchemeSwitchPrecom() const noexcept {
+        return m_scheme && m_scheme->HasSchemeSwitchPrecom();
+    }
+
     /**
     * @brief Serializes either all EvalMult keys (if keyTag is empty) or the EvalMult keys for keyTag
     *

src/pke/include/cryptocontextfactory.h

@@ -35,6 +35,7 @@
 #include "cryptocontext-fwd.h"
 #include "lattice/lat-hal.h"
 #include "scheme/scheme-id.h"
+#include "utils/memory.h"
 
 #include <memory>
 #include <string>
@@ -55,19 +56,75 @@ class CryptoParametersBase;
 template <typename Element>
 class CryptoContextFactory {
     static std::vector<CryptoContext<Element>> AllContexts;
+    // When AutoReleaseMode is true, newly built contexts are NOT tracked in
+    // AllContexts. Dropping the caller's last handle then runs the context
+    // destructor immediately, which releases every scheme-level cache and
+    // static eval-key ref without needing a manual ReleaseAllContexts call
+    // (issue #533). Default is false to preserve pre-fix semantics.
+    static bool s_autoReleaseMode;
 
 protected:
     static CryptoContext<Element> FindContext(std::shared_ptr<CryptoParametersBase<Element>> params,
                                               std::shared_ptr<SchemeBase<Element>> scheme);
     static void AddContext(CryptoContext<Element>);
 
 public:
+    /**
+     * Enable or disable AutoRelease mode. When enabled, the factory no longer
+     * holds a strong reference to newly created contexts; the caller becomes
+     * the sole owner and cache memory is reclaimed automatically when they
+     * drop the last reference. Existing tracked contexts are released when
+     * the switch is flipped so the behavior is consistent (issue #533).
+     *
+     * Trade-off: FindContext / GetContext deduplication only inspects the
+     * tracked list, so two identical parameter sets created in AutoRelease
+     * mode will produce two distinct contexts instead of aliasing. Enable
+     * this only if you don't rely on factory-level deduplication.
+     */
+    static void SetAutoReleaseMode(bool enabled) {
+        if (enabled && !AllContexts.empty()) {
+            // Flush tracked contexts so the semantics are consistent: from
+            // now on, no context is tracked by the factory.
+            ReleaseAllContexts();
+        }
+        s_autoReleaseMode = enabled;
+    }
+
+    static bool IsAutoReleaseMode() noexcept {
+        return s_autoReleaseMode;
+    }
+
+
     static void ReleaseAllContexts() {
+        // Drop scheme-level caches (bootstrap + scheme-switch) on every live
+        // context before clearing the static maps. This ensures memory held
+        // by FHECKKSRNS::m_bootPrecomMap and SWITCHCKKSRNS state is freed
+        // even when users still hold CryptoContext shared pointers after
+        // this call returns (issue #533).
+        for (auto& cc : AllContexts) {
+            if (cc)
+                cc->ClearAllCKKSCaches();
+        }
         if (AllContexts.size() > 0)
             AllContexts[0]->ClearStaticMapsAndVectors();
         AllContexts.clear();
     }
 
+    /**
+     * Release every live context's scheme-level caches, clear the static key
+     * maps, and ask the C allocator to return freed pages to the OS. Behaves
+     * like ReleaseAllContexts() followed by TrimAllocator() (issue #533).
+     *
+     * Use this when heap usage drops after clearing but RSS doesn't — the C
+     * allocator is holding freed pages in its arena. On platforms without a
+     * trim primitive (non-glibc Linux, FreeBSD, WASM), this degrades into the
+     * plain ReleaseAllContexts() path.
+     */
+    static void ReleaseAllContextsAndTrim() {
+        ReleaseAllContexts();
+        TrimAllocator();
+    }
+
     static int GetContextCount() {
         return AllContexts.size();
     }
@@ -89,6 +146,9 @@ class CryptoContextFactory {
 template <>
 std::vector<CryptoContext<DCRTPoly>> CryptoContextFactory<DCRTPoly>::AllContexts;
 
+template <>
+bool CryptoContextFactory<DCRTPoly>::s_autoReleaseMode;
+
 }  // namespace lbcrypto
 
 #endif