From dc078281c2fd83db9fc4860e8b24979ab2a1e407 Mon Sep 17 00:00:00 2001
From: Kenneth Heafield <github@kheafield.com>
Date: Thu, 16 Aug 2012 17:02:56 -0400
Subject: KenLM update.  Remove a couple of segfaults for weird input.  Other
 oddball stuff.

---
 klm/lm/trie_sort.cc | 24 ++++++++++++++++--------
 1 file changed, 16 insertions(+), 8 deletions(-)

(limited to 'klm/lm/trie_sort.cc')

diff --git a/klm/lm/trie_sort.cc b/klm/lm/trie_sort.cc
index b80fed02..0d83221e 100644
--- a/klm/lm/trie_sort.cc
+++ b/klm/lm/trie_sort.cc
@@ -148,13 +148,17 @@ template <class Combine> FILE *MergeSortedFiles(FILE *first_file, FILE *second_f
 } // namespace
 
 void RecordReader::Init(FILE *file, std::size_t entry_size) {
-  rewind(file);
-  file_ = file;
+  entry_size_ = entry_size;
   data_.reset(malloc(entry_size));
   UTIL_THROW_IF(!data_.get(), util::ErrnoException, "Failed to malloc read buffer");
-  remains_ = true;
-  entry_size_ = entry_size;
-  ++*this;
+  file_ = file;
+  if (file) {
+    rewind(file);
+    remains_ = true;
+    ++*this;
+  } else {
+    remains_ = false;
+  }
 }
 
 void RecordReader::Overwrite(const void *start, std::size_t amount) {
@@ -169,9 +173,13 @@ void RecordReader::Overwrite(const void *start, std::size_t amount) {
 }
 
 void RecordReader::Rewind() {
-  rewind(file_);
-  remains_ = true;
-  ++*this;
+  if (file_) {
+    rewind(file_);
+    remains_ = true;
+    ++*this;
+  } else {
+    remains_ = false;
+  }
 }
 
 SortedFiles::SortedFiles(const Config &config, util::FilePiece &f, std::vector<uint64_t> &counts, size_t buffer, const std::string &file_prefix, SortedVocabulary &vocab) {
-- 
cgit v1.2.3