Review comments.

.github/workflows/hax.yml

@@ -50,5 +50,9 @@ jobs:
       - name: 🏃🏻‍♀️ Run hax extraction
         run: |
           eval $(opam env)
-          python3 extract_to_fstar.py --kyber-reference
-          python3 extract_to_fstar.py --crate specs/kyber
+          ./extract_to_fstar.py --kyber-reference
+          # Extract the functions in the compress module individually to test
+          # the function-extraction code.
+          # Extract functions from the remaining modules to test the
+          # module-extraction code.
+          ./extract_to_fstar.py --crate-path specs/kyber --functions compress::compress compress::decompress compress::compress_d compress::decompress_d --modules ind_cpa root matrix ntt parameters sampling serialize

extract_to_fstar.py

@@ -1,18 +1,22 @@
+#! /usr/bin/env python3
+
 import os
 import argparse
 import subprocess
 import sys
 
 
-def shell(command, expect=0, cwd=None, format_function_string=False):
+def shell(command, expect=0, cwd=None, format_selection_string=False):
     subprocess_stdout = subprocess.DEVNULL
 
-    command_formatted = command.copy()
-    if format_function_string:
-        command_formatted[4] = "'{}'".format(command_formatted[4])
+    print("Command: ", end="")
+    for i, word in enumerate(command):
+        if i == 4:
+            print("'{}' ".format(word), end="")
+        else:
+            print("{} ".format(word), end="")
 
-    print("Command: {}".format(" ".join(command_formatted)))
-    print("Directory: {}".format(cwd))
+    print("\nDirectory: {}".format(cwd))
 
     ret = subprocess.run(
         command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, cwd=cwd
@@ -30,9 +34,9 @@ def shell(command, expect=0, cwd=None, format_function_string=False):
     help="Extract the Kyber reference implementation.",
 )
 parser.add_argument(
-    "--crate",
+    "--crate-path",
     type=str,
-    dest="crate",
+    dest="crate_path",
     nargs="?",
     default=".",
     help="Path to the crate from which to extract the code (default is path to libcrux).",
@@ -41,31 +45,60 @@ def shell(command, expect=0, cwd=None, format_function_string=False):
     "--functions",
     type=str,
     nargs="*",
+    dest="functions",
+    default="",
+    help="Space-separated list of functions to extract. The functions must be fully qualified from the crate root.",
+)
+
+parser.add_argument(
+    "--modules",
+    type=str,
+    dest="modules",
+    nargs="*",
     default="",
-    help="Space-separated list of functions to extract",
+    help='Space-separated list of modules to extract. The modules must be fully qualified from the crate root. The special argument"root" can be used to extract the lib.rs file.',
 )
-args = parser.parse_args()
 
-if args.functions:
-    args.functions = " ".join(["+" + function for function in args.functions])
+options = parser.parse_args()
+
+if options.modules or options.functions:
+    if options.modules:
+        options.modules = " ".join(
+            [
+                "+::*" if module == "root" else "+" + module + "::*"
+                for module in options.modules
+            ]
+        )
+        options.modules = " {}".format(options.modules)
+
+    if options.functions:
+        options.functions = " ".join(["+" + function for function in options.functions])
+        options.functions = " {}".format(options.functions)
+
     shell(
-        ["cargo", "hax", "into", "-i", "-** {}".format(args.functions), "fstar"],
-        cwd=args.crate,
-        format_function_string=True,
+        [
+            "cargo",
+            "hax",
+            "into",
+            "-i",
+            "-**{}{}".format(options.functions, options.modules),
+            "fstar",
+        ],
+        cwd=options.crate_path,
+        format_selection_string=True,
+    )
+elif options.kyber_reference:
+    shell(
+        [
+            "cargo",
+            "hax",
+            "into",
+            "-i",
+            "-** +kem::kyber::** -kem::kyber::arithmetic::mutable_operations::**",
+            "fstar",
+        ],
+        cwd=".",
+        format_selection_string=True,
     )
 else:
-    if args.kyber_reference:
-        shell(
-            [
-                "cargo",
-                "hax",
-                "into",
-                "-i",
-                "-** +kem::kyber::** -kem::kyber::arithmetic::mutable_operations::**",
-                "fstar",
-            ],
-            cwd=".",
-            format_function_string=True,
-        )
-    else:
-        shell(["cargo", "hax", "into", "fstar"], cwd=args.crate)
+    shell(["cargo", "hax", "into", "fstar"], cwd=options.crate_path)