1 files changed, 36 insertions, 14 deletions
diff --git a/tools/releasetools/sparse_img.py b/tools/releasetools/sparse_img.py
index 78f3057..7019f00 100644
--- a/tools/releasetools/sparse_img.py
+++ b/tools/releasetools/sparse_img.py
@@ -21,10 +21,17 @@ import rangelib
 
 
 class SparseImage(object):
-  """Wraps a sparse image file (and optional file map) into an image
-  object suitable for passing to BlockImageDiff."""
+  """Wraps a sparse image file into an image object.
 
-  def __init__(self, simg_fn, file_map_fn=None):
+  Wraps a sparse image file (and optional file map and clobbered_blocks) into
+  an image object suitable for passing to BlockImageDiff. file_map contains
+  the mapping between files and their blocks. clobbered_blocks contains the set
+  of blocks that should be always written to the target regardless of the old
+  contents (i.e. copying instead of patching). clobbered_blocks should be in
+  the form of a string like "0" or "0 1-5 8".
+  """
+
+  def __init__(self, simg_fn, file_map_fn=None, clobbered_blocks=None):
     self.simg_f = f = open(simg_fn, "rb")
 
     header_bin = f.read(28)
@@ -57,6 +64,7 @@ class SparseImage(object):
     pos = 0   # in blocks
     care_data = []
     self.offset_map = offset_map = []
+    self.clobbered_blocks = rangelib.RangeSet(data=clobbered_blocks)
 
     for i in range(total_chunks):
       header_bin = f.read(12)
@@ -103,17 +111,23 @@ class SparseImage(object):
     self.offset_index = [i[0] for i in offset_map]
 
     if file_map_fn:
-      self.LoadFileBlockMap(file_map_fn)
+      self.LoadFileBlockMap(file_map_fn, self.clobbered_blocks)
     else:
       self.file_map = {"__DATA": self.care_map}
 
   def ReadRangeSet(self, ranges):
     return [d for d in self._GetRangeData(ranges)]
 
-  def TotalSha1(self):
-    """Return the SHA-1 hash of all data in the 'care' regions of this image."""
+  def TotalSha1(self, include_clobbered_blocks=False):
+    """Return the SHA-1 hash of all data in the 'care' regions.
+
+    If include_clobbered_blocks is True, it returns the hash including the
+    clobbered_blocks."""
+    ranges = self.care_map
+    if not include_clobbered_blocks:
+      ranges = ranges.subtract(self.clobbered_blocks)
     h = sha1()
-    for d in self._GetRangeData(self.care_map):
+    for d in self._GetRangeData(ranges):
       h.update(d)
     return h.hexdigest()
 
@@ -156,7 +170,7 @@ class SparseImage(object):
           yield fill_data * (this_read * (self.blocksize >> 2))
         to_read -= this_read
 
-  def LoadFileBlockMap(self, fn):
+  def LoadFileBlockMap(self, fn, clobbered_blocks):
     remaining = self.care_map
     self.file_map = out = {}
 
@@ -166,14 +180,20 @@ class SparseImage(object):
         ranges = rangelib.RangeSet.parse(ranges)
         out[fn] = ranges
         assert ranges.size() == ranges.intersect(remaining).size()
+
+        # Currently we assume that blocks in clobbered_blocks are not part of
+        # any file.
+        assert not clobbered_blocks.overlaps(ranges)
         remaining = remaining.subtract(ranges)
 
+    remaining = remaining.subtract(clobbered_blocks)
+
     # For all the remaining blocks in the care_map (ie, those that
-    # aren't part of the data for any file), divide them into blocks
-    # that are all zero and blocks that aren't.  (Zero blocks are
-    # handled specially because (1) there are usually a lot of them
-    # and (2) bsdiff handles files with long sequences of repeated
-    # bytes especially poorly.)
+    # aren't part of the data for any file nor part of the clobbered_blocks),
+    # divide them into blocks that are all zero and blocks that aren't.
+    # (Zero blocks are handled specially because (1) there are usually
+    # a lot of them and (2) bsdiff handles files with long sequences of
+    # repeated bytes especially poorly.)
 
     zero_blocks = []
     nonzero_blocks = []
@@ -201,12 +221,14 @@ class SparseImage(object):
           nonzero_blocks.append(b)
           nonzero_blocks.append(b+1)
 
-    assert zero_blocks or nonzero_blocks
+    assert zero_blocks or nonzero_blocks or clobbered_blocks
 
     if zero_blocks:
       out["__ZERO"] = rangelib.RangeSet(data=zero_blocks)
     if nonzero_blocks:
       out["__NONZERO"] = rangelib.RangeSet(data=nonzero_blocks)
+    if clobbered_blocks:
+      out["__COPY"] = clobbered_blocks
 
   def ResetFileMap(self):
     """Throw away the file map and treat the entire image as