1
use crate::ast::{
2
    DeclarationKind, DomainOpError, Expression, FuncAttr, Literal, Metadata, Moo, PartitionAttr,
3
    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
6961672
    fn from(v: T) -> Self {
39
6961672
        IntVal::Const(v.into())
40
6961672
    }
41
}
42

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

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

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

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

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

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

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

            
88
2172
    pub fn new_ref(re: &Reference) -> Result<IntVal, DomainOpError> {
89
2172
        match re.ptr.kind().deref() {
90
1880
            DeclarationKind::ValueLetting(expr, _)
91
1880
            | DeclarationKind::TemporaryValueLetting(expr) => match expr.return_type() {
92
1880
                ReturnType::Int => Ok(IntVal::Reference(re.clone())),
93
                _ => Err(DomainOpError::WrongType),
94
            },
95
212
            DeclarationKind::Given(dom) => match dom.return_type() {
96
212
                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
80
            DeclarationKind::Quantified(inner) => match inner.domain().return_type() {
104
80
                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) => match var.return_type() {
115
                ReturnType::Int => Ok(IntVal::Reference(re.clone())),
116
                _ => Err(DomainOpError::WrongType),
117
            },
118
            DeclarationKind::DomainLetting(_) => Err(DomainOpError::WrongType),
119
        }
120
2172
    }
121

            
122
1216
    pub fn new_expr(value: Moo<Expression>) -> Result<IntVal, DomainOpError> {
123
1216
        if value.return_type() != ReturnType::Int {
124
            return Err(DomainOpError::WrongType);
125
1216
        }
126
1216
        Ok(IntVal::Expr(value))
127
1216
    }
128

            
129
8143832
    pub fn resolve(&self) -> Result<Int, DomainOpError> {
130
8143832
        match self {
131
7205896
            IntVal::Const(value) => (*value).try_into().map_err(|_| DomainOpError::OutOfBounds),
132
71560
            IntVal::Expr(expr) => eval_expr_to_int(expr).ok_or(DomainOpError::NotGround),
133
866376
            IntVal::Reference(re) => match re.ptr.kind().deref() {
134
859980
                DeclarationKind::ValueLetting(expr, _)
135
4800
                | DeclarationKind::TemporaryValueLetting(expr) => {
136
864780
                    eval_expr_to_int(expr).ok_or(DomainOpError::NotGround)
137
                }
138
                // If this is an int given we will be able to resolve it eventually, but not yet
139
36
                DeclarationKind::Given(_) => Err(DomainOpError::NotGround),
140
200
                DeclarationKind::Quantified(inner) => {
141
200
                    if let Some(generator) = inner.generator()
142
                        && let Some(expr) = generator.as_value_letting()
143
                    {
144
                        eval_expr_to_int(&expr).ok_or(DomainOpError::NotGround)
145
                    } else {
146
200
                        Err(DomainOpError::NotGround)
147
                    }
148
                }
149
                // Decision variables inside domains are unresolved until solving.
150
1360
                DeclarationKind::Find(_) => Err(DomainOpError::NotGround),
151
                DeclarationKind::DomainLetting(_) | DeclarationKind::QuantifiedExpr(_) => bug!(
152
                    "Expected integer expression, given, or letting inside int domain; Got: {re}"
153
                ),
154
            },
155
        }
156
8143832
    }
157

            
158
    pub fn try_add<T: Into<Expression>>(self, rhs: T) -> Result<IntVal, DomainOpError> {
159
        let sum = Expression::Sum(
160
            Metadata::new(),
161
            Moo::new(matrix_expr!(self.try_into()?, rhs.into())),
162
        );
163
        Ok(IntVal::Expr(Moo::new(sum)))
164
    }
165

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

            
176
// ------------------------------------------
177
// ------ Expression-related miscellanea ----
178
// ------------------------------------------
179

            
180
impl Range<IntVal> {
181
    /// Generates the expression to compute the size of this range
182
    pub fn len_expr(self) -> Result<Expression, DomainOpError> {
183
        match self {
184
            Range::Single(a) => Ok(a.try_into()?),
185
            Range::Bounded(a, b) => {
186
                let neg_b = Expression::Neg(Metadata::new(), b.try_into()?);
187
                let sum_matr = into_matrix_expr!(vec![a.try_into()?, neg_b]);
188
                Ok(Expression::Sum(Metadata::new(), sum_matr.into()))
189
            }
190
            _ => Err(DomainOpError::Unbounded),
191
        }
192
    }
193

            
194
    /// Generates the expression to compute the size of a list of ranges
195
    pub fn len_expr_of(rngs: &[Range<IntVal>]) -> Result<Expression, DomainOpError> {
196
        let mut rng_sizes = Vec::with_capacity(rngs.len());
197
        for rng in rngs {
198
            rng_sizes.push(rng.clone().len_expr()?);
199
        }
200
        let rng_sizes = into_matrix_expr!(rng_sizes);
201
        Ok(Expression::Sum(Metadata::new(), rng_sizes.into()))
202
    }
203
}
204

            
205
936340
fn eval_expr_to_int(expr: &Expression) -> Option<Int> {
206
936340
    match eval_constant(expr)? {
207
929540
        Literal::Int(v) => Some(v),
208
        _ => bug!("Expected integer expression, got: {expr}"),
209
    }
210
936340
}
211

            
212
impl TryFrom<IntVal> for Moo<Expression> {
213
    type Error = DomainOpError;
214

            
215
1360
    fn try_from(value: IntVal) -> Result<Self, Self::Error> {
216
1360
        match value {
217
320
            IntVal::Const(val) => {
218
320
                let val: Int = val.try_into().map_err(|_| DomainOpError::OutOfBounds)?;
219
320
                Ok(Moo::new(val.into()))
220
            }
221
360
            IntVal::Reference(re) => Ok(Moo::new(re.into())),
222
680
            IntVal::Expr(expr) => Ok(expr),
223
        }
224
1360
    }
225
}
226

            
227
impl TryFrom<IntVal> for Expression {
228
    type Error = DomainOpError;
229
1360
    fn try_from(value: IntVal) -> Result<Self, Self::Error> {
230
1360
        Ok(Moo::unwrap_or_clone(value.try_into()?))
231
1360
    }
232
}
233

            
234
// --------------------------------------------------------------
235
// ------ Derive into / resolve for container types by macro ----
236
// --------------------------------------------------------------
237

            
238
macro_rules! impl_int_conversions {
239
    ($container:ident) => {
240
        impl From<$container<Int>> for $container<IntVal> {
241
3533844
            fn from(val: $container<Int>) -> Self {
242
3533844
                val.func_map(IntVal::from)
243
3533844
            }
244
        }
245

            
246
        impl TryFrom<$container<IntVal>> for $container<Int> {
247
            type Error = DomainOpError;
248

            
249
73262
            fn try_from(val: $container<IntVal>) -> Result<Self, Self::Error> {
250
73262
                val.try_func_map(IntVal::try_into)
251
73262
            }
252
        }
253

            
254
        impl $container<IntVal> {
255
            // All inner types are either i64 or pointers so cloning should be relatively cheap;
256
            // so, for ergonomics, we pretend that `resolve` methods take a reference :)
257
926672
            pub fn resolve(&self) -> Result<$container<Int>, DomainOpError> {
258
1825640
                self.clone().try_func_map(|x| IntVal::resolve(&x))
259
926672
            }
260
        }
261
    };
262
}
263

            
264
macro_rules! impl_int_conversions_for {
265
    ($($container:ident),+ $(,)?) => {
266
        $(impl_int_conversions!($container);)+
267
    };
268
}
269

            
270
// To add a new type in the future:
271
// 1. Add #[derive(FuncMap, TryFuncMap)] to the container type or impl the traits yourself
272
// 2. Add the container type to the list below
273
impl_int_conversions_for!(
274
    Range,
275
    SetAttr,
276
    MSetAttr,
277
    FuncAttr,
278
    SequenceAttr,
279
    PartitionAttr,
280
    RelAttr
281
);