1
use crate::ast::{
2
    DeclarationKind, DomainOpError, Expression, FuncAttr, Literal, Metadata, Moo, PartitionAttr,
3
    PermutationAttr, Reference, RelAttr, ReturnType, SequenceAttr, Typeable,
4
    domains::{Int, MSetAttr, Range, SetAttr},
5
    eval_constant,
6
};
7
use crate::{bug, into_matrix_expr, matrix_expr};
8
use funcmap::{FuncMap, TryFuncMap};
9
use polyquine::Quine;
10
use serde::{Deserialize, Serialize};
11
use std::fmt::{Display, Formatter};
12
use std::ops::Deref;
13
use uniplate::Uniplate;
14

            
15
/// A variable or expression appearing inside an int range of an unresolved domain;
16
/// E.g `int(1..x)`, `int(2, 4..(2*y))`, `set (minSize x) of int(1..5)`, etc
17
#[derive(Clone, Debug, PartialEq, Eq, Hash, Serialize, Deserialize, Quine, Uniplate)]
18
#[path_prefix(conjure_cp::ast)]
19
#[biplate(to=Expression)]
20
#[biplate(to=Reference)]
21
pub enum IntVal {
22
    // For ergonomics, we use a type bigger than both Int and UInt so both fit;
23
    // Overflows are handled when resolving
24
    Const(i64),
25
    #[polyquine_skip]
26
    Reference(Reference),
27
    Expr(Moo<Expression>),
28
}
29

            
30
// ------------------------------------
31
// ------ Trait impls for IntVal ------
32
// ------------------------------------
33

            
34
impl<T> From<T> for IntVal
35
where
36
    T: Into<i64>,
37
{
38
13897975
    fn from(v: T) -> Self {
39
13897975
        IntVal::Const(v.into())
40
13897975
    }
41
}
42

            
43
impl TryFrom<IntVal> for Int {
44
    type Error = DomainOpError;
45

            
46
1263849
    fn try_from(value: IntVal) -> Result<Int, Self::Error> {
47
1263849
        match value {
48
1023208
            IntVal::Const(val) => val.try_into().map_err(|_| DomainOpError::OutOfBounds),
49
240641
            _ => Err(DomainOpError::NotGround),
50
        }
51
1263849
    }
52
}
53

            
54
impl Display for IntVal {
55
1321565
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
56
1321565
        match self {
57
620130
            IntVal::Const(val) => write!(f, "{val}"),
58
565722
            IntVal::Reference(re) => write!(f, "{re}"),
59
135713
            IntVal::Expr(expr) => write!(f, "({expr})"),
60
        }
61
1321565
    }
62
}
63

            
64
impl std::ops::Neg for IntVal {
65
    type Output = IntVal;
66

            
67
378519
    fn neg(self) -> Self::Output {
68
378519
        match self {
69
378519
            IntVal::Const(val) => IntVal::Const(-val),
70
            IntVal::Reference(re) => IntVal::Expr(Moo::new(Expression::Neg(
71
                Metadata::new(),
72
                Moo::new(re.into()),
73
            ))),
74
            IntVal::Expr(expr) => IntVal::Expr(Moo::new(Expression::Neg(Metadata::new(), expr))),
75
        }
76
378519
    }
77
}
78

            
79
// ----------------------------------------
80
// ------ Core IntVal implementation ------
81
// ----------------------------------------
82

            
83
impl IntVal {
84
5409110
    pub fn new_const(val: Int) -> IntVal {
85
5409110
        IntVal::Const(val as i64)
86
5409110
    }
87

            
88
195594
    pub fn new_ref(re: &Reference) -> Result<IntVal, DomainOpError> {
89
195594
        match re.ptr.kind().deref() {
90
            DeclarationKind::ValueLetting(expr, _)
91
86462
            | DeclarationKind::TemporaryValueLetting(expr) => match expr.return_type() {
92
86462
                ReturnType::Int => Ok(IntVal::Reference(re.clone())),
93
                _ => Err(DomainOpError::WrongType),
94
            },
95
100287
            DeclarationKind::Given(dom) => match dom.return_type() {
96
100287
                ReturnType::Int => Ok(IntVal::Reference(re.clone())),
97
                _ => Err(DomainOpError::WrongType),
98
            },
99
            // TODO: I'm not sure if this is correct, see discussion #1890
100
            // Assume that for `y : int(1..x)`, where x is an induction variable, to be valid,
101
            // x must be an integer (e.g x : int(1..3)). I don't know how the generator expression
102
            // fits into this.
103
348
            DeclarationKind::Quantified(inner) => match inner.domain().return_type() {
104
348
                ReturnType::Int => Ok(IntVal::Reference(re.clone())),
105
                _ => Err(DomainOpError::WrongType),
106
            },
107
            // TODO: I'm not sure if this is correct, see discussion #1890
108
            // Assume that for `y : int(1..x)`, where `x <- expr`, to be valid,
109
            // x must be a collection of integers; E.g, `x <- {1, 2, 3}`.
110
            DeclarationKind::QuantifiedExpr(expr) => match expr.return_type().elem_type() {
111
                Some(ReturnType::Int) => Ok(IntVal::Reference(re.clone())),
112
                _ => Err(DomainOpError::WrongType),
113
            },
114
            DeclarationKind::Find(var) | DeclarationKind::FindAuxiliary(var) => {
115
8497
                match var.return_type() {
116
8497
                    ReturnType::Int => Ok(IntVal::Reference(re.clone())),
117
                    _ => Err(DomainOpError::WrongType),
118
                }
119
            }
120
            DeclarationKind::DomainLetting(_)
121
            | DeclarationKind::EnumeratedType(_)
122
            | DeclarationKind::UnnamedType(_) => Err(DomainOpError::WrongType),
123
        }
124
195594
    }
125

            
126
70996
    pub fn new_expr(value: Moo<Expression>) -> Result<IntVal, DomainOpError> {
127
70996
        if value.return_type() != ReturnType::Int {
128
            return Err(DomainOpError::WrongType);
129
70996
        }
130
70996
        Ok(IntVal::Expr(value))
131
70996
    }
132

            
133
102875241
    pub fn resolve(&self) -> Result<Int, DomainOpError> {
134
102875241
        match self {
135
53130644
            IntVal::Const(value) => (*value).try_into().map_err(|_| DomainOpError::OutOfBounds),
136
19685669
            IntVal::Expr(expr) => eval_expr_to_int(expr).ok_or(DomainOpError::NotGround),
137
30058928
            IntVal::Reference(re) => match re.ptr.kind().deref() {
138
                DeclarationKind::ValueLetting(expr, _)
139
                | DeclarationKind::TemporaryValueLetting(expr) => {
140
29554040
                    eval_expr_to_int(expr).ok_or(DomainOpError::NotGround)
141
                }
142
                // If this is an int given we will be able to resolve it eventually, but not yet
143
45992
                DeclarationKind::Given(_) => Err(DomainOpError::NotGround),
144
                DeclarationKind::Quantified(inner) => {
145
145
                    if let Some(generator) = inner.generator()
146
                        && let Some(expr) = generator.as_value_letting()
147
                    {
148
                        eval_expr_to_int(&expr).ok_or(DomainOpError::NotGround)
149
                    } else {
150
145
                        Err(DomainOpError::NotGround)
151
                    }
152
                }
153
                // Decision variables inside domains are unresolved until solving.
154
                DeclarationKind::Find(_) | DeclarationKind::FindAuxiliary(_) => {
155
458751
                    Err(DomainOpError::NotGround)
156
                }
157
                DeclarationKind::DomainLetting(_)
158
                | DeclarationKind::EnumeratedType(_)
159
                | DeclarationKind::UnnamedType(_)
160
                | DeclarationKind::QuantifiedExpr(_) => bug!(
161
                    "Expected integer expression, given, or letting inside int domain; Got: {re}"
162
                ),
163
            },
164
        }
165
102875241
    }
166

            
167
    pub fn try_add<T: Into<Expression>>(self, rhs: T) -> Result<IntVal, DomainOpError> {
168
        let sum = Expression::Sum(
169
            Metadata::new(),
170
            Moo::new(matrix_expr!(self.try_into()?, rhs.into())),
171
        );
172
        Ok(IntVal::Expr(Moo::new(sum)))
173
    }
174

            
175
    pub fn try_sub<T: Into<Expression>>(self, rhs: T) -> Result<IntVal, DomainOpError> {
176
        let rhs_neg = Expression::Neg(Metadata::new(), Moo::new(rhs.into()));
177
        let sum = Expression::Sum(
178
            Metadata::new(),
179
            Moo::new(matrix_expr!(self.try_into()?, rhs_neg)),
180
        );
181
        Ok(IntVal::Expr(Moo::new(sum)))
182
    }
183
}
184

            
185
// ------------------------------------------
186
// ------ Expression-related miscellanea ----
187
// ------------------------------------------
188

            
189
impl Range<IntVal> {
190
    /// Generates the expression to compute the size of this range
191
    pub fn len_expr(self) -> Result<Expression, DomainOpError> {
192
        match self {
193
            Range::Single(a) => Ok(a.try_into()?),
194
            Range::Bounded(a, b) => {
195
                let neg_b = Expression::Neg(Metadata::new(), b.try_into()?);
196
                let sum_matr = into_matrix_expr!(vec![a.try_into()?, neg_b]);
197
                Ok(Expression::Sum(Metadata::new(), sum_matr.into()))
198
            }
199
            _ => Err(DomainOpError::Unbounded),
200
        }
201
    }
202

            
203
    /// Generates the expression to compute the size of a list of ranges
204
    pub fn len_expr_of(rngs: &[Range<IntVal>]) -> Result<Expression, DomainOpError> {
205
        let mut rng_sizes = Vec::with_capacity(rngs.len());
206
        for rng in rngs {
207
            rng_sizes.push(rng.clone().len_expr()?);
208
        }
209
        let rng_sizes = into_matrix_expr!(rng_sizes);
210
        Ok(Expression::Sum(Metadata::new(), rng_sizes.into()))
211
    }
212
}
213

            
214
49239709
fn eval_expr_to_int(expr: &Expression) -> Option<Int> {
215
49239709
    match eval_constant(expr)? {
216
48974619
        Literal::Int(v) => Some(v),
217
        _ => bug!("Expected integer expression, got: {expr}"),
218
    }
219
49239709
}
220

            
221
impl TryFrom<IntVal> for Moo<Expression> {
222
    type Error = DomainOpError;
223

            
224
25346
    fn try_from(value: IntVal) -> Result<Self, Self::Error> {
225
25346
        match value {
226
            IntVal::Const(val) => {
227
8497
                let val: Int = val.try_into().map_err(|_| DomainOpError::OutOfBounds)?;
228
8497
                Ok(Moo::new(val.into()))
229
            }
230
10556
            IntVal::Reference(re) => Ok(Moo::new(re.into())),
231
6293
            IntVal::Expr(expr) => Ok(expr),
232
        }
233
25346
    }
234
}
235

            
236
impl TryFrom<IntVal> for Expression {
237
    type Error = DomainOpError;
238
25346
    fn try_from(value: IntVal) -> Result<Self, Self::Error> {
239
25346
        Ok(Moo::unwrap_or_clone(value.try_into()?))
240
25346
    }
241
}
242

            
243
// --------------------------------------------------------------
244
// ------ Derive into / resolve for container types by macro ----
245
// --------------------------------------------------------------
246

            
247
macro_rules! impl_int_conversions {
248
    ($container:ident) => {
249
        impl From<$container<Int>> for $container<IntVal> {
250
10204859
            fn from(val: $container<Int>) -> Self {
251
10204859
                val.func_map(IntVal::from)
252
10204859
            }
253
        }
254

            
255
        impl TryFrom<$container<IntVal>> for $container<Int> {
256
            type Error = DomainOpError;
257

            
258
944612
            fn try_from(val: $container<IntVal>) -> Result<Self, Self::Error> {
259
944612
                val.try_func_map(IntVal::try_into)
260
944612
            }
261
        }
262

            
263
        impl $container<IntVal> {
264
            // All inner types are either i64 or pointers so cloning should be relatively cheap;
265
            // so, for ergonomics, we pretend that `resolve` methods take a reference :)
266
52534893
            pub fn resolve(&self) -> Result<$container<Int>, DomainOpError> {
267
102673048
                self.clone().try_func_map(|x| IntVal::resolve(&x))
268
52534893
            }
269
        }
270
    };
271
}
272

            
273
macro_rules! impl_int_conversions_for {
274
    ($($container:ident),+ $(,)?) => {
275
        $(impl_int_conversions!($container);)+
276
    };
277
}
278

            
279
// To add a new type in the future:
280
// 1. Add #[derive(FuncMap, TryFuncMap)] to the container type or impl the traits yourself
281
// 2. Add the container type to the list below
282
impl_int_conversions_for!(
283
    Range,
284
    SetAttr,
285
    MSetAttr,
286
    FuncAttr,
287
    SequenceAttr,
288
    PartitionAttr,
289
    PermutationAttr,
290
    RelAttr
291
);