Optimize PERM_LINESWAP for full statements

README.md

@@ -42,7 +42,7 @@ The .c file may be modified with any of the following macros which affect manual
 - `PERM_GENERAL(a, b, ...)` expands to any of `a`, `b`, ...
 - `PERM_VAR(a, b)` sets the meta-variable `a` to `b`, `PERM_VAR(a)` expands to the meta-variable `a`.
 - `PERM_RANDOMIZE(code)` expands to `code`, but allows randomization within that region. Multiple regions may be specified.
-- `PERM_LINESWAP(lines)` expands to a permutation of the ordered set of non-whitespace lines (split by `\n`). Warning: this gets slow with more than 5 lines or so!
+- `PERM_LINESWAP(lines)` expands to a permutation of the ordered set of non-whitespace lines (split by `\n`). Each line must contain zero or more complete C statements. (For incomplete statements use `PERM_LINESWAP_TEXT`, which is slower because it has to repeatedly parse C code.)
 - `PERM_INT(lo, hi)` expands to an integer between `lo` and `hi` (which must be constants).
 - `PERM_IGNORE(code)` expands to `code`, without passing it through the C parser library (pycparser)/randomizer. This can be used to avoid parse errors for non-standard C, e.g. `asm` blocks.
 - `PERM_PRETEND(code)` expands to `code` for the purpose of the C parser/randomizer, but gets removed afterwards. This can be used together with `PERM_IGNORE` to enable the permuter to deal with input it isn't designed for (e.g. inline functions, C++, non-code).

run-tests.sh

@@ -1,2 +1,3 @@
 #!/bin/sh
 python3 -m unittest discover -s test/
+# python3 -m pytest test/

src/candidate.py

@@ -10,6 +10,7 @@ from .compiler import Compiler
 from .randomizer import Randomizer
 from .scorer import Scorer
 from .perm.perm import EvalState, Perm
+from .perm.ast import apply_ast_perms
 from .helpers import try_remove
 from .profiler import Profiler
 from . import perm
@@ -38,7 +39,6 @@ class Candidate:
 
     ast: ca.FileAST
 
-    orig_fn: ca.FuncDef
     fn_index: int
     rng_seed: int
     randomizer: Randomizer
@@ -58,7 +58,9 @@ class Candidate:
         return orig_fn, fn_index, ast
 
     @staticmethod
-    def from_source(source: str, fn_name: str, rng_seed: int) -> "Candidate":
+    def from_source(
+        source: str, eval_state: EvalState, fn_name: str, rng_seed: int
+    ) -> "Candidate":
         # Use the same AST for all instances of the same original source, but
         # with the target function deeply copied. Since we never change the
         # AST outside of the target function, this is fine, and it saves us
@@ -66,10 +68,11 @@ class Candidate:
         orig_fn, fn_index, ast = Candidate._cached_shared_ast(source, fn_name)
         ast = copy.copy(ast)
         ast.ext = copy.copy(ast.ext)
-        ast.ext[fn_index] = copy.deepcopy(orig_fn)
+        fn_copy = copy.deepcopy(orig_fn)
+        ast.ext[fn_index] = fn_copy
+        apply_ast_perms(fn_copy, eval_state)
         return Candidate(
             ast=ast,
-            orig_fn=orig_fn,
             fn_index=fn_index,
             rng_seed=rng_seed,
             randomizer=Randomizer(rng_seed),

src/perm/ast.py

@@ -0,0 +1,58 @@
+from typing import List, Optional, Tuple
+
+from .. import ast_util
+from .perm import EvalState, Perm
+from ..ast_util import Block, Statement, Expression
+from ..error import CandidateConstructionFailure
+
+from pycparser import c_ast as ca
+
+
+class _Done(Exception):
+    pass
+
+
+def _apply_perm(fn: ca.FuncDef, perm_id: int, perm: Perm, seed: int) -> None:
+    """Find and apply a single late perm macro in the AST."""
+    # Currently we search for statement macros only.
+    wanted_pragma = f"_permuter ast_perm {perm_id}"
+    Loc = Tuple[List[Statement], int]
+
+    def try_handle_block(block: ca.Node, where: Optional[Loc]) -> None:
+        if not isinstance(block, ca.Compound) or not block.block_items:
+            return
+        pragma = block.block_items[0]
+        if not isinstance(pragma, ca.Pragma) or pragma.string != wanted_pragma:
+            return
+
+        args: List[Statement] = block.block_items[1:]
+        stmts = perm.eval_statement_ast(args, seed)
+
+        if where:
+            where[0][where[1] : where[1] + 1] = stmts
+        else:
+            block.block_items = stmts
+        raise _Done
+
+    def rec(block: Block) -> None:
+        # if (x) { _Pragma(...); inputs } -> if (x) { outputs }
+        try_handle_block(block, None)
+        stmts = ast_util.get_block_stmts(block, False)
+        for i, stmt in enumerate(stmts):
+            # { ... { _Pragma(...); inputs } ... } -> { ... outputs ... }
+            try_handle_block(stmt, (stmts, i))
+            ast_util.for_nested_blocks(stmt, rec)
+
+    try:
+        rec(fn.body)
+        raise CandidateConstructionFailure("Failed to find PERM macro in AST.")
+    except _Done:
+        pass
+
+
+def apply_ast_perms(fn: ca.FuncDef, eval_state: EvalState) -> None:
+    """Find all late perm macros in the AST and apply them."""
+    # Nested perms will have smaller IDs, so apply the perms from lowest ID to
+    # highest to ensure that all arguments to perms have already been evaluated.
+    for perm_id, (perm, seed) in enumerate(eval_state.ast_perms):
+        _apply_perm(fn, perm_id, perm, seed)

src/perm/perm.py

@@ -1,10 +1,16 @@
 from base64 import b64encode
 from collections import defaultdict
 from dataclasses import dataclass, field
-from typing import Dict, List, Optional
+from typing import Dict, List, Tuple, TypeVar, Optional
 import math
 import itertools
 
+from pycparser import c_ast as ca
+
+from ..ast_util import Statement
+
+T = TypeVar("T")
+
 
 @dataclass
 class PreprocessState:
@@ -17,6 +23,24 @@ class PreprocessState:
 class EvalState:
     vars: Dict[str, str] = field(default_factory=dict)
     once_choices: Dict[str, "Perm"] = field(default_factory=dict)
+    ast_perms: List[Tuple["Perm", int]] = field(default_factory=list)
+
+    def register_ast_perm(self, perm: "Perm", seed: int) -> int:
+        ret = len(self.ast_perms)
+        self.ast_perms.append((perm, seed))
+        return ret
+
+    def gen_ast_statement_perm(
+        self, perm: "Perm", seed: int, *, statements: List[str]
+    ) -> str:
+        perm_id = self.register_ast_perm(perm, seed)
+        lines = [
+            "{",
+            f"#pragma _permuter ast_perm {perm_id}",
+            *["{" + stmt + "}" for stmt in statements],
+            "}",
+        ]
+        return "\n".join(lines)
 
 
 class Perm:
@@ -35,6 +59,9 @@ class Perm:
     def evaluate(self, seed: int, state: EvalState) -> str:
         return ""
 
+    def eval_statement_ast(self, args: List[Statement], seed: int) -> List[Statement]:
+        raise NotImplementedError
+
     def preprocess(self, state: PreprocessState) -> None:
         for p in self.children:
             p.preprocess(state)
@@ -71,6 +98,17 @@ def _count_either(perms: List[Perm]) -> int:
     return sum(p.perm_count for p in perms)
 
 
+def _shuffle(items: List[T], seed: int) -> List[T]:
+    items = items[:]
+    output = []
+    while items:
+        ind = seed % len(items)
+        seed //= len(items)
+        output.append(items[ind])
+        del items[ind]
+    return output
+
+
 class RootPerm(Perm):
     def __init__(self, inner: Perm) -> None:
         self.children = [inner]
@@ -219,13 +257,26 @@ class LineSwapPerm(Perm):
     def evaluate(self, seed: int, state: EvalState) -> str:
         sub_seed, variation = divmod(seed, self.own_count)
         texts = _eval_all(sub_seed, self.children, state)
-        output = []
-        while texts:
-            ind = variation % len(texts)
-            variation //= len(texts)
-            output.append(texts[ind])
-            del texts[ind]
-        return "\n".join(output)
+        return "\n".join(_shuffle(texts, variation))
+
+
+class LineSwapAstPerm(Perm):
+    def __init__(self, lines: List[Perm]) -> None:
+        self.children = lines
+        self.own_count = math.factorial(len(lines))
+        self.perm_count = self.own_count * _count_all(self.children)
+
+    def evaluate(self, seed: int, state: EvalState) -> str:
+        sub_seed, variation = divmod(seed, self.own_count)
+        texts = _eval_all(sub_seed, self.children, state)
+        return state.gen_ast_statement_perm(self, variation, statements=texts)
+
+    def eval_statement_ast(self, args: List[Statement], seed: int) -> List[Statement]:
+        ret = []
+        for item in _shuffle(args, seed):
+            assert isinstance(item, ca.Compound)
+            ret.extend(item.block_items or [])
+        return ret
 
 
 class IntPerm(Perm):

src/perm/perm_eval.py

@@ -1,5 +1,5 @@
 import random
-from typing import List, Iterable, Set
+from typing import List, Iterable, Set, Tuple
 
 from .perm import Perm, EvalState
 
@@ -31,5 +31,6 @@ def perm_gen_all_seeds(perm: Perm) -> Iterable[int]:
             break
 
 
-def perm_evaluate_one(perm: Perm) -> str:
-    return perm.evaluate(0, EvalState())
+def perm_evaluate_one(perm: Perm) -> Tuple[str, EvalState]:
+    eval_state = EvalState()
+    return perm.evaluate(0, eval_state), eval_state

src/perm/perm_gen.py

@@ -7,6 +7,7 @@ from .perm import (
     IgnorePerm,
     IntPerm,
     LineSwapPerm,
+    LineSwapAstPerm,
     OncePerm,
     Perm,
     PretendPerm,
@@ -78,7 +79,8 @@ PERM_FACTORIES: Dict[str, Callable[[str], Perm]] = {
     "PERM_ONCE": lambda text: make_once_perm(text),
     "PERM_RANDOMIZE": lambda text: RandomizerPerm(rec_perm_gen(text)),
     "PERM_VAR": lambda text: make_var_perm(text),
-    "PERM_LINESWAP": lambda text: LineSwapPerm(split_args_newline(text)),
+    "PERM_LINESWAP_TEXT": lambda text: LineSwapPerm(split_args_newline(text)),
+    "PERM_LINESWAP": lambda text: LineSwapAstPerm(split_args_newline(text)),
     "PERM_INT": lambda text: IntPerm(*map(int, split_args_text(text))),
     "PERM_IGNORE": lambda text: IgnorePerm(rec_perm_gen(text)),
     "PERM_PRETEND": lambda text: PretendPerm(rec_perm_gen(text)),
@@ -137,4 +139,6 @@ def perm_gen(text: str) -> Perm:
     ret = RootPerm(ret)
     if not ret.is_random():
         print(f"Will run for {ret.perm_count} iterations.")
+    else:
+        print(f"Will try {ret.perm_count} different base sources.")
     return ret

src/permuter.py

@@ -101,8 +101,10 @@ class Permuter:
         self._last_score: Optional[int] = None
 
     def _create_and_score_base(self) -> Tuple[int, str, str]:
-        base_source = perm_eval.perm_evaluate_one(self._permutations)
-        base_cand = Candidate.from_source(base_source, self.fn_name, rng_seed=0)
+        base_source, eval_state = perm_eval.perm_evaluate_one(self._permutations)
+        base_cand = Candidate.from_source(
+            base_source, eval_state, self.fn_name, rng_seed=0
+        )
         o_file = base_cand.compile(self.compiler, show_errors=True)
         if not o_file:
             raise CandidateConstructionFailure(f"Unable to compile {self.source_file}")
@@ -137,11 +139,12 @@ class Permuter:
         # This means we're not guaranteed to test all seeds, but it doesn't really matter since
         # we're randomizing anyway.
         if not self._cur_cand or not keep:
-            cand_c = self._permutations.evaluate(seed, EvalState())
+            eval_state = EvalState()
+            cand_c = self._permutations.evaluate(seed, eval_state)
             rng_seed = self._force_rng_seed or random.randrange(1, 10 ** 20)
             self._cur_seed = (seed, rng_seed)
             self._cur_cand = Candidate.from_source(
-                cand_c, self.fn_name, rng_seed=rng_seed
+                cand_c, eval_state, self.fn_name, rng_seed=rng_seed
             )
 
         # Randomize the candidate

src/randomizer.py

@@ -18,7 +18,7 @@ from typing import (
     Union,
 )
 
-from pycparser import c_ast as ca, c_parser, c_generator
+from pycparser import c_ast as ca
 
 from . import ast_util
 from .ast_util import Block, Indices, Statement, Expression

test/test_perm.py

@@ -184,6 +184,36 @@ class TestPermMacros(unittest.TestCase):
         )
         self.assertEqual(score, 0)
 
+    def test_lineswap(self) -> None:
+        score = self.go(
+            "void a(); void b(); void c(); void test(void) {",
+            "}",
+            """
+            PERM_LINESWAP(
+                a();
+                b();
+                c();
+            )
+            """,
+            "b(); a(); c();",
+        )
+        self.assertEqual(score, 0)
+
+    def test_lineswap_text(self) -> None:
+        score = self.go(
+            "void a(); void b(); void c(); void test(void) {",
+            "}",
+            """
+            PERM_LINESWAP_TEXT(
+                a();
+                b();
+                c();
+            )
+            """,
+            "b(); a(); c();",
+        )
+        self.assertEqual(score, 0)
+
     def test_randomizer(self) -> None:
         score = self.go(
             "void foo(); void bar(); void test(void) {",