Generate: fix count & add valid filter

kaist-plrg · Mar 6, 2024 · 000b5e7 · 000b5e7
1 parent ebd4441
commit 000b5e7
Show file tree

Hide file tree

Showing 2 changed files with 14 additions and 6 deletions.
diff --git 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
@@ -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