1 files changed, 33 insertions, 11 deletions
diff --git a/src/passes/SafeHeap.cpp b/src/passes/SafeHeap.cpp
index b2f7db567..7baeb365e 100644
--- a/src/passes/SafeHeap.cpp
+++ b/src/passes/SafeHeap.cpp
@@ -288,6 +288,7 @@ struct SafeHeap : public Pass {
     auto func = Builder::makeFunction(name, funcSig, {indexType});
     Builder builder(*module);
     auto* block = builder.makeBlock();
+    // stash the sum of the pointer (0) and the size (1) in a local (2)
     block->list.push_back(builder.makeLocalSet(
       2,
       builder.makeBinary(memory->is64() ? AddInt64 : AddInt32,
@@ -296,6 +297,7 @@ struct SafeHeap : public Pass {
     // check for reading past valid memory: if pointer + offset + bytes
     block->list.push_back(makeBoundsCheck(style.type,
                                           builder,
+                                          0,
                                           2,
                                           style.bytes,
                                           module,
@@ -346,6 +348,7 @@ struct SafeHeap : public Pass {
     // check for reading past valid memory: if pointer + offset + bytes
     block->list.push_back(makeBoundsCheck(style.valueType,
                                           builder,
+                                          0,
                                           3,
                                           style.bytes,
                                           module,
@@ -386,9 +389,14 @@ struct SafeHeap : public Pass {
       builder.makeCall(alignfault, {}, Type::none));
   }
 
+  // Constructs a bounds check. This receives the indexes of two locals: the
+  // pointer local, which contains the pointer we are checking, and the sum
+  // local which contains the pointer added to the number of bytes being
+  // accessed.
   Expression* makeBoundsCheck(Type type,
                               Builder& builder,
-                              Index local,
+                              Index ptrLocal,
+                              Index sumLocal,
                               Index bytes,
                               Module* module,
                               Type indexType,
@@ -415,19 +423,33 @@ struct SafeHeap : public Pass {
     }
     auto gtuOp = is64 ? GtUInt64 : GtUInt32;
     auto addOp = is64 ? AddInt64 : AddInt32;
+    auto* upperCheck =
+      builder.makeBinary(upperOp,
+                         builder.makeLocalGet(sumLocal, indexType),
+                         builder.makeConstPtr(upperBound, indexType));
+    auto* lowerCheck = builder.makeBinary(
+      gtuOp,
+      builder.makeBinary(addOp,
+                         builder.makeLocalGet(sumLocal, indexType),
+                         builder.makeConstPtr(bytes, indexType)),
+      brkLocation);
+    // Check for an overflow when adding the pointer and the size, using the
+    // rule that for any unsigned x and y,
+    //    x + y < x    <=>   x + y overflows
+    auto* overflowCheck =
+      builder.makeBinary(is64 ? LtUInt64 : LtUInt32,
+                         builder.makeLocalGet(sumLocal, indexType),
+                         builder.makeLocalGet(ptrLocal, indexType));
+    // Add an unreachable right after the call to segfault for performance
+    // reasons: the call never returns, and this helps optimizations benefit
+    // from that.
     return builder.makeIf(
       builder.makeBinary(
         OrInt32,
-        builder.makeBinary(upperOp,
-                           builder.makeLocalGet(local, indexType),
-                           builder.makeConstPtr(upperBound, indexType)),
-        builder.makeBinary(
-          gtuOp,
-          builder.makeBinary(addOp,
-                             builder.makeLocalGet(local, indexType),
-                             builder.makeConstPtr(bytes, indexType)),
-          brkLocation)),
-      builder.makeCall(segfault, {}, Type::none));
+        upperCheck,
+        builder.makeBinary(OrInt32, lowerCheck, overflowCheck)),
+      builder.makeSequence(builder.makeCall(segfault, {}, Type::none),
+                           builder.makeUnreachable()));
   }
 };