From f51a547a2548a67b9c096660715edee236f32695 Mon Sep 17 00:00:00 2001
From: Manik Jain <mjain@jumptrading.com>
Date: Thu, 19 Dec 2024 18:04:22 +0000
Subject: [PATCH] Added rekeying + add/remove functionality to
 regenerate-fixtures

---
 commands.md                      |  41 +++++++------
 src/test_suite/features_utils.py |   6 ++
 src/test_suite/test_suite.py     | 101 +++++++++++++++++++++++++------
 3 files changed, 113 insertions(+), 35 deletions(-)

diff --git a/commands.md b/commands.md
index a24494e..db6ecf0 100644
--- a/commands.md
+++ b/commands.md
@@ -24,8 +24,8 @@ $ solana-test-suite [OPTIONS] COMMAND [ARGS]...
 * `execute`: Execute Context or Fixture message(s) and...
 * `fix-to-ctx`: Extract Context messages from Fixtures.
 * `list-harness-types`: List harness types available for use.
-* `regenerate-all-fixtures`: Regenerate all fixtures in provided...
-* `regenerate-fixtures`: Regenerate Fixture messages by checking...
+* `mass-regenerate-fixtures`: Regenerate features for fixtures in...
+* `regenerate-fixtures`: Regenerate features in fixture messages.
 * `run-tests`: Run tests on a set of targets with a...
 
 ## `solana-test-suite create-env`
@@ -40,9 +40,9 @@ $ solana-test-suite create-env [OPTIONS]
 
 **Options**:
 
-* `-s, --solana-target PATH`: Solana (or ground truth) shared object (.so) target file path  [default: impl/lib/libsolfuzz_agave_v2.0.so]
+* `-s, --solana-target PATH`: Solana (or ground truth) shared object (.so) target file path  [default: /data/mjain/repos/solfuzz-agave/target/release/libsolfuzz_agave.so]
 * `-h, --default-harness-type TEXT`: Harness type to use for Context protobufs  [default: InstrHarness]
-* `-t, --target PATH`: Shared object (.so) target file paths (pairs with --keep-passing). Targets must have required function entrypoints defined  [default: impl/lib/libsolfuzz_firedancer.so]
+* `-t, --target PATH`: Shared object (.so) target file paths (pairs with --keep-passing). Targets must have required function entrypoints defined  [default: /data/mjain/repos/firedancer/build/native/clang/lib/libfd_exec_sol_compat.so]
 * `-o, --output-dir PATH`: Output directory for messages  [default: debug_mismatch]
 * `-u, --repro-urls TEXT`: Comma-delimited list of FuzzCorp mismatch links
 * `-n, --section-names TEXT`: Comma-delimited list of FuzzCorp section names
@@ -72,7 +72,7 @@ $ solana-test-suite create-fixtures [OPTIONS]
 
 * `-i, --input PATH`: Input protobuf file or directory of protobuf files  [default: corpus8]
 * `-h, --default-harness-type TEXT`: Harness type to use for Context protobufs  [default: InstrHarness]
-* `-s, --solana-target PATH`: Solana (or ground truth) shared object (.so) target file path  [default: impl/lib/libsolfuzz_agave_v2.0.so]
+* `-s, --solana-target PATH`: Solana (or ground truth) shared object (.so) target file path  [default: /data/mjain/repos/solfuzz-agave/target/release/libsolfuzz_agave.so]
 * `-t, --target PATH`: Shared object (.so) target file paths (pairs with --keep-passing). Targets must have required function entrypoints defined
 * `-o, --output-dir PATH`: Output directory for fixtures  [default: test_fixtures]
 * `-p, --num-processes INTEGER`: Number of processes to use  [default: 4]
@@ -96,9 +96,9 @@ $ solana-test-suite debug-mismatches [OPTIONS]
 
 **Options**:
 
-* `-s, --solana-target PATH`: Solana (or ground truth) shared object (.so) target file path  [default: impl/lib/libsolfuzz_agave_v2.0.so]
+* `-s, --solana-target PATH`: Solana (or ground truth) shared object (.so) target file path  [default: /data/mjain/repos/solfuzz-agave/target/release/libsolfuzz_agave.so]
 * `-h, --default-harness-type TEXT`: Harness type to use for Context protobufs  [default: InstrHarness]
-* `-t, --target PATH`: Shared object (.so) target file paths (pairs with --keep-passing). Targets must have required function entrypoints defined  [default: impl/lib/libsolfuzz_firedancer.so]
+* `-t, --target PATH`: Shared object (.so) target file paths (pairs with --keep-passing). Targets must have required function entrypoints defined  [default: /data/mjain/repos/firedancer/build/native/clang/lib/libfd_exec_sol_compat.so]
 * `-o, --output-dir PATH`: Output directory for messages  [default: debug_mismatch]
 * `-u, --repro-urls TEXT`: Comma-delimited list of FuzzCorp mismatch links
 * `-n, --section-names TEXT`: Comma-delimited list of FuzzCorp section names
@@ -197,27 +197,31 @@ $ solana-test-suite list-harness-types [OPTIONS]
 
 * `--help`: Show this message and exit.
 
-## `solana-test-suite regenerate-all-fixtures`
+## `solana-test-suite mass-regenerate-fixtures`
 
-Regenerate all fixtures in provided test-vectors folder
+Regenerate features for fixtures in provided test-vectors folder.
 
 **Usage**:
 
 ```console
-$ solana-test-suite regenerate-all-fixtures [OPTIONS]
+$ solana-test-suite mass-regenerate-fixtures [OPTIONS]
 ```
 
 **Options**:
 
 * `-i, --input PATH`: Input test-vectors directory  [default: corpus8]
 * `-o, --output-dir PATH`: Output directory for regenerated fixtures  [default: /tmp/regenerated_fixtures]
-* `-t, --target PATH`: Shared object (.so) target file path to execute  [default: impl/lib/libsolfuzz_agave_v2.0.so]
-* `-s, --stubbed-target PATH`: Stubbed shared object (.so) target file path to execute  [default: impl/lib/libsolfuzz_firedancer.so]
+* `-t, --target PATH`: Shared object (.so) target file path to execute  [default: /data/mjain/repos/solfuzz-agave/target/release/libsolfuzz_agave.so]
+* `-s, --stubbed-target PATH`: Stubbed shared object (.so) target file path to execute  [default: /data/mjain/repos/solfuzz-agave/target/debug/libsolfuzz_agave_stubbed.so]
+* `-f, --add-feature TEXT`: List of feature pubkeys to force add to the fixtures.
+* `-r, --remove-feature TEXT`: List of feature pubkeys to force remove from the fixtures.
+* `-k, --rekey-feature TEXT`: List of feature pubkeys to rekey in the fixtures, formatted 'old/new' (e.g. `--rekey-feature old/new`).
+* `-m, --merge-with-latest`: Merge with the latest cleaned-up and supported featureset pulled from the target.
 * `--help`: Show this message and exit.
 
 ## `solana-test-suite regenerate-fixtures`
 
-Regenerate Fixture messages by checking FeatureSet compatibility with the target shared library. 
+Regenerate features in fixture messages.
 
 **Usage**:
 
@@ -228,10 +232,13 @@ $ solana-test-suite regenerate-fixtures [OPTIONS]
 **Options**:
 
 * `-i, --input PATH`: Either a file or directory containing messages  [default: corpus8]
-* `-t, --target PATH`: Shared object (.so) target file path to execute  [default: impl/lib/libsolfuzz_agave_v2.0.so]
+* `-t, --target PATH`: Shared object (.so) target file path to execute  [default: /data/mjain/repos/solfuzz-agave/target/release/libsolfuzz_agave.so]
 * `-o, --output-dir PATH`: Output directory for regenerated fixtures  [default: regenerated_fixtures]
 * `-d, --dry-run`: Only print the fixtures that would be regenerated
-* `-a, --all-fixtures`: Regenerate all fixtures, regardless of FeatureSet compatibility. Will apply minimum compatible features.
+* `-f, --add-feature TEXT`: List of feature pubkeys to force add to the fixtures.
+* `-r, --remove-feature TEXT`: List of feature pubkeys to force remove from the fixtures.
+* `-k, --rekey-feature TEXT`: List of feature pubkeys to rekey in the fixtures, formatted 'old/new' (e.g. `--rekey-feature old/new`).
+* `-m, --merge-with-latest`: Merge with the latest cleaned-up and supported featureset pulled from the target.
 * `-l, --log-level INTEGER`: FD logging level  [default: 5]
 * `--help`: Show this message and exit.
 
@@ -251,8 +258,8 @@ $ solana-test-suite run-tests [OPTIONS]
 
 * `-i, --input PATH`: Input protobuf file or directory of protobuf files  [default: corpus8]
 * `-h, --default-harness-type TEXT`: Harness type to use for Context protobufs  [default: InstrHarness]
-* `-s, --solana-target PATH`: Solana (or ground truth) shared object (.so) target file path  [default: impl/lib/libsolfuzz_agave_v2.0.so]
-* `-t, --target PATH`: Shared object (.so) target file paths  [default: impl/lib/libsolfuzz_firedancer.so]
+* `-s, --solana-target PATH`: Solana (or ground truth) shared object (.so) target file path  [default: /data/mjain/repos/solfuzz-agave/target/release/libsolfuzz_agave.so]
+* `-t, --target PATH`: Shared object (.so) target file paths  [default: /data/mjain/repos/firedancer/build/native/clang/lib/libfd_exec_sol_compat.so]
 * `-o, --output-dir PATH`: Output directory for test results  [default: test_results]
 * `-p, --num-processes INTEGER`: Number of processes to use  [default: 4]
 * `-r, --randomize-output-buffer`: Randomizes bytes in output buffer before shared library execution
diff --git a/src/test_suite/features_utils.py b/src/test_suite/features_utils.py
index d7e1e39..600a1fb 100644
--- a/src/test_suite/features_utils.py
+++ b/src/test_suite/features_utils.py
@@ -1,5 +1,7 @@
 from ctypes import *
 from dataclasses import dataclass, field
+import struct
+import fd58
 
 
 @dataclass
@@ -100,3 +102,7 @@ def print_featureset_compatibility_report(
     else:
         print("Unsupported features:")
         print(context_features.difference(target.union_features))
+
+
+def feature_bytes_to_ulong(feature_pubkey: str) -> int:
+    return struct.unpack("<Q", fd58.dec32(feature_pubkey.encode())[:8])[0]
diff --git a/src/test_suite/test_suite.py b/src/test_suite/test_suite.py
index bc7ea22..b9eb3f1 100644
--- a/src/test_suite/test_suite.py
+++ b/src/test_suite/test_suite.py
@@ -734,7 +734,7 @@ def debug_mismatches(
 
 @app.command(
     help=f"""
-        Regenerate Fixture messages by checking FeatureSet compatibility with the target shared library. 
+        Regenerate features in fixture messages.
     """
 )
 def regenerate_fixtures(
@@ -762,11 +762,29 @@ def regenerate_fixtures(
         "-d",
         help="Only print the fixtures that would be regenerated",
     ),
-    all_fixtures: bool = typer.Option(
+    add_features: List[str] = typer.Option(
+        [],
+        "--add-feature",
+        "-f",
+        help="List of feature pubkeys to force add to the fixtures.",
+    ),
+    remove_features: List[str] = typer.Option(
+        [],
+        "--remove-feature",
+        "-r",
+        help="List of feature pubkeys to force remove from the fixtures.",
+    ),
+    rekeyed_features: List[str] = typer.Option(
+        [],
+        "--rekey-feature",
+        "-k",
+        help="List of feature pubkeys to rekey in the fixtures, formatted 'old/new' (e.g. `--rekey-feature old/new`).",
+    ),
+    merge_with_latest: bool = typer.Option(
         False,
-        "--all-fixtures",
-        "-a",
-        help="Regenerate all fixtures, regardless of FeatureSet compatibility. Will apply minimum compatible features.",
+        "--merge-with-latest",
+        "-m",
+        help="Merge with the latest cleaned-up and supported featureset pulled from the target.",
     ),
     log_level: int = typer.Option(
         5,
@@ -790,6 +808,16 @@ def regenerate_fixtures(
     test_cases = list(input.iterdir()) if input.is_dir() else [input]
     num_regenerated = 0
 
+    features_to_add = set(map(features_utils.feature_bytes_to_ulong, add_features))
+    features_to_remove = set(
+        map(features_utils.feature_bytes_to_ulong, remove_features)
+    )
+
+    rekey_features = list(
+        tuple(map(features_utils.feature_bytes_to_ulong, feature.split("/")))
+        for feature in rekeyed_features
+    )
+
     for file in test_cases:
         fixture = read_fixture(file)
         harness_ctx = ENTRYPOINT_HARNESS_MAP[fixture.metadata.fn_entrypoint]
@@ -804,13 +832,20 @@ def regenerate_fixtures(
         features_path = features_path[0]
 
         features = pb_utils.access_nested_field_safe(fixture.input, features_path)
-        feature_set = set(features.features) if features else set()
+        original_feature_set = set(features.features) if features else set()
+        new_feature_set = (original_feature_set | features_to_add) - features_to_remove
 
-        regenerate = True
-        if not all_fixtures:
-            # Skip regeneration if the features are already compatible with the target
-            if features_utils.is_featureset_compatible(target_features, feature_set):
-                regenerate = False
+        for old_feature, new_feature in rekey_features:
+            if old_feature in new_feature_set:
+                new_feature_set.remove(old_feature)
+                new_feature_set.add(new_feature)
+
+        if merge_with_latest:
+            new_feature_set = features_utils.min_compatible_featureset(
+                target_features, new_feature_set
+            )
+
+        regenerate = new_feature_set != original_feature_set
 
         if regenerate:
             num_regenerated += 1
@@ -820,9 +855,7 @@ def regenerate_fixtures(
                 print(f"Regenerating {file}")
                 # Apply minimum compatible features
                 if features is not None:
-                    features.features[:] = features_utils.min_compatible_featureset(
-                        target_features, feature_set
-                    )
+                    features.features[:] = new_feature_set
                 regenerated_fixture = create_fixture_from_context(
                     harness_ctx, fixture.input
                 )
@@ -838,10 +871,10 @@ def regenerate_fixtures(
 
 @app.command(
     help=f"""
-        Regenerate all fixtures in provided test-vectors folder
+        Regenerate features for fixtures in provided test-vectors folder.
     """
 )
-def regenerate_all_fixtures(
+def mass_regenerate_fixtures(
     test_vectors: Path = typer.Option(
         Path("corpus8"),
         "--input",
@@ -866,6 +899,30 @@ def regenerate_all_fixtures(
         "-s",
         help="Stubbed shared object (.so) target file path to execute",
     ),
+    add_features: List[str] = typer.Option(
+        [],
+        "--add-feature",
+        "-f",
+        help="List of feature pubkeys to force add to the fixtures.",
+    ),
+    remove_features: List[str] = typer.Option(
+        [],
+        "--remove-feature",
+        "-r",
+        help="List of feature pubkeys to force remove from the fixtures.",
+    ),
+    rekeyed_features: List[str] = typer.Option(
+        [],
+        "--rekey-feature",
+        "-k",
+        help="List of feature pubkeys to rekey in the fixtures, formatted 'old/new' (e.g. `--rekey-feature old/new`).",
+    ),
+    merge_with_latest: bool = typer.Option(
+        False,
+        "--merge-with-latest",
+        "-m",
+        help="Merge with the latest cleaned-up and supported featureset pulled from the target.",
+    ),
 ):
     globals.output_dir = output_dir
 
@@ -918,7 +975,11 @@ def get_harness_type_for_folder(src, regenerate_folder):
                 shared_library=stubbed_shared_library,
                 output_dir=Path(output_folder),
                 dry_run=False,
-                all_fixtures=True,
+                add_features=add_features,
+                remove_features=remove_features,
+                rekeyed_features=rekeyed_features,
+                merge_with_latest=merge_with_latest,
+                log_level=5,
             )
         elif folder_harness_type in ["ElfLoaderHarness"]:
             shutil.copytree(source_folder, output_folder, dirs_exist_ok=True)
@@ -928,7 +989,11 @@ def get_harness_type_for_folder(src, regenerate_folder):
                 shared_library=shared_library,
                 output_dir=Path(output_folder),
                 dry_run=False,
-                all_fixtures=True,
+                add_features=add_features,
+                remove_features=remove_features,
+                rekeyed_features=rekeyed_features,
+                merge_with_latest=merge_with_latest,
+                log_level=5,
             )
 
     print(f"Regenerated fixtures from {test_vectors} to {output_dir}")