Adds WrenchesCache

JuliaRobotics · Oct 22, 2024 · 9e47e3d · 9e47e3d
1 parent 7cd36b9
commit 9e47e3d
Show file tree

Hide file tree

Showing 2 changed files with 27 additions and 1 deletion.
diff --git a/src/RigidBodyDynamics.jl b/src/RigidBodyDynamics.jl
@@ -21,7 +21,8 @@ export
     DynamicsResult,
     StateCache,
     DynamicsResultCache,
-    SegmentedVectorCache
+    SegmentedVectorCache,
+    WrenchesCache
 
 # specific joint types
 export

diff --git a/src/caches.jl b/src/caches.jl
@@ -107,3 +107,28 @@ end
 @inline function makevalue(c::SegmentedVectorCache{K, KeyRange}, ::Type{T}) where {K, T, KeyRange}
     SegmentedVector{K, T, KeyRange}(Vector{T}(undef, c.length), c.ranges)
 end
+
+"""
+$(TYPEDEF)
+
+A container that manages the creation and storage of `BodyDict{Wrench{T}}`.
+Similar to [`StateCache`](@ref).
+"""
+struct WrenchesCache{M} <: AbstractTypeDict
+    mechanism::Mechanism{M}
+    keys::Vector{UInt}
+    values::Vector  # Can also be `Any`
+end
+
+Base.show(io::IO, ::WrenchesCache{M}) where {M} = print(io, "WrenchesCache{$M}(…)")
+
+WrenchesCache(mechanism::Mechanism{M}) where {M} = WrenchesCache{M}(mechanism, [], [])
+
+@inline function valuetype(::Type{WrenchesCache{M}}, ::Type{T}) where {M,T}
+    BodyDict{Wrench{T}}
+end
+
+@inline function makevalue(c::WrenchesCache, ::Type{T}) where {T}
+    rootframe = root_frame(c.mechanism)
+    BodyDict{Wrench{T}}(b => zero(Wrench{T}, rootframe) for b in bodies(c.mechanism))
+end