evolving_opcode_swarm.py

import time
import random
from execution_fitness import ExecutionFitness
from adversarial_tester import AdversarialTester
from self_patcher import SelfPatcher
from repair_memory import RepairMemory

class EvolvingOpcodeSwarm:
    """
    Genetic evolution system for opcode implementations.

    Pipeline:
        generate → mutate → validate → score → select → repeat
    """

    def __init__(self, swarm):
        self.swarm = swarm

        # (opcode, language) -> list of candidates
        self.population = {}

        # best survivors
        self.elite = {}
        self.patcher = SelfPatcher(self.swarm.llm_router)
        # executor for fitness evaluation
        self.executor = ExecutionFitness(timeout=2)
        self.tester = AdversarialTester(self.swarm.llm_router)  # Initialize the adversarial tester
        self.repair_memory = RepairMemory()
    # =====================================================
    # FITNESS FUNCTION
    # =====================================================
    def fitness(self, code, language, executor, expected_output=None):

        if language != "python":
            # fallback for now (you can extend later per-language sandboxes)
            return 0.0

        result = executor.run_python(code)

        return executor.score(result, expected_output)

    # =====================================================
    # INITIAL GENERATION
    # =====================================================
    def seed(self, opcode, language, n=3):

        key = (opcode, language)

        self.population[key] = []

        for _ in range(n):
            code = self.swarm.generate(opcode, language)

            mutated = self.swarm.mutate(code, language)
            validation = self.swarm.validate(mutated, language)

            self.population[key].append({
                "code": mutated,
                "validation": validation,
                "score": self.fitness(mutated, validation)
            })

    # =====================================================
    # SELECTION (KEEP BEST)
    # =====================================================
    def select(self, opcode, language, top_k=2):

        key = (opcode, language)

        pop = self.population.get(key, [])

        pop = sorted(pop, key=lambda x: x["score"], reverse=True)

        survivors = pop[:top_k]

        self.elite[key] = survivors[0] if survivors else None

        return survivors

    # =====================================================
    # MUTATION STEP
    # =====================================================
    def evolve_step(self, opcode, language):

        key = (opcode, language)

        survivors = self.select(opcode, language)

        new_population = []

        code = self.swarm.generate(opcode, language)
        code = self.swarm.mutate(code, language)

        tests = self.tester.generate_tests(opcode, language, code)

        results = self.executor.run_with_tests(code, tests)

        # -----------------------------------------------------
        # DETECT FAILURE
        # -----------------------------------------------------
        execution = self.executor.run_with_tests(code, tests)

        score = self.executor.score(execution)

        error_log = "\n".join(
            r["stderr"] for r in execution if r["stderr"]
        )

        # -----------------------------------------------------
        # STEP 1: CHECK REPAIR MEMORY FIRST
        # -----------------------------------------------------
        cached_fix = self.repair_memory.lookup(error_log, language)

        if cached_fix:
            repaired = cached_fix

        else:
            # STEP 2: LLM REPAIR
            repaired = self.patcher.repair(
                code,
                language,
                error_log
            )

        # -----------------------------------------------------
        # RE-TEST REPAIR
        # -----------------------------------------------------
        new_execution = self.executor.run_with_tests(repaired, tests)

        new_score = self.executor.score(new_execution)

        # -----------------------------------------------------
        # STORE NEW KNOWLEDGE
        # -----------------------------------------------------
        if new_score > score:

            self.repair_memory.record(
                error_log,
                code,
                repaired,
                language
            )

            code = repaired
            score = new_score
                
        # keep elite
        for s in survivors:
            new_population.append(s)

            # mutate elite
            mutated = self.swarm.mutate(s["code"], language)
            validation = self.swarm.validate(mutated, language)

            new_population.append({
                "code": mutated,
                "validation": validation,
                "score": self.fitness(mutated, validation)
            })

        # refill population if needed
        while len(new_population) < 5:
            code = self.swarm.generate(opcode, language)
            mutated = self.swarm.mutate(code, language)
            validation = self.swarm.validate(mutated, language)

            new_population.append({
                "code": mutated,
                "validation": validation,
                "score": self.fitness(
                    mutated,
                    language,
                    self.executor
                )
            })

        self.population[key] = new_population

    # =====================================================
    # FULL EVOLUTION LOOP
    # =====================================================
    def run_generations(self, opcodes, languages, generations=3):

        for gen in range(generations):

            for opcode in opcodes:
                for lang in languages:

                    key = (opcode, lang)

                    # seed if first run
                    if key not in self.population:
                        self.seed(opcode, lang)

                    # evolve population
                    self.evolve_step(opcode, lang)

            print(f"Generation {gen} complete")

        return self.elite