1  import glob, re, os.path, textwrap
  2  from colorama import Fore, Back, Style
  3  
  4  def target_prefix(dir):
  5      if dir.startswith("net"):
  6          return "net"
  7      elif dir.startswith("serial"):
  8          return "serial"
  9      elif dir.startswith("util"):
 10          return "util"
 11      elif dir.startswith("runtime"):
 12          return "runtime"
 13      elif dir.startswith("zk"):
 14          return "zk"
 15      elif dir.startswith("raft"):
 16          return "raft"
 17      elif dir.startswith("sdk/src/crypto"):
 18          return "sdk::crypto"
 19      elif dir.startswith("sdk"):
 20          return "sdk"
 21      elif dir.startswith("contract/dao"):
 22          return "dao"
 23      elif dir.startswith("contract/money"):
 24          return "money"
 25      elif dir.startswith("rpc"):
 26          return "rpc"
 27      elif dir.startswith("system"):
 28          return "system"
 29      elif dir.startswith("dht"):
 30          return "dht"
 31      elif dir.startswith("consensus"):
 32          return "consensus"
 33      elif dir.startswith("zkas"):
 34          return "zkas"
 35      elif dir.startswith("blockchain"):
 36          return "blockchain"
 37      elif dir.startswith("wallet"):
 38          return "wallet"
 39      else:
 40          assert not dir or dir == "tx"
 41          return ""
 42  
 43  def target_suffix(prefix, base):
 44      if prefix in ("dao", "money"):
 45          # Just shorten the target to simply "dao" or "money"
 46          # We don't need the fine grained details
 47          return ""
 48      elif base in ("mod.rs", "lib.rs"):
 49          # Just use the module name as the target with these files
 50          return ""
 51      # Otherwise just use the filename as the target suffix
 52      return base.removesuffix(".rs")
 53  
 54  def log_target(fname):
 55      dir, base = os.path.dirname(fname), os.path.basename(fname)
 56      prefix = target_prefix(dir)
 57      suffix = target_suffix(prefix, base)
 58      # you don't need :: when the suffix is empty
 59      if not suffix and not prefix:
 60          return ""
 61      if not suffix:
 62          return prefix
 63      if not prefix:
 64          return suffix
 65      return f"{prefix}::{suffix}"
 66  
 67  def replace(fname, contents):
 68      target = log_target(fname)
 69  
 70      # You can debug like this:
 71      #if target != "net::protocol_seed":
 72      #    return None
 73  
 74      print(f"Replacing {target}" + " "*(40 - len(target)) + f"[{fname}]")
 75  
 76      result = []
 77      lines = contents.split("\n")
 78      i = 0
 79      while i < len(lines):
 80          line = lines[i]
 81          # Line number
 82          ln = lambda: i + 1
 83  
 84          # only used for debug output
 85          old_text = None
 86          new_text = None
 87          # This is used as a debug goto
 88          is_modified = False
 89  
 90          log_level = None
 91          if "trace!(" in line:
 92              log_level = "trace"
 93          elif "debug!(" in line:
 94              log_level = "debug"
 95          elif "info!(" in line:
 96              log_level = "info"
 97          elif "warn!(" in line:
 98              log_level = "warn"
 99          elif "error!(" in line:
100              log_level = "error"
101  
102          if log_level is not None:
103              # Walk backwards to find the function name
104              # Range is (i, 0]
105              function_name = None
106              for j in range(i - 1, -1, -1):
107                  past_line = lines[j]
108                  if (match := re.search(
109                      r"fn ([\w]+)(<[\w: \+']+>)?\(",
110                      past_line
111                  )):
112                      function_name = match.group(1)
113                      break
114              assert function_name is not None
115  
116              local_target = target
117              if target in (
118                  "consensus::protocol_proposal",
119                  "consensus::protocol_sync",
120                  "consensus::protocol_sync_consensus",
121                  "consensus::protocol_tx",
122                  "runtime::db",
123                  "net::hosts",
124                  "net::protocol_address",
125                  "net::protocol_ping",
126                  "net::protocol_seed",
127                  "net::protocol_version",
128                  "net::p2p",
129                  "net::message_subscriber",
130                  "net::channel",
131              ):
132                  local_target += f"::{function_name}()"
133  
134              # No target exists for this file at all. Just ignore these
135              # We would normally delete any target set for these files
136              # but so far we have none of them, so just ignore them.
137              if not target:
138                  print(
139                      "    "
140                      + Back.RED + "Skip [no target]:" + Style.RESET_ALL
141                      + f" {line}"
142                  )
143              # Single lines with a target that's a constant or string
144              elif re.search(rf'{log_level}!\(target: ([\w]+|"[\w:\-\(\)]+"),', line):
145                  old_text = f"{ln()}: {line}"
146  
147                  line = re.sub(
148                      r'target: ([\w]+|"[\w:\-\(\)]+"),',
149                      f'target: "{local_target}",',
150                      line
151                  )
152  
153                  is_modified = True
154                  new_text = f"{ln()}: {line}"
155              # Normal single lines with no target set
156              elif f'{log_level}!("' in line:
157                  old_text = f"{ln()}: {line}"
158  
159                  #print(f"    No target: {line}")
160                  line = line.replace(f'{log_level}!(',
161                                      f'{log_level}!(target: "{local_target}", ')
162  
163                  is_modified = True
164                  new_text = f"{ln()}: {line}"
165              # Multiline logs
166              # We read the next line and check if there's a target set or not
167              else:
168                  assert re.search(rf"{log_level}!\($", line)
169  
170                  old_text = f"{ln()}: {line}"
171                  new_text = f"{ln()}: {line}"
172  
173                  result.append(line)
174                  i += 1
175                  assert i < len(lines)
176                  line = lines[i]
177  
178                  old_text += f"\n{ln()}: {line}"
179  
180                  # Constants or target strings set
181                  if re.search(r'target: ([\w]+|"[\w:\-\(\)]+"),', line):
182                      line = re.sub(
183                          r'target: ([\w]+|"[\w:\-\(\)]+"),',
184                          f'target: "{local_target}",',
185                          line
186                      )
187  
188                      new_text += f"\n{ln()}: {line}"
189                  # Multi-line logs with no target set
190                  # Insert an extra line with the target
191                  else:
192                      assert re.search('^"', line)
193  
194                      leading_space = lambda line: len(line) - len(line.lstrip())
195  
196                      added_line = (" "*leading_space(line)
197                                    + f'target: "{local_target}",')
198                      result.append(added_line)
199  
200                      new_text += f"\n{ln()}: {added_line}\n{ln() + 1}: {line}"
201  
202                  is_modified = True
203  
204          if is_modified:
205              assert old_text is not None and new_text is not None
206              print(
207                  Fore.RED
208                  + textwrap.indent(old_text, "    < ")
209                  + Style.RESET_ALL
210              )
211              print(
212                  Fore.GREEN
213                  + textwrap.indent(new_text, "    > ")
214                  + Style.RESET_ALL
215              )
216              print()
217  
218          result.append(line)
219          i += 1
220      return "\n".join(result)
221  
222  def main():
223      for fname in glob.glob("**/*.rs", root_dir="src/", recursive=True):
224          with open(f"src/{fname}", "r") as f:
225              contents = f.read()
226  
227          if (contents := replace(fname, contents)) is None:
228              # Skip this file
229              continue
230  
231          # Uncomment this to apply the changes
232          with open(f"src/{fname}", "w") as f:
233              f.write(contents)
234  
235  if __name__ == "__main__":
236      main()
237