1 files changed, 37 insertions, 7 deletions

diff --git a/klm/util/mmap.cc b/klm/util/mmap.cc
index 648b5d0a..8685170f 100644
--- a/klm/util/mmap.cc
+++ b/klm/util/mmap.cc
@@ -53,10 +53,8 @@ void *MapOrThrow(std::size_t size, bool for_write, int flags, bool prefault, int
   if (prefault) {
     flags |= MAP_POPULATE;
   }
-  int protect = for_write ? (PROT_READ | PROT_WRITE) : PROT_READ;
-#else
-  int protect = for_write ? (PROT_READ | PROT_WRITE) : PROT_READ;
 #endif
+  int protect = for_write ? (PROT_READ | PROT_WRITE) : PROT_READ;
   void *ret = mmap(NULL, size, protect, flags, fd, offset);
   if (ret == MAP_FAILED) {
     UTIL_THROW(ErrnoException, "mmap failed for size " << size << " at offset " << offset);
@@ -64,8 +62,40 @@ void *MapOrThrow(std::size_t size, bool for_write, int flags, bool prefault, int
   return ret;
 }
 
-void *MapForRead(std::size_t size, bool prefault, int fd, off_t offset) {
-  return MapOrThrow(size, false, MAP_FILE | MAP_PRIVATE, prefault, fd, offset);
+namespace {
+void ReadAll(int fd, void *to_void, std::size_t amount) {
+  uint8_t *to = static_cast<uint8_t*>(to_void);
+  while (amount) {
+    ssize_t ret = read(fd, to, amount);
+    if (ret == -1) UTIL_THROW(ErrnoException, "Reading " << amount << " from fd " << fd << " failed.");
+    if (ret == 0) UTIL_THROW(Exception, "Hit EOF in fd " << fd << " but there should be " << amount << " more bytes to read.");
+    amount -= ret;
+    to += ret;
+  }
+}
+} // namespace
+
+void MapRead(LoadMethod method, int fd, off_t offset, std::size_t size, scoped_memory &out) {
+  switch (method) {
+    case LAZY:
+      out.reset(MapOrThrow(size, false, MAP_FILE | MAP_SHARED, false, fd, offset), size, scoped_memory::MMAP_ALLOCATED);
+      break;
+    case POPULATE_OR_LAZY:
+#ifdef MAP_POPULATE
+    case POPULATE_OR_READ:
+#endif
+      out.reset(MapOrThrow(size, false, MAP_FILE | MAP_SHARED, true, fd, offset), size, scoped_memory::MMAP_ALLOCATED);
+      break;
+#ifndef MAP_POPULATE
+    case POPULATE_OR_READ:
+#endif
+    case READ:
+      out.reset(malloc(size), size, scoped_memory::MALLOC_ALLOCATED);
+      if (!out.get()) UTIL_THROW(util::ErrnoException, "Allocating " << size << " bytes with malloc");
+      if (-1 == lseek(fd, offset, SEEK_SET)) UTIL_THROW(ErrnoException, "lseek to " << offset << " in fd " << fd << " failed.");
+      ReadAll(fd, out.get(), size);
+      break;
+  }
 }
 
 void *MapAnonymous(std::size_t size) {
@@ -78,14 +108,14 @@ void *MapAnonymous(std::size_t size) {
       | MAP_PRIVATE, false, -1, 0);
 }
 
-void MapZeroedWrite(const char *name, std::size_t size, scoped_fd &file, scoped_mmap &mem) {
+void *MapZeroedWrite(const char *name, std::size_t size, scoped_fd &file) {
   file.reset(open(name, O_CREAT | O_RDWR | O_TRUNC, S_IRUSR | S_IWUSR | S_IRGRP | S_IROTH));
   if (-1 == file.get())
     UTIL_THROW(ErrnoException, "Failed to open " << name << " for writing");
   if (-1 == ftruncate(file.get(), size))
     UTIL_THROW(ErrnoException, "ftruncate on " << name << " to " << size << " failed");
   try {
-    mem.reset(MapOrThrow(size, true, MAP_FILE | MAP_SHARED, false, file.get(), 0), size);
+    return MapOrThrow(size, true, MAP_FILE | MAP_SHARED, false, file.get(), 0);
   } catch (ErrnoException &e) {
     e << " in file " << name;
     throw;