pyk.ktool.kfuzz: Add support for custom substitution functions in fuzz (#4792)

Adds an optional `subst_func` parameter to the `fuzz` function, allowing
users to customize how substitutions are applied to the template.

By default, substitutions are applied using a bottom-up traversal on the
whole template pattern. Custom substitution functions can be used to
target specific parts of the configuration, enabling more optimized
fuzzing.

Also includes a test case demonstrating the use of a custom substitution
function.

---------

Co-authored-by: Tamás Tóth <[email protected]>

Author: bbyalcinkaya

pyk/src/pyk/kore/manip.py

@@ -5,7 +5,7 @@
 from .syntax import And, App, EVar, MLQuant, Top
 
 if TYPE_CHECKING:
-    from collections.abc import Collection
+    from collections.abc import Collection, Mapping
 
     from .syntax import Pattern
 
@@ -57,6 +57,24 @@ def add_symbol(pattern: Pattern) -> None:
     return res
 
 
+def substitute_vars(pattern: Pattern, subst_map: Mapping[EVar, Pattern]) -> Pattern:
+    """Substitute variables in a pattern using a bottom-up traversal.
+
+    Args:
+        pattern: The pattern containing variables to be substituted.
+        subst_map: A mapping from variables to their replacement patterns.
+    """
+
+    def subst(pattern: Pattern) -> Pattern:
+        match pattern:
+            case EVar() as var:
+                return subst_map.get(var, var)
+            case _:
+                return pattern
+
+    return pattern.bottom_up(subst)
+
+
 def elim_aliases(pattern: Pattern) -> Pattern:
     r"""Eliminate subpatterns of the form ``\and{S}(p, X : S)``.
 
@@ -66,19 +84,12 @@ def elim_aliases(pattern: Pattern) -> Pattern:
 
     def inline_aliases(pattern: Pattern) -> Pattern:
         match pattern:
-            case And(_, (p, EVar(name))):
-                aliases[name] = p
+            case And(_, (p, EVar() as var)):
+                aliases[var] = p
                 return p
             case _:
                 return pattern
 
-    def substitute_vars(pattern: Pattern) -> Pattern:
-        match pattern:
-            case EVar(name) as var:
-                return aliases.get(name, var)
-            case _:
-                return pattern
-
     pattern = pattern.bottom_up(inline_aliases)
-    pattern = pattern.bottom_up(substitute_vars)
+    pattern = substitute_vars(pattern, aliases)
     return pattern

pyk/src/pyk/ktool/kfuzz.py

@@ -6,6 +6,7 @@
 from hypothesis import Phase, given, settings
 from hypothesis.strategies import fixed_dictionaries, integers
 
+from ..kore.manip import substitute_vars
 from ..kore.parser import KoreParser
 from ..kore.prelude import inj
 from ..kore.syntax import DV, EVar, SortApp, String
@@ -54,10 +55,17 @@ class KFuzz:
 
     definition_dir: Path
     handler: KFuzzHandler
+    subst_func: Callable[[Pattern, Mapping[EVar, Pattern]], Pattern]
 
-    def __init__(self, definition_dir: Path, handler: KFuzzHandler = _DEFAULT_HANDLER) -> None:
+    def __init__(
+        self,
+        definition_dir: Path,
+        handler: KFuzzHandler = _DEFAULT_HANDLER,
+        subst_func: Callable[[Pattern, Mapping[EVar, Pattern]], Pattern] = substitute_vars,
+    ) -> None:
         self.definition_dir = definition_dir
         self.handler = handler
+        self.subst_func = subst_func
 
     def fuzz_with_check(
         self,
@@ -76,6 +84,7 @@ def fuzz_with_check(
             subst_strategy,
             check_func=check_func,
             handler=self.handler,
+            subst_func=self.subst_func,
             **hypothesis_args,
         )
 
@@ -95,6 +104,7 @@ def fuzz_with_exit_code(
             subst_strategy,
             check_exit_code=True,
             handler=self.handler,
+            subst_func=self.subst_func,
             **hypothesis_args,
         )
 
@@ -133,6 +143,7 @@ def fuzz(
     check_func: Callable[[Pattern], Any] | None = None,
     check_exit_code: bool = False,
     handler: KFuzzHandler = _DEFAULT_HANDLER,
+    subst_func: Callable[[Pattern, Mapping[EVar, Pattern]], Pattern] = substitute_vars,
     **hypothesis_args: Any,
 ) -> None:
     """Fuzz a property test with concrete execution over a K term.
@@ -148,6 +159,9 @@ def fuzz(
           An exit code of 0 indicates a passing test.
           A RuntimeError will be thrown if this is True and check_func is also passed as an argument.
         handler: An instance of a `KFuzzHandler` implementing custom behavior while fuzzing.
+        subst_func: A function used to apply the substitution to the template.
+          Takes the original template and a substitution map, and returns the substituted pattern.
+          Defaults to a bottom-up traversal that replaces variables with the corresponding values in the substitution map.
         hypothesis_args: Keyword arguments that will be passed as settings for the hypothesis test. Defaults:
 
           deadline: 5000
@@ -169,7 +183,7 @@ def sub(p: Pattern) -> Pattern:
                 return p
 
         handler.handle_test(subst_case)
-        test_pattern = template.bottom_up(sub)
+        test_pattern = subst_func(template, subst_case)
         res = llvm_interpret_raw(definition_dir, test_pattern.text, check=False)
 
         try:

pyk/src/tests/integration/ktool/test_fuzz.py

@@ -5,8 +5,9 @@
 
 import pytest
 
+from pyk.kore.manip import substitute_vars
 from pyk.kore.parser import KoreParser
-from pyk.kore.prelude import inj, top_cell_initializer
+from pyk.kore.prelude import generated_counter, generated_top, inj, int_dv, k, map_pattern, top_cell_initializer
 from pyk.kore.syntax import DV, App, EVar, SortApp, String
 from pyk.ktool.kfuzz import KFuzz, kintegers
 from pyk.ktool.kprint import _kast
@@ -15,6 +16,7 @@
 from ..utils import K_FILES, TEST_DATA_DIR
 
 if TYPE_CHECKING:
+    from collections.abc import Mapping
     from pathlib import Path
     from typing import Final
 
@@ -28,6 +30,37 @@
 VAR_Y = EVar(name='VarY', sort=SortApp('SortInt'))
 
 
+def t(k_cell: Pattern, state_cell: Pattern) -> App:
+    return App("Lbl'-LT-'T'-GT-'", (), (k_cell, state_cell))
+
+
+def state(pattern: Pattern) -> App:
+    return App("Lbl'-LT-'state'-GT-'", (), (pattern,))
+
+
+def subst_on_k_cell(template: Pattern, subst_case: Mapping[EVar, Pattern]) -> Pattern:
+    """A substitution function that only applies substitutions within the K cell.
+    Used to test custom substitution option in fuzzing and optimize fuzzers by
+    restricting changes to relevant parts of the configuration.
+
+    Args:
+        template: The template configuration containing variables in the K cell.
+        subst_case: A mapping from variables to their replacement patterns.
+    """
+    match template:
+        case App(
+            "Lbl'-LT-'generatedTop'-GT-'",
+            args=(
+                App("Lbl'-LT-'T'-GT-'", args=(k_cell, state_cell)),
+                generated_counter_cell,
+            ),
+        ):
+            k_cell = substitute_vars(k_cell, subst_case)
+            return generated_top((t(k_cell, state_cell), generated_counter_cell))
+
+    raise ValueError(template)
+
+
 class TestImpFuzz(KompiledTest):
     KOMPILE_MAIN_FILE = K_FILES / 'imp.k'
     KOMPILE_BACKEND = 'llvm'
@@ -37,6 +70,10 @@ class TestImpFuzz(KompiledTest):
     def kfuzz(self, definition_dir: Path) -> KFuzz:
         return KFuzz(definition_dir)
 
+    @pytest.fixture
+    def kfuzz_with_subst(self, definition_dir: Path) -> KFuzz:
+        return KFuzz(definition_dir, subst_func=subst_on_k_cell)
+
     @staticmethod
     def check(p: Pattern) -> None:
         def check_inner(p: Pattern) -> Pattern:
@@ -76,6 +113,25 @@ def replace_var_ids(p: Pattern) -> Pattern:
 
         return init_pattern
 
+    @staticmethod
+    def setup_program_with_config(definition_dir: Path, text: str) -> Pattern:
+
+        kore_text = _kast(definition_dir=definition_dir, input='program', output='kore', expression=text).stdout
+
+        program_pattern = KoreParser(kore_text).pattern()
+
+        def replace_var_ids(p: Pattern) -> Pattern:
+            match p:
+                case App('inj', _, (DV(_, String('varx')),)):
+                    return VAR_X
+                case App('inj', _, (DV(_, String('vary')),)):
+                    return VAR_Y
+            return p
+
+        program_pattern = program_pattern.top_down(replace_var_ids)
+
+        return generated_top((t(k(program_pattern), state(map_pattern())), generated_counter(int_dv(0))))
+
     def test_fuzz(
         self,
         definition_dir: Path,
@@ -124,3 +180,28 @@ def test_fuzz_fail(
         # Then
         with pytest.raises(AssertionError):
             kfuzz.fuzz_with_check(init_pattern, self.SUBSTS, self.check)
+
+    def test_fuzz_with_subst(
+        self,
+        definition_dir: Path,
+        kfuzz_with_subst: KFuzz,
+    ) -> None:
+        # Given
+        program_text = """
+            // Checks the commutativity of addition
+            int x, y, a, b, res;
+            x = varx;
+            y = vary;
+            a = x + y;
+            b = y + x;
+            if ((a <= b) && (b <= a)) { // a == b
+                res = 0;
+            } else {
+                res = 1;
+            }
+            """
+
+        init_pattern = self.setup_program_with_config(definition_dir, program_text)
+
+        # Then
+        kfuzz_with_subst.fuzz_with_check(init_pattern, self.SUBSTS, self.check)