Skip to content

Commit e9214eb

Browse files
committed
Eigen caster: fix lifetime issue when casting std::optional<Eigen::Ref<const T>> with conversion
The conversion process initially produces an Eigen::Ref that owns its storage, but copying or moving it produces an Eigen::Ref that refers non-owningly to the storage owned by the first Ref. The first Ref is destroyed before the result can be used, leading to a dangling reference. Fix by having the Eigen caster return a type that encapsulates the constructor argument of Ref rather than a Ref itself, so that the owning Ref can be constructed in-place in its final location.
1 parent d4b245a commit e9214eb

File tree

4 files changed

+81
-5
lines changed

4 files changed

+81
-5
lines changed

include/nanobind/eigen/dense.h

Lines changed: 57 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -381,6 +381,7 @@ struct type_caster<Eigen::Ref<T, Options, StrideType>,
381381
// Extended version taking arbitrary strides
382382
using DMap = Eigen::Map<const T, Options, DStride>;
383383
using DMapCaster = make_caster<DMap>;
384+
static constexpr bool DMapIsDifferent = !std::is_same_v<Map, DMap>;
384385

385386
/**
386387
* The constructor of ``Ref<const T>`` uses one of two strategies
@@ -391,14 +392,66 @@ struct type_caster<Eigen::Ref<T, Options, StrideType>,
391392
*
392393
* When the value below is ``true``, then it is guaranteed that
393394
* ``Ref(<DMap instance>)`` owns the underlying data.
395+
* Note that this ownership does *not* propagate to any further ``Ref``
396+
* that might be constructed from it; the ``ToRef`` wrapper below
397+
* attempts to avoid this hazard by ensuring there is no need to copy
398+
* or move the initial ``Ref``.
394399
*/
395400
static constexpr bool DMapConstructorOwnsData =
396401
!Eigen::internal::traits<Ref>::template match<DMap>::type::value;
397402

403+
/**
404+
* Result type of our cast operator, which can be used to construct a
405+
* ``Ref`` from either ``Map`` or ``DMap``. Unfortunately, ``Ref`` does
406+
* not model any sort of shared ownership; if we wind up creating a ``Ref``
407+
* that owns the data, then copying it will produce a second ``Ref`` that
408+
* refers to the original ``Ref``, which is likely to go out of scope
409+
* soon. Eigen also has very minimal support for C++11 features such as
410+
* moves, and in particular moving a ``Ref`` is equivalent to copying it
411+
* in all released versions of Eigen as of this writing.
412+
* This creates problems for nested type casters such as
413+
* ``std::optional<Eigen::Ref<..>>``, which can't directly construct the
414+
* cast operator's return value inside their data structure.
415+
*/
416+
struct ToRef {
417+
// DMap has maximally dynamic strides, so it won't contain
418+
// fewer fields than Map (which might have some stride info
419+
// known at compile time).
420+
static_assert(alignof(DMap) >= alignof(Map) &&
421+
sizeof(DMap) >= sizeof(Map));
422+
alignas(DMap) char storage[sizeof(DMap)];
423+
bool holds_dmap;
424+
static constexpr bool MightHoldDMap = MaybeConvert && DMapIsDifferent;
425+
426+
ToRef(Map m) : holds_dmap(false) { new(&storage) Map(m); }
427+
template <bool Enable = MightHoldDMap, enable_if_t<Enable> = 0>
428+
ToRef(DMap m) : holds_dmap(true) { new(&storage) DMap(m); }
429+
ToRef(const ToRef&) = delete;
430+
ToRef& operator=(const ToRef&) = delete;
431+
~ToRef() {
432+
if constexpr (MightHoldDMap) {
433+
if (holds_dmap) {
434+
((DMap *) storage)->~DMap();
435+
return;
436+
}
437+
}
438+
((Map *) storage)->~Map();
439+
}
440+
441+
operator Ref() const {
442+
if constexpr (MightHoldDMap) {
443+
if (holds_dmap) {
444+
return Ref(*(DMap *) storage);
445+
}
446+
}
447+
return Ref(*(Map *) storage);
448+
}
449+
};
450+
398451
static constexpr auto Name =
399452
const_name<MaybeConvert>(DMapCaster::Name, MapCaster::Name);
400453

401-
template <typename T_> using Cast = Ref;
454+
template <typename T_> using Cast = ToRef;
402455
template <typename T_> static constexpr bool can_cast() { return true; }
403456

404457
MapCaster caster;
@@ -459,13 +512,13 @@ struct type_caster<Eigen::Ref<T, Options, StrideType>,
459512
cleanup);
460513
}
461514

462-
operator Ref() {
515+
operator ToRef() {
463516
if constexpr (MaybeConvert) {
464517
if (dcaster.caster.value.is_valid())
465-
return Ref(dcaster.operator DMap());
518+
return ToRef(dcaster.operator DMap());
466519
}
467520

468-
return Ref(caster.operator Map());
521+
return ToRef(caster.operator Map());
469522
}
470523
};
471524

tests/test_eigen.cpp

Lines changed: 17 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,5 @@
11
#include <nanobind/stl/complex.h>
2+
#include <nanobind/stl/optional.h>
23
#include <nanobind/eigen/dense.h>
34
#include <nanobind/eigen/sparse.h>
45
#include <nanobind/trampoline.h>
@@ -144,6 +145,22 @@ NB_MODULE(test_eigen_ext, m) {
144145
m.def("updateRefVXi", [](Eigen::Ref<Eigen::VectorXi> a) { a[2] = 123; });
145146
m.def("updateRefVXi_nc", [](Eigen::Ref<Eigen::VectorXi> a) { a[2] = 123; }, nb::arg().noconvert());
146147

148+
m.def("addOptRefVXi",
149+
[](std::optional<Eigen::Ref<const Eigen::VectorXi>> r1,
150+
std::optional<Eigen::Ref<const Eigen::VectorXi>> r2)
151+
-> std::optional<Eigen::VectorXi> {
152+
if (!r1.has_value() && !r2.has_value()) {
153+
return std::nullopt;
154+
}
155+
if (!r1.has_value()) {
156+
return *r2;
157+
}
158+
if (!r2.has_value()) {
159+
return *r1;
160+
}
161+
return *r1 + *r2;
162+
}, "r1"_a.none(), "r2"_a.none());
163+
147164
using SparseMatrixR = Eigen::SparseMatrix<float, Eigen::RowMajor>;
148165
using SparseMatrixC = Eigen::SparseMatrix<float>;
149166
Eigen::MatrixXf mat(5, 6);

tests/test_eigen.py

Lines changed: 6 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -77,6 +77,12 @@ def test02_vector_dynamic():
7777
# Try with a big array. This will move the result to avoid a copy
7878
assert_array_equal(t.addVXi(x, x), 2*x)
7979

80+
# Regression test for lifetime issues casting std::optional<Ref<VectorXi>>
81+
assert_array_equal(t.addOptRefVXi(af, b), c)
82+
assert_array_equal(t.addOptRefVXi(af, None), af)
83+
assert_array_equal(t.addOptRefVXi(None, b), b)
84+
assert t.addOptRefVXi(None, None) is None
85+
8086

8187
@needs_numpy_and_eigen
8288
def test03_update_map():

0 commit comments

Comments
 (0)