releasetools: Fix an issue in image size computation.

When building a system image with system_root_image enabled, the size
computation should include files under both of in_dir (i.e. /system
files) and root (pointed by 'root_dir'). Because files from both
locations will end up into the built image. The files under root are
usually only a few MiBs, but should be accounted for especially in the
context of logical partitions (where the partition size will be
allocated based on the actual need). Note that we will still need some
"reserved space" (defined via BOARD_*_PARTITION_RESERVED_SIZE) to cover
the cost for filesystem and/or verity metadata.

This CL moves the combination of the two dirs up, before parsing and
computing other properties. This doesn't affect anything for a
successful image building path. It may however increase the time to
error out in certain error path, since it copies the files earlier now.

Test: python -m unittest test_build_image
Test: `make dist`
Test: Setup a target with PRODUCT_USE_LOGICAL_PARTITIONS == true and
      BOARD_SYSTEMIMAGE_PARTITION_RESERVED_SIZE == 20MiB. Build system
      image successfully.
Test: Setup a target with PRODUCT_USE_LOGICAL_PARTITIONS == true and
      BOARD_SYSTEMIMAGE_PARTITION_RESERVED_SIZE == 20MiB. Write a large
      file to root dir (PRODUCT_OUT_ROOT). The image building fails, but
      reporting a size that accounts for both of /system and root.
Change-Id: Idfb26b8e259626ba57ec3bd4f85d357c30e56163
diff --git a/tools/releasetools/build_image.py b/tools/releasetools/build_image.py
index 8abc442..b7e114d 100755
--- a/tools/releasetools/build_image.py
+++ b/tools/releasetools/build_image.py
@@ -425,6 +425,52 @@
   return base_fs_file if exit_code == 0 else None
 
 
+def SetUpInDirAndFsConfig(origin_in, prop_dict):
+  """Returns the in_dir and fs_config that should be used for image building.
+
+  If the target uses system_root_image and it's building system.img, it creates
+  and returns a staged dir that combines the contents of /system (i.e. in the
+  given in_dir) and root.
+
+  Args:
+    origin_in: Path to the input directory.
+    prop_dict: A property dict that contains info like partition size. Values
+        may be updated.
+
+  Returns:
+    A tuple of in_dir and fs_config that should be used to build the image.
+  """
+  fs_config = prop_dict.get("fs_config")
+  if (prop_dict.get("system_root_image") != "true" or
+      prop_dict["mount_point"] != "system"):
+    return origin_in, fs_config
+
+  # Construct a staging directory of the root file system.
+  in_dir = common.MakeTempDir()
+  root_dir = prop_dict.get("root_dir")
+  if root_dir:
+    shutil.rmtree(in_dir)
+    shutil.copytree(root_dir, in_dir, symlinks=True)
+  in_dir_system = os.path.join(in_dir, "system")
+  shutil.rmtree(in_dir_system, ignore_errors=True)
+  shutil.copytree(origin_in, in_dir_system, symlinks=True)
+
+  # Change the mount point to "/".
+  prop_dict["mount_point"] = "/"
+  if fs_config:
+    # We need to merge the fs_config files of system and root.
+    merged_fs_config = common.MakeTempFile(
+        prefix="merged_fs_config", suffix=".txt")
+    with open(merged_fs_config, "w") as fw:
+      if "root_fs_config" in prop_dict:
+        with open(prop_dict["root_fs_config"]) as fr:
+          fw.writelines(fr.readlines())
+      with open(fs_config) as fr:
+        fw.writelines(fr.readlines())
+    fs_config = merged_fs_config
+  return in_dir, fs_config
+
+
 def CheckHeadroom(ext4fs_output, prop_dict):
   """Checks if there's enough headroom space available.
 
@@ -468,40 +514,25 @@
 
 
 def BuildImage(in_dir, prop_dict, out_file, target_out=None):
-  """Build an image to out_file from in_dir with property prop_dict.
-  After the function call, values in prop_dict is updated with
-  computed values.
+  """Builds an image for the files under in_dir and writes it to out_file.
+
+  When using system_root_image, it will additionally look for the files under
+  root (specified by 'root_dir') and builds an image that contains both sources.
 
   Args:
-    in_dir: path of input directory.
-    prop_dict: property dictionary.
-    out_file: path of the output image file.
-    target_out: path of the product out directory to read device specific FS
-        config files.
+    in_dir: Path to input directory.
+    prop_dict: A property dict that contains info like partition size. Values
+        will be updated with computed values.
+    out_file: The output image file.
+    target_out: Path to the TARGET_OUT directory as in Makefile. It actually
+        points to the /system directory under PRODUCT_OUT. fs_config (the one
+        under system/core/libcutils) reads device specific FS config files from
+        there.
 
   Returns:
     True iff the image is built successfully.
   """
-  # system_root_image=true: build a system.img that combines the contents of
-  # /system and root, which should be mounted at the root of the file system.
-  origin_in = in_dir
-  fs_config = prop_dict.get("fs_config")
-  if (prop_dict.get("system_root_image") == "true" and
-      prop_dict["mount_point"] == "system"):
-    in_dir = common.MakeTempDir()
-    # Change the mount point to "/".
-    prop_dict["mount_point"] = "/"
-    if fs_config:
-      # We need to merge the fs_config files of system and root.
-      merged_fs_config = common.MakeTempFile(prefix="merged_fs_config",
-                                             suffix=".txt")
-      with open(merged_fs_config, "w") as fw:
-        if "root_fs_config" in prop_dict:
-          with open(prop_dict["root_fs_config"]) as fr:
-            fw.writelines(fr.readlines())
-        with open(fs_config) as fr:
-          fw.writelines(fr.readlines())
-      fs_config = merged_fs_config
+  in_dir, fs_config = SetUpInDirAndFsConfig(in_dir, prop_dict)
 
   build_command = []
   fs_type = prop_dict.get("fs_type", "")
@@ -518,11 +549,11 @@
   if (prop_dict.get("use_logical_partitions") == "true" and
       "partition_size" not in prop_dict):
     # if partition_size is not defined, use output of `du' + reserved_size
-    success, size = GetDiskUsage(origin_in)
+    success, size = GetDiskUsage(in_dir)
     if not success:
       return False
     if OPTIONS.verbose:
-      print("The tree size of %s is %d MB." % (origin_in, size // BYTES_IN_MB))
+      print("The tree size of %s is %d MB." % (in_dir, size // BYTES_IN_MB))
     size += int(prop_dict.get("partition_reserved_size", 0))
     # Round this up to a multiple of 4K so that avbtool works
     size = common.RoundUpTo4K(size)
@@ -643,27 +674,17 @@
     print("Error: unknown filesystem type '%s'" % (fs_type))
     return False
 
-  if in_dir != origin_in:
-    # Construct a staging directory of the root file system.
-    root_dir = prop_dict.get("root_dir")
-    if root_dir:
-      shutil.rmtree(in_dir)
-      shutil.copytree(root_dir, in_dir, symlinks=True)
-    staging_system = os.path.join(in_dir, "system")
-    shutil.rmtree(staging_system, ignore_errors=True)
-    shutil.copytree(origin_in, staging_system, symlinks=True)
-
   (mkfs_output, exit_code) = RunCommand(build_command)
   if exit_code != 0:
     print("Error: '%s' failed with exit code %d:\n%s" % (
         build_command, exit_code, mkfs_output))
-    success, du = GetDiskUsage(origin_in)
+    success, du = GetDiskUsage(in_dir)
     du_str = ("%d bytes (%d MB)" % (du, du // BYTES_IN_MB)
              ) if success else "unknown"
     print(
         "Out of space? The tree size of {} is {}, with reserved space of {} "
         "bytes ({} MB).".format(
-            origin_in, du_str,
+            in_dir, du_str,
             int(prop_dict.get("partition_reserved_size", 0)),
             int(prop_dict.get("partition_reserved_size", 0)) // BYTES_IN_MB))
     if "original_partition_size" in prop_dict: