<commit_before>// Copyright 2017 <NAME>. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#pragma once

#include <type_traits>

#include <rs/backreference.h>
#include <rs/element_count.h>
#include <rs/publisher.h>
#include <rs/subscriber.h>
#include <rs/subscription.h>

namespace shk {
namespace detail {

template <typename InnerSubscriberType, typename Predicate>
class TakeWhileSubscriber : public Subscriber {
 public:
  TakeWhileSubscriber(
      InnerSubscriberType &&inner_subscriber,
      const Predicate &predicate)
      : inner_subscriber_(std::move(inner_subscriber)),
        predicate_(predicate) {}

  void TakeSubscription(Backreference<AnySubscription> &&subscription) {
    subscription_ = std::move(subscription);
  }

  template <typename T, class = RequireRvalue<T>>
  void OnNext(T &&t) {
    if (cancelled_) {
      return;
    }

    bool predicate_match = false;
    try {
      predicate_match = predicate_(static_cast<const T &>(t));
    } catch (...) {
      Cancel();
      inner_subscriber_.OnError(std::current_exception());
    }
    if (!cancelled_) { // If predicate_ threw, cancelled_ could be true here
      if (predicate_match) {
        inner_subscriber_.OnNext(std::forward<T>(t));
      } else {
        inner_subscriber_.OnComplete();
        Cancel();
      }
    }
  }

  void OnError(std::exception_ptr &&error) {
    if (!cancelled_) {
      inner_subscriber_.OnError(std::move(error));
    }
  }

  void OnComplete() {
    if (!cancelled_) {
      inner_subscriber_.OnComplete();
    }
  }

  void Cancel() {
    if (subscription_) {
      // TODO(peck): It's wrong that subscription_ can be null here; when that
      // happens we still actually need to be able to cancel the subscription.
      //
      // I think one approach to fix this is to change so that destroying the
      // Subscription implies cancellation.
      subscription_->Cancel();
    }
    cancelled_ = true;
  }

 private:
  bool cancelled_ = false;
  InnerSubscriberType inner_subscriber_;
  Backreference<AnySubscription> subscription_;
  Predicate predicate_;
};

}  // namespace detail

template <typename Predicate>
auto TakeWhile(Predicate &&predicate) {
  // Return an operator (it takes a Publisher and returns a Publisher)
  return [predicate = std::forward<Predicate>(predicate)](auto source) {
    // Return a Publisher
    return MakePublisher([predicate, source = std::move(source)](
        auto &&subscriber) {
      using TakeWhileSubscriberT = detail::TakeWhileSubscriber<
          typename std::decay<decltype(subscriber)>::type,
          typename std::decay<Predicate>::type>;

      Backreference<TakeWhileSubscriberT> take_while_ref;
      auto take_while_subscriber = WithBackreference(
          TakeWhileSubscriberT(
              std::forward<decltype(subscriber)>(subscriber),
              predicate),
          &take_while_ref);

      Backreference<AnySubscription> sub_ref;
      auto sub = WithBackreference(
          AnySubscription(source.Subscribe(std::move(take_while_subscriber))),
          &sub_ref);

      if (take_while_ref) {  // TODO(peck): Test what happens if it's is empty
        take_while_ref->TakeSubscription(std::move(sub_ref));
      }

      return sub;
    });
  };
}

}  // namespace shk<commit_msg>Avoid unnecessary heap allocation in TakeWhile<commit_after>// Copyright 2017 <NAME>. All Rights Reserved.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#pragma once

#include <type_traits>

#include <rs/backreference.h>
#include <rs/element_count.h>
#include <rs/publisher.h>
#include <rs/subscriber.h>
#include <rs/subscription.h>

namespace shk {
namespace detail {

template <typename InnerSubscriberType, typename Predicate>
class TakeWhileSubscriber : public Subscriber {
 public:
  TakeWhileSubscriber(
      InnerSubscriberType &&inner_subscriber,
      const Predicate &predicate)
      : inner_subscriber_(std::move(inner_subscriber)),
        predicate_(predicate) {}

  void TakeSubscription(Backreference<PureVirtualSubscription> &&subscription) {
    subscription_ = std::move(subscription);
  }

  template <typename T, class = RequireRvalue<T>>
  void OnNext(T &&t) {
    if (cancelled_) {
      return;
    }

    bool predicate_match = false;
    try {
      predicate_match = predicate_(static_cast<const T &>(t));
    } catch (...) {
      Cancel();
      inner_subscriber_.OnError(std::current_exception());
    }
    if (!cancelled_) { // If predicate_ threw, cancelled_ could be true here
      if (predicate_match) {
        inner_subscriber_.OnNext(std::forward<T>(t));
      } else {
        inner_subscriber_.OnComplete();
        Cancel();
      }
    }
  }

  void OnError(std::exception_ptr &&error) {
    if (!cancelled_) {
      inner_subscriber_.OnError(std::move(error));
    }
  }

  void OnComplete() {
    if (!cancelled_) {
      inner_subscriber_.OnComplete();
    }
  }

  void Cancel() {
    if (subscription_) {
      // TODO(peck): It's wrong that subscription_ can be null here; when that
      // happens we still actually need to be able to cancel the subscription.
      //
      // I think one approach to fix this is to change so that destroying the
      // Subscription implies cancellation.
      subscription_->Cancel();
    }
    cancelled_ = true;
  }

 private:
  bool cancelled_ = false;
  InnerSubscriberType inner_subscriber_;
  Backreference<PureVirtualSubscription> subscription_;
  Predicate predicate_;
};

}  // namespace detail

template <typename Predicate>
auto TakeWhile(Predicate &&predicate) {
  // Return an operator (it takes a Publisher and returns a Publisher)
  return [predicate = std::forward<Predicate>(predicate)](auto source) {
    // Return a Publisher
    return MakePublisher([predicate, source = std::move(source)](
        auto &&subscriber) {
      using TakeWhileSubscriberT = detail::TakeWhileSubscriber<
          typename std::decay<decltype(subscriber)>::type,
          typename std::decay<Predicate>::type>;

      Backreference<TakeWhileSubscriberT> take_while_ref;
      auto take_while_subscriber = WithBackreference(
          TakeWhileSubscriberT(
              std::forward<decltype(subscriber)>(subscriber),
              predicate),
          &take_while_ref);

      Backreference<PureVirtualSubscription> sub_ref;
      auto sub = WithBackreference(
          MakeVirtualSubscription(
              source.Subscribe(std::move(take_while_subscriber))),
          &sub_ref);

      if (take_while_ref) {  // TODO(peck): Test what happens if it's is empty
        take_while_ref->TakeSubscription(std::move(sub_ref));
      }

      return sub;
    });
  };
}

}  // namespace shk
{'fraction_non_alphanumeric': 0.10667012717363093, 'fraction_numerical': 0.002076304178562159, 'mean_word_length': 3.232289950576606, 'pattern_counts': {'":': 0, '<': 47, '<?xml version=': 0, '>': 51, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5875934', 'n_tokens_mistral': 2229, 'n_tokens_neox': 2182, 'n_words': 609}