folly/experimental/StringGen-inl.h

   1 /*
   2  * Copyright 2012 Facebook, Inc.
   3  *
   4  * Licensed under the Apache License, Version 2.0 (the "License");
   5  * you may not use this file except in compliance with the License.
   6  * You may obtain a copy of the License at
   7  *
   8  *   http://www.apache.org/licenses/LICENSE-2.0
   9  *
  10  * Unless required by applicable law or agreed to in writing, software
  11  * distributed under the License is distributed on an "AS IS" BASIS,
  12  * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  13  * See the License for the specific language governing permissions and
  14  * limitations under the License.
  15  */
  16
  17 #ifndef FOLLY_STRINGGEN_H_
  18 #error This file may only be included from folly/experimental/StringGen.h
  19 #endif
  20
  21 #include "folly/experimental/io/IOBuf.h"
  22
  23 namespace folly {
  24 namespace gen {
  25 namespace detail {
  26
  27 inline bool splitPrefix(StringPiece& in, StringPiece& prefix, char delimiter) {
  28   auto p = static_cast<const char*>(memchr(in.data(), delimiter, in.size()));
  29   if (p) {
  30     prefix.assign(in.data(), p);
  31     in.assign(p + 1, in.end());
  32     return true;
  33   }
  34   prefix.clear();
  35   return false;
  36 }
  37
  38 inline const char* ch(const unsigned char* p) {
  39   return reinterpret_cast<const char*>(p);
  40 }
  41
  42 class StringResplitter : public Operator<StringResplitter> {
  43   char delimiter_;
  44  public:
  45   explicit StringResplitter(char delimiter) : delimiter_(delimiter) { }
  46
  47   template <class Source>
  48   class Generator : public GenImpl<StringPiece, Generator<Source>> {
  49     Source source_;
  50     char delimiter_;
  51     static constexpr size_t kDefaultLineSize = 256;
  52    public:
  53     Generator(Source source, char delimiter)
  54       : source_(std::move(source)), delimiter_(delimiter) { }
  55
  56     template <class Body>
  57     bool apply(Body&& body) const {
  58       std::unique_ptr<IOBuf> buffer;
  59
  60       auto fn = [&](StringPiece in) -> bool {
  61         StringPiece prefix;
  62         bool found = splitPrefix(in, prefix, this->delimiter_);
  63         if (found && buffer && buffer->length() != 0) {
  64           // Append to end of buffer, return line
  65           if (!prefix.empty()) {
  66             buffer->reserve(0, prefix.size());
  67             memcpy(buffer->writableTail(), prefix.data(), prefix.size());
  68             buffer->append(prefix.size());
  69           }
  70           if (!body(StringPiece(ch(buffer->data()), buffer->length()))) {
  71             return false;
  72           }
  73           buffer->clear();
  74           found = splitPrefix(in, prefix, this->delimiter_);
  75         }
  76         // Buffer is empty, return lines directly from input (no buffer)
  77         while (found) {
  78           if (!body(prefix)) {
  79             return false;
  80           }
  81           found = splitPrefix(in, prefix, this->delimiter_);
  82         }
  83         if (!in.empty()) {
  84           // Incomplete line left, append to buffer
  85           if (!buffer) {
  86             // Arbitrarily assume that we have half a line and get enough
  87             // room for twice that.
  88             buffer = IOBuf::create(std::max(kDefaultLineSize, 2 * in.size()));
  89           }
  90           buffer->reserve(0, in.size());
  91           memcpy(buffer->writableTail(), in.data(), in.size());
  92           buffer->append(in.size());
  93         }
  94         return true;
  95       };
  96
  97       // Iterate
  98       if (!source_.apply(std::move(fn))) {
  99         return false;
 100       }
 101
 102       // Incomplete last line
 103       if (buffer && buffer->length() != 0) {
 104         if (!body(StringPiece(ch(buffer->data()), buffer->length()))) {
 105           return false;
 106         }
 107       }
 108       return true;
 109     }
 110   };
 111
 112   template<class Source,
 113            class Value,
 114            class Gen = Generator<Source>>
 115   Gen compose(GenImpl<Value, Source>&& source) const {
 116     return Gen(std::move(source.self()), delimiter_);
 117   }
 118
 119   template<class Source,
 120            class Value,
 121            class Gen = Generator<Source>>
 122   Gen compose(const GenImpl<Value, Source>& source) const {
 123     return Gen(source.self(), delimiter_);
 124   }
 125 };
 126
 127 class SplitStringSource : public GenImpl<StringPiece, SplitStringSource> {
 128   StringPiece source_;
 129   char delimiter_;
 130  public:
 131   SplitStringSource(const StringPiece& source,
 132                     char delimiter)
 133     : source_(source)
 134     , delimiter_(delimiter) { }
 135
 136   template <class Body>
 137   bool apply(Body&& body) const {
 138     StringPiece rest(source_);
 139     StringPiece prefix;
 140     while (splitPrefix(rest, prefix, this->delimiter_)) {
 141       if (!body(prefix)) {
 142         return false;
 143       }
 144     }
 145     if (!rest.empty()) {
 146       if (!body(rest)) {
 147         return false;
 148       }
 149     }
 150     return true;
 151   }
 152 };
 153
 154
 155 }  // namespace detail
 156 }  // namespace gen
 157 }  // namespace folly
 158