Add script to refactor crate:: to super:: references

This will be run on crates that are to be moved into the
base crate.
The script is run in the follow up CL.

Note: This is a one-off script. It just has to work one,
not for the general case.

BUG=b:223206469
TEST=python tools/contrib/refactor_use_references.py

Change-Id: Ic92109572649c9130784b986c67fddc8f9b838d6
Reviewed-on: https://chromium-review.googlesource.com/c/chromiumos/platform/crosvm/+/3530497
Reviewed-by: Anton Romanov <romanton@google.com>
Tested-by: kokoro <noreply+kokoro@google.com>

tools/contrib/refactor_use_references.py

@@ -0,0 +1,145 @@
+# Copyright 2022 The Chromium OS Authors. All rights reserved.
+# Use of this source code is governed by a BSD-style license that can be
+# found in the LICENSE file.
+
+# Tools for refactoring references in rust code.
+#
+# Contains the last run refactoring for reference. Don't run this script, it'll
+# fail, but use it as a foundation for other refactorings.
+
+from contextlib import contextmanager
+import os
+import re
+import subprocess
+from pathlib import Path
+from typing import Callable, NamedTuple, Union
+
+SearchPattern = Union[str, re.Pattern[str]]
+
+
+class Token(NamedTuple):
+    token: str
+    start: int
+    end: int
+
+
+def tokenize(source: str):
+    "Split source by whitespace with start/end indices annotated."
+    start = 0
+    for i in range(len(source)):
+        if source[i] in (" ", "\n", "\t") and i - start > 0:
+            token = source[start:i].strip()
+            if token:
+                yield Token(token, start, i)
+            start = i
+
+
+def parse_module_chunks(source: str):
+    """Terrible parser to split code by `mod foo { ... }` statements. Please don't judge me.
+
+    Returns the original source split with module names anntated as ('module name', 'source')
+    """
+    tokens = list(tokenize(source))
+    prev = 0
+    for i in range(len(tokens) - 2):
+        if tokens[i].token == "mod" and tokens[i + 2].token == "{":
+            brackets = 1
+            for j in range(i + 3, len(tokens)):
+                if "{" not in tokens[j].token or "}" not in tokens[j].token:
+                    if "{" in tokens[j].token:
+                        brackets += 1
+                    elif "}" in tokens[j].token:
+                        brackets -= 1
+                if brackets == 0:
+                    start = tokens[i + 2].end
+                    end = tokens[j].start
+                    yield ("", source[prev:start])
+                    yield (tokens[i + 1].token, source[start:end])
+                    prev = end
+                    break
+    if prev != len(source):
+        yield ("", source[prev:])
+
+
+def replace_use_references(file_path: Path, callback: Callable[[list[str], str], str]):
+    """Calls 'callback' for each foo::bar reference in `file_path`.
+
+    The callback is called with the reference as an argument and is expected to return the rewritten
+    reference.
+    Additionally, the absolute path in the module tree is provided, taking into account the file
+    path as well as modules defined in the source itself.
+
+    eg.
+    src/foo.rs:
+    ```
+    mod tests {
+        use crate::baz;
+    }
+    ```
+    will call `callback(['foo', 'tests'], 'crate::baz')`
+    """
+    module_parts = list(file_path.parts[:-1])
+    if file_path.stem not in ("mod", "lib"):
+        module_parts.append(file_path.stem)
+
+    with open(file_path, "r") as file:
+        contents = file.read()
+    chunks: list[str] = []
+    for module, source in parse_module_chunks(contents):
+        if module:
+            full_module_parts = module_parts + [module]
+        else:
+            full_module_parts = module_parts
+        chunks.append(
+            re.sub(
+                r"([\w\*\_\$]+\:\:)+[\w\*\_]+",
+                lambda m: callback(full_module_parts, m.group(0)),
+                source,
+            )
+        )
+    with open(file_path, "w") as file:
+        file.write("".join(chunks))
+
+
+@contextmanager
+def chdir(path: Union[Path, str]):
+    origin = Path().absolute()
+    try:
+        os.chdir(path)
+        yield
+    finally:
+        os.chdir(origin)
+
+
+def use_super_instead_of_crate(root: Path):
+    """Expects to be run directly on the src directory and assumes
+    that directory to be the module crate:: refers to."""
+
+    def replace(module: list[str], use: str):
+        if use.startswith("crate::"):
+            new_use = use.replace("crate::", "super::" * len(module))
+            print("::".join(module), use, "->", new_use)
+            return new_use
+        return use
+
+    with chdir(root):
+        for file in Path().glob("**/*.rs"):
+            replace_use_references(file, replace)
+
+
+def main():
+    for crate in ("sys_util", "win_sys_util", "sys_util_core", "cros_async"):
+        path = Path("common") / crate / "src"
+        subprocess.check_call(["git", "checkout", "-f", str(path)])
+
+        # Use rustfmt to re-format use statements to be one per line.
+        subprocess.check_call(
+            ["rustfmt", "+nightly", "--config=imports_granularity=item", f"{path}/lib.rs"]
+        )
+        use_super_instead_of_crate(path)
+        subprocess.check_call(
+            ["rustfmt", "+nightly", "--config=imports_granularity=crate", f"{path}/lib.rs"]
+        )
+
+
+main()