[rom_ctrl,util] Make collision check faster in scramble_image.py

Update the collision check implementation using a dict to traverse the
chunks only once.

Signed-off-by: Emmanuel Blot <[email protected]>

hw/ip/rom_ctrl/util/mem.py

@@ -323,24 +323,16 @@ def add_ecc32(self) -> None:
             chunk.add_ecc32(self.config)
         self.width = 39
 
-    def collisions(self) -> List[Tuple[int, int]]:
-        '''Find "collisions" in the scrambled memory
-
-        This looks at all pairs of words in the memory, looking for addresses
-        where the words are equal. Returns a list of pairs (addr0, addr1) of
-        such addresses where addr0 < addr1 and in ascending order of addr0.
+    def first_collision(self) -> Optional[Tuple[int, int]]:
+        '''Return the address of the first pair of colliding addresses
 
+           If there is no such pair (which is hopefully the case), return None.
         '''
-        ret = []
-        for idx0, chunk0 in enumerate(self.chunks):
-            for off0, word0 in enumerate(chunk0.words):
-                for diff_idx, chunk1 in enumerate(self.chunks[idx0:]):
-                    first_off1 = 0 if diff_idx else off0 + 1
-                    for diff_off, word1 in enumerate(chunk1.words[first_off1:]):
-                        off1 = first_off1 + diff_off
-
-                        if word0 == word1:
-                            addr0 = chunk0.base_addr + off0
-                            addr1 = chunk1.base_addr + off1
-                            ret.append((addr0, addr1))
-        return ret
+        known = {}  # type: Dict[int, int]
+        for chunk in self.chunks:
+            addr = chunk.base_addr
+            for off, word in enumerate(chunk.words):
+                if word in known:
+                    return known[word], addr + off
+                known[word] = addr + off
+        return None

hw/ip/rom_ctrl/util/scramble_image.py

@@ -391,14 +391,14 @@ def add_hash(self, scr_mem: MemFile) -> None:
         num_digest_words = 256 // 32
 
         # Read out the scrambled data in logical address order
-        to_hash = b''
+        to_hash = bytearray()
         for log_addr in range(self.rom_size_words - num_digest_words):
             phy_addr = self.addr_sp_enc(log_addr)
             scr_word = scr_chunk.words[phy_addr]
             # Note that a scrambled word with ECC amounts to 39bit. The
             # expression (39 + 7) // 8 calculates the amount of bytes that are
             # required to store these bits.
-            to_hash += scr_word.to_bytes((39 + 7) // 8, byteorder='little')
+            to_hash.extend(scr_word.to_bytes((39 + 7) // 8, byteorder='little'))
 
         # Hash it
         hash_obj = cSHAKE256.new(data=to_hash,
@@ -465,8 +465,9 @@ def main() -> int:
     scrambler.add_hash(scr_mem)
 
     # Check for collisions
-    collisions = scr_mem.collisions()
-    if collisions:
+    collision = scr_mem.first_collision()
+    if collision:
+        addr0, addr1 = collision
         print(
             'ERROR: This combination of ROM contents and scrambling\n'
             '       key results in one or more collisions where\n'
@@ -475,11 +476,9 @@ def main() -> int:
             '       Looks like we\'ve been (very) unlucky with the\n'
             '       birthday problem. As a work-around, try again after\n'
             '       generating some different RndCnst* parameters.\n',
+            '\n'
+            f'First colliding addresses: {addr0:#010x}, {addr1:#010x}',
             file=sys.stderr)
-        print('{} colliding addresses:'.format(len(collisions)),
-              file=sys.stderr)
-        for addr0, addr1 in collisions:
-            print('  {:#010x}, {:#010x}'.format(addr0, addr1), file=sys.stderr)
         return 1
 
     scr_mem.write_vmem(args.outfile)

util/design/secded_gen.py

@@ -409,7 +409,7 @@ def ecc_encode(config: Dict[str, Any], codetype: str, k: int, dataword: int) ->
 def ecc_encode_some(config: Dict[str, Any],
                     codetype: str,
                     k: int,
-                    datawords: int) -> Tuple[List[int], int]:
+                    datawords: List[int]) -> Tuple[List[int], int]:
     m, bitmasks, invert = _ecc_pick_code(config, codetype, k)
     codewords = [int(_ecc_encode(k, m, bitmasks, invert, w), 2)
                  for w in datawords]