Implemented compress which removes constant part in programs

Author: Theomat

examples/pbe/dataset_generator_unique.py

@@ -9,6 +9,7 @@
 
 from synth import Dataset, PBE
 from synth.pbe.task_generator import TaskGenerator
+from synth.semantic import DSLEvaluator
 from synth.utils import chrono
 from synth.syntax import CFG, Type, Program
 
@@ -156,6 +157,8 @@ def generate_programs_and_samples_for(
         prog, unique = task_generator.generate_program(tr)
         while not unique:
             prog, unique = task_generator.generate_program(tr)
+        if isinstance(task_generator.evaluator, DSLEvaluator):
+            prog = task_generator.evaluator.compress(prog)
         programs.add(prog)
     # Phase 2
     samples, equiv = generate_samples_for(
@@ -173,6 +176,8 @@ def generate_programs_and_samples_for(
         prog, unique = task_generator.generate_program(tr)
         while not unique:
             prog, unique = task_generator.generate_program(tr)
+        if isinstance(task_generator.evaluator, DSLEvaluator):
+            prog = task_generator.evaluator.compress(prog)
         # Compute semantic hash
         cl = None
         has_none = False

synth/semantic/evaluator.py

@@ -36,6 +36,26 @@ def __init__(self, semantics: Dict[Primitive, Any], use_cache: bool = True) -> N
         self._total_requests = 0
         self._cache_hits = 0
 
+    def compress(self, program: Program) -> Program:
+        """
+        Return a semantically equivalent version of the program by evaluating constant expressions.
+        Note for data saving/loading purposes, partial applications are left untouched.
+        """
+        if isinstance(program, Function):
+            args = [self.compress(p) for p in program.arguments]
+            if len(program.type.returns().arguments()) == 0 and all(
+                a.is_constant() for a in args
+            ):
+                before = self.use_cache
+                self.use_cache = False
+                value = self.eval(program, [])
+                self.use_cache = before
+                return Constant(program.type.returns(), value, True)
+            else:
+                return Function(program.function, args)
+        else:
+            return program
+
     def eval(self, program: Program, input: List) -> Any:
         key = __tuplify__(input)
         if self.use_cache and key not in self._cache: