Coverage for datesat / symbolic_bitvector / alpha_beta_bv.py: 82.7%
193 statements
« prev ^ index » next coverage.py v7.13.4, created at 2026-02-10 23:47 +0000
1"""
2Alpha-beta DateSAT implementation using epoch-based conversion.
4This module implements the alpha-beta approach where dates are represented
5as (months, days) since an epoch.
6"""
8from typing import List, Tuple, Union
10from z3 import (
11 And,
12 BitVec,
13 BitVecRef,
14 BitVecVal,
15 BoolRef,
16 CheckSatResult,
17 If,
18 ModelRef,
19 Not,
20 Optimize,
21 Or,
22 Solver,
23 sat,
24 unsat,
25)
27from ..core import Date, Period
28from .naive_bv import (
29 eom_clamp,
30 days_in_month
31)
32from .epoch_days_bv import (
33 days_since_epoch_from_ymd,
34 ymd_from_days_since_epoch,
35)
36from .bitwidths import LEGACY_BITS
38# -------------------------------
39# Alpha (months-since-epoch) helpers
40# Epoch month: 2000-03 (alpha = 0)
41# alpha = 12*y + m - (12*2000 + 3)
42# Inverse: let k = alpha + (12*2000 + 3), then
43# y = (k - 1) / 12
44# m = k - 12*y
45# -------------------------------
46# Python int epoch constants (for arithmetic outside Z3)
47_EPOCH_YEAR = 2000
48_EPOCH_MONTH = 3
49# Z3 epoch constants
50_EPOCH_LINEAR = _EPOCH_YEAR * 12 + _EPOCH_MONTH # 12*2000 + 3
51# Alpha bounds constants (months since epoch)
52_ALPHA_MIN = (1900 - _EPOCH_YEAR) * 12 + (3 - _EPOCH_MONTH) # -1200
53_ALPHA_MAX = (2100 - _EPOCH_YEAR) * 12 + (2 - _EPOCH_MONTH) # 1199
56def months_since_epoch_from_ym(y, m) -> BitVecRef:
57 """Z3-pure: compute months-since-epoch (alpha) from year/month."""
58 return (y * BitVecVal(12, LEGACY_BITS) + m) - _EPOCH_LINEAR
61def ym_from_months_since_epoch(alpha) -> Tuple[BitVecRef, BitVecRef]:
62 """Z3-pure inverse: decode (year, month) from alpha months-since-epoch."""
63 k = alpha + _EPOCH_LINEAR
64 y = (k - BitVecVal(1, LEGACY_BITS)) / BitVecVal(12, LEGACY_BITS)
65 m = k - y * BitVecVal(12, LEGACY_BITS)
66 return y, m
69class DateVar:
70 """Symbolic date variable using alpha-beta representation.
72 alpha (months_var): months since epoch month 2000-03 (March 2000 = 0)
73 beta (beta_var): extra days within that month (0-based), so DOM = 1+beta
74 """
76 def __init__(self, name: str):
77 """Create a symbolic date variable."""
78 self.name = name
79 # Alpha: Z3 bitvector variable for months since epoch-month
80 self.months_var = BitVec(f"{name}_months", LEGACY_BITS)
81 # Beta: Z3 bitvector variable for extra days (0-based) within month
82 self.beta_var = BitVec(f"{name}_beta", LEGACY_BITS)
83 # Solver reference for adding bounds to intermediate dates (set after creation if needed)
84 self._solver = None
86 def __str__(self) -> str:
87 return f"DateVar({self.name})"
89 @property
90 def year(self) -> BitVecRef:
91 """Get symbolic year component (decodes from months_var)."""
92 y, _ = ym_from_months_since_epoch(self.months_var)
93 return y
95 @property
96 def month(self) -> BitVecRef:
97 """Get symbolic month component (decodes from months_var)."""
98 _, m = ym_from_months_since_epoch(self.months_var)
99 return m
101 @property
102 def day(self) -> BitVecRef:
103 """Get symbolic day component (beta_var + 1, since beta is 0-based)."""
104 return self.beta_var + BitVecVal(1, LEGACY_BITS)
106 def to_concrete_date(self, model: ModelRef) -> Date:
107 """Convert Z3 model to concrete Date using (alpha, beta)."""
108 alpha_val = model.evaluate(
109 self.months_var, model_completion=True
110 ).as_signed_long()
111 beta_val = model.evaluate(self.beta_var, model_completion=True).as_signed_long()
112 k = alpha_val + (_EPOCH_YEAR * 12 + _EPOCH_MONTH)
113 year = (k - 1) // 12
114 month = k - year * 12
115 day = beta_val + 1
116 try:
117 return Date(year, month, day)
118 except ValueError:
119 # Intermediate result went out of bounds - use unbounded date
120 return Date(year, month, day, bounded=False)
122 def __ge__(self, other) -> BoolRef:
123 """Support x >= date comparison."""
124 if isinstance(other, Date):
125 # Convert Date to bitvector values if needed
126 alpha_o = months_since_epoch_from_ym(
127 BitVecVal(other.year, LEGACY_BITS), BitVecVal(other.month, LEGACY_BITS)
128 )
129 beta_o = BitVecVal(other.day - 1, LEGACY_BITS)
131 return Or(
132 self.months_var > alpha_o,
133 And(self.months_var == alpha_o, self.beta_var >= beta_o),
134 )
135 elif isinstance(other, DateVar): 135 ↛ 143line 135 didn't jump to line 143 because the condition on line 135 was always true
136 return Or(
137 self.months_var > other.months_var,
138 And(
139 self.months_var == other.months_var, self.beta_var >= other.beta_var
140 ),
141 )
142 else:
143 raise TypeError(f"Cannot compare DateVar with {type(other)}")
145 def __le__(self, other) -> BoolRef:
146 """Support x <= date comparison."""
147 if isinstance(other, Date):
148 alpha_o = months_since_epoch_from_ym(
149 BitVecVal(other.year, LEGACY_BITS), BitVecVal(other.month, LEGACY_BITS)
150 )
151 beta_o = BitVecVal(other.day - 1, LEGACY_BITS)
153 return Or(
154 self.months_var < alpha_o,
155 And(self.months_var == alpha_o, self.beta_var <= beta_o),
156 )
157 elif isinstance(other, DateVar): 157 ↛ 165line 157 didn't jump to line 165 because the condition on line 157 was always true
158 return Or(
159 self.months_var < other.months_var,
160 And(
161 self.months_var == other.months_var, self.beta_var <= other.beta_var
162 ),
163 )
164 else:
165 raise TypeError(f"Cannot compare DateVar with {type(other)}")
167 def __lt__(self, other) -> BoolRef:
168 """Support x < date comparison."""
169 if isinstance(other, (Date, DateVar)): 169 ↛ 172line 169 didn't jump to line 172 because the condition on line 169 was always true
170 return Not(self.__ge__(other))
171 else:
172 raise TypeError(f"Cannot compare DateVar with {type(other)}")
174 def __gt__(self, other) -> BoolRef:
175 """Support x > date comparison."""
176 if isinstance(other, (Date, DateVar)): 176 ↛ 179line 176 didn't jump to line 179 because the condition on line 176 was always true
177 return Not(self.__le__(other))
178 else:
179 raise TypeError(f"Cannot compare DateVar with {type(other)}")
181 def __eq__(self, other) -> BoolRef:
182 """Support x == date comparison."""
183 if isinstance(other, Date):
184 alpha_o = months_since_epoch_from_ym(
185 BitVecVal(other.year, LEGACY_BITS), BitVecVal(other.month, LEGACY_BITS)
186 )
187 beta_o = BitVecVal(other.day - 1, LEGACY_BITS)
189 return And(self.months_var == alpha_o, self.beta_var == beta_o)
190 elif isinstance(other, DateVar): 190 ↛ 195line 190 didn't jump to line 195 because the condition on line 190 was always true
191 return And(
192 self.months_var == other.months_var, self.beta_var == other.beta_var
193 )
194 else:
195 raise TypeError(f"Cannot compare DateVar with {type(other)}")
197 def __ne__(self, other) -> BoolRef:
198 """Support x != date comparison using ordinal arithmetic."""
199 if isinstance(other, (Date, DateVar)): 199 ↛ 202line 199 didn't jump to line 202 because the condition on line 199 was always true
200 return Not(self.__eq__(other))
201 else:
202 raise TypeError(f"Cannot compare DateVar with {type(other)}")
204 def _add_bounds(self) -> None:
205 """Add date validation bounds to this DateVar if solver is available."""
206 if self._solver is None: 206 ↛ 207line 206 didn't jump to line 207 because the condition on line 206 was never true
207 return
209 # Alpha bounds: months since 2000-03
210 # 1900-03 => -1200, 2100-02 => 1199
211 self._solver.add(self.months_var >= BitVecVal(_ALPHA_MIN, LEGACY_BITS))
212 self._solver.add(self.months_var <= BitVecVal(_ALPHA_MAX, LEGACY_BITS))
214 # Beta bounds depend on month length: 0 <= beta < days_in_month(y,m)
215 y, m = ym_from_months_since_epoch(self.months_var)
216 self._solver.add(self.beta_var >= BitVecVal(0, LEGACY_BITS))
217 self._solver.add(self.beta_var < days_in_month(y, m))
219 def __add__(self, other) -> "DateVar":
220 """DateVar + Period using alpha for Y/M and beta for D.
221 Steps:
222 - Fast path 1: If days-only period and result stays within month, simple addition.
223 - Fast path 2: If days-only period but crosses month boundary, normalize via epoch-days.
224 - Full path: Add months to alpha, EOM clamp beta, add days to beta.
225 If result stays within month, simple addition; otherwise normalize via epoch-days.
226 """
227 if isinstance(other, Period): 227 ↛ 312line 227 didn't jump to line 312 because the condition on line 227 was always true
228 result = DateVar(f"{self.name}_plus")
229 months_delta = BitVecVal(other.years * 12 + other.months, LEGACY_BITS)
230 days_delta = BitVecVal(other.days, LEGACY_BITS)
232 # Fast path: days-only period
233 if other.years == 0 and other.months == 0:
234 # Get current month info for within-month check
235 y0, m0 = ym_from_months_since_epoch(self.months_var)
236 dim0 = days_in_month(y0, m0)
238 # Add days directly to beta
239 new_beta = self.beta_var + days_delta
241 # Check if result stays within same month
242 stays_in_month = And(new_beta >= BitVecVal(0, LEGACY_BITS), new_beta < dim0)
244 # Within-month fast path: simple addition
245 alpha_within = self.months_var
246 beta_within = new_beta
248 # Fallback: normalize via epoch-days conversion (handles month overflow)
249 d0 = new_beta + BitVecVal(1, LEGACY_BITS) # Convert 0-based beta to 1-based day
250 epoch_days = days_since_epoch_from_ymd(y0, m0, d0)
251 y2, m2, d2 = ymd_from_days_since_epoch(epoch_days)
252 months_ordinal = months_since_epoch_from_ym(y2, m2)
253 beta_ordinal = d2 - BitVecVal(1, LEGACY_BITS)
255 # Select result based on within-month condition
256 result.months_var = If(stays_in_month, alpha_within, months_ordinal)
257 result.beta_var = If(stays_in_month, beta_within, beta_ordinal)
259 # Add bounds to intermediate result
260 result._solver = self._solver
261 result._add_bounds()
262 return result
264 # Full path: Add to alpha and beta directly, then normalize
265 # Step 1: Add months to alpha directly
266 alpha1 = self.months_var + months_delta
267 y1, m1 = ym_from_months_since_epoch(alpha1)
268 dim1 = days_in_month(y1, m1)
270 # Step 2: EOM clamp beta (needed when adding months - e.g., Jan 31 + 1 month = Feb 28/29)
271 d1 = eom_clamp(y1, m1, self.beta_var + BitVecVal(1, LEGACY_BITS))
272 beta1 = d1 - BitVecVal(1, LEGACY_BITS) # Convert back to 0-based beta
274 # Fast path: years/months-only period (no days)
275 if other.days == 0:
276 # No day addition needed - we're done!
277 result.months_var = alpha1
278 result.beta_var = beta1
280 result._solver = self._solver
281 result._add_bounds()
282 return result
284 else:
285 # Step 3: Add days to beta directly
286 new_beta = beta1 + days_delta
288 # Check if result stays within same month
289 stays_in_month = And(new_beta >= BitVecVal(0, LEGACY_BITS), new_beta < dim1)
291 # Within-month fast path: simple addition
292 alpha_within = alpha1
293 beta_within = new_beta
295 # Fallback: Normalize via epoch-days conversion (handles all overflow cases)
296 d_temp = new_beta + BitVecVal(1, LEGACY_BITS) # Convert 0-based beta to 1-based day
297 epoch_days = days_since_epoch_from_ymd(y1, m1, d_temp)
298 y2, m2, d2 = ymd_from_days_since_epoch(epoch_days)
300 months_ordinal = months_since_epoch_from_ym(y2, m2)
301 beta_ordinal = d2 - BitVecVal(1, LEGACY_BITS)
303 # Select result based on within-month condition
304 result.months_var = If(stays_in_month, alpha_within, months_ordinal)
305 result.beta_var = If(stays_in_month, beta_within, beta_ordinal)
307 # Add bounds to intermediate result
308 result._solver = self._solver
309 result._add_bounds()
310 return result
311 else:
312 raise TypeError(f"Cannot add {type(other)} to DateVar")
314 def __sub__(self, other) -> "DateVar":
315 """DateVar - Period implemented as DateVar + (-Period)."""
316 if isinstance(other, Period): 316 ↛ 320line 316 didn't jump to line 320 because the condition on line 316 was always true
317 neg = Period(-other.years, -other.months, -other.days)
318 return self.__add__(neg)
319 else:
320 raise TypeError(f"Cannot subtract {type(other)} from DateVar")
323class AlphaBetaSolver:
324 """Alpha-beta date constraint solver using epoch-based conversion."""
326 def __init__(self, timeout_ms=600000, use_maxsat=False):
327 """Initialize the solver with timeout.
329 Args:
330 timeout_ms: Timeout in milliseconds (default: 60 seconds)
331 use_maxsat: If True, use MaxSAT optimization with soft constraints
332 """
333 self.use_maxsat = use_maxsat
334 if use_maxsat: 334 ↛ 335line 334 didn't jump to line 335 because the condition on line 334 was never true
335 self.solver = Optimize()
336 else:
337 self.solver = Solver()
338 self.solver.set("timeout", timeout_ms)
339 self.date_vars = {}
340 self.constraints = []
341 self.timeout_ms = timeout_ms
343 def add_date_var(self, name: str) -> DateVar:
344 """Add a symbolic date variable with basic constraints."""
345 date_var = DateVar(name)
346 date_var._solver = self.solver
347 self.date_vars[name] = date_var
349 # Add bounds using _add_bounds method
350 date_var._add_bounds()
351 return date_var
353 def add_constraint(self, constraint: BoolRef) -> None:
354 """Add a constraint to the solver."""
355 self.constraints.append(constraint)
356 self.solver.add(constraint)
358 def check(self) -> CheckSatResult:
359 """Check if constraints are satisfiable."""
360 return self.solver.check()
362 def model(self) -> ModelRef:
363 """Get the model if satisfiable."""
364 return self.solver.model()
366 def get_concrete_dates(self, model: ModelRef) -> dict:
367 """Get concrete dates from the model."""
368 return {
369 name: var.to_concrete_date(model) for name, var in self.date_vars.items()
370 }
372 def solve(self) -> Union[bool, dict]:
373 """Solve the constraints."""
374 # Add MaxSAT soft constraints if enabled
375 if self.use_maxsat: 375 ↛ 376line 375 didn't jump to line 376 because the condition on line 375 was never true
376 from datetime import date
378 today = date.today()
379 # Calculate months since epoch for today
380 today_months = (today.year - _EPOCH_YEAR) * 12 + (
381 today.month - _EPOCH_MONTH
382 )
384 # Convert years to months
385 months_50_years = 50 * 12 # 600 months
386 months_10_years = 10 * 12 # 120 months
388 # Add soft constraints for each date variable
389 for name, date_var in self.date_vars.items():
390 # High weight: today ± 50 years
391 within_50_years = And(
392 date_var.months_var
393 >= BitVecVal(today_months - months_50_years, LEGACY_BITS),
394 date_var.months_var
395 <= BitVecVal(today_months + months_50_years, LEGACY_BITS),
396 )
397 self.solver.add_soft(within_50_years, weight=100)
399 # Low weight: today ± 10 years
400 within_10_years = And(
401 date_var.months_var
402 >= BitVecVal(today_months - months_10_years, LEGACY_BITS),
403 date_var.months_var
404 <= BitVecVal(today_months + months_10_years, LEGACY_BITS),
405 )
406 self.solver.add_soft(within_10_years, weight=10)
408 result = self.check()
409 if result == sat:
410 model = self.model()
411 return {
412 "status": "sat",
413 "dates": self.get_concrete_dates(model),
414 }
415 elif result == unsat: 415 ↛ 419line 415 didn't jump to line 419 because the condition on line 415 was always true
416 return {"status": "unsat", "dates": {}}
417 else:
418 # result == unknown (timeout or resource limit)
419 return {"status": "timeout", "dates": {}}
421 def to_smt2(self) -> str:
422 """Return the current problem in SMT-LIB v2 format."""
423 return self.solver.to_smt2()
425 def get_assertions(self) -> List[BoolRef]:
426 """Return the list of current Z3 assertions (BoolRef)."""
427 return list(self.solver.assertions())