From 000b5e7f6efdf225fed33f6b4d4265edc03b861a Mon Sep 17 00:00:00 2001
From: Hyerin Park <hyerin.park@kaist.ac.kr>
Date: Wed, 6 Mar 2024 07:57:29 +0000
Subject: [PATCH] Generate: fix count & add valid filter

---
 src/main/scala/fhetest/Generate/ValidFilter.scala |  8 ++++++++
 src/main/scala/fhetest/Phase/Generate.scala       | 12 ++++++------
 2 files changed, 14 insertions(+), 6 deletions(-)

diff --git a/src/main/scala/fhetest/Generate/ValidFilter.scala b/src/main/scala/fhetest/Generate/ValidFilter.scala
index ca40fd6..51b067a 100644
--- a/src/main/scala/fhetest/Generate/ValidFilter.scala
+++ b/src/main/scala/fhetest/Generate/ValidFilter.scala
@@ -31,3 +31,11 @@ def boundIsLessThanPowerOfModSize(
   case intBound: Int       => intBound < math.pow(2, firstModSize)
   case doubleBound: Double => doubleBound < math.pow(2, firstModSize)
 }
+
+def boundIsLessThanPlainMod(
+  bound: Int | Double,
+  plainMod: Int,
+): Boolean = bound match {
+  case intBound: Int => intBound < plainMod
+  case _: Double     => true
+}
diff --git a/src/main/scala/fhetest/Phase/Generate.scala b/src/main/scala/fhetest/Phase/Generate.scala
index d631fef..bea83cd 100644
--- a/src/main/scala/fhetest/Phase/Generate.scala
+++ b/src/main/scala/fhetest/Phase/Generate.scala
@@ -38,12 +38,8 @@ case class Generate(
 
   def apply(nOpt: Option[Int]): LazyList[T2Program] = {
     println(s"Genrating Strategy: $strategy")
-    val absPrograms = nOpt match {
-      case Some(n) => allAbsPrograms.take(n)
-      case None    => allAbsPrograms
-    }
     val adjustedAbsPrograms: LazyList[AbsProgram] = for {
-      absProgram <- absPrograms
+      absProgram <- allAbsPrograms
     } yield {
       val assigned = absProgram.assignRandValues()
       val adjusted = assigned.adjustScale(encType)
@@ -52,7 +48,11 @@ case class Generate(
     val resultAbsPrograms = if (checkValid) {
       adjustedAbsPrograms.filter(_.isValid)
     } else { adjustedAbsPrograms.filterNot(_.isValid) }
-    resultAbsPrograms.map(toT2Program)
+    val takenResultAbsPrograms = nOpt match {
+      case Some(n) => resultAbsPrograms.take(n)
+      case None    => resultAbsPrograms
+    }
+    takenResultAbsPrograms.map(toT2Program)
   }
 
   // This is for testing purpose