Expose minimal asset cache API.

PiperOrigin-RevId: 803126923
Change-Id: Iabf63802572f7cb86864b5ece85f44b3f15a4ec0

doc/APIreference/functions.rst

@@ -1289,6 +1289,12 @@ Add file to VFS. The directory argument is optional and can be NULL or empty. Re
 
 *Nullable:* ``directory``
 
+
+.. Assetcache:
+
+The asset cache is a mechanism for caching assets (e.g. textures, meshes, etc.) to avoid repeated slow recompilation.
+The following methods provide way to control the capacity of the cache or to disable it altogether.
+
 .. _mj_addBufferVFS:
 
 `mj_addBufferVFS <#mj_addBufferVFS>`__

doc/APIreference/functions_override.rst

@@ -25,6 +25,12 @@ Add file to VFS. The directory argument is optional and can be NULL or empty. Re
 
 *Nullable:* ``directory``
 
+
+.. Assetcache:
+
+The asset cache is a mechanism for caching assets (e.g. textures, meshes, etc.) to avoid repeated slow recompilation.
+The following methods provide way to control the capacity of the cache or to disable it altogether.
+
 .. _Parseandcompile:
 
 The key function here is :ref:`mj_loadXML`. It invokes the built-in parser and compiler, and either returns a pointer to

doc/includes/references.h

@@ -789,6 +789,10 @@ struct mjLROpt_ {                 // options for mj_setLengthRange()
   mjtNum tolrange;                // convergence tolerance (relative to range)
 };
 typedef struct mjLROpt_ mjLROpt;
+struct mjCache_ {                 // asset cache used by the compiler
+  void* impl_;                    // internal pointer to cache
+};
+typedef struct mjCache_ mjCache;
 struct mjVFS_ {                   // virtual file system for loading from memory
   void* impl_;                    // internal pointer to VFS memory
 };
@@ -2999,6 +3003,11 @@ int mj_addFileVFS(mjVFS* vfs, const char* directory, const char* filename);
 int mj_addBufferVFS(mjVFS* vfs, const char* name, const void* buffer, int nbuffer);
 int mj_deleteFileVFS(mjVFS* vfs, const char* filename);
 void mj_deleteVFS(mjVFS* vfs);
+size_t mj_getCacheSize(const mjCache* cache);
+size_t mj_getCacheCapacity(const mjCache* cache);
+size_t mj_setCacheCapacity(mjCache* cache, size_t size);
+mjCache* mj_getCache(void);
+void mj_clearCache(mjCache* cache);
 mjModel* mj_loadXML(const char* filename, const mjVFS* vfs, char* error, int error_sz);
 mjSpec* mj_parseXML(const char* filename, const mjVFS* vfs, char* error, int error_sz);
 mjSpec* mj_parseXMLString(const char* xml, const mjVFS* vfs, char* error, int error_sz);

include/mujoco/mjmodel.h

@@ -450,6 +450,12 @@ struct mjLROpt_ {                 // options for mj_setLengthRange()
 };
 typedef struct mjLROpt_ mjLROpt;
 
+//---------------------------------- mjCache -------------------------------------------------------
+
+struct mjCache_ {                 // asset cache used by the compiler
+  void* impl_;                    // internal pointer to cache
+};
+typedef struct mjCache_ mjCache;
 
 //---------------------------------- mjVFS ---------------------------------------------------------
 

include/mujoco/mujoco.h

@@ -90,6 +90,22 @@ MJAPI int mj_deleteFileVFS(mjVFS* vfs, const char* filename);
 // Delete all files from VFS and deallocates VFS internal memory.
 MJAPI void mj_deleteVFS(mjVFS* vfs);
 
+//------------------------------------ Asset cache -------------------------------------------------
+
+// Get the current size of the asset cache in bytes.
+MJAPI size_t mj_getCacheSize(const mjCache* cache);
+
+// Get the capacity of the asset cache in bytes.
+MJAPI size_t mj_getCacheCapacity(const mjCache* cache);
+
+// 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 the internal asset cache used by the compiler.
+MJAPI mjCache* mj_getCache(void);
+
+// Clear the asset cache.
+MJAPI void mj_clearCache(mjCache* cache);
 
 //---------------------------------- Parse and compile ---------------------------------------------
 

python/mujoco/introspect/functions.py

@@ -130,6 +130,75 @@ FUNCTIONS: Mapping[str, FunctionDecl] = dict([
          ),
          doc='Delete all files from VFS and deallocates VFS internal memory.',
      )),
+    ('mj_getCacheSize',
+     FunctionDecl(
+         name='mj_getCacheSize',
+         return_type=ValueType(name='size_t'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='cache',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjCache', is_const=True),
+                 ),
+             ),
+         ),
+         doc='Get the current size of the asset cache in bytes.',
+     )),
+    ('mj_getCacheCapacity',
+     FunctionDecl(
+         name='mj_getCacheCapacity',
+         return_type=ValueType(name='size_t'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='cache',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjCache', is_const=True),
+                 ),
+             ),
+         ),
+         doc='Get the capacity of the asset cache in bytes.',
+     )),
+    ('mj_setCacheCapacity',
+     FunctionDecl(
+         name='mj_setCacheCapacity',
+         return_type=ValueType(name='size_t'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='cache',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjCache'),
+                 ),
+             ),
+             FunctionParameterDecl(
+                 name='size',
+                 type=ValueType(name='size_t'),
+             ),
+         ),
+         doc='Set the capacity of the asset cache in bytes (0 to disable); returns the new capacity.',  # pylint: disable=line-too-long
+     )),
+    ('mj_getCache',
+     FunctionDecl(
+         name='mj_getCache',
+         return_type=PointerType(
+             inner_type=ValueType(name='mjCache'),
+         ),
+         parameters=(),
+         doc='Get the internal asset cache used by the compiler.',
+     )),
+    ('mj_clearCache',
+     FunctionDecl(
+         name='mj_clearCache',
+         return_type=ValueType(name='void'),
+         parameters=(
+             FunctionParameterDecl(
+                 name='cache',
+                 type=PointerType(
+                     inner_type=ValueType(name='mjCache'),
+                 ),
+             ),
+         ),
+         doc='Clear the asset cache.',
+     )),
     ('mj_loadXML',
      FunctionDecl(
          name='mj_loadXML',

python/mujoco/introspect/structs.py

@@ -85,6 +85,20 @@ STRUCTS: Mapping[str, StructDecl] = dict([
              ),
          ),
      )),
+    ('mjCache',
+     StructDecl(
+         name='mjCache',
+         declname='struct mjCache_',
+         fields=(
+             StructFieldDecl(
+                 name='impl_',
+                 type=PointerType(
+                     inner_type=ValueType(name='void'),
+                 ),
+                 doc='internal pointer to cache',
+             ),
+         ),
+     )),
     ('mjVFS',
      StructDecl(
          name='mjVFS',

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) {

test/user/user_cache_test.cc

@@ -106,7 +106,7 @@ TEST(CacheTest, InsertReplace) {
 // Trim cache based off of access count
 TEST(CacheTest, Limit1) {
   mjCCache cache(kMaxSize);
-  EXPECT_THAT(cache.MaxSize(), kMaxSize);
+  EXPECT_THAT(cache.Capacity(), kMaxSize);
 
   CacheText(cache, "fil.xml", "foo.obj", kText);
   CacheText(cache, "fil.xml", "bar.obj", kText);
@@ -117,7 +117,7 @@ TEST(CacheTest, Limit1) {
   GetCachedText(cache, "file.xml", "bar.obj", kText);
 
   // make max size so cache can hold only one asset
-  cache.SetMaxSize(12);
+  cache.SetCapacity(12);
 
   // foo should still be in cache
   EXPECT_THAT(cache.HasAsset("foo.obj"), NotNull());
@@ -138,7 +138,7 @@ TEST(CacheTest, Limit2) {
   GetCachedText(cache, "file.xml", "bar.obj", kText);
 
   // make max size so cache can hold only one asset
-  cache.SetMaxSize(12);
+  cache.SetCapacity(12);
 
   // foo should be gone because it's older
   EXPECT_THAT(cache.HasAsset("foo.obj"), IsNull());

unity/Runtime/Bindings/MjBindings.cs

@@ -5075,6 +5075,11 @@ public unsafe struct mjLROpt_ {
   public double tolrange;
 }
 
+[StructLayout(LayoutKind.Sequential)]
+public unsafe struct mjCache_ {
+  public void* impl_;
+}
+
 [StructLayout(LayoutKind.Sequential)]
 public unsafe struct _mjVFS
 {
@@ -6334,6 +6339,21 @@ public static unsafe extern int mj_deleteFileVFS(void* vfs, [MarshalAs(Unmanaged
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern void mj_deleteVFS(void* vfs);
 
+[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
+public static unsafe extern UIntPtr mj_getCacheSize(mjCache_* cache);
+
+[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
+public static unsafe extern UIntPtr mj_getCacheCapacity(mjCache_* cache);
+
+[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
+public static unsafe extern UIntPtr mj_setCacheCapacity(mjCache_* cache, UIntPtr size);
+
+[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
+public static unsafe extern mjCache_* mj_getCache();
+
+[DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
+public static unsafe extern void mj_clearCache(mjCache_* cache);
+
 [DllImport("mujoco", CallingConvention = CallingConvention.Cdecl)]
 public static unsafe extern mjModel_* mj_loadXML([MarshalAs(UnmanagedType.LPStr)]string filename, void* vfs, StringBuilder error, int error_sz);