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: