diff --git a/src/tests/util/list.cpp b/src/tests/util/list.cpp index e28386c71..00c3312b8 100644 --- a/src/tests/util/list.cpp +++ b/src/tests/util/list.cpp @@ -161,6 +161,23 @@ static void tst12() { lean_assert(sum == 30); } +static void tst13() { + list l({1, 2, 3, 4, 5, 6, 7}); + list l2 = map_filter(l, [](int in, int & out) { + if (in % 2 == 0) { + out = in + 10; + return true; + } else { + return false; + } + }); + std::cout << l2 << "\n"; + list l3({12, 14, 16}); + lean_assert_eq(l2, l3); + lean_assert(empty(map_filter(l, [](int, int &) { return false; }))); + lean_assert(empty(map_filter(list(), [](int in, int & out) { out = in; return true; }))); +} + int main() { tst1(); tst2(); @@ -174,5 +191,6 @@ int main() { tst10(1000, 5); tst11(1000, 5); tst12(); + tst13(); return has_violations() ? 1 : 0; } diff --git a/src/util/list_fn.h b/src/util/list_fn.h index a2f616d1d..fd7f2e01d 100644 --- a/src/util/list_fn.h +++ b/src/util/list_fn.h @@ -156,6 +156,32 @@ list map(list const & l, F && f) { } } +/** + \brief Similar to \c map but \c f has signature + + bool f(T const & in, T & out) + + If \c out becomes part of the result iff \c f returns true. +*/ +template +list map_filter(list const & l, F && f) { + if (is_nil(l)) { + return l; + } else { + buffer::cell*> tmp; + to_buffer(l, tmp); + unsigned i = tmp.size(); + list r; + while (i > 0) { + --i; + T out; + if (f(tmp[i]->head(), out)) + r = cons(out, r); + } + return r; + } +} + /** \brief Semantically equivalent to \c map, but it tries to reuse list cells. The elements are compared using the predicate \c eq.