Merge pull request #1108 from pq-code-package/color_tests

lint, autogen, format: Improve and unify UX

Author: mkannwischer

scripts/autogen

@@ -31,10 +31,44 @@ montgomery_factor = pow(2, 16, modulus)
 # - header guards
 # - #undef's for CU-local macros
 
+# Standard color definitions
+GREEN = "\033[32m"
+RED = "\033[31m"
+BLUE = "\033[94m"
+BOLD = "\033[1m"
+NORMAL = "\033[0m"
+
+
+def clear_status_line():
+    """Clear any existing status line by overwriting with spaces and returning to start of line"""
+    print(f"\r{' ' * 160}", end="", flush=True)
+
 
 def status_update(task, msg):
-    print(f"\r{'':<160}", end="", flush=True)
-    print(f"\r[{task}]: {msg} ...", end="", flush=True)
+    clear_status_line()
+    print(f"\r{BLUE}[{task}]{NORMAL}: {msg} ...", end="", flush=True)
+
+
+def high_level_status(msg):
+    clear_status_line()
+    print(
+        f"\r{GREEN}✓{NORMAL} {msg}"
+    )  # This will end with a newline, clearing the status line
+
+
+def info(msg):
+    clear_status_line()
+    print(f"\r{GREEN}info{NORMAL} {msg}")
+
+
+def error(msg):
+    clear_status_line()
+    print(f"\r{RED}error{NORMAL} {msg}")
+
+
+def file_updated(filename):
+    clear_status_line()
+    print(f"\r{BOLD}updated {filename}{NORMAL}")
 
 
 def gen_header():
@@ -240,7 +274,7 @@ class CondParser:
         return CondParser.print_exp(self.parse_condition(exp))
 
 
-def adjust_preprocessor_comments_for_filename(content, source_file):
+def adjust_preprocessor_comments_for_filename(content, source_file, show_status=False):
     """Automatically add comments to large `#if ... #else ... #endif`
     blocks indicating the guarding conditions.
 
@@ -270,7 +304,8 @@ def adjust_preprocessor_comments_for_filename(content, source_file):
 
     ```
     """
-    status_update("if-else", source_file)
+    if show_status:
+        status_update("if-else", source_file)
 
     content = content.split("\n")
     new_content = []
@@ -390,7 +425,9 @@ def adjust_preprocessor_comments_for_filename(content, source_file):
 def gen_preprocessor_comments_for(source_file, dry_run=False):
     with open(source_file, "r") as f:
         content = f.read()
-    new_content = adjust_preprocessor_comments_for_filename(content, source_file)
+    new_content = adjust_preprocessor_comments_for_filename(
+        content, source_file, show_status=True
+    )
     update_file(source_file, new_content, dry_run=dry_run)
 
 
@@ -402,10 +439,17 @@ def gen_preprocessor_comments(dry_run=False):
         )
 
 
-def update_file(filename, content, dry_run=False, force_format=False):
+def update_file(
+    filename,
+    content,
+    dry_run=False,
+    force_format=False,
+    skip_preprocessor_comments=False,
+):
 
     if force_format is True or filename.endswith((".c", ".h", ".i")):
-        content = adjust_preprocessor_comments_for_filename(content, filename)
+        if skip_preprocessor_comments is False:
+            content = adjust_preprocessor_comments_for_filename(content, filename)
         content = format_content(content)
 
     if os.path.exists(filename) is True:
@@ -418,15 +462,15 @@ def update_file(filename, content, dry_run=False, force_format=False):
         return
 
     if dry_run is False:
+        file_updated(filename)
         with open(filename, "w+") as f:
             f.write(content)
     else:
         filename_new = f"{filename}.new"
-        print(
-            f"Autogenerated file {filename} needs updating. Have you called scripts/autogen?",
-            file=sys.stderr,
+        error(
+            f"Autogenerated file {filename} needs updating. Have you called scripts/autogen?"
         )
-        print(f"Writing new version to {filename_new}", file=sys.stderr)
+        info(f"Writing new version to {filename_new}")
         with open(filename_new, "w") as f:
             f.write(content)
         # If the file exists, print diff between old and new version for debugging
@@ -2000,6 +2044,7 @@ def synchronize_backends(
         update_via_copy(
             f"dev/aarch64_{ty}/meta.h",
             "mlkem/src/native/aarch64/meta.h",
+            dry_run=dry_run,
             transform=lambda c: adjust_header_guard_for_filename(
                 c, "mlkem/src/native/aarch64/meta.h"
             ),
@@ -2439,10 +2484,13 @@ def gen_c_citations_for(filename, bibliography, dry_run=False):
         # Remember this usage of the bibliography entry
         entry.register_usages(uses)
 
-    update_file(filename, "\n".join(content), dry_run=dry_run)
+    update_file(
+        filename, "\n".join(content), dry_run=dry_run, skip_preprocessor_comments=True
+    )
 
 
 def gen_citations_for(filename, bibliography, dry_run=False):
+    status_update("citations", filename)
     if filename.endswith(".md"):
         gen_markdown_citations_for(filename, bibliography, dry_run=dry_run)
     elif filename.endswith((".c", ".h", ".S")):
@@ -2557,13 +2605,18 @@ def _main():
     os.chdir(os.path.join(os.path.dirname(__file__), ".."))
 
     gen_citations(args.dry_run)
+    high_level_status("Generated citations")
 
     if args.slothy == []:
         args.slothy = slothy_choices
-    gen_slothy(args.slothy, args.dry_run)
+    if args.slothy is not None:
+        gen_slothy(args.slothy, args.dry_run)
+        high_level_status("Generated SLOTHY optimized assembly")
 
     check_asm_register_aliases()
+    high_level_status("Checked assembly register aliases")
     check_asm_loop_labels()
+    high_level_status("Checked assembly loop labels")
 
     gen_c_zeta_file(args.dry_run)
     gen_aarch64_hol_light_zeta_file(args.dry_run)
@@ -2573,21 +2626,29 @@ def _main():
     gen_avx2_zeta_file(args.dry_run)
     gen_avx2_rej_uniform_table(args.dry_run)
     gen_avx2_mulcache_twiddles_file(args.dry_run)
+    high_level_status("Generated zeta and lookup tables")
 
     if platform.machine().lower() in ["arm64", "aarch64"]:
         gen_hol_light_asm(args.dry_run)
+        high_level_status("Generated HOL Light assembly")
 
     synchronize_backends(
         dry_run=args.dry_run,
         clean=args.aarch64_clean,
         no_simplify=args.no_simplify,
         force_cross=args.force_cross,
     )
+    high_level_status("Synchronized backends")
+
     gen_header_guards(args.dry_run)
+    high_level_status("Generated header guards")
     gen_preprocessor_comments(args.dry_run)
+    high_level_status("Generated preprocessor comments")
     gen_monolithic_source_file(args.dry_run)
     gen_monolithic_asm_file(args.dry_run)
+    high_level_status("Generated monolithic source files")
     gen_undefs(args.dry_run)
+    high_level_status("Generated undefs")
 
     synchronize_backends(
         dry_run=args.dry_run,
@@ -2596,9 +2657,10 @@ def _main():
         force_cross=args.force_cross,
         no_simplify=args.no_simplify,
     )
+    high_level_status("Completed final backend synchronization")
 
     check_macro_typos()
-    print()
+    high_level_status("Checked macro typos")
 
 
 if __name__ == "__main__":

scripts/format

@@ -10,40 +10,49 @@ set -o pipefail
 # consts
 ROOT="$(realpath "$(dirname "$0")"/../)"
 
-GREEN="$(tput setaf 2)"
-NORMAL="$(tput sgr0)"
+# Standard color definitions
+GREEN="\033[32m"
+RED="\033[31m"
+BLUE="\033[94m"
+BOLD="\033[1m"
+NORMAL="\033[0m"
 
 # utility
 info()
 {
-  printf "%s %b\n" "${GREEN}info" "${NORMAL}${*}"
+  printf "%b %b\n" "${GREEN}info" "${NORMAL}${*}"
+}
+
+error()
+{
+  printf "%b %b\n" "${RED}error" "${NORMAL}${*}"
 }
 
 info "Formatting nix files"
 if ! command -v nixpkgs-fmt 2>&1 >/dev/null; then
-  echo "nixpkgs-fmt not found. Are you running in a nix shell? See BUILDING.md."
+  error "nixpkgs-fmt not found. Are you running in a nix shell? See BUILDING.md."
   exit 1
 fi
 
 nixpkgs-fmt "$ROOT"
 
 info "Formatting shell scripts"
 if ! command -v shfmt 2>&1 >/dev/null; then
-  echo "shfmt not found. Are you running in a nix shell? See BUILDING.md."
+  error "shfmt not found. Are you running in a nix shell? See BUILDING.md."
   exit 1
 fi
 shfmt -s -w -l -i 2 -ci -fn $(shfmt -f $(git grep -l '' :/))
 
 info "Formatting python scripts"
 if ! command -v black 2>&1 >/dev/null; then
-  echo "black not found. Are you running in a nix shell? See BUILDING.md."
+  error "black not found. Are you running in a nix shell? See BUILDING.md."
   exit 1
 fi
 black --include "(scripts/tests|scripts/simpasm|scripts/autogen|scripts/check-namespace|\.py$)" "$ROOT"
 
 info "Formatting c files"
 if ! command -v clang-format 2>&1 >/dev/null; then
-  echo "clang-format not found. Are you running in a nix shell? See BUILDING.md."
+  error "clang-format not found. Are you running in a nix shell? See BUILDING.md."
   exit 1
 fi