Make collect_model_values work on a list of SymData

dgpv · Nov 10, 2023 · 1af33b1 · 1af33b1
1 parent e5b8dcc
commit 1af33b1
Showing 1 changed file with 26 additions and 25 deletions.
diff --git a/bsst/__init__.py b/bsst/__init__.py
@@ -2378,6 +2378,31 @@ def getval(v: SymData) -> int | bytes:
         a.set_static(getval(b))
 
 
+def collect_model_values(
+    values: Iterable['SymData'],
+    cb: Callable[[Optional[dict[str, 'ConstrainedValue']]],
+                 Optional['z3.BoolRef']]
+) -> None:
+    z3_push_context()
+
+    mvdict_req: dict[str, tuple[str, SymDataRType]] = {}
+    mvnamemap: dict[str, 'SymData'] = {}
+
+    for v in values:
+        v.update_model_values_request_dict(mvdict_req, mvnamemap)
+
+    mvdict: Optional[dict[str, 'ConstrainedValue']] = None
+
+    while cb(mvdict):
+        try:
+            mvdict = z3check(force_check=True,
+                             model_values_to_retrieve=mvdict_req)
+        except ScriptFailure:
+            break
+
+    z3_pop_context()
+
+
 def is_cond_possible(  # noqa
     cond: Union[bool, 'z3.BoolRef'], sd: 'SymData',
     *, name: str = '', fail_msg: str = '',
@@ -5207,30 +5232,6 @@ def set_known_bool(self, value: bool, set_size: bool = False) -> None:
 
         cur_context().known_bool_values[self._unique_name] = value
 
-    def collect_model_values(
-        self, cb: Callable[[Optional[dict[str, 'ConstrainedValue']]],
-                           Optional['z3.BoolRef']]
-    ) -> None:
-        z3_push_context()
-
-        mvdict_req: dict[str, tuple[str, SymDataRType]] = {}
-        mvnamemap: dict[str, 'SymData'] = {}
-
-        self.update_model_values_request_dict(mvdict_req, mvnamemap)
-
-        mvdict: Optional[dict[str, 'ConstrainedValue']] = None
-
-        try:
-            while cb(mvdict):
-                mvdict = z3check(force_check=True,
-                                 model_values_to_retrieve=mvdict_req)
-        except ScriptFailure:
-            return
-
-        z3_pop_context()
-
-        return
-
     def collect_integer_model_values(self, max_count: int) -> list[int]:
 
         result: list[int] = []
@@ -5259,7 +5260,7 @@ def collect(mvdict: Optional[dict[str, 'ConstrainedValue']]) -> bool:
 
             return True
 
-        self.collect_model_values(collect)
+        collect_model_values([self], collect)
 
         return result