[Utils][NFC] Clean up update_mc_test_checks.py.

Refine the code a bit to make it easier to comprehend the logic.

Author: kosarev

llvm/utils/update_mc_test_checks.py

@@ -222,9 +222,6 @@ def update_test(ti: common.TestInfo):
     testlines = list(dict.fromkeys(testlines))
     common.debug("Valid test line found: ", len(testlines))
 
-    run_list_size = len(run_list)
-    testnum = len(testlines)
-
     raw_output = []
     raw_prefixes = []
     for (
@@ -266,60 +263,47 @@ def update_test(ti: common.TestInfo):
     prefix_set = set([prefix for p in run_list for prefix in p[0]])
     common.debug("Rewriting FileCheck prefixes:", str(prefix_set))
 
-    for test_id in range(testnum):
-        input_line = testlines[test_id]
-
+    for test_id, input_line in enumerate(testlines):
         # a {prefix : output, [runid] } dict
         # insert output to a prefix-key dict, and do a max sorting
         # to select the most-used prefix which share the same output string
         p_dict = {}
-        for run_id in range(run_list_size):
+        for run_id in range(len(run_list)):
             out = raw_output[run_id][test_id]
 
             if hasErr(out):
                 o = getErrString(out)
             else:
                 o = getOutputString(out)
 
-            prefixes = raw_prefixes[run_id]
-
-            for p in prefixes:
+            for p in raw_prefixes[run_id]:
                 if p not in p_dict:
                     p_dict[p] = o, [run_id]
-                else:
-                    if p_dict[p] == (None, []):
-                        continue
+                    continue
 
-                    prev_o, run_ids = p_dict[p]
-                    if o == prev_o:
-                        run_ids.append(run_id)
-                        p_dict[p] = o, run_ids
-                    else:
-                        # conflict, discard
-                        p_dict[p] = None, []
+                if p_dict[p] == (None, []):
+                    continue
 
-        p_dict_sorted = dict(sorted(p_dict.items(), key=lambda item: -len(item[1][1])))
+                prev_o, run_ids = p_dict[p]
+                if o == prev_o:
+                    run_ids.append(run_id)
+                    p_dict[p] = o, run_ids
+                else:
+                    # conflict, discard
+                    p_dict[p] = None, []
 
         # prefix is selected and generated with most shared output lines
         # each run_id can only be used once
-        used_runid = set()
-
+        used_run_ids = set()
         selected_prefixes = set()
-        for prefix, tup in p_dict_sorted.items():
-            o, run_ids = tup
-
-            if len(run_ids) == 0:
-                continue
-
-            skip = False
-            for i in run_ids:
-                if i in used_runid:
-                    skip = True
-                else:
-                    used_runid.add(i)
-            if not skip:
+        get_num_runs = lambda item: len(item[1][1])
+        p_dict_sorted = sorted(p_dict.items(), key=get_num_runs, reverse=True)
+        for prefix, (o, run_ids) in p_dict_sorted:
+            if run_ids and used_run_ids.isdisjoint(run_ids):
                 selected_prefixes.add(prefix)
 
+            used_run_ids.update(run_ids)
+
         # Generate check lines in alphabetical order.
         check_lines = []
         for prefix in sorted(selected_prefixes):