folly/experimental/FileGen-inl.h

   1 /*
   2  * Copyright 2013 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;
  41       do {
  42         n = ::read(file_.fd(), buffer_->writableTail(), buffer_->capacity());
  43       } while (n == -1 && errno == EINTR);
  44       if (n == -1) {
  45         throw std::system_error(errno, std::system_category(), "read failed");
  46       }
  47       if (n == 0) {
  48         return true;
  49       }
  50       if (!body(ByteRange(buffer_->tail(), n))) {
  51         return false;
  52       }
  53     }
  54   }
  55  private:
  56   File file_;
  57   std::unique_ptr<IOBuf> buffer_;
  58 };
  59
  60 class FileWriter : public Operator<FileWriter> {
  61  public:
  62   FileWriter(File file, std::unique_ptr<IOBuf> buffer)
  63     : file_(std::move(file)),
  64       buffer_(std::move(buffer)) {
  65     if (buffer_) {
  66       buffer_->clear();
  67     }
  68   }
  69
  70   template <class Source, class Value>
  71   void compose(const GenImpl<Value, Source>& source) const {
  72     auto fn = [&](ByteRange v) {
  73       if (!this->buffer_ || v.size() >= this->buffer_->capacity()) {
  74         this->flushBuffer();
  75         this->write(v);
  76       } else {
  77         if (v.size() > this->buffer_->tailroom()) {
  78           this->flushBuffer();
  79         }
  80         memcpy(this->buffer_->writableTail(), v.data(), v.size());
  81         this->buffer_->append(v.size());
  82       }
  83     };
  84
  85     // Iterate
  86     source.foreach(std::move(fn));
  87
  88     flushBuffer();
  89     file_.close();
  90   }
  91
  92  private:
  93   void write(ByteRange v) const {
  94     ssize_t n;
  95     while (!v.empty()) {
  96       do {
  97         n = ::write(file_.fd(), v.data(), v.size());
  98       } while (n == -1 && errno == EINTR);
  99       if (n == -1) {
 100         throw std::system_error(errno, std::system_category(),
 101                                 "write() failed");
 102       }
 103       v.advance(n);
 104     }
 105   }
 106
 107   void flushBuffer() const {
 108     if (buffer_ && buffer_->length() != 0) {
 109       write(ByteRange(buffer_->data(), buffer_->length()));
 110       buffer_->clear();
 111     }
 112   }
 113
 114   mutable File file_;
 115   std::unique_ptr<IOBuf> buffer_;
 116 };
 117
 118 }  // namespace detail
 119
 120 inline auto byLine(File file, char delim='\n') ->
 121 decltype(fromFile(std::move(file)) | eachAs<StringPiece>() | resplit(delim)) {
 122   return fromFile(std::move(file)) | eachAs<StringPiece>() | resplit(delim);
 123 }
 124
 125 }  // namespace gen
 126 }  // namespace folly
 127