[Zucchini] Fix underflow / overflow for extended forward-projection.

Forward-projection is how Zucchini uses the equivalence map to create
estimated "new" targets from "old" targets. Extended forward-projection
is defined to transform non-covered offsets: Given an offset, it finds
the equivalence unit with nearest "old" block, then applies the
"old"-to-"new" displacement to the offset. However, this makes it
possible to map an "old" offset to an offset outside "new" image.

Another issue is that Zucchini uses "dangling targets" that use "fake
offsets" outside the image file to represent .bss data. These targets
also undergo forward-projection, and should be properly handled.

This CL fixes the existing behavior, where underflow / overflow go
unchecked (although these values are rendered benign downstream, since
the nearest actual "new" target is found). The updated extended
forward-projection specifies:
- For "old" targets with real offsets: Take nearest equivalence unit,
  clamp output to be inside [0, "new" image size).
- For "old" dangling targets with fake offsets: Use difference in
  file size as displacement.

The main impact w.r.t. patch is to reduce possible variance in patch
sizes -- dangling targets are now handled better. Extensive unit tests
are also added.

Bug: 832572
Change-Id: I41fea175e4c13585d14a97a712a191afc2fcc6d6
Reviewed-on: https://chromium-review.googlesource.com/1111467
Reviewed-by: Samuel Huang <huangs@chromium.org>
Reviewed-by: Greg Thompson <grt@chromium.org>
Commit-Queue: Samuel Huang <huangs@chromium.org>
Cr-Commit-Position: refs/heads/master@{#570401}
NOKEYCHECK=True
GitOrigin-RevId: ad7a5c086f00de62997714b84d6d6b5817ccc9d8

6 files changed, 310 insertions, 50 deletions

diff --git a/equivalence_map.cc b/equivalence_map.cc
index b3181ab..40071c0 100644
--- a/equivalence_map.cc
+++ b/equivalence_map.cc
@@ -8,6 +8,7 @@
 #include <utility>
 
 #include "base/logging.h"
+#include "base/numerics/safe_conversions.h"
 #include "components/zucchini/encoded_view.h"
 #include "components/zucchini/patch_reader.h"
 #include "components/zucchini/suffix_array.h"
@@ -191,15 +192,25 @@ EquivalenceCandidate VisitEquivalenceSeed(
 
 /******** OffsetMapper ********/
 
-OffsetMapper::OffsetMapper(std::vector<Equivalence>&& equivalences)
-    : equivalences_(std::move(equivalences)) {
+OffsetMapper::OffsetMapper(std::vector<Equivalence>&& equivalences,
+                           size_t old_image_size,
+                           size_t new_image_size)
+    : equivalences_(std::move(equivalences)),
+      old_image_size_(old_image_size),
+      new_image_size_(new_image_size) {
+  DCHECK_GT(new_image_size_, 0U);
   DCHECK(std::is_sorted(equivalences_.begin(), equivalences_.end(),
                         [](const Equivalence& a, const Equivalence& b) {
                           return a.src_offset < b.src_offset;
                         }));
+  // This is for testing. Assume pruned.
 }
 
-OffsetMapper::OffsetMapper(EquivalenceSource&& equivalence_source) {
+OffsetMapper::OffsetMapper(EquivalenceSource&& equivalence_source,
+                           size_t old_image_size,
+                           size_t new_image_size)
+    : old_image_size_(old_image_size), new_image_size_(new_image_size) {
+  DCHECK_GT(new_image_size_, 0U);
   for (auto e = equivalence_source.GetNext(); e.has_value();
        e = equivalence_source.GetNext()) {
     equivalences_.push_back(*e);
@@ -207,8 +218,13 @@ OffsetMapper::OffsetMapper(EquivalenceSource&& equivalence_source) {
   PruneEquivalencesAndSortBySource(&equivalences_);
 }
 
-OffsetMapper::OffsetMapper(const EquivalenceMap& equivalence_map)
-    : equivalences_(equivalence_map.size()) {
+OffsetMapper::OffsetMapper(const EquivalenceMap& equivalence_map,
+                           size_t old_image_size,
+                           size_t new_image_size)
+    : equivalences_(equivalence_map.size()),
+      old_image_size_(old_image_size),
+      new_image_size_(new_image_size) {
+  DCHECK_GT(new_image_size_, 0U);
   std::transform(equivalence_map.begin(), equivalence_map.end(),
                  equivalences_.begin(),
                  [](const EquivalenceCandidate& c) { return c.eq; });
@@ -217,17 +233,40 @@ OffsetMapper::OffsetMapper(const EquivalenceMap& equivalence_map)
 
 OffsetMapper::~OffsetMapper() = default;
 
-offset_t OffsetMapper::ForwardProject(offset_t offset) const {
-  auto pos = std::upper_bound(
-      equivalences_.begin(), equivalences_.end(), offset,
-      [](offset_t a, const Equivalence& b) { return a < b.src_offset; });
-  if (pos != equivalences_.begin()) {
-    if (pos == equivalences_.end() || offset < pos[-1].src_end() ||
-        offset - pos[-1].src_end() < pos->src_offset - offset) {
+// Safely evaluates |offset - unit.src_offset + unit.dst_offset| with signed
+// arithmetic, then clips the result to |[0, new_image_size_)|.
+offset_t OffsetMapper::NaiveExtendedForwardProject(const Equivalence& unit,
+                                                   offset_t offset) const {
+  int64_t old_offset64 = offset;
+  int64_t src_offset64 = unit.src_offset;
+  int64_t dst_offset64 = unit.dst_offset;
+  uint64_t new_offset64 = std::min<uint64_t>(
+      std::max<int64_t>(0LL, old_offset64 - src_offset64 + dst_offset64),
+      new_image_size_ - 1);
+  return base::checked_cast<offset_t>(new_offset64);
+}
+
+offset_t OffsetMapper::ExtendedForwardProject(offset_t offset) const {
+  DCHECK(!equivalences_.empty());
+  if (offset < old_image_size_) {
+    // Finds the equivalence unit whose "old" block is nearest to |offset|,
+    // favoring the block with lower offset in case of a tie.
+    auto pos = std::upper_bound(
+        equivalences_.begin(), equivalences_.end(), offset,
+        [](offset_t a, const Equivalence& b) { return a < b.src_offset; });
+    // For tiebreaking: |offset - pos[-1].src_end()| is actually 1 less than
+    // |offset|'s distance to "old" block of |pos[-1]|. Therefore "<" is used.
+    if (pos != equivalences_.begin() &&
+        (pos == equivalences_.end() || offset < pos[-1].src_end() ||
+         offset - pos[-1].src_end() < pos->src_offset - offset)) {
       --pos;
     }
+    return NaiveExtendedForwardProject(*pos, offset);
   }
-  return offset - pos->src_offset + pos->dst_offset;
+  // Fake offsets.
+  offset_t delta = offset - old_image_size_;
+  return delta < kOffsetBound - new_image_size_ ? new_image_size_ + delta
+                                                : kOffsetBound - 1;
 }
 
 void OffsetMapper::ForwardProjectAll(std::vector<offset_t>* offsets) const {
diff --git a/equivalence_map.h b/equivalence_map.h
index 91b215c..2ae8801 100644
--- a/equivalence_map.h
+++ b/equivalence_map.h
@@ -84,26 +84,46 @@ class OffsetMapper {
  public:
   using const_iterator = std::vector<Equivalence>::const_iterator;
 
-  // Constructors for various data sources.
+  // Constructors for various data sources. "Old" and "new" image sizes are
+  // needed for bounds checks and to handle dangling targets.
   // - From a list of |equivalences|, already sorted (by |src_offset|) and
   //   pruned, useful for tests.
-  explicit OffsetMapper(std::vector<Equivalence>&& equivalences);
+  OffsetMapper(std::vector<Equivalence>&& equivalences,
+               size_t old_image_size,
+               size_t new_image_size);
   // - From a generator, useful for Zucchini-apply.
-  explicit OffsetMapper(EquivalenceSource&& equivalence_source);
+  OffsetMapper(EquivalenceSource&& equivalence_source,
+               size_t old_image_size,
+               size_t new_image_size);
   // - From an EquivalenceMap that needs to be processed, useful for
   //   Zucchini-gen.
-  explicit OffsetMapper(const EquivalenceMap& equivalence_map);
+  OffsetMapper(const EquivalenceMap& equivalence_map,
+               size_t old_image_size,
+               size_t new_image_size);
   ~OffsetMapper();
 
   size_t size() const { return equivalences_.size(); }
   const_iterator begin() const { return equivalences_.begin(); }
   const_iterator end() const { return equivalences_.end(); }
 
+  // Returns naive extended forward-projection of "old" |offset| that follows
+  // |eq|'s delta. |eq| needs not cover |offset|.
+  // - Averts underflow / overflow by clamping to |[0, new_image_size_)|.
+  // - However, |offset| is *not* restricted to |[0, old_image_size_)|; the
+  //   caller must to make the check (hence "naive").
+  offset_t NaiveExtendedForwardProject(const Equivalence& unit,
+                                       offset_t offset) const;
+
   // Returns an offset in |new_image| corresponding to |offset| in |old_image|.
-  // If |offset| is not part of an equivalence, the equivalence nearest to
-  // |offset| is used as if it contained |offset|. This assumes |equivalences_|
-  // is not empty.
-  offset_t ForwardProject(offset_t offset) const;
+  // Assumes |equivalences_| to be non-empty. Cases:
+  // - If |offset| is covered (i.e., in an "old" block), then use the delta of
+  //   the (unique) equivalence unit that covers |offset|.
+  // - If |offset| is non-covered, but in |[0, old_image_size_)|, then find the
+  //   nearest "old" block, use its delta, and avert underflow / overflow by
+  //   clamping the result to |[0, new_image_size_)|.
+  // - If |offset| is >= |new_image_size_| (a "fake offset"), then use
+  //   |new_image_size_ - old_image_size_| as the delta.
+  offset_t ExtendedForwardProject(offset_t offset) const;
 
   // Given sorted |offsets|, applies a projection in-place of all offsets that
   // are part of a pruned equivalence from |old_image| to |new_image|. Other
@@ -123,7 +143,11 @@ class OffsetMapper {
       std::vector<Equivalence>* equivalences);
 
  private:
-  std::vector<Equivalence> equivalences_;
+  // |equivalences_| is pruned, i.e., no "old" blocks overlap (and no "new"
+  // block overlaps). Also, it is sorted by "old" offsets.
+  std::vector<Equivalence> equivalences_;  // P
+  const size_t old_image_size_;
+  const size_t new_image_size_;
 };
 
 // Container of equivalences between |old_image_index| and |new_image_index|,
diff --git a/equivalence_map_unittest.cc b/equivalence_map_unittest.cc
index 9c4166f..b3a4ea4 100644
--- a/equivalence_map_unittest.cc
+++ b/equivalence_map_unittest.cc
@@ -5,6 +5,9 @@
 #include "components/zucchini/equivalence_map.h"
 
 #include <cstring>
+#include <deque>
+#include <map>
+#include <string>
 #include <utility>
 #include <vector>
 
@@ -296,7 +299,209 @@ TEST(EquivalenceMapTest, PruneEquivalencesAndSortBySource) {
             }()));
 }
 
-TEST(EquivalenceMapTest, ForwardProject) {
+TEST(EquivalenceMapTest, NaiveExtendedForwardProject) {
+  constexpr size_t kOldImageSize = 1000U;
+  constexpr size_t kNewImageSize = 1000U;
+  OffsetMapper offset_mapper(std::vector<Equivalence>(), kOldImageSize,
+                             kNewImageSize);
+
+  // Convenience function to declutter.
+  auto project = [&offset_mapper](const Equivalence& eq, offset_t offset) {
+    return offset_mapper.NaiveExtendedForwardProject(eq, offset);
+  };
+
+  // Equivalence with delta = 0.
+  Equivalence eq_stay = {10, 10, +5};  // [10,15) -> [10,15).
+  for (offset_t offset = 0U; offset < 1000U; ++offset) {
+    EXPECT_EQ(offset, project(eq_stay, offset));
+  }
+  // Saturate since result would overflow "new".
+  EXPECT_EQ(999U, project(eq_stay, 1000U));
+  EXPECT_EQ(999U, project(eq_stay, 2000U));
+  EXPECT_EQ(999U, project(eq_stay, kOffsetBound - 1));
+
+  // Equivalence with delta = -10.
+  Equivalence eq_dec = {20, 10, +12};  // [20,32) --> [10,22).
+  // Offsets in "old" block.
+  EXPECT_EQ(10U, project(eq_dec, 20U));
+  EXPECT_EQ(11U, project(eq_dec, 21U));
+  EXPECT_EQ(21U, project(eq_dec, 31U));
+  // Offsets before "old" block, no underflow
+  EXPECT_EQ(9U, project(eq_dec, 19U));
+  EXPECT_EQ(1U, project(eq_dec, 11U));
+  EXPECT_EQ(0U, project(eq_dec, 10U));
+  // Offsets before "old" block, underflow (possible since delta < 0).
+  EXPECT_EQ(0U, project(eq_dec, 9U));
+  EXPECT_EQ(0U, project(eq_dec, 5U));
+  EXPECT_EQ(0U, project(eq_dec, 0U));
+  // Offsets after "old" block, no overflow.
+  EXPECT_EQ(20U, project(eq_dec, 30U));
+  EXPECT_EQ(64U, project(eq_dec, 74U));
+  EXPECT_EQ(90U, project(eq_dec, 100U));
+  EXPECT_EQ(490U, project(eq_dec, 500U));
+  EXPECT_EQ(999U, project(eq_dec, 1009U));
+  // Offsets after "old" block, overflow.
+  EXPECT_EQ(999U, project(eq_dec, 1010U));
+  EXPECT_EQ(999U, project(eq_dec, 2000U));
+  EXPECT_EQ(999U, project(eq_dec, kOffsetBound - 1));
+
+  // Equivalence with delta = +10.
+  Equivalence eq_inc = {7, 17, +80};  // [7,87) --> [17,97).
+  // Offsets in "old" block.
+  EXPECT_EQ(17U, project(eq_inc, 7U));
+  EXPECT_EQ(60U, project(eq_inc, 50U));
+  EXPECT_EQ(96U, project(eq_inc, 86U));
+  // Offsets before "old" block, underflow impossible since delta >= 0.
+  EXPECT_EQ(16U, project(eq_inc, 6U));
+  EXPECT_EQ(10U, project(eq_inc, 0U));
+  // Offsets after "old" block, no overflow.
+  EXPECT_EQ(97U, project(eq_inc, 87U));
+  EXPECT_EQ(510U, project(eq_inc, 500U));
+  EXPECT_EQ(999U, project(eq_inc, 989U));
+  // Offsets after "old" block, overflow.
+  EXPECT_EQ(999U, project(eq_inc, 990U));
+  EXPECT_EQ(999U, project(eq_inc, 2000U));
+  EXPECT_EQ(999U, project(eq_inc, kOffsetBound - 1));
+}
+
+TEST(EquivalenceMapTest, ExtendedForwardProject) {
+  // EquivalenceMaps provided must be sorted by "old" offset, and pruned.
+  // [0,2) --> [10,12), [2,3) --> [13,14), [4,6) --> [16,18).
+  OffsetMapper offset_mapper1({{0, 10, +2}, {2, 13, +1}, {4, 16, +2}}, 20U,
+                              25U);
+  EXPECT_EQ(10U, offset_mapper1.ExtendedForwardProject(0U));
+  EXPECT_EQ(11U, offset_mapper1.ExtendedForwardProject(1U));
+  EXPECT_EQ(13U, offset_mapper1.ExtendedForwardProject(2U));
+  EXPECT_EQ(14U, offset_mapper1.ExtendedForwardProject(3U));  // Previous equiv.
+  EXPECT_EQ(16U, offset_mapper1.ExtendedForwardProject(4U));
+  EXPECT_EQ(17U, offset_mapper1.ExtendedForwardProject(5U));
+  EXPECT_EQ(18U, offset_mapper1.ExtendedForwardProject(6U));  // Previous equiv.
+  // Fake offsets.
+  EXPECT_EQ(25U, offset_mapper1.ExtendedForwardProject(20U));
+  EXPECT_EQ(26U, offset_mapper1.ExtendedForwardProject(21U));
+  EXPECT_EQ(1005U, offset_mapper1.ExtendedForwardProject(1000U));
+  EXPECT_EQ(kOffsetBound - 1,
+            offset_mapper1.ExtendedForwardProject(kOffsetBound - 1));
+
+  // [0,2) --> [10,12), [13,14) --> [2,3), [16,18) --> [4,6).
+  OffsetMapper offset_mapper2({{0, 10, +2}, {13, 2, +1}, {16, 4, +2}}, 25U,
+                              20U);
+  EXPECT_EQ(10U, offset_mapper2.ExtendedForwardProject(0U));
+  EXPECT_EQ(11U, offset_mapper2.ExtendedForwardProject(1U));
+  EXPECT_EQ(2U, offset_mapper2.ExtendedForwardProject(13U));
+  EXPECT_EQ(3U, offset_mapper2.ExtendedForwardProject(14U));  // Previous equiv.
+  EXPECT_EQ(4U, offset_mapper2.ExtendedForwardProject(16U));
+  EXPECT_EQ(5U, offset_mapper2.ExtendedForwardProject(17U));
+  EXPECT_EQ(6U, offset_mapper2.ExtendedForwardProject(18U));  // Previous equiv.
+  // Fake offsets.
+  EXPECT_EQ(20U, offset_mapper2.ExtendedForwardProject(25U));
+  EXPECT_EQ(21U, offset_mapper2.ExtendedForwardProject(26U));
+  EXPECT_EQ(995U, offset_mapper2.ExtendedForwardProject(1000U));
+  EXPECT_EQ(kOffsetBound - 1 - 5,
+            offset_mapper2.ExtendedForwardProject(kOffsetBound - 1));
+}
+
+TEST(EquivalenceMapTest, ExtendedForwardProjectEncoding) {
+  // Tests OffsetMapper::ExtendedForwardProject(), which maps every "old" offset
+  // to a "new" offset, with possible overlap (even though blocks don't
+  // overlap). Not testing real offsets only (no fake offsets).
+  // |old_spec| is a string like "<<aaAAaabbBBbcCCc>>":
+  // - Upper case letters are covered "old" offsets.
+  // - Lower case letters are non-covered offsets that are properly mapped using
+  //   nearest "old" block.
+  // - '<' denotes underflow (clamped to 0).
+  // - '>' denotes overflow (clampled to "new" size - 1).
+  // |new_spec| is a string like "aaAA(ab)(ab)BBb..cCCc":
+  // - Upper and lower case letters are mapped "new" targets, occurring in the
+  //   order that they appear in |old_spec|.
+  // - '.' are "new" offsets that appear as output.
+  // - '(' and ')' surround a single "new" location that are repeated as output.
+  int case_no = 0;
+  auto run_test = [&case_no](std::vector<Equivalence>&& equivalences,
+                             const std::string& old_spec,
+                             const std::string& new_spec) {
+    const size_t old_size = old_spec.length();
+    // Build expected "new" offsets, queue up for each letter.
+    std::map<char, std::deque<offset_t>> expected;
+    offset_t cur_new_offset = 0;
+    char state = ')';  // ')' = increase offset, '(' = stay.
+    for (char ch : new_spec) {
+      if (ch == '(' || ch == ')')
+        state = ch;
+      else
+        expected[ch].push_back(cur_new_offset);
+      cur_new_offset += (state == ')') ? 1 : 0;
+    }
+    const size_t new_size = cur_new_offset;
+    // Forward-project for each "old" index, pull from queue from matching
+    // letter, and compare.
+    OffsetMapper offset_mapper(std::move(equivalences), old_size, new_size);
+    for (offset_t old_offset = 0; old_offset < old_size; ++old_offset) {
+      offset_t new_offset = offset_mapper.ExtendedForwardProject(old_offset);
+      char ch = old_spec[old_offset];
+      if (ch == '<') {  // Special case: Underflow.
+        EXPECT_EQ(0U, new_offset) << "in case " << case_no;
+      } else if (ch == '>') {  // Special case: Overflow.
+        EXPECT_EQ(static_cast<offset_t>(new_size - 1), new_offset)
+            << "in case " << case_no;
+      } else {
+        std::deque<offset_t>& q = expected[ch];
+        ASSERT_FALSE(q.empty());
+        EXPECT_EQ(q.front(), new_offset) << "in case " << case_no;
+        q.pop_front();
+        if (q.empty())
+          expected.erase(ch);
+      }
+    }
+    // Clear useless '.', and ensure everything is consumed.
+    expected.erase('.');
+    EXPECT_TRUE(expected.empty()) << "in case " << case_no;
+    ++case_no;
+  };
+
+  // Trivial: [5,9) --> [5,9).
+  run_test({{5, 5, +4}}, "aaaaaAAAAaaaaa", "aaaaaAAAAaaaaa");
+  // Swap: [0,4) --> [6,10), [4,10) --> [0,6).
+  run_test({{0, 6, +4}, {4, 0, +6}}, "AAAABBBBBB", "BBBBBBAAAA");
+  // Overlap: [0,4) --> [2,6), [4,10) --> [3,9).
+  run_test({{0, 2, +4}, {4, 3, +6}}, "AAAABBBBBB", "..A(AB)(AB)(AB)BBB.");
+  // Converge: [1,3) --> [2,4), [7,8) --> [6,7).
+  run_test({{1, 2, +2}, {7, 6, +1}}, "aAAaabbBbb", ".aAA(ab)(ab)Bbb.");
+  // Converge with tie-breaker: [1,3) --> [2,4), [8,9) --> [7,8).
+  run_test({{1, 2, +2}, {8, 7, +1}}, "aAAaaabbBb", ".aAAa(ab)(ab)Bb.");
+  // Shift left: [6,8) --> [2,4): Underflow occurs.
+  run_test({{6, 2, +2}}, "<<<<aaAAaa", "aaAAaa....");
+  // Shift right: [2,5) --> [6,9): Overflow occurs.
+  run_test({{2, 6, +3}}, "aaAAAa>>>>", "....aaAAAa");
+  // Diverge: [3,5) --> [1,3], [7,9) --> [9,11).
+  run_test({{3, 1, +2}, {7, 9, +2}}, "<<aAAabBBb>>", "aAAa....bBBb");
+  // Pile-up: [0,2) --> [7,9), [9,11) --> [9,11), [18,20) --> [11,13).
+  run_test({{0, 7, +2}, {9, 9, +2}, {18, 11, +2}}, "AAaaaabbbBBbbbbcccCC",
+           "......b(Ab)(Abc)(Bac)(Bac)(Cab)(Cab)bb.....");
+  // Inverse pile-up: [7,9) --> [0,2), [9,11) --> [9,11), [13,15) --> [18,20).
+  run_test({{7, 0, +2}, {9, 9, +2}, {11, 18, +2}}, "<<<<<<<AABBCC>>>>>>>",
+           "AA.......BB.......CC");
+  // Sparse rotate: [3,4) -> [10,11), [10,11) --> [17,18), [17,18) --> [3,4).
+  run_test({{3, 10, +1}, {10, 17, +1}, {17, 3, +1}}, "aaaAaaabbbBbbbcccCccc",
+           "cccCcccaaaAaaabbbBbbb");
+  // Messy swap: [2,4) --> [10,12), [12,16) --> [3,7).
+  run_test({{2, 10, +2}, {12, 3, +4}}, "aaAAaa>><bbbBBBBbb",
+           "bbbBBBBb(ab)aAAaa");
+  // Messy expand: [6,8) --> [3,5), [10,11) -> [11,12), [14,17) --> [16,19).
+  run_test({{6, 3, +2}, {10, 11, +1}, {14, 16, +3}}, "<<<aaaAAabBbbcCCCc>>>>>",
+           "aaaAAa....bBbb.cCCCc");
+  // Interleave: [1,2) --> [0,1), [5,6) --> [10,11), [6,8) --> [3,5),
+  //             [11,13) --> [12,14), [14,16) --> [6,8), [17,18) --> [17,18).
+  run_test({{1, 0, +1},
+            {5, 10, +1},
+            {6, 3, +2},
+            {11, 12, +2},
+            {14, 6, +2},
+            {17, 17, +1}},
+           "<AaabBCCccdDDdEEeFf>", "AaaCCc(Ec)EebBdDDd..Ff");
+}
+
+TEST(EquivalenceMapTest, ForwardProjectAll) {
   auto ForwardProjectAllTest = [](const OffsetMapper& offset_mapper,
                                   std::initializer_list<offset_t> offsets) {
     OffsetVector offsets_vec(offsets);
@@ -304,7 +509,9 @@ TEST(EquivalenceMapTest, ForwardProject) {
     return offsets_vec;
   };
 
-  OffsetMapper offset_mapper1({{0, 10, 2}, {2, 13, 1}, {4, 16, 2}});
+  // [0,2) --> [10,12), [2,3) --> [13,14), [4,6) --> [16,18).
+  OffsetMapper offset_mapper1({{0, 10, +2}, {2, 13, +1}, {4, 16, +2}}, 100U,
+                              100U);
   EXPECT_EQ(OffsetVector({10}), ForwardProjectAllTest(offset_mapper1, {0}));
   EXPECT_EQ(OffsetVector({13}), ForwardProjectAllTest(offset_mapper1, {2}));
   EXPECT_EQ(OffsetVector({}), ForwardProjectAllTest(offset_mapper1, {3}));
@@ -317,7 +524,9 @@ TEST(EquivalenceMapTest, ForwardProject) {
   EXPECT_EQ(OffsetVector({10, 11, 13, 16, 17}),
             ForwardProjectAllTest(offset_mapper1, {0, 1, 2, 3, 4, 5, 6}));
 
-  OffsetMapper offset_mapper2({{0, 10, 2}, {13, 2, 1}, {16, 4, 2}});
+  // [0,2) --> [10,12), [13,14) --> [2,3), [16,18) --> [4,6).
+  OffsetMapper offset_mapper2({{0, 10, +2}, {13, 2, +1}, {16, 4, +2}}, 100U,
+                              100U);
   EXPECT_EQ(OffsetVector({2}), ForwardProjectAllTest(offset_mapper2, {13}));
   EXPECT_EQ(OffsetVector({10, 2}),
             ForwardProjectAllTest(offset_mapper2, {0, 13}));
@@ -329,26 +538,6 @@ TEST(EquivalenceMapTest, ForwardProject) {
             ForwardProjectAllTest(offset_mapper2, {0, 1, 13, 14, 16, 17, 18}));
 }
 
-TEST(EquivalenceMapTest, ProjectOffset) {
-  OffsetMapper offset_mapper1({{0, 10, 2}, {2, 13, 1}, {4, 16, 2}});
-  EXPECT_EQ(10U, offset_mapper1.ForwardProject(0));
-  EXPECT_EQ(11U, offset_mapper1.ForwardProject(1));
-  EXPECT_EQ(13U, offset_mapper1.ForwardProject(2));
-  EXPECT_EQ(14U, offset_mapper1.ForwardProject(3));  // Previous equivalence.
-  EXPECT_EQ(16U, offset_mapper1.ForwardProject(4));
-  EXPECT_EQ(17U, offset_mapper1.ForwardProject(5));
-  EXPECT_EQ(18U, offset_mapper1.ForwardProject(6));  // Previous equivalence.
-
-  OffsetMapper offset_mapper2({{0, 10, 2}, {13, 2, 1}, {16, 4, 2}});
-  EXPECT_EQ(10U, offset_mapper2.ForwardProject(0));
-  EXPECT_EQ(11U, offset_mapper2.ForwardProject(1));
-  EXPECT_EQ(2U, offset_mapper2.ForwardProject(13));
-  EXPECT_EQ(3U, offset_mapper2.ForwardProject(14));  // Previous equivalence.
-  EXPECT_EQ(4U, offset_mapper2.ForwardProject(16));
-  EXPECT_EQ(5U, offset_mapper2.ForwardProject(17));
-  EXPECT_EQ(6U, offset_mapper2.ForwardProject(18));  // Previous equivalence.
-}
-
 TEST(EquivalenceMapTest, Build) {
   auto test_build_equivalence = [](const ImageIndex old_index,
                                    const ImageIndex new_index,
diff --git a/zucchini_apply.cc b/zucchini_apply.cc
index f6226bc..e202085 100644
--- a/zucchini_apply.cc
+++ b/zucchini_apply.cc
@@ -112,7 +112,8 @@ bool ApplyReferencesCorrection(ExecutableType exe_type,
   for (const auto& ref_group : old_disasm->MakeReferenceGroups())
     pool_groups[ref_group.pool_tag()].push_back(ref_group);
 
-  OffsetMapper offset_mapper(patch.GetEquivalenceSource());
+  OffsetMapper offset_mapper(patch.GetEquivalenceSource(), old_image.size(),
+                             new_image.size());
 
   std::vector<ReferenceGroup> new_groups = new_disasm->MakeReferenceGroups();
   for (const auto& pool_and_sub_groups : pool_groups) {
@@ -149,7 +150,8 @@ bool ApplyReferencesCorrection(ExecutableType exe_type,
           DCHECK_GE(ref->location, equivalence->src_offset);
           DCHECK_LT(ref->location, equivalence->src_end());
 
-          offset_t projected_target = offset_mapper.ForwardProject(ref->target);
+          offset_t projected_target =
+              offset_mapper.ExtendedForwardProject(ref->target);
           offset_t expected_key = targets.KeyForNearestOffset(projected_target);
           auto delta = ref_delta_source.GetNext();
           if (!delta.has_value()) {
diff --git a/zucchini_gen.cc b/zucchini_gen.cc
index 19da0af..0bb4658 100644
--- a/zucchini_gen.cc
+++ b/zucchini_gen.cc
@@ -211,7 +211,8 @@ bool GenerateReferencesDelta(const ReferenceSet& src_refs,
                 dst_ref->location - equiv.dst_offset);
       offset_t old_offset =
           src_refs.target_pool().OffsetForKey(src_ref->target_key);
-      offset_t new_estimated_offset = offset_mapper.ForwardProject(old_offset);
+      offset_t new_estimated_offset =
+          offset_mapper.ExtendedForwardProject(old_offset);
       offset_t new_estimated_key =
           projected_target_pool.KeyForNearestOffset(new_estimated_offset);
       offset_t new_offset =
@@ -288,7 +289,7 @@ bool GenerateExecutableElement(ExecutableType exe_type,
   EquivalenceMap equivalences =
       CreateEquivalenceMap(old_image_index, new_image_index,
                            new_disasm->num_equivalence_iterations());
-  OffsetMapper offset_mapper(equivalences);
+  OffsetMapper offset_mapper(equivalences, old_image.size(), new_image.size());
 
   ReferenceDeltaSink reference_delta_sink;
   for (const auto& old_targets : old_image_index.target_pools()) {
diff --git a/zucchini_gen_unittest.cc b/zucchini_gen_unittest.cc
index fb16cdd..97b223e 100644
--- a/zucchini_gen_unittest.cc
+++ b/zucchini_gen_unittest.cc
@@ -33,6 +33,11 @@ std::vector<int32_t> GenerateReferencesDeltaTest(
     std::vector<offset_t>&& exp_old_targets,
     std::vector<offset_t>&& exp_projected_old_targets,
     EquivalenceMap&& equivalence_map) {
+  // OffsetMapper needs image sizes for forward-projection overflow check. These
+  // are tested elsewhere, so just use arbitrary large value.
+  constexpr size_t kOldImageSize = 1000000;
+  constexpr size_t kNewImageSize = 1001000;
+
   ReferenceDeltaSink reference_delta_sink;
 
   TargetPool old_targets;
@@ -46,7 +51,7 @@ std::vector<int32_t> GenerateReferencesDeltaTest(
   ReferenceSet new_refs({1, TypeTag(0), PoolTag(0)}, new_targets);
   new_refs.InitReferences(new_references);
 
-  OffsetMapper offset_mapper(equivalence_map);
+  OffsetMapper offset_mapper(equivalence_map, kOldImageSize, kNewImageSize);
   TargetPool projected_old_targets = old_targets;
   projected_old_targets.FilterAndProject(offset_mapper);