From c080bc4734f2062241624ab6cdca8a0598b410c9 Mon Sep 17 00:00:00 2001
From: Theodoros Theodoridis <theodoros.theodoridis@inf.ethz.ch>
Date: Thu, 25 Apr 2024 17:24:25 +0200
Subject: [PATCH] Add a yarpgen wrapper

---
 .pre-commit-config.yaml |  4 +--
 diopter/generator.py    | 80 +++++++++++++++++++++++++++++++++++++++++
 setup.cfg               |  2 +-
 3 files changed, 83 insertions(+), 3 deletions(-)

diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
index d29636c..ff8819c 100644
--- a/.pre-commit-config.yaml
+++ b/.pre-commit-config.yaml
@@ -1,6 +1,6 @@
 repos:
 -   repo: https://github.com/psf/black
-    rev: 24.2.0
+    rev: 24.4.1
     hooks:
     -   id: black
 -   repo: https://github.com/PyCQA/isort
@@ -13,7 +13,7 @@ repos:
     hooks:
     -   id: flake8
 -   repo: https://github.com/pre-commit/mirrors-mypy
-    rev: v1.8.0
+    rev: v1.10.0
     hooks:
     -   id: mypy
         additional_dependencies: [pytest >= 7.0.0]
diff --git a/diopter/generator.py b/diopter/generator.py
index 1bdb0e9..b8aa8e3 100644
--- a/diopter/generator.py
+++ b/diopter/generator.py
@@ -11,6 +11,7 @@
 
 from diopter.compiler import Language, SourceProgram
 from diopter.sanitizer import Sanitizer
+from diopter.utils import TempDirEnv
 
 
 def dummy_func(generator: Generator) -> SourceProgram:
@@ -187,3 +188,82 @@ def filter_program(self, program: SourceProgram) -> bool:
         ):
             return False
         return bool(self.sanitizer.sanitize(program))
+
+
+class YarpGen(Generator):
+    def __init__(
+        self,
+        sanitizer: Sanitizer,
+        yarpgen: str | None = None,
+        language: Language = Language.C,
+        additional_flags: tuple[str, ...] = (),
+        minimum_length: int = 10000,
+        maximum_length: int = 50000,
+    ):
+        """
+        Args:
+            sanitizer (Sanitizer):
+                used to sanitize and discard generated code
+            yarpgen (str | None):
+                Path to yarpgen executable, if empty "csmith" will be used
+            language (Language):
+                The language of the generated code (--std flag).
+            additional_flags (tuple[str, ...]):
+                Additional flags to pass to yarpgen.
+            minimum_length (int):
+                The minimum length of a generated test case in characters.
+            maximum_length (int):
+                The maximum length of a generated test case in characters.
+        """
+        super().__init__(sanitizer)
+        self.minimum_length = minimum_length
+        self.maximum_length = maximum_length
+        self.yarpgen = yarpgen if yarpgen else "yarpgen"
+        self.language = language
+        self.additional_flags = additional_flags
+
+        if not which(self.yarpgen):
+            raise ValueError(f"Invalid yarpgen executable: {self.yarpgen}")
+
+    def generate_program_impl(self) -> SourceProgram:
+        """Generate random code with yarpgen.
+
+        Returns:
+            SourceProgram: yarpgen generated program.
+        """
+
+        lang_flag = {
+            Language.C: "--std=c",
+            Language.CPP: "--std=c++",
+        }[self.language]
+        with TempDirEnv() as temp_dir:
+            cmd = [self.yarpgen, lang_flag, "-o", str(temp_dir.resolve())] + list(
+                self.additional_flags
+            )
+            result = subprocess.run(
+                cmd, stdout=subprocess.PIPE, stderr=subprocess.STDOUT
+            )
+            assert result.returncode == 0, result.stdout.decode("utf-8")
+            driver_str = (temp_dir / ("driver" + self.language.to_suffix())).read_text()
+            func_str = (
+                (temp_dir / ("func" + self.language.to_suffix()))
+                .read_text()
+                .replace('#include "init.h"', "")
+            )
+
+            return SourceProgram(
+                code=driver_str + "\n" + func_str,
+                language=self.language,
+                defined_macros=(),
+                include_paths=(),
+                system_include_paths=(),
+                flags=(),
+            )
+
+    def filter_program(self, program: SourceProgram) -> bool:
+        if (
+            len(program.code) < self.minimum_length
+            or len(program.code) > self.maximum_length
+        ):
+            return False
+        return bool(self.sanitizer.sanitize(program))
diff --git a/setup.cfg b/setup.cfg
index 9e6e9ca..a4c8e95 100644
--- a/setup.cfg
+++ b/setup.cfg
@@ -1,6 +1,6 @@
 [metadata]
 name = diopter
-version = 0.0.27
+version = 0.0.28
 author = Theodoros Theodoridis, Yann Girsberger
 author_email = theodort@inf.ethz.ch
 description = A library for building differential tests