Use correct sat variables for model counting and model enumeration

booleworks · Apr 15, 2024 · dd2d495 · dd2d495
1 parent 0a87c70
commit dd2d495
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 3 deletions.
diff --git a/src/main/kotlin/com/booleworks/boolerules/computations/generic/FeatureFunctions.kt b/src/main/kotlin/com/booleworks/boolerules/computations/generic/FeatureFunctions.kt
@@ -15,7 +15,8 @@ fun computeRelevantVars(
     features: List<String>,
 ): SortedSet<Variable> =
     if (features.isEmpty()) {
-        (info.knownVariables + info.encodingContext.relevantSatVariables).toSortedSet()
+        val intSatVariables = info.integerVariables.flatMap { info.encodingContext.variableMap[it.variable]?.values ?: emptySet() }
+        (info.knownVariables + intSatVariables).toSortedSet()
     } else {
         val vars = TreeSet<Variable>()
         info.integerVariables

diff --git a/src/main/kotlin/com/booleworks/boolerules/computations/modelcount/ModelCountComputation.kt b/src/main/kotlin/com/booleworks/boolerules/computations/modelcount/ModelCountComputation.kt
@@ -65,18 +65,21 @@ internal object ModelCountComputation :
     ): ModelCountInternalResult {
         // currently no projected model counting
 //        val variables = computeRelevantVars(f, info, listOf())
-        val variables = (info.knownVariables + info.intPredicateMapping.values + info.encodingContext.relevantSatVariables).toSortedSet()
+        val intSatVariables = info.integerVariables.flatMap { info.encodingContext.variableMap[it.variable]?.values ?: emptySet() }
+        val variables = (info.knownVariables + info.intPredicateMapping.values + intSatVariables).toSortedSet()
         val formulas = info.propositions.map { it.formula() }
         val additionalDefinitions =
             request.additionalConstraints.map { constraint -> processConstraint(f, constraint, model, info, status) }
         val additionalConstraints = additionalDefinitions.map { it!!.first.formula() }
         val additionalVarDefinitions = additionalDefinitions.flatMap { it!!.second }.toSet().map { info.integerEncodings.getEncoding(it)!! }
+        val allFormulas = formulas + additionalConstraints + additionalVarDefinitions
+        val allVariables = (allFormulas.flatMap { it.variables(f) } + variables).toSortedSet()
         if (!status.successful()) {
             return ModelCountInternalResult(slice, BigInteger.ZERO)
         }
         return ModelCountInternalResult(
             slice,
-            ModelCounter.count(f, formulas + additionalConstraints + additionalVarDefinitions, variables)
+            ModelCounter.count(f, formulas + additionalConstraints + additionalVarDefinitions, allVariables)
         )
     }