folly/experimental/FileGen-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_FILEGEN_H_
  18 #error This file may only be included from folly/experimental/FileGen.h
  19 #endif
  20
  21 #include <system_error>
  22
  23 #include "folly/experimental/StringGen.h"
  24
  25 namespace folly {
  26 namespace gen {
  27 namespace detail {
  28
  29 class FileReader : public GenImpl<ByteRange, FileReader> {
  30  public:
  31   FileReader(File file, std::unique_ptr<IOBuf> buffer)
  32     : file_(std::move(file)),
  33       buffer_(std::move(buffer)) {
  34     buffer_->clear();
  35   }
  36
  37   template <class Body>
  38   bool apply(Body&& body) const {
  39     for (;;) {
  40       ssize_t n = ::read(file_.fd(), buffer_->writableTail(),
  41                          buffer_->capacity());
  42       if (n == -1) {
  43         throw std::system_error(errno, std::system_category(), "read failed");
  44       }
  45       if (n == 0) {
  46         return true;
  47       }
  48       if (!body(ByteRange(buffer_->tail(), n))) {
  49         return false;
  50       }
  51     }
  52   }
  53  private:
  54   File file_;
  55   std::unique_ptr<IOBuf> buffer_;
  56 };
  57
  58 class FileWriter : public Operator<FileWriter> {
  59  public:
  60   FileWriter(File file, std::unique_ptr<IOBuf> buffer)
  61     : file_(std::move(file)),
  62       buffer_(std::move(buffer)) {
  63     if (buffer_) {
  64       buffer_->clear();
  65     }
  66   }
  67
  68   template <class Source, class Value>
  69   void compose(const GenImpl<Value, Source>& source) const {
  70     auto fn = [&](ByteRange v) {
  71       if (!this->buffer_ || v.size() >= this->buffer_->capacity()) {
  72         this->flushBuffer();
  73         this->write(v);
  74       } else {
  75         if (v.size() > this->buffer_->tailroom()) {
  76           this->flushBuffer();
  77         }
  78         memcpy(this->buffer_->writableTail(), v.data(), v.size());
  79         this->buffer_->append(v.size());
  80       }
  81     };
  82
  83     // Iterate
  84     source.foreach(std::move(fn));
  85
  86     flushBuffer();
  87     file_.close();
  88   }
  89
  90  private:
  91   void write(ByteRange v) const {
  92     ssize_t n;
  93     while (!v.empty()) {
  94       n = ::write(file_.fd(), v.data(), v.size());
  95       if (n == -1) {
  96         if (errno == EINTR) {
  97           continue;
  98         }
  99         throw std::system_error(errno, std::system_category(),
 100                                 "write() failed");
 101       }
 102       v.advance(n);
 103     }
 104   }
 105
 106   void flushBuffer() const {
 107     if (buffer_ && buffer_->length() != 0) {
 108       write(ByteRange(buffer_->data(), buffer_->length()));
 109       buffer_->clear();
 110     }
 111   }
 112
 113   mutable File file_;
 114   std::unique_ptr<IOBuf> buffer_;
 115 };
 116
 117 }  // namespace detail
 118
 119 auto byLine(File file, char delim='\n') ->
 120 decltype(fromFile(std::move(file)) | eachAs<StringPiece>() | resplit(delim)) {
 121   return fromFile(std::move(file)) | eachAs<StringPiece>() | resplit(delim);
 122 }
 123
 124 }  // namespace gen
 125 }  // namespace folly
 126