1 files changed, 24 insertions, 9 deletions
diff --git a/klm/util/file_piece.cc b/klm/util/file_piece.cc
index 81eb9bb9..67681f7e 100644
--- a/klm/util/file_piece.cc
+++ b/klm/util/file_piece.cc
@@ -79,22 +79,22 @@ FilePiece::~FilePiece() {
 }
 
 StringPiece FilePiece::ReadLine(char delim) throw (GZException, EndOfFileException) {
-  const char *start = position_;
-  do {
-    for (const char *i = start; i < position_end_; ++i) {
+  size_t skip = 0;
+  while (true) {
+    for (const char *i = position_ + skip; i < position_end_; ++i) {
       if (*i == delim) {
         StringPiece ret(position_, i - position_);
         position_ = i + 1;
         return ret;
       }
     }
-    size_t skip = position_end_ - position_;
+    if (at_end_) {
+      if (position_ == position_end_) Shift();
+      return Consume(position_end_);
+    }
+    skip = position_end_ - position_;
     Shift();
-    start = position_ + skip;
-  } while (!at_end_);
-  StringPiece ret(position_, position_end_ - position_);
-  position_ = position_end_;
-  return ret;
+  }
 }
 
 float FilePiece::ReadFloat() throw(GZException, EndOfFileException, ParseNumberException) {
@@ -186,6 +186,21 @@ template <class T> T FilePiece::ReadNumber() throw(GZException, EndOfFileExcepti
   return ret;
 }
 
+const char *FilePiece::FindDelimiterOrEOF(const bool *delim) throw (GZException, EndOfFileException) {
+  size_t skip = 0;
+  while (true) {
+    for (const char *i = position_ + skip; i < position_end_; ++i) {
+      if (delim[static_cast<unsigned char>(*i)]) return i;
+    }
+    if (at_end_) {
+      if (position_ == position_end_) Shift();
+      return position_end_;
+    }
+    skip = position_end_ - position_;
+    Shift();
+  }
+}
+
 void FilePiece::Shift() throw(GZException, EndOfFileException) {
   if (at_end_) {
     progress_.Finished();