Skip to main content

theory_core/
ir.rs

1//! Constraint Intermediate Representation (IR)
2//!
3//! Provides a typed, versioned intermediate representation for constraints
4//! that can be consumed by both neural networks and SMT solvers.
5
6use indexmap::IndexMap;
7use serde::{Deserialize, Serialize};
8use smallvec::SmallVec;
9use std::fmt;
10
11/// Versioned constraint IR to enable evolution
12#[derive(Debug, Clone, Serialize, Deserialize)]
13pub struct ConstraintIR {
14    pub version: u32,
15    pub variables: IndexMap<VarId, Variable>,
16    pub constraints: Vec<Constraint>,
17    pub objective: Option<Objective>,
18    pub theory_tags: SmallVec<[TheoryTag; 4]>,
19}
20
21impl ConstraintIR {
22    pub fn new() -> Self {
23        Self {
24            version: 1,
25            variables: IndexMap::new(),
26            constraints: Vec::new(),
27            objective: None,
28            theory_tags: SmallVec::new(),
29        }
30    }
31
32    pub fn add_variable(&mut self, var: Variable) -> VarId {
33        let id = VarId(self.variables.len());
34        self.variables.insert(id, var);
35        id
36    }
37
38    pub fn add_constraint(&mut self, constraint: Constraint) {
39        self.constraints.push(constraint);
40    }
41
42    pub fn set_objective(&mut self, objective: Objective) {
43        self.objective = Some(objective);
44    }
45
46    pub fn add_theory_tag(&mut self, tag: TheoryTag) {
47        if !self.theory_tags.contains(&tag) {
48            self.theory_tags.push(tag);
49        }
50    }
51
52    /// Create a simple test constraint IR for demonstration purposes
53    pub fn new_simple_test() -> Self {
54        let mut ir = Self::new();
55
56        // Add variables
57        let x = ir.add_variable(Variable {
58            name: "x".to_string(),
59            domain: Domain::Integer {
60                min: Some(0),
61                max: Some(100),
62            },
63            metadata: VariableMetadata::default(),
64        });
65
66        let y = ir.add_variable(Variable {
67            name: "y".to_string(),
68            domain: Domain::Integer {
69                min: Some(0),
70                max: Some(100),
71            },
72            metadata: VariableMetadata::default(),
73        });
74
75        // Add constraints: x + y = 10
76        ir.add_constraint(Constraint::Equal {
77            lhs: Expr::Binary {
78                op: BinaryOp::Add,
79                lhs: Box::new(Expr::Var(x)),
80                rhs: Box::new(Expr::Var(y)),
81            },
82            rhs: Expr::Const(ConstValue::Integer(10)),
83        });
84
85        // Add constraint: x >= 2*y (represented as NOT(x < 2*y))
86        ir.add_constraint(Constraint::Not {
87            constraint: Box::new(Constraint::Less {
88                lhs: Expr::Var(x),
89                rhs: Expr::Binary {
90                    op: BinaryOp::Mul,
91                    lhs: Box::new(Expr::Const(ConstValue::Integer(2))),
92                    rhs: Box::new(Expr::Var(y)),
93                },
94            }),
95        });
96
97        // Add theory tag for integer linear programming
98        ir.add_theory_tag(TheoryTag::LIA);
99
100        ir
101    }
102}
103
104/// Variable identifier
105#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize)]
106pub struct VarId(pub usize);
107
108impl fmt::Display for VarId {
109    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
110        write!(f, "v{}", self.0)
111    }
112}
113
114/// Variable in the constraint system
115#[derive(Debug, Clone, Serialize, Deserialize)]
116pub struct Variable {
117    pub name: String,
118    pub domain: Domain,
119    pub metadata: VariableMetadata,
120}
121
122/// Domain specification for variables
123#[derive(Debug, Clone, Serialize, Deserialize)]
124pub enum Domain {
125    /// Boolean variable (true/false)
126    Boolean,
127    /// Integer with optional bounds
128    Integer { min: Option<i64>, max: Option<i64> },
129    /// Real number with optional bounds
130    Real { min: Option<f64>, max: Option<f64> },
131    /// Bit-vector of specified width
132    BitVector { width: u32 },
133    /// String with optional length bounds
134    String { max_length: Option<usize> },
135    /// Array of elements
136    Array {
137        element_type: Box<Domain>,
138        size: Option<usize>,
139    },
140    /// Enumeration of possible values
141    Enum { values: Vec<String> },
142}
143
144/// Metadata for variables (optimization hints, etc.)
145#[derive(Debug, Clone, Default, Serialize, Deserialize)]
146pub struct VariableMetadata {
147    pub description: Option<String>,
148    pub importance: f32,       // 0.0 to 1.0
149    pub computation_cost: f32, // estimated cost to compute
150}
151
152/// Constraint in the system
153#[derive(Debug, Clone, Serialize, Deserialize)]
154pub enum Constraint {
155    /// Equality constraint: lhs = rhs
156    Equal { lhs: Expr, rhs: Expr },
157    /// Inequality constraint: lhs < rhs
158    Less { lhs: Expr, rhs: Expr },
159    /// Less than or equal: lhs <= rhs
160    LessEqual { lhs: Expr, rhs: Expr },
161    /// Logical AND of sub-constraints
162    And { constraints: Vec<Constraint> },
163    /// Logical OR of sub-constraints
164    Or { constraints: Vec<Constraint> },
165    /// Negation of constraint
166    Not { constraint: Box<Constraint> },
167    /// Implication: if antecedent then consequent
168    Implies {
169        antecedent: Box<Constraint>,
170        consequent: Box<Constraint>,
171    },
172    /// Custom constraint with solver-specific encoding
173    Custom {
174        name: String,
175        params: IndexMap<String, Expr>,
176    },
177}
178
179/// Expression in constraints
180#[derive(Debug, Clone, Serialize, Deserialize)]
181pub enum Expr {
182    /// Variable reference
183    Var(VarId),
184    /// Constant values
185    Const(ConstValue),
186    /// Binary operations
187    Binary {
188        op: BinaryOp,
189        lhs: Box<Expr>,
190        rhs: Box<Expr>,
191    },
192    /// Unary operations
193    Unary { op: UnaryOp, operand: Box<Expr> },
194    /// Function call
195    Call { function: String, args: Vec<Expr> },
196}
197
198/// Constant values in expressions
199#[derive(Debug, Clone, Serialize, Deserialize)]
200pub enum ConstValue {
201    Bool(bool),
202    Integer(i64),
203    Real(f64),
204    String(String),
205    BitVector { value: u64, width: u32 },
206}
207
208/// Binary operators
209#[derive(Debug, Clone, Serialize, Deserialize)]
210pub enum BinaryOp {
211    Add,
212    Sub,
213    Mul,
214    Div,
215    Mod,
216    Power, // Exponentiation: a^b
217    And,
218    Or,
219    Xor,
220    Equal,
221    NotEqual,
222    Less,
223    LessEqual,
224    Greater,
225    GreaterEqual,
226}
227
228/// Unary operators
229#[derive(Debug, Clone, Serialize, Deserialize)]
230pub enum UnaryOp {
231    Not,
232    Negate, // Arithmetic negation (-x)
233    Sqrt,   // Square root
234    Abs,    // Absolute value
235}
236
237/// Optimization objective
238#[derive(Debug, Clone, Serialize, Deserialize)]
239pub struct Objective {
240    pub goal: OptimizationGoal,
241    pub expr: Expr,
242    pub weight: f64,
243}
244
245/// Optimization goals
246#[derive(Debug, Clone, Serialize, Deserialize)]
247pub enum OptimizationGoal {
248    Minimize,
249    Maximize,
250    Satisfy, // Just find a satisfying assignment
251}
252
253/// Theory tags to route to appropriate solvers
254#[derive(Debug, Clone, PartialEq, Eq, Hash, Serialize, Deserialize)]
255pub enum TheoryTag {
256    /// Linear integer arithmetic
257    LIA,
258    /// Linear real arithmetic  
259    LRA,
260    /// Non-linear arithmetic
261    NLA,
262    /// Bit-vectors
263    BV,
264    /// Arrays
265    Array,
266    /// Strings
267    String,
268    /// Floating-point arithmetic
269    FP,
270    /// Uninterpreted functions
271    UF,
272    /// Quantified formulas
273    Quantified,
274}
275
276impl Default for ConstraintIR {
277    fn default() -> Self {
278        Self::new()
279    }
280}