Expose minimal asset cache API.

PiperOrigin-RevId: 803126923
Change-Id: Iabf63802572f7cb86864b5ece85f44b3f15a4ec0

src/user/user_api.cc

@@ -1730,22 +1730,63 @@ const void* mjs_getPluginAttributes(const mjsPlugin* plugin) {
 
 // -------------------------- GLOBAL ASSET CACHE -------------------------------
 
-void mj_setCacheSize(mjCache cache, std::size_t size) {
-  mjCCache* ccache = reinterpret_cast<mjCCache*>(cache);
-  if (ccache) {
-    ccache->SetMaxSize(size);
+// get the capacity of the asset cache in bytes
+size_t mj_getCacheCapacity(const mjCache* cache) {
+  if (cache) {
+    const mjCCache* ccache = reinterpret_cast<const mjCCache*>(cache->impl_);
+    if (ccache) {
+      return ccache->Capacity();
+    }
   }
+  return 0;
 }
 
 
+// set the capacity of the asset cache in bytes (0 to disable)
+size_t mj_setCacheCapacity(mjCache* cache, size_t size) {
+  if (cache) {
+    mjCCache* ccache = reinterpret_cast<mjCCache*>(cache->impl_);
+    if (ccache) {
+      ccache->SetCapacity(size);
+      return ccache->Capacity();
+    }
+  }
+  return 0;
+}
 
-mjCache mj_globalCache() {
+
+// get the current size of the asset cache in bytes
+size_t mj_getCacheSize(const mjCache* cache) {
+  if (cache) {
+    const mjCCache* ccache = reinterpret_cast<const mjCCache*>(cache->impl_);
+    if (ccache) {
+      return ccache->Size();
+    }
+  }
+  return 0;
+}
+
+
+// clear the asset cache
+void mj_clearCache(mjCache* cache) {
+  if (cache) {
+    mjCCache* ccache = reinterpret_cast<mjCCache*>(cache->impl_);
+    if (ccache) {
+      ccache->Reset();
+    }
+  }
+}
+
+// get the internal asset cache used by the compiler
+mjCache* mj_getCache() {
   // mjCCache is not trivially destructible and so the global cache needs to
   // allocated on the heap
   if constexpr (kGlobalCacheSize != 0) {
     static mjCCache* cache = new(std::nothrow) mjCCache(kGlobalCacheSize);
-    return (mjCache) cache;
+    static mjCache cache_cwrapper;
+    cache_cwrapper.impl_ = cache;
+    return cache->Capacity() > 0 ? &cache_cwrapper : nullptr;
   } else {
-    return NULL;
+    return nullptr;
   }
 }

src/user/user_api.h

@@ -505,13 +505,20 @@ MJAPI void mjs_defaultPlugin(mjsPlugin* plugin);
 
 //---------------------------------- Compiler cache ------------------------------------------------
 
-typedef struct mjCache_* mjCache;
+// Get the capacity of the asset cache in bytes.
+MJAPI size_t mj_getCacheCapacity(const mjCache* cache);
 
-// Set the size of the cache in bytes.
-MJAPI void mj_setCacheSize(mjCache cache, size_t size);
+// Set the capacity of the asset cache in bytes (0 to disable); returns the new capacity.
+MJAPI size_t mj_setCacheCapacity(mjCache* cache, size_t size);
 
-// Get internal global cache context.
-MJAPI mjCache mj_globalCache(void);
+// Get the current size of the asset cache in bytes.
+MJAPI size_t mj_getCacheSize(const mjCache* cache);
+
+// Clear the asset cache.
+MJAPI void mj_clearCache(mjCache* cache);
+
+// Get the internal asset cache used by the compiler.
+MJAPI mjCache* mj_getCache(void);
 
 #ifdef __cplusplus
 }  // extern "C"

src/user/user_cache.cc

@@ -41,9 +41,9 @@ mjCAsset mjCAsset::Copy(const mjCAsset& other) {
 // sets the total maximum size of the cache in bytes
 // low-priority cached assets will be dropped to make the new memory
 // requirement
-void mjCCache::SetMaxSize(std::size_t size) {
+void mjCCache::SetCapacity(std::size_t size) {
   std::lock_guard<std::mutex> lock(mutex_);
-  max_size_ = size;
+  capacity_ = size;
   Trim();
 }
 
@@ -69,7 +69,7 @@ bool mjCCache::Insert(const std::string& modelname, const mjResource *resource,
   std::lock_guard<std::mutex> lock(mutex_);
 
   // check if asset is too large to fit in the cache
-  if ((size_ + size > max_size_) &&
+  if ((size_ + size > capacity_) &&
       lookup_.find(resource->name) == lookup_.end()) {
     return false;
   }
@@ -78,7 +78,7 @@ bool mjCCache::Insert(const std::string& modelname, const mjResource *resource,
   mjCAsset* asset_ptr = &(it->second);
 
   if (!inserted) {
-    if (size_ - asset_ptr->BytesCount() + size > max_size_) {
+    if (size_ - asset_ptr->BytesCount() + size > capacity_) {
       return false;
     }
     models_[modelname].insert(asset_ptr);  // add it for the model
@@ -164,9 +164,9 @@ void mjCCache::Reset() {
 
 
 
-std::size_t mjCCache::MaxSize() const {
+std::size_t mjCCache::Capacity() const {
   std::lock_guard<std::mutex> lock(mutex_);
-  return max_size_;
+  return capacity_;
 }
 
 
@@ -218,7 +218,7 @@ void mjCCache::Delete(mjCAsset* asset, const std::string& skip) {
 
 // trims out data to meet memory requirements
 void mjCCache::Trim() {
-  while (size_ > max_size_) {
+  while (size_ > capacity_) {
     Delete(*entries_.begin());
   }
 }

src/user/user_cache.h

@@ -117,7 +117,7 @@ struct mjCAssetCompare {
 // the class container for a thread-safe asset cache
 class mjCCache {
  public:
-  explicit mjCCache(std::size_t size)  :  max_size_(size) {}
+  explicit mjCCache(std::size_t size)  :  capacity_(size) {}
 
   // move only
   mjCCache(mjCCache&& other) = delete;
@@ -125,10 +125,10 @@ class mjCCache {
   mjCCache(const mjCCache& other) = delete;
   mjCCache& operator=(const mjCCache& other) = delete;
 
-  // sets the total maximum size of the cache in bytes
+  // sets the capacity of the cache in bytes
   // low-priority cached assets will be dropped to make the new memory
   // requirement
-  void SetMaxSize(std::size_t size);
+  void SetCapacity(std::size_t size);
 
   // returns the corresponding timestamp, if the given asset is stored in
   // the cache
@@ -156,7 +156,7 @@ class mjCCache {
   void Reset();
 
   // accessors
-  std::size_t MaxSize() const;
+  std::size_t Capacity() const;
   std::size_t Size() const;
 
  private:
@@ -170,7 +170,7 @@ class mjCCache {
   mutable std::mutex mutex_;
   std::size_t insert_num_ = 0;  // a running counter of assets being inserted
   std::size_t size_ = 0;        // current size of the cache in bytes
-  std::size_t max_size_ = 0;    // max size of the cache in bytes
+  std::size_t capacity_ = 0;    // capacity of the cache in bytes
 
   // internal constant look up table for assets
   std::unordered_map<std::string, mjCAsset> lookup_;

src/user/user_mesh.cc

@@ -688,7 +688,7 @@ void mjCMesh::TryCompile(const mjVFS* vfs) {
   bool fromCache = false;
   CopyFromSpec();
   visual_ = true;
-  mjCCache *cache = reinterpret_cast<mjCCache*>(mj_globalCache());
+  mjCCache *cache = reinterpret_cast<mjCCache*>(mj_getCache()->impl_);
 
   // load file
   if (!file_.empty()) {

src/user/user_objects.cc

@@ -77,7 +77,7 @@ PNGImage PNGImage::Load(const mjCBase* obj, mjResource* resource,
                         LodePNGColorType color_type) {
   PNGImage image;
   image.color_type_ = color_type;
-  mjCCache *cache = reinterpret_cast<mjCCache*>(mj_globalCache());
+  mjCCache *cache = reinterpret_cast<mjCCache*>(mj_getCache()->impl_);
 
   // cache callback
   auto callback = [&image](const void* data) {