--- a/memory/replace/logalloc/replay/Replay.cpp
+++ b/memory/replace/logalloc/replay/Replay.cpp
@@ -332,96 +332,103 @@ public:
 #endif
   }
 
   MemSlot& operator[] (size_t index) const
   {
     return mSlots[index];
   }
 
-  void malloc(MemSlot& aSlot, Buffer& aArgs)
+  void malloc(Buffer& aArgs, Buffer& aResult)
   {
+    MemSlot& aSlot = SlotForResult(aResult);
     mOps++;
     size_t size = parseNumber(aArgs);
     aSlot.mPtr = ::malloc_impl(size);
   }
 
-  void posix_memalign(MemSlot& aSlot, Buffer& aArgs)
+  void posix_memalign(Buffer& aArgs, Buffer& aResult)
   {
+    MemSlot& aSlot = SlotForResult(aResult);
     mOps++;
     size_t alignment = parseNumber(aArgs.SplitChar(','));
     size_t size = parseNumber(aArgs);
     void* ptr;
     if (::posix_memalign_impl(&ptr, alignment, size) == 0) {
       aSlot.mPtr = ptr;
     } else {
       aSlot.mPtr = nullptr;
     }
   }
 
-  void aligned_alloc(MemSlot& aSlot, Buffer& aArgs)
+  void aligned_alloc(Buffer& aArgs, Buffer& aResult)
   {
+    MemSlot& aSlot = SlotForResult(aResult);
     mOps++;
     size_t alignment = parseNumber(aArgs.SplitChar(','));
     size_t size = parseNumber(aArgs);
     aSlot.mPtr = ::aligned_alloc_impl(alignment, size);
   }
 
-  void calloc(MemSlot& aSlot, Buffer& aArgs)
+  void calloc(Buffer& aArgs, Buffer& aResult)
   {
+    MemSlot& aSlot = SlotForResult(aResult);
     mOps++;
     size_t num = parseNumber(aArgs.SplitChar(','));
     size_t size = parseNumber(aArgs);
     aSlot.mPtr = ::calloc_impl(num, size);
   }
 
-  void realloc(MemSlot& aSlot, Buffer& aArgs)
+  void realloc(Buffer& aArgs, Buffer& aResult)
   {
+    MemSlot& aSlot = SlotForResult(aResult);
     mOps++;
     Buffer dummy = aArgs.SplitChar('#');
     if (dummy) {
       die("Malformed input");
     }
     size_t slot_id = parseNumber(aArgs.SplitChar(','));
     size_t size = parseNumber(aArgs);
     MemSlot& old_slot = (*this)[slot_id];
     void* old_ptr = old_slot.mPtr;
     old_slot.mPtr = nullptr;
     aSlot.mPtr = ::realloc_impl(old_ptr, size);
   }
 
-  void free(Buffer& aArgs)
+  void free(Buffer& aArgs, Buffer& aResult)
   {
     mOps++;
     Buffer dummy = aArgs.SplitChar('#');
     if (dummy) {
       die("Malformed input");
     }
     size_t slot_id = parseNumber(aArgs);
     MemSlot& slot = (*this)[slot_id];
     ::free_impl(slot.mPtr);
     slot.mPtr = nullptr;
   }
 
-  void memalign(MemSlot& aSlot, Buffer& aArgs)
+  void memalign(Buffer& aArgs, Buffer& aResult)
   {
+    MemSlot& aSlot = SlotForResult(aResult);
     mOps++;
     size_t alignment = parseNumber(aArgs.SplitChar(','));
     size_t size = parseNumber(aArgs);
     aSlot.mPtr = ::memalign_impl(alignment, size);
   }
 
-  void valloc(MemSlot& aSlot, Buffer& aArgs)
+  void valloc(Buffer& aArgs, Buffer& aResult)
   {
+    MemSlot& aSlot = SlotForResult(aResult);
     mOps++;
     size_t size = parseNumber(aArgs);
     aSlot.mPtr = ::valloc_impl(size);
   }
 
-  void jemalloc_stats(Buffer& aArgs)
+  void jemalloc_stats(Buffer& aArgs, Buffer& aResult)
   {
     if (aArgs) {
       die("Malformed input");
     }
     mOps++;
     jemalloc_stats_t stats;
     ::jemalloc_stats(&stats);
     FdPrintf(mStdErr,
@@ -429,16 +436,29 @@ public:
              "bookkeep: %zu; binunused: %zu\n", mOps, stats.mapped,
              stats.allocated, stats.waste, stats.page_cache,
              stats.bookkeeping, stats.bin_unused);
     /* TODO: Add more data, like actual RSS as measured by OS, but compensated
      * for the replay internal data. */
   }
 
 private:
+  MemSlot& SlotForResult(Buffer& aResult)
+  {
+    /* Parse result value and get the corresponding slot. */
+    Buffer dummy = aResult.SplitChar('=');
+    Buffer dummy2 = aResult.SplitChar('#');
+    if (dummy || dummy2) {
+      die("Malformed input");
+    }
+
+    size_t slot_id = parseNumber(aResult);
+    return mSlots[slot_id];
+  }
+
   intptr_t mStdErr;
   size_t mOps;
   MemSlotList mSlots;
 };
 
 
 int
 main()
@@ -479,49 +499,33 @@ main()
 
     /* The log contains thread ids for manual analysis, but we just ignore them
      * for now. */
     parseNumber(line.SplitChar(' '));
 
     Buffer func = line.SplitChar('(');
     Buffer args = line.SplitChar(')');
 
-    /* jemalloc_stats and free are functions with no result. */
     if (func == Buffer("jemalloc_stats")) {
-      replay.jemalloc_stats(args);
-      continue;
-    }
-    if (func == Buffer("free")) {
-      replay.free(args);
-      continue;
-    }
-
-    /* Parse result value and get the corresponding slot. */
-    Buffer dummy = line.SplitChar('=');
-    Buffer dummy2 = line.SplitChar('#');
-    if (dummy || dummy2) {
-      die("Malformed input");
-    }
-
-    size_t slot_id = parseNumber(line);
-    MemSlot& slot = replay[slot_id];
-
-    if (func == Buffer("malloc")) {
-      replay.malloc(slot, args);
+      replay.jemalloc_stats(args, line);
+    } else if (func == Buffer("free")) {
+      replay.free(args, line);
+    } else if (func == Buffer("malloc")) {
+      replay.malloc(args, line);
     } else if (func == Buffer("posix_memalign")) {
-      replay.posix_memalign(slot, args);
+      replay.posix_memalign(args, line);
     } else if (func == Buffer("aligned_alloc")) {
-      replay.aligned_alloc(slot, args);
+      replay.aligned_alloc(args, line);
     } else if (func == Buffer("calloc")) {
-      replay.calloc(slot, args);
+      replay.calloc(args, line);
     } else if (func == Buffer("realloc")) {
-      replay.realloc(slot, args);
+      replay.realloc(args, line);
     } else if (func == Buffer("memalign")) {
-      replay.memalign(slot, args);
+      replay.memalign(args, line);
     } else if (func == Buffer("valloc")) {
-      replay.valloc(slot, args);
+      replay.valloc(args, line);
     } else {
       die("Malformed input");
     }
   }
 
   return 0;
 }