diff --git a/doc/APIreference/functions.rst b/doc/APIreference/functions.rst index 9606e94d..e13ad19d 100644 --- a/doc/APIreference/functions.rst +++ b/doc/APIreference/functions.rst @@ -4459,6 +4459,17 @@ Transform body into a frame. Set user payload, overriding the existing value for the specified key if present. +.. _mjs_setUserValueWithCleanup: + +`mjs_setUserValueWithCleanup <#mjs_setUserValueWithCleanup>`__ +~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ + +.. mujoco-include:: mjs_setUserValueWithCleanup + +Set user payload, overriding the existing value for the specified key if +present. This version differs from mjs_setUserValue in that it takes a +cleanup function that will be called when the user payload is deleted. + .. _mjs_getUserValue: `mjs_getUserValue <#mjs_getUserValue>`__ diff --git a/doc/includes/references.h b/doc/includes/references.h index 6ebc759a..0689e150 100644 --- a/doc/includes/references.h +++ b/doc/includes/references.h @@ -3732,6 +3732,9 @@ const char* mjs_resolveOrientation(double quat[4], mjtByte degree, const char* s const mjsOrientation* orientation); mjsFrame* mjs_bodyToFrame(mjsBody** body); void mjs_setUserValue(mjsElement* element, const char* key, const void* data); +void mjs_setUserValueWithCleanup(mjsElement* element, const char* key, + const void* data, + void (*cleanup)(const void*)); const void* mjs_getUserValue(mjsElement* element, const char* key); void mjs_deleteUserValue(mjsElement* element, const char* key); void mjs_defaultSpec(mjSpec* spec); diff --git a/include/mujoco/mujoco.h b/include/mujoco/mujoco.h index 63bfafcc..5bca71ea 100644 --- a/include/mujoco/mujoco.h +++ b/include/mujoco/mujoco.h @@ -1644,6 +1644,13 @@ MJAPI mjsFrame* mjs_bodyToFrame(mjsBody** body); // Set user payload, overriding the existing value for the specified key if present. MJAPI void mjs_setUserValue(mjsElement* element, const char* key, const void* data); +// Set user payload, overriding the existing value for the specified key if +// present. This version differs from mjs_setUserValue in that it takes a +// cleanup function that will be called when the user payload is deleted. +MJAPI void mjs_setUserValueWithCleanup(mjsElement* element, const char* key, + const void* data, + void (*cleanup)(const void*)); + // Return user payload or NULL if none found. MJAPI const void* mjs_getUserValue(mjsElement* element, const char* key); diff --git a/python/mujoco/introspect/codegen/generate_functions.py b/python/mujoco/introspect/codegen/generate_functions.py index 1966e328..06beeaf2 100644 --- a/python/mujoco/introspect/codegen/generate_functions.py +++ b/python/mujoco/introspect/codegen/generate_functions.py @@ -99,8 +99,13 @@ class MjFunctionVisitor: return ''.join(strings) def visit(self, node: ClangJsonNode) -> None: - if (node.get('kind') == 'FunctionDecl' and - node.get('name', '').startswith('mj')): + # Skip mjs_setUserValueWithCleanup as it's only useful for heap allocated + # objects and doesn't need a python wrapper. + if ( + node.get('kind') == 'FunctionDecl' + and node.get('name', '').startswith('mj') + and node.get('name', '') != 'mjs_setUserValueWithCleanup' + ): func_decl = self._make_function(node) self._functions[func_decl.name] = func_decl diff --git a/src/user/user_api.cc b/src/user/user_api.cc index e103aa8e..41cba2d2 100644 --- a/src/user/user_api.cc +++ b/src/user/user_api.cc @@ -860,15 +860,17 @@ mjsFrame* mjs_bodyToFrame(mjsBody** body) { return &frameC->spec; } - - -// set user payload void mjs_setUserValue(mjsElement* element, const char* key, const void* data) { - mjCBase* baseC = static_cast(element); - baseC->SetUserValue(key, data); + mjs_setUserValueWithCleanup(element, key, data, nullptr); } - +// set user payload +void mjs_setUserValueWithCleanup(mjsElement* element, const char* key, + const void* data, + void (*cleanup)(const void*)) { + mjCBase* baseC = static_cast(element); + baseC->SetUserValue(key, data, cleanup); +} // return user payload or NULL if none found const void* mjs_getUserValue(mjsElement* element, const char* key) { diff --git a/src/user/user_api.h b/src/user/user_api.h index e1b73d86..cd228147 100644 --- a/src/user/user_api.h +++ b/src/user/user_api.h @@ -374,6 +374,11 @@ MJAPI mjsFrame* mjs_bodyToFrame(mjsBody** body); // Set user payload. MJAPI void mjs_setUserValue(mjsElement* element, const char* key, const void* data); +// Set user payload. +MJAPI void mjs_setUserValueWithCleanup(mjsElement* element, const char* key, + const void* data, + void (*cleanup)(const void*)); + // Return user payload or NULL if none found. MJAPI const void* mjs_getUserValue(mjsElement* element, const char* key); diff --git a/src/user/user_objects.cc b/src/user/user_objects.cc index 82853307..51c0f63a 100644 --- a/src/user/user_objects.cc +++ b/src/user/user_objects.cc @@ -784,15 +784,14 @@ void mjCBase::SetFrame(mjCFrame* _frame) { frame = _frame; } - -void mjCBase::SetUserValue(std::string_view key, const void* data) { - user_payload_[std::string(key)] = data; +void mjCBase::SetUserValue(std::string_view key, const void* data, + void (*cleanup)(const void*)) { + user_payload_[std::string(key)] = UserValue(data, cleanup); } - const void* mjCBase::GetUserValue(std::string_view key) { auto found = user_payload_.find(std::string(key)); - return found != user_payload_.end() ? found->second : nullptr; + return found != user_payload_.end() ? found->second.value : nullptr; } diff --git a/src/user/user_objects.h b/src/user/user_objects.h index bfec0797..9cdae417 100644 --- a/src/user/user_objects.h +++ b/src/user/user_objects.h @@ -281,7 +281,8 @@ class mjCBase : public mjCBase_ { } // Set and get user payload - void SetUserValue(std::string_view key, const void* data); + void SetUserValue(std::string_view key, const void* data, + void (*cleanup)(const void*)); const void* GetUserValue(std::string_view key); void DeleteUserValue(std::string_view key); @@ -292,8 +293,44 @@ class mjCBase : public mjCBase_ { // reference count for allowing deleting an attached object int refcount = 1; + // Arbitrary user value that cleans up the data when destroyed. + struct UserValue { + const void* value = nullptr; + void (*cleanup)(const void*) = nullptr; + + UserValue() {} + UserValue(const void* value, void (*cleanup)(const void*)) + : value(value), cleanup(cleanup) {} + UserValue(const UserValue& other) = delete; + UserValue& operator=(const UserValue& other) = delete; + + UserValue(UserValue&& other) : value(other.value), cleanup(other.cleanup) { + other.value = nullptr; + other.cleanup = nullptr; + } + + UserValue& operator=(UserValue&& other) { + if (this != &other) { + if (cleanup && value) { + cleanup(value); + } + value = other.value; + cleanup = other.cleanup; + other.value = nullptr; + other.cleanup = nullptr; + } + return *this; + } + + ~UserValue() { + if (cleanup && value) { + cleanup(value); + } + } + }; + // user payload - std::unordered_map user_payload_; + std::unordered_map user_payload_; }; diff --git a/test/user/user_api_test.cc b/test/user/user_api_test.cc index 4db9f25b..1669ea20 100644 --- a/test/user/user_api_test.cc +++ b/test/user/user_api_test.cc @@ -2845,6 +2845,14 @@ TEST_F(MujocoTest, UserValue) { EXPECT_STREQ(static_cast(payload), data.c_str()); mjs_deleteUserValue(body->element, "key"); EXPECT_THAT(mjs_getUserValue(body->element, "key"), IsNull()); + + std::string* heap_data = new std::string("heap_data"); + mjs_setUserValueWithCleanup( + body->element, "key", heap_data, + [](const void* data) { delete static_cast(data); }); + payload = mjs_getUserValue(body->element, "key"); + EXPECT_STREQ(static_cast(payload)->c_str(), + heap_data->c_str()); mj_deleteSpec(spec); }