Compiler Explorer

Source code

#include <ranges>

namespace std::ranges {
namespace __detail {
template<class R, class F>
concept flat_mappable = input_range<R> &&  // exposition only
  regular_invocable<F&, range_reference_t<R>> &&
  input_range<invoke_result_t<F&, range_reference_t<R>>>;
}

template<input_range V, move_constructible F>
  requires view<V> && is_object_v<F> && __detail::flat_mappable<V, F>
class flat_map_view : public view_interface<flat_map_view<V, F>> {
 private:
  using MappedRng = invoke_result_t<F&, range_reference_t<V>>;  // exposition only

// [range.flat.map.iterator], class template flat_map_view::iterator
  template<bool Const>
  struct iterator {
    friend flat_map_view;

private:
    using Parent = __detail::__maybe_const_t<Const, flat_map_view>;  // exposition only
    using Base = __detail::__maybe_const_t<Const, V>;                // exposition only
    using MappedBase =                                               // exposition only
      invoke_result_t<__detail::__maybe_const_t<Const, F>&, range_reference_t<Base>>;
    using BaseIter = iterator_t<Base>;          // exposition only
    using MappedIter = iterator_t<MappedBase>;  // exposition only

static constexpr bool mapped_is_glvalue =  // exposition only
      is_reference_v<MappedBase>;

BaseIter current_ = BaseIter();   // exposition only, present only
                                      // if Base models forward_range
    optional<MappedIter> mapped_it_;  // exposition only
    Parent* parent_ = nullptr;        // exposition only

__detail::__maybe_present_t<mapped_is_glvalue,
                                add_pointer_t<MappedBase>>
      mapped_;  // exposition only, present only
                // if mapped_is_glvalue is true

constexpr BaseIter&
    current() {
      if constexpr (forward_range<Base>)
        return current_;
      else
        return *parent_->current_;
    }

constexpr const BaseIter&
    current() const {
      if constexpr (forward_range<Base>)
        return current_;
      else
        return *parent_->current_;
    }

constexpr auto&
    update_mapped() {
      if constexpr (mapped_is_glvalue) {
        mapped_ = addressof(__detail::__as_lvalue(std::__invoke(*parent_->fun_, *current())));
        return *mapped_;
      } else {
        class map_iterator {
          friend iterator;

__detail::__maybe_const_t<Const, F>& fun_;
          const BaseIter& current_;
          constexpr map_iterator(__detail::__maybe_const_t<Const, F>& fun, const BaseIter& current)
            : fun_(fun), current_(current) { }

public:
          constexpr auto
          operator*() const noexcept(noexcept(std::__invoke(fun_, *current_))) {
            return std::__invoke(fun_, *current_);
          }
        };
        return parent_->mapped_._M_emplace_deref(map_iterator(*parent_->fun_, current()));
      }
    }

constexpr auto&
    get_mapped() {
      if constexpr (mapped_is_glvalue)
        return *mapped_;
      else
        return *parent_->mapped_;
    }

constexpr void
    satisfy() {
      for (; current() != ranges::end(parent_->base_); ++current()) {
        auto& mapped = update_mapped();
        mapped_it_ = ranges::begin(mapped);
        if (*mapped_it_ != ranges::end(mapped))
          return;
      }
      if constexpr (mapped_is_glvalue)
        mapped_it_.reset();
    }

constexpr iterator(Parent& parent, BaseIter current)
      requires forward_range<Base>
      : current_(std::move(current)), parent_(addressof(parent)) {
      satisfy();
    }
    constexpr explicit iterator(Parent& parent)
      requires(!forward_range<Base>)
      : parent_(addressof(parent)) {
      satisfy();
    }

public:
    using iterator_concept = decltype([] {
      if constexpr (mapped_is_glvalue && bidirectional_range<Base> &&
                    __detail::__bidirectional_common<MappedBase>)
        return bidirectional_iterator_tag{};
      else if constexpr (mapped_is_glvalue && forward_range<Base> && forward_range<MappedBase>)
        return forward_iterator_tag{};
      else
        return input_iterator_tag{};
    }());
    using iterator_category = decltype([] {
      using BASEC = typename iterator_traits<iterator_t<Base>>::iterator_category;
      using MAPPEDC = typename iterator_traits<iterator_t<MappedBase>>::iterator_category;
      if constexpr (derived_from<BASEC, bidirectional_iterator_tag> &&
                    derived_from<MAPPEDC, bidirectional_iterator_tag> && common_range<MappedBase>)
        return bidirectional_iterator_tag{};
      else if constexpr (derived_from<BASEC, forward_iterator_tag> &&
                         derived_from<MAPPEDC, forward_iterator_tag>)
        return forward_iterator_tag{};
      else
        return input_iterator_tag{};
    }());
    using value_type = range_value_t<MappedBase>;
    using difference_type = common_type_t<range_difference_t<Base>, range_difference_t<MappedBase>>;

iterator() = default;

template<class MappedRng = invoke_result_t<F&, range_reference_t<V>>>
    constexpr iterator(iterator<!Const> i)
      requires Const && is_reference_v<MappedRng> && convertible_to<iterator_t<V>, BaseIter> &&
                 convertible_to<iterator_t<MappedRng>, MappedIter> &&
                 convertible_to<add_pointer_t<MappedRng>, add_pointer_t<MappedBase>>
      : current_(std::move(i.current_)),
        mapped_it_(std::move(i.mapped_it_)),
        parent_(i.parent_),
        mapped_(i.mapped_) { }

constexpr decltype(auto)
    operator*() const {
      return **mapped_it_;
    }

constexpr MappedIter
    operator->() const
      requires __detail::__has_arrow<MappedIter> && copyable<MappedIter>
    {
      return *mapped_it_;
    }

constexpr iterator&
    operator++() {
      if (++*mapped_it_ == ranges::end(get_mapped())) {
        ++current();
        satisfy();
      }
      return *this;
    }

constexpr void
    operator++(int) {
      ++*this;
    }

constexpr iterator
    operator++(int)
      requires mapped_is_glvalue && forward_range<Base> && forward_range<MappedBase>
    {
      auto tmp = *this;
      ++*this;
      return tmp;
    }

constexpr iterator&
    operator--()
      requires mapped_is_glvalue
      && bidirectional_range<Base> && __detail::__bidirectional_common<MappedBase>
    {
      if (current_ == ranges::end(parent_->base_)) {
        --current_;
        mapped_it_ = ranges::end(update_mapped());
      }

while (mapped_it_ == ranges::begin(get_mapped())) {
        --current_;
        mapped_it_ = ranges::end(update_mapped());
      }
      --*mapped_it_;
      return *this;
    }

constexpr iterator
    operator--(int)
      requires mapped_is_glvalue
      && bidirectional_range<Base> && __detail::__bidirectional_common<MappedBase>
    {
      auto tmp = *this;
      --*this;
      return tmp;
    }

friend constexpr bool
    operator==(const iterator& x, const iterator& y)
      requires mapped_is_glvalue
      && forward_range<Base> && equality_comparable<iterator_t<MappedBase>>
    {
      return x.current_ == y.current_ && x.mapped_it_ == y.mapped_it_;
    }

friend constexpr decltype(auto)
    iter_move(const iterator& i) noexcept(noexcept(ranges::iter_move(*i.mapped_it_))) {
      return ranges::iter_move(*i.mapped_it_);
    }

friend constexpr void
    iter_swap(const iterator& x,
              const iterator& y) noexcept(noexcept(ranges::iter_swap(*x.mapped_it_, *y.mapped_it_)))
      requires indirectly_swappable<MappedIter>
    {
      ranges::iter_swap(*x.mapped_it_, *y.mapped_it_);
    }
  };  // exposition only

// [range.flat.map.sentinel], class template flat_map_view::sentinel
  template<bool Const>
  struct sentinel {
    template<bool>
    friend class iterator;

private:
    using Parent = __detail::__maybe_const_t<Const, flat_map_view>;  // exposition only
    using Base = __detail::__maybe_const_t<Const, V>;                // exposition only
    sentinel_t<Base> end_ = sentinel_t<Base>();                      // exposition only

public:
    sentinel() = default;

constexpr explicit sentinel(Parent& parent) : end_(ranges::end(parent.base_)) { }
    constexpr sentinel(sentinel<!Const> s)
      requires Const && convertible_to<sentinel_t<V>, sentinel_t<Base>>
      : end_(std::move(end_)) { }

template<bool OtherConst>
      requires sentinel_for<sentinel_t<Base>, iterator_t<__detail::__maybe_const_t<OtherConst, V>>>
    friend constexpr bool
    operator==(const iterator<OtherConst>& x, const sentinel& y) {
      return x.current() == y.end_;
    }
  };  // exposition only

V base_ = V();            // exposition only
  __detail::__box<F> fun_;  // exposition only

__detail::__maybe_present_t<!forward_range<V>, __detail::__non_propagating_cache<iterator_t<V>>>
    current_;  // exposition only, present only
               // when !forward_range<V>

__detail::__maybe_present_t<!is_reference_v<MappedRng>,
                              __detail::__non_propagating_cache<MappedRng>>
    mapped_;  // exposition only, present only
              // if is_reference_v<MappedRng> is false
 public:
  flat_map_view()
    requires default_initializable<V> && default_initializable<F>
  = default;
  constexpr explicit flat_map_view(V base, F fun) : base_(std::move(base)), fun_(std::move(fun)) { }

constexpr V
  base() const&
    requires copy_constructible<V>
  {
    return base_;
  }
  constexpr V
  base() && {
    return std::move(base_);
  }

constexpr auto
  begin() {
    if constexpr (forward_range<V>) {
      constexpr bool use_const = __detail::__simple_view<V> && is_reference_v<MappedRng>;
      return iterator<use_const>{*this, ranges::begin(base_)};
    } else {
      current_ = ranges::begin(base_);
      return iterator<false>{*this};
    }
  }

constexpr auto
  begin() const
    requires forward_range<const V> && __detail::flat_mappable<const V, const F> &&
    is_reference_v<invoke_result_t<const F&, range_reference_t<const V>>>
  {
    return iterator<true>{*this, ranges::begin(base_)};
  }

constexpr auto
  end() {
    if constexpr (forward_range<V> && is_reference_v<MappedRng> && forward_range<MappedRng> &&
                  common_range<V> && common_range<MappedRng>)
      return iterator<__detail::__simple_view<V>>{*this, ranges::end(base_)};
    else
      return sentinel<__detail::__simple_view<V>>{*this};
  }

constexpr auto
  end() const
    requires forward_range<const V> && __detail::flat_mappable<const V, const F> &&
    is_reference_v<invoke_result_t<const F&, range_reference_t<const V>>>
  {
    using MappedConstRng = invoke_result_t<const F&, range_reference_t<const V>>;
    if constexpr (forward_range<MappedConstRng> && common_range<const V> &&
                  common_range<MappedConstRng>)
      return iterator<true>{*this, ranges::end(base_)};
    else
      return sentinel<true>{*this};
  }
};

template<class R, class F>
flat_map_view(R&&, F) -> flat_map_view<views::all_t<R>, F>;

namespace views {
struct FlatMap : __adaptor::_RangeAdaptor<FlatMap> {
  template<viewable_range R, class F>
    requires requires { flat_map_view(std::declval<R>(), std::declval<F>()); }
  constexpr auto
  operator()(R&& r, F&& fun) const {
    return flat_map_view(std::forward<R>(r), std::forward<F>(fun));
  }

using _RangeAdaptor<FlatMap>::operator();
  static constexpr int _S_arity = 2;
  static constexpr bool _S_has_simple_extra_args = true;
};

inline constexpr FlatMap flat_map;
}  // namespace views

}  // namespace std::ranges

#include <print>
int main() {
  auto ints = {0, 1, 2};
  auto flatten1 = ints | std::views::flat_map([](int i) { return std::views::repeat(i, 3); });
  std::println("{}", flatten1);

auto outer = std::views::iota(1, 5);
  auto inner = std::views::iota(1, 10);
  auto flatten2 = std::views::flat_map(outer, [&](int i) -> auto& { return inner; });
  std::println("{}", flatten2 | std::views::reverse);
}