merge_dtbs: Use logging and argparse

Use the logging module to log important information to various log
levels: info, debug, warning and error.
- Modify __str__ representations of classes DeviceTree and
  InnerMergedDeviceTree for readability in debug logs
- Make parse_dt_files() return instead of yield so as to order all the
  filename parsing at one go after the preceding log messages

Use the argparse module to specify script arguments cleanly. Also make
the corresponding change to merge_dtbs.sh to invoke this script with the
correct flags.

Change-Id: Id09d9d7b34cde60eff99876f0c0c4f6643852aee
Signed-off-by: Guru Das Srinagesh <quic_gurus@quicinc.com>
Signed-off-by: Jis G Jacob <studiokeys@blissroms.org>
diff --git a/build/tasks/kernel.mk b/build/tasks/kernel.mk
index 2bf7270..644a0bd 100644
--- a/build/tasks/kernel.mk
+++ b/build/tasks/kernel.mk
@@ -634,7 +634,7 @@
 	$(hide) find $(DTBS_BASE) -type f -name "*.dtb*" | xargs rm -f
 	$(hide) find $(DTBS_OUT) -type f -name "*.dtb*" | xargs rm -f
 	mv $(DTB_OUT)/arch/$(KERNEL_ARCH)/boot/dts/vendor/*/*.dtb $(DTB_OUT)/arch/$(KERNEL_ARCH)/boot/dts/vendor/*/*.dtbo $(DTBS_BASE)/
-	PATH=$(abspath $(HOST_OUT_EXECUTABLES)):$${PATH} python3 $(BUILD_TOP)/vendor/bliss/build/tools/merge_dtbs.py $(DTBS_BASE) $(DTB_OUT)/arch/$(KERNEL_ARCH)/boot/dts/vendor/qcom $(DTBS_OUT)
+	PATH=$(abspath $(HOST_OUT_EXECUTABLES)):$${PATH} python3 $(BUILD_TOP)/vendor/bliss/build/tools/merge_dtbs.py --base $(DTBS_BASE) --techpack $(DTB_OUT)/arch/$(KERNEL_ARCH)/boot/dts/vendor/qcom --out $(DTBS_OUT)
 	cat $(shell find $(DTB_OUT)/out -type f -name "${TARGET_MERGE_DTBS_WILDCARD}.dtb" | sort) > $@
 else
 	cat $(shell find $(DTB_OUT)/arch/$(KERNEL_ARCH)/boot/dts -type f -name "*.dtb" | sort) > $@
diff --git a/build/tools/merge_dtbs.py b/build/tools/merge_dtbs.py
index fe5cd08..17778af 100755
--- a/build/tools/merge_dtbs.py
+++ b/build/tools/merge_dtbs.py
@@ -1,6 +1,7 @@
 #! /usr/bin/env python3
 
 # Copyright (c) 2020-2021, The Linux Foundation. All rights reserved.
+# Copyright (c) 2023 Qualcomm Innovation Center, Inc. All rights reserved.
 #
 # Redistribution and use in source and binary forms, with or without
 # modification, are permitted provided that the following conditions are
@@ -33,6 +34,8 @@
 import subprocess
 import shutil
 from itertools import product, combinations_with_replacement, chain
+import logging
+import argparse
 
 def split_array(array, cells):
 	"""
@@ -230,6 +233,7 @@
 class DeviceTree(DeviceTreeInfo):
 	def __init__(self, filename):
 		self.filename = filename
+		logging.debug('Initializing new DeviceTree: {}'.format(os.path.basename(filename)))
 		msm_id = split_array(self.get_prop('/', 'qcom,msm-id', check_output=False), 2)
 		board_id = split_array(self.get_prop('/', 'qcom,board-id', check_output=False), 2)
 		# default pmic-id-size is 4
@@ -239,7 +243,7 @@
 		super().__init__(msm_id, board_id, pmic_id, miboard_id)
 
 		if not self.has_any_properties():
-			print('WARNING! {} has no properties and may match with any other devicetree'.format(os.path.basename(self.filename)))
+			logging.warning('{} has no properties and may match with any other devicetree'.format(os.path.basename(self.filename)))
 
 	def get_prop(self, node, property, prop_type='i', check_output=True):
 		r = subprocess.run(["fdtget", "-t", prop_type, self.filename, node, property],
@@ -264,7 +268,7 @@
 		return out
 
 	def __str__(self):
-		return "{} [{}]".format(super().__str__(), os.path.basename(self.filename))
+		return "[{}] {}".format(os.path.basename(self.filename), super().__str__())
 
 class InnerMergedDeviceTree(DeviceTreeInfo):
 	"""
@@ -274,6 +278,7 @@
 	"""
 	def __init__(self, filename, plat_id, board_id, pmic_id, miboard_id, techpacks=None):
 		self.base = filename
+		# All inner merged device trees start with zero techpacks
 		self.techpacks = techpacks or []
 		super().__init__(plat_id, board_id, pmic_id, miboard_id)
 
@@ -283,6 +288,7 @@
 		intersection = techpack & self
 		if intersection in self:
 			self.techpacks.append(intersection)
+			logging.debug('Appended after intersection: {}'.format(repr(self)))
 			return True
 		return False
 
@@ -290,12 +296,13 @@
 		if name is None:
 			name = self.get_name()
 
+		logging.info('Saving to: {}'.format(name))
 		out_file = os.path.join(out_dir, name)
 		ext = os.path.splitext(os.path.basename(self.base))[1]
 
 		# This check might fail in future if we get into an edge case
 		# when splitting the base devicetree into multiple merged DTs
-		assert not os.path.exists(out_file)
+		assert not os.path.exists(out_file), "Cannot overwrite: {}".format(out_file)
 
 		if len(self.techpacks) == 0:
 			cmd = ['cp', self.base, out_file]
@@ -308,25 +315,25 @@
 			cmd.extend([tp.filename for tp in self.techpacks])
 			cmd.extend(['-o', out_file])
 
-		print(' {}'.format(' '.join(cmd)))
+		logging.debug(' {}'.format(' '.join(cmd)))
 		subprocess.run(cmd, check=True)
 
 		if self.plat_id:
 			plat_iter = self.plat_id if isinstance(self.plat_id, tuple) else chain.from_iterable(self.plat_id)
 			cmd = ['fdtput', '-t', 'i', out_file, '/', 'qcom,msm-id'] + list(map(str, plat_iter))
-			print('  {}'.format(' '.join(cmd)))
+			logging.debug('  {}'.format(' '.join(cmd)))
 			subprocess.run(cmd, check=True)
 
 		if self.board_id:
 			board_iter = self.board_id if isinstance(self.board_id, tuple) else chain.from_iterable(self.board_id)
 			cmd = ['fdtput', '-t', 'i', out_file, '/', 'qcom,board-id'] + list(map(str, board_iter))
-			print('  {}'.format(' '.join(cmd)))
+			logging.debug('  {}'.format(' '.join(cmd)))
 			subprocess.run(cmd, check=True)
 
 		if self.pmic_id:
 			pmic_iter = self.pmic_id if isinstance(self.pmic_id, tuple) else chain.from_iterable(self.pmic_id)
 			cmd = ['fdtput', '-t', 'i', out_file, '/', 'qcom,pmic-id'] + list(map(str, pmic_iter))
-			print('  {}'.format(' '.join(cmd)))
+			logging.debug('  {}'.format(' '.join(cmd)))
 			subprocess.run(cmd, check=True)
 
 		if self.miboard_id:
@@ -350,15 +357,16 @@
 		return [part for part in parts if part not in ignored_parts]
 
 	def __str__(self):
-		return "{} [{} + {{{}}}]".format(super().__str__(), os.path.basename(self.base), " ".join(os.path.basename(t.filename) for t in self.techpacks))
+		return "[{} + {{{}}}] {}".format(os.path.basename(self.base), " ".join(os.path.basename(t.filename) for t in self.techpacks), super().__str__())
 
 class MergedDeviceTree(object):
 	def __init__(self, other):
 		self.merged_devicetrees = {InnerMergedDeviceTree(other.filename, other.plat_id, other.board_id, other.pmic_id, other.miboard_id)}
 
-	def try_add(self, techpack):
+	def merged_dt_try_add(self, techpack):
 		did_add = False
 		for mdt in self.merged_devicetrees.copy():
+			logging.debug('{}'.format(repr(mdt)))
 			# techpack and kernel devicetree need only to overlap in order to merge,
 			# and not match exactly. Think: venn diagram.
 			# Need 2 things: The devicetree part that applies to
@@ -372,14 +380,16 @@
 			intersection = techpack & mdt
 			if intersection not in mdt:
 				continue
+			logging.debug('Found intersection!')
+
 			# mdt may apply to a superset of devices the techpack DT applies to
 			# (mdt - intersection) splits mdt into appropriate number of devicetrees
 			# such that we can apply techpack onto one of the resulting DTs in the
 			# difference
 			difference = mdt - intersection
 			if len(difference) > 1:
-				print('Splitting {}'.format(mdt))
-				print(' because  {}'.format(techpack))
+				logging.debug('Splitting {}'.format(mdt))
+				logging.debug('because {}'.format(techpack))
 				self.merged_devicetrees.remove(mdt)
 				self.merged_devicetrees.update(difference)
 
@@ -396,47 +406,61 @@
 		else:
 			name = None
 		for mdt in self.merged_devicetrees:
-			print()
 			yield mdt.save(name, out_dir)
 
 def parse_dt_files(dt_folder):
+	devicetrees = []
 	for root, dirs, files in os.walk(dt_folder):
 		for filename in files:
 			if os.path.splitext(filename)[1] not in ['.dtb', '.dtbo']:
 				continue
 			filepath = os.path.join(root, filename)
-			yield DeviceTree(filepath)
+			devicetrees.append(DeviceTree(filepath))
+	return devicetrees
 
 def main():
-	if len(sys.argv) != 4:
-		print("Usage: {} <base dtb folder> <techpack dtb folder> <output folder>"
-		      .format(sys.argv[0]))
-		sys.exit(1)
 
-	# 1. Parse the devicetrees -- extract the device info (msm-id, board-id, pmic-id, miboard-id)
-	bases = parse_dt_files(sys.argv[1])
-	techpacks = parse_dt_files(sys.argv[2])
+	parser = argparse.ArgumentParser(description='Merge devicetree blobs of techpacks with Kernel Platform SoCs')
+	parser.add_argument('-b', '--base', required=True, help="Folder containing base DTBs from Kernel Platform output")
+	parser.add_argument('-t', '--techpack', required=True, help="Folder containing techpack DLKM DTBOs")
+	parser.add_argument('-o', '--out', required=True, help="Output folder where merged DTBs will be saved")
+	parser.add_argument('--loglevel', choices=['debug', 'info', 'warn', 'error'], default='info', help="Set loglevel to see debug messages")
+
+	args = parser.parse_args()
+
+	logging.basicConfig(level=args.loglevel.upper(), format='%(levelname)s: %(message)s'.format(os.path.basename(sys.argv[0])))
+
+	# 1. Parse the devicetrees -- extract the device info (msm-id, board-id, pmic-id)
+	logging.info('Parsing base dtb files from {}'.format(args.base))
+	bases = parse_dt_files(args.base)
+	all_bases = '\n'.join(list(map(lambda x: str(x), bases)))
+	logging.info('Parsed bases: \n{}'.format(all_bases))
+
+	logging.info('Parsing techpack dtb files from {}'.format(args.techpack))
+	techpacks = parse_dt_files(args.techpack)
+	all_techpacks = '\n'.join(list(map(lambda x: str(x), techpacks)))
+	logging.info('Parsed techpacks: \n{}'.format(all_techpacks))
 
 	# 2.1: Create an intermediate representation of the merged devicetrees, starting with the base
 	merged_devicetrees = list(map(lambda dt: MergedDeviceTree(dt), bases))
 	# 2.2: Try to add techpack devicetree to each base DT
 	for techpack in techpacks:
+		logging.debug('Trying to add techpack: {}'.format(techpack))
 		did_add = False
 		for dt in merged_devicetrees:
-			if dt.try_add(techpack):
+			if dt.merged_dt_try_add(techpack):
 				did_add = True
 		if not did_add:
-			print('WARNING! Could not apply {} to any devicetrees'.format(techpack))
+			logging.warning('Could not apply {} to any devicetrees'.format(techpack))
 
-	print()
-	print('==================================')
+	final_inner_merged_dts = '\n'.join(list(str(mdt.merged_devicetrees) for mdt in merged_devicetrees))
+	logging.info('Final Inner Merged Devicetrees: \n{}'.format(final_inner_merged_dts))
+
 	created = []
 	# 3. Save the deviectrees to real files
 	for dt in merged_devicetrees:
-		created.extend(dt.save(sys.argv[3]))
+		created.extend(dt.save(args.out))
 
-	print()
-	print('==================================')
 	# 4. Try to apply merged DTBOs onto merged DTBs, when appropriate
 	#    This checks that DTBOs and DTBs generated by merge_dtbs.py can be merged by bootloader
 	#    at runtime.
@@ -446,7 +470,7 @@
 		# See DeviceTreeInfo.__gt__; this checks whether dtbo is more specific than the base
 		if dtbo > base:
 			cmd = ['ufdt_apply_overlay', base.filename, dtbo.filename, '/dev/null']
-			print(' '.join(cmd))
+			logging.debug(' '.join(cmd))
 			subprocess.run(cmd, check=True)