z3_solver.py•9.54 kB
"""
Z3 SMT solver implementation for constraint satisfaction problems.
"""
from typing import cast, Dict, List, Optional, Tuple
import time
import z3
from returns.maybe import Maybe, Nothing, Some
from returns.result import Failure, Result, Success
from ..models.z3_models import (
Z3Constraint,
Z3Problem,
Z3Ref,
Z3Solution,
Z3Value,
Z3Variable,
Z3VariableType,
)
def create_variable(variable: Z3Variable) -> Result[Tuple[str, Z3Ref], str]:
"""Takes a Variable model and converts it into a Z3 variable representation.
Args:
variable: The Variable model to convert.
Returns:
Result: A Result containing either a successful tuple of (variable_name, z3_variable)
or a failure with error details.
"""
try:
name = variable.name
match variable.type:
case Z3VariableType.INTEGER:
z3_var = cast(Z3Ref, z3.Int(name))
case Z3VariableType.REAL:
z3_var = cast(Z3Ref, z3.Real(name))
case Z3VariableType.BOOLEAN:
z3_var = cast(Z3Ref, z3.Bool(name))
case Z3VariableType.STRING:
z3_var = cast(Z3Ref, z3.String(name))
case _:
return Failure(f"Unsupported variable type: {variable.type}")
# Add bounds if specified
if variable.lower_bound is not None:
z3_var = z3_var >= variable.lower_bound
if variable.upper_bound is not None:
z3_var = z3_var <= variable.upper_bound
return Success((name, z3_var))
except Exception as e:
return Failure(f"Error creating variable {variable.name}: {e!s}")
def create_variables(variables: List[Z3Variable]) -> Result[Dict[str, Z3Ref], str]:
"""Processes a list of variable definitions and creates corresponding Z3 variables.
Args:
variables: List of Variable definitions to process.
Returns:
Result: A Result containing either a dictionary mapping variable names to Z3 variables
or a failure with error details.
"""
result_dict = {}
for var in variables:
result = create_variable(var)
match result:
case Success(value):
name, z3_var = value
result_dict[name] = z3_var
case Failure(_):
return result
return Success(result_dict)
def parse_constraint(
constraint: Z3Constraint, variables: Dict[str, Z3Ref]
) -> Result[z3.BoolRef, str]:
"""Converts a constraint expression into a Z3 constraint using the provided variables.
Args:
constraint: The constraint definition to parse.
variables: Dictionary of variable names to Z3 variables.
Returns:
Result: A Result containing either a Z3 boolean reference or a failure with error details.
"""
try:
local_dict = {
**variables,
**{
"And": z3.And,
"Or": z3.Or,
"Not": z3.Not,
"Implies": z3.Implies,
"ForAll": z3.ForAll,
"Exists": z3.Exists,
"If": z3.If,
"Distinct": z3.Distinct,
"true": True,
"false": False,
},
}
z3_constraint = eval(constraint.expression, {"__builtins__": {}}, local_dict)
return Success(z3_constraint)
except Exception as e:
return Failure(f"Error parsing constraint '{constraint.expression}': {e!s}")
def create_constraints(
constraints: List[Z3Constraint], variables: Dict[str, Z3Ref]
) -> Result[List[z3.BoolRef], str]:
"""Transforms a list of constraint definitions into Z3 constraints.
Args:
constraints: List of constraint definitions to transform.
variables: Dictionary of variable names to Z3 variables.
Returns:
Result: A Result containing either a list of Z3 boolean references or a failure with error details.
"""
z3_constraints = []
for constraint in constraints:
result = parse_constraint(constraint, variables)
match result:
case Success(value):
z3_constraints.append(value)
case Failure(_):
return result
return Success(z3_constraints)
def solve(
variables: Dict[str, Z3Ref], constraints: List[z3.BoolRef], timeout: Optional[int] = None
) -> Result[Tuple[z3.CheckSatResult, Optional[z3.ModelRef]], str]:
"""Attempts to solve a Z3 problem with the given variables and constraints.
Args:
variables: Dictionary of variable names to Z3 variables.
constraints: List of Z3 constraints to solve.
timeout: Optional timeout in milliseconds.
Returns:
Result: A Result containing either a tuple of (satisfiability_result, model)
or a failure with error details.
"""
try:
solver = z3.Solver()
# Set timeout if specified
if timeout is not None:
solver.set("timeout", timeout)
solver.add(constraints)
result = solver.check()
model = None
if result == z3.sat:
model = solver.model()
return Success((result, model))
except Exception as e:
return Failure(f"Error solving constraints: {e!s}")
def get_z3_value(model: z3.ModelRef, z3_var: Z3Ref) -> Maybe[Z3Value]:
"""Extracts a value from a Z3 model for a given variable using Maybe monad.
Args:
model: The Z3 model to extract from.
z3_var: The Z3 variable to extract.
Returns:
Maybe: A Maybe containing either the extracted value or Nothing if extraction fails.
"""
try:
z3_value = model[z3_var]
if z3_value is None:
return Nothing
str_value = str(z3_value)
match z3_var:
case _ if z3.is_bool(z3_var):
return Some(z3.is_true(z3_value))
case _ if z3.is_int(z3_var):
return Some(int(str_value))
case _ if z3.is_real(z3_var):
if "/" in str_value:
num_str, den_str = str_value.split("/")
return Some(float(int(num_str)) / float(int(den_str)))
else:
return Some(float(str_value))
case _:
return Some(str_value)
except Exception:
return Nothing
def extract_solution(
result: z3.CheckSatResult,
model: Optional[z3.ModelRef],
variables: Dict[str, Z3Ref],
solve_time: Optional[float] = None,
) -> Result[Z3Solution, str]:
"""Processes a Z3 solution and extracts variable values from the model.
Args:
result: The satisfiability result from Z3.
model: The Z3 model if satisfiable, None otherwise.
variables: Dictionary of variable names to Z3 variables.
solve_time: Optional solve time in seconds.
Returns:
Result: A Result containing either a Z3Solution or a failure with error details.
"""
try:
match result:
case z3.sat:
if model is None:
return Failure("Model is None but result is sat")
values = {
name: maybe_value.unwrap()
for name, z3_var in variables.items()
for maybe_value in [get_z3_value(model, z3_var)]
if isinstance(maybe_value, Some)
}
return Success(
Z3Solution(
values=values,
is_satisfiable=True,
status=str(result),
model=model,
solve_time=solve_time,
)
)
case _:
return Success(
Z3Solution(
values={},
is_satisfiable=False,
status=str(result),
solve_time=solve_time,
)
)
except Exception as e:
return Failure(f"Error extracting solution: {e!s}")
def solve_z3_problem(problem: Z3Problem) -> Result[Z3Solution, str]:
"""Orchestrates the complete process of solving a Z3 problem.
Args:
problem: The problem definition containing variables and constraints.
Returns:
Result: A Result containing either a Solution object or a failure with error details.
"""
try:
# Validate problem
if not problem.validate():
return Failure("Invalid problem definition")
start_time = time.time()
vars_result = create_variables(problem.variables)
if isinstance(vars_result, Failure):
return vars_result
variables = vars_result.unwrap()
constraints_result = create_constraints(problem.constraints, variables)
if isinstance(constraints_result, Failure):
return constraints_result
z3_constraints = constraints_result.unwrap()
solve_result = solve(variables, z3_constraints, problem.timeout)
if isinstance(solve_result, Failure):
return solve_result
result_tuple = solve_result.unwrap()
solve_time = time.time() - start_time
return extract_solution(result_tuple[0], result_tuple[1], variables, solve_time)
except Exception as e:
return Failure(f"Error solving problem: {e!s}")