From 1e261dea7c0455c8034d75c5acec08cbd04b26e7 Mon Sep 17 00:00:00 2001
From: Quentin Bourgerie <bourgerie.quentin@gmail.com>
Date: Fri, 24 Jun 2022 15:27:48 +0200
Subject: [PATCH] feat(test): add option to provide manual crypto parameters
 for tests

---
 compiler/tests/fixture/EndToEndFixture.cpp    | 23 +++++++++++++++++++
 compiler/tests/fixture/EndToEndFixture.h      |  2 ++
 compiler/tests/unittest/end_to_end_jit_fhe.cc | 10 +++++++-
 3 files changed, 34 insertions(+), 1 deletion(-)

diff --git a/compiler/tests/fixture/EndToEndFixture.cpp b/compiler/tests/fixture/EndToEndFixture.cpp
index 52650c349..c3bc187b8 100644
--- a/compiler/tests/fixture/EndToEndFixture.cpp
+++ b/compiler/tests/fixture/EndToEndFixture.cpp
@@ -181,6 +181,29 @@ template <> struct llvm::yaml::MappingTraits<EndToEndDesc> {
     io.mapRequired("description", desc.description);
     io.mapRequired("program", desc.program);
     io.mapRequired("tests", desc.tests);
+    std::vector<int64_t> v0parameter;
+    io.mapOptional("v0-parameter", v0parameter);
+    if (!v0parameter.empty()) {
+      if (v0parameter.size() != 7) {
+        io.setError("v0-parameter expect to be a list 7 elemnts "
+                    "[glweDimension, logPolynomialSize, nSmall, brLevel, "
+                    "brLobBase, ksLevel, ksLogBase]");
+      }
+      desc.v0Parameter = mlir::concretelang::V0Parameter(
+          v0parameter[0], v0parameter[1], v0parameter[2], v0parameter[3],
+          v0parameter[4], v0parameter[5], v0parameter[6]);
+    }
+    std::vector<int64_t> v0constraint;
+    io.mapOptional("v0-constraint", v0constraint);
+    if (!v0constraint.empty()) {
+      if (v0constraint.size() != 2) {
+        io.setError("v0-constraint expect to be a list 2 elemnts "
+                    "[p, norm2]");
+      }
+      desc.v0Constraint = mlir::concretelang::V0FHEConstraint();
+      desc.v0Constraint->p = v0constraint[0];
+      desc.v0Constraint->norm2 = v0constraint[1];
+    }
   }
 };
 
diff --git a/compiler/tests/fixture/EndToEndFixture.h b/compiler/tests/fixture/EndToEndFixture.h
index 1305b87ba..bc115b5d0 100644
--- a/compiler/tests/fixture/EndToEndFixture.h
+++ b/compiler/tests/fixture/EndToEndFixture.h
@@ -41,6 +41,8 @@ struct EndToEndDesc {
   std::string description;
   std::string program;
   std::vector<TestDescription> tests;
+  llvm::Optional<mlir::concretelang::V0Parameter> v0Parameter;
+  llvm::Optional<mlir::concretelang::V0FHEConstraint> v0Constraint;
 };
 
 llvm::Expected<mlir::concretelang::LambdaArgument *>
diff --git a/compiler/tests/unittest/end_to_end_jit_fhe.cc b/compiler/tests/unittest/end_to_end_jit_fhe.cc
index c095888e8..43712a2b0 100644
--- a/compiler/tests/unittest/end_to_end_jit_fhe.cc
+++ b/compiler/tests/unittest/end_to_end_jit_fhe.cc
@@ -12,8 +12,16 @@
 template <typename LambdaSupport>
 void compile_and_run(EndToEndDesc desc, LambdaSupport support) {
 
+  mlir::concretelang::CompilationOptions options("main");
+  if (desc.v0Constraint.hasValue()) {
+    options.v0FHEConstraints = *desc.v0Constraint;
+  }
+  if (desc.v0Parameter.hasValue()) {
+    options.v0Parameter = *desc.v0Parameter;
+  }
+
   /* 1 - Compile the program */
-  auto compilationResult = support.compile(desc.program);
+  auto compilationResult = support.compile(desc.program, options);
   ASSERT_EXPECTED_SUCCESS(compilationResult);
 
   /* 2 - Load the client parameters and build the keySet */