diff --git a/tools/releasetools/blockimgdiff.py b/tools/releasetools/blockimgdiff.py
index 42007eb..31d9645 100644
--- a/tools/releasetools/blockimgdiff.py
+++ b/tools/releasetools/blockimgdiff.py
@@ -338,8 +338,8 @@
           sid = next_stash_id
           next_stash_id += 1
         stashes[s] = sid
-        stashed_blocks += sr.size()
         if self.version == 2:
+          stashed_blocks += sr.size()
           out.append("stash %d %s\n" % (sid, sr.to_string_raw()))
         else:
           sh = self.HashBlocks(self.src, sr)
@@ -347,12 +347,14 @@
             stashes[sh] += 1
           else:
             stashes[sh] = 1
+            stashed_blocks += sr.size()
             out.append("stash %s %s\n" % (sh, sr.to_string_raw()))
 
       if stashed_blocks > max_stashed_blocks:
         max_stashed_blocks = stashed_blocks
 
       free_string = []
+      free_size = 0
 
       if self.version == 1:
         src_str = xf.src_ranges.to_string_raw()
@@ -371,7 +373,6 @@
         mapped_stashes = []
         for s, sr in xf.use_stash:
           sid = stashes.pop(s)
-          stashed_blocks -= sr.size()
           unstashed_src_ranges = unstashed_src_ranges.subtract(sr)
           sh = self.HashBlocks(self.src, sr)
           sr = xf.src_ranges.map_within(sr)
@@ -384,11 +385,13 @@
             # and lead to OTA failures.
             # Bug: 23119955
             free_string.append("free %d\n" % (sid,))
+            free_size += sr.size()
           else:
             assert sh in stashes
             src_str.append("%s:%s" % (sh, sr.to_string_raw()))
             stashes[sh] -= 1
             if stashes[sh] == 0:
+              free_size += sr.size()
               free_string.append("free %s\n" % (sh))
               stashes.pop(sh)
           heapq.heappush(free_stash_ids, sid)
@@ -492,6 +495,7 @@
 
       if free_string:
         out.append("".join(free_string))
+        stashed_blocks -= free_size
 
       if self.version >= 2 and common.OPTIONS.cache_size is not None:
         # Sanity check: abort if we're going to need more stash space than