1
use std::fmt;
2
use std::ops::Deref;
3

            
4
use ahash::AHashSet;
5
use delegate::delegate;
6

            
7
use merc_aterm::ATerm;
8
use merc_aterm::ATermArgs;
9
use merc_aterm::ATermIndex;
10
use merc_aterm::ATermRef;
11
use merc_aterm::ATermString;
12
use merc_aterm::Markable;
13
use merc_aterm::Symb;
14
use merc_aterm::SymbolRef;
15
use merc_aterm::Term;
16
use merc_aterm::TermBuilder;
17
use merc_aterm::TermIterator;
18
use merc_aterm::Transmutable;
19
use merc_aterm::Yield;
20
use merc_aterm::storage::Marker;
21
use merc_aterm::storage::THREAD_TERM_POOL;
22
use merc_macros::merc_derive_terms;
23
use merc_macros::merc_ignore;
24
use merc_macros::merc_term;
25

            
26
use crate::DATA_SYMBOLS;
27
use crate::SortExpression;
28
use crate::SortExpressionRef;
29
use crate::is_data_application;
30
use crate::is_data_expression;
31
use crate::is_data_function_symbol;
32
use crate::is_data_machine_number;
33
use crate::is_data_variable;
34

            
35
// This module is only used internally to run the proc macro.
36
#[merc_derive_terms]
37
mod inner {
38

            
39
    use std::iter;
40

            
41
    use merc_aterm::ATermStringRef;
42
    use merc_utilities::MercError;
43

            
44
    use super::*;
45

            
46
    /// A data expression is an [merc_aterm::ATerm] with additional structure.
47
    ///
48
    /// # Details
49
    ///
50
    /// A data expression can be any of:
51
    ///     - a variable
52
    ///     - a function symbol, i.e. f without arguments.
53
    ///     - a term applied to a number of arguments, i.e., t_0(t1, ..., tn).
54
    ///     - an abstraction lambda x: Sort . e, or forall and exists.
55
    ///     - machine number, a value [0, ..., 2^64-1].
56
    ///
57
    /// Not supported:
58
    ///     - a where clause "e where [x := f, ...]"
59
    ///     - set enumeration
60
    ///     - bag enumeration
61
    ///
62
    #[merc_term(is_data_expression)]
63
    pub struct DataExpression {
64
        term: ATerm,
65
    }
66

            
67
    impl DataExpression {
68
        /// Returns the head symbol a data expression
69
        ///     - function symbol                  f -> f
70
        ///     - application       f(t_0, ..., t_n) -> f
71
161035746
        pub fn data_function_symbol(&self) -> DataFunctionSymbolRef<'_> {
72
161035746
            if is_data_application(&self.term) {
73
132238305
                self.term.arg(0).into()
74
28797441
            } else if is_data_function_symbol(&self.term) {
75
28797441
                self.term.copy().into()
76
            } else {
77
                // This can only happen if the term is an incorrect data expression.
78
                panic!("data_function_symbol not implemented for {self}");
79
            }
80
161035746
        }
81

            
82
        /// Returns the arguments of a data expression
83
        ///     - function symbol                  f -> []
84
        ///     - application       f(t_0, ..., t_n) -> [t_0, ..., t_n]
85
        #[merc_ignore]
86
607369
        pub fn data_arguments(&self) -> impl ExactSizeIterator<Item = DataExpressionRef<'_>> + use<'_> {
87
607369
            let mut result = self.term.arguments();
88
607369
            if is_data_application(&self.term) {
89
527937
                result.next();
90
527937
            } else if is_data_function_symbol(&self.term) || is_data_variable(&self.term) {
91
79432
                result.next();
92
79432
                result.next();
93
79432
            } else {
94
                // This can only happen if the term is an incorrect data expression.
95
                panic!("data_arguments not implemented for {self}");
96
            }
97

            
98
635397
            result.map(|t| t.into())
99
607369
        }
100

            
101
        /// Creates a closed [DataExpression] from a string, i.e., has no free variables.
102
        #[merc_ignore]
103
2513
        pub fn from_string(text: &str) -> Result<DataExpression, MercError> {
104
2513
            Ok(to_untyped_data_expression(ATerm::from_string(text)?, None))
105
2513
        }
106

            
107
        /// Creates a [DataExpression] from a string with free untyped variables indicated by the set of names.
108
        #[merc_ignore]
109
24
        pub fn from_string_untyped(text: &str, variables: &AHashSet<String>) -> Result<DataExpression, MercError> {
110
24
            Ok(to_untyped_data_expression(ATerm::from_string(text)?, Some(variables)))
111
24
        }
112

            
113
        /// Returns the ith argument of a data application.
114
        #[merc_ignore]
115
2
        pub fn data_arg(&self, index: usize) -> DataExpressionRef<'_> {
116
2
            debug_assert!(is_data_application(self), "Term {self:?} is not a data application");
117
2
            debug_assert!(
118
2
                index + 1 < self.get_head_symbol().arity(),
119
                "data_arg({index}) is not defined for term {self:?}"
120
            );
121

            
122
2
            self.term.arg(index + 1).into()
123
2
        }
124

            
125
        /// Returns the arguments of a data expression
126
        ///     - function symbol                  f -> []
127
        ///     - application       f(t_0, ..., t_n) -> [t_0, ..., t_n]
128
        pub fn data_sort(&self) -> SortExpression {
129
            if is_data_function_symbol(&self.term) {
130
                DataFunctionSymbolRef::from(self.term.copy()).sort().protect()
131
            } else if is_data_variable(&self.term) {
132
                DataVariableRef::from(self.term.copy()).sort().protect()
133
            } else {
134
                panic!("data_sort not implemented for {self}");
135
            }
136
        }
137
    }
138

            
139
    impl fmt::Display for DataExpression {
140
9065
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
141
9065
            if is_data_function_symbol(&self.term) {
142
5001
                write!(f, "{}", DataFunctionSymbolRef::from(self.term.copy()))
143
4064
            } else if is_data_application(&self.term) {
144
3464
                write!(f, "{}", DataApplicationRef::from(self.term.copy()))
145
600
            } else if is_data_variable(&self.term) {
146
                write!(f, "{}", DataVariableRef::from(self.term.copy()))
147
600
            } else if is_data_machine_number(&self.term) {
148
600
                write!(f, "{}", MachineNumberRef::from(self.term.copy()))
149
            } else {
150
                write!(f, "{}", self.term)
151
            }
152
9065
        }
153
    }
154

            
155
    #[merc_term(is_data_function_symbol)]
156
    pub struct DataFunctionSymbol {
157
        term: ATerm,
158
    }
159

            
160
    impl DataFunctionSymbol {
161
        #[merc_ignore]
162
69091
        pub fn new<N>(name: N) -> DataFunctionSymbol
163
69091
        where
164
69091
            N: Into<ATermString> + AsRef<str>,
165
        {
166
69091
            DATA_SYMBOLS.with_borrow(|ds| DataFunctionSymbol {
167
69091
                term: ATerm::with_args(
168
69091
                    ds.data_function_symbol.deref(),
169
69091
                    &[Into::<ATerm>::into(name.into()), SortExpression::unknown_sort().into()],
170
69091
                )
171
69091
                .protect(),
172
69091
            })
173
69091
        }
174

            
175
        /// Returns the name of the function symbol
176
8469
        pub fn name(&self) -> ATermStringRef<'_> {
177
8469
            ATermStringRef::from(self.term.arg(0))
178
8469
        }
179

            
180
        /// Returns the sort of the function symbol.
181
        pub fn sort(&self) -> SortExpressionRef<'_> {
182
            self.term.arg(1).into()
183
        }
184

            
185
        /// Returns the internal operation id (a unique number) for the data::function_symbol.
186
101925488
        pub fn operation_id(&self) -> usize {
187
101925488
            self.term.index()
188
101925488
        }
189
    }
190

            
191
    impl fmt::Display for DataFunctionSymbol {
192
8467
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
193
8467
            write!(f, "{}", self.name())
194
8467
        }
195
    }
196

            
197
    #[merc_term(is_data_variable)]
198
    pub struct DataVariable {
199
        term: ATerm,
200
    }
201

            
202
    impl DataVariable {
203
        /// Create a new untyped variable with the given name.
204
        #[merc_ignore]
205
20656
        pub fn new<N: Into<ATermString>>(name: N) -> DataVariable {
206
20656
            DATA_SYMBOLS.with_borrow(|ds| {
207
                // TODO: Storing terms temporarily is not optimal.
208
20656
                let t = name.into();
209
20656
                let args: &[ATerm] = &[t.into(), SortExpression::unknown_sort().into()];
210

            
211
20656
                DataVariable {
212
20656
                    term: ATerm::with_args(ds.data_variable.deref(), args).protect(),
213
20656
                }
214
20656
            })
215
20656
        }
216

            
217
        /// Create a variable with the given sort and name.
218
        pub fn with_sort<N: Into<ATermString>>(name: N, sort: SortExpressionRef<'_>) -> DataVariable {
219
            DATA_SYMBOLS.with_borrow(|ds| {
220
                // TODO: Storing terms temporarily is not optimal.
221
                let t = name.into();
222
                let args: &[ATermRef<'_>] = &[t.copy().into(), sort.into()];
223

            
224
                DataVariable {
225
                    term: ATerm::with_args(ds.data_variable.deref(), args).protect(),
226
                }
227
            })
228
        }
229

            
230
        /// Returns the name of the variable.
231
4480
        pub fn name(&self) -> &str {
232
            // We only change the lifetime, but that is fine since it is derived from the current term.
233
4480
            self.term.arg(0).get_head_symbol().name()
234
4480
        }
235

            
236
        /// Returns the sort of the variable.
237
        pub fn sort(&self) -> SortExpressionRef<'_> {
238
            self.term.arg(1).into()
239
        }
240
    }
241

            
242
    impl fmt::Display for DataVariable {
243
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
244
            write!(f, "{}", self.name())
245
        }
246
    }
247

            
248
    #[merc_term(is_data_application)]
249
    pub struct DataApplication {
250
        term: ATerm,
251
    }
252

            
253
    impl DataApplication {
254
        /// Create a new data application with the given head and arguments.
255
        #[merc_ignore]
256
1090
        pub fn with_args<'a, 'b, H: Term<'a, 'b>, T: Term<'a, 'b>>(head: &'b H, arguments: &'b [T]) -> DataApplication {
257
1090
            DATA_SYMBOLS.with_borrow_mut(|ds| {
258
1090
                let symbol = ds.get_data_application_symbol(arguments.len() + 1).copy();
259

            
260
1973
                let args = iter::once(head.copy()).chain(arguments.iter().map(|t| t.copy()));
261
1090
                let term = ATerm::with_iter(&symbol, args);
262

            
263
1090
                DataApplication { term }
264
1090
            })
265
1090
        }
266

            
267
        /// Create a new data application with the given head and arguments.
268
        ///
269
        /// arity must be equal to the number of arguments + 1.
270
        #[merc_ignore]
271
4325067
        pub fn with_iter<'a, 'b, 'c, 'd, T, H, I>(head: &'b H, arity: usize, arguments: I) -> DataApplication
272
4325067
        where
273
4325067
            I: Iterator<Item = T>,
274
4325067
            T: Term<'c, 'd>,
275
4325067
            H: Term<'a, 'b>,
276
        {
277
4325067
            DATA_SYMBOLS.with_borrow_mut(|ds| {
278
4325067
                let symbol = ds.get_data_application_symbol(arity + 1).copy();
279

            
280
4325067
                let term = ATerm::with_iter_head(&symbol, head, arguments);
281

            
282
4325067
                DataApplication { term }
283
4325067
            })
284
4325067
        }
285

            
286
        /// Returns the head symbol a data application
287
3465
        pub fn data_function_symbol(&self) -> DataFunctionSymbolRef<'_> {
288
3465
            self.term.arg(0).into()
289
3465
        }
290

            
291
        /// Returns the arguments of a data application
292
3466
        pub fn data_arguments(&self) -> ATermArgs<'_> {
293
3466
            let mut result = self.term.arguments();
294
3466
            result.next();
295
3466
            result
296
3466
        }
297

            
298
        /// Returns the ith argument of a data application.
299
        pub fn data_arg(&self, index: usize) -> DataExpressionRef<'_> {
300
            debug_assert!(
301
                index + 1 < self.get_head_symbol().arity(),
302
                "data_arg({index}) is not defined for term {self:?}"
303
            );
304

            
305
            self.term.arg(index + 1).into()
306
        }
307

            
308
        /// Returns the sort of a data application.
309
        pub fn sort(&self) -> SortExpressionRef<'_> {
310
            // We only change the lifetime, but that is fine since it is derived from the current term.
311
            SortExpressionRef::from(self.term.arg(0))
312
        }
313
    }
314

            
315
    impl fmt::Display for DataApplication {
316
3465
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
317
3465
            write!(f, "{}", self.data_function_symbol())?;
318

            
319
3465
            let mut first = true;
320
6329
            for arg in self.data_arguments() {
321
6329
                if !first {
322
2864
                    write!(f, ", ")?;
323
                } else {
324
3465
                    write!(f, "(")?;
325
                }
326

            
327
6329
                write!(f, "{}", DataExpressionRef::from(arg.copy()))?;
328
6329
                first = false;
329
            }
330

            
331
3465
            if !first {
332
3465
                write!(f, ")")?;
333
            }
334

            
335
3465
            Ok(())
336
3465
        }
337
    }
338

            
339
    #[merc_term(is_data_machine_number)]
340
    struct MachineNumber {
341
        pub term: ATerm,
342
    }
343

            
344
    impl MachineNumber {
345
        /// Obtain the underlying value of a machine number.
346
600
        pub fn value(&self) -> u64 {
347
600
            self.term
348
600
                .copy()
349
600
                .annotation()
350
600
                .expect("MachineNumber must have an integer annotation") as u64
351
600
        }
352
    }
353

            
354
    impl fmt::Display for MachineNumber {
355
600
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
356
600
            write!(f, "{}", self.value())
357
600
        }
358
    }
359

            
360
    /// Conversions to `DataExpression`
361
    #[merc_ignore]
362
    impl From<DataFunctionSymbol> for DataExpression {
363
7610480
        fn from(value: DataFunctionSymbol) -> Self {
364
7610480
            value.term.into()
365
7610480
        }
366
    }
367

            
368
    #[merc_ignore]
369
    impl From<DataApplication> for DataExpression {
370
34238657
        fn from(value: DataApplication) -> Self {
371
34238657
            value.term.into()
372
34238657
        }
373
    }
374

            
375
    #[merc_ignore]
376
    impl From<DataVariable> for DataExpression {
377
2400
        fn from(value: DataVariable) -> Self {
378
2400
            value.term.into()
379
2400
        }
380
    }
381

            
382
    #[merc_ignore]
383
    impl From<DataExpression> for DataFunctionSymbol {
384
20000
        fn from(value: DataExpression) -> Self {
385
20000
            value.term.into()
386
20000
        }
387
    }
388

            
389
    #[merc_ignore]
390
    impl From<DataExpression> for DataVariable {
391
230768
        fn from(value: DataExpression) -> Self {
392
230768
            value.term.into()
393
230768
        }
394
    }
395

            
396
    #[merc_ignore]
397
    impl<'a> From<DataExpressionRef<'a>> for DataVariableRef<'a> {
398
289368
        fn from(value: DataExpressionRef<'a>) -> Self {
399
289368
            value.term.into()
400
289368
        }
401
    }
402
}
403

            
404
pub use inner::*;
405

            
406
impl<'a> DataExpressionRef<'a> {
407
20148680
    pub fn data_arguments(&self) -> impl ExactSizeIterator<Item = DataExpressionRef<'a>> + use<'a> {
408
20148680
        let mut result = self.term.arguments();
409
20148680
        if is_data_application(&self.term) {
410
15492256
            result.next();
411
15492256
        } else if is_data_function_symbol(&self.term) || is_data_variable(&self.term) {
412
4656424
            result.next();
413
4656424
            result.next();
414
4656424
        } else {
415
            // This can only happen if the term is not a data expression.
416
            panic!("data_arguments not implemented for {self}");
417
        }
418

            
419
20148680
        result.map(|t| t.into())
420
20148680
    }
421

            
422
    /// Returns the ith argument of a data application.
423
130863913
    pub fn data_arg(&self, index: usize) -> DataExpressionRef<'a> {
424
130863913
        debug_assert!(is_data_application(self), "Term {self:?} is not a data application");
425
130863913
        debug_assert!(
426
130863913
            index + 1 < self.get_head_symbol().arity(),
427
            "data_arg({index}) is not defined for term {self:?}"
428
        );
429

            
430
130863913
        self.term.arg(index + 1).into()
431
130863913
    }
432
}
433

            
434
/// Converts an [ATerm] to an untyped data expression.
435
17721
pub fn to_untyped_data_expression(t: ATerm, variables: Option<&AHashSet<String>>) -> DataExpression {
436
17721
    let mut builder = TermBuilder::<ATerm, ATerm>::new();
437
17721
    THREAD_TERM_POOL.with_borrow(|tp| {
438
17721
        builder
439
17721
            .evaluate(
440
17721
                tp,
441
17721
                t,
442
84285
                |_tp, args, t| {
443
84285
                    if variables.is_some_and(|v| v.contains(t.get_head_symbol().name())) {
444
                        // Convert a constant variable, for example 'x', into an untyped variable.
445
15224
                        Ok(Yield::Term(DataVariable::new(t.get_head_symbol().name()).into()))
446
69061
                    } else if t.get_head_symbol().arity() == 0 {
447
22739
                        Ok(Yield::Term(DataFunctionSymbol::new(t.get_head_symbol().name()).into()))
448
                    } else {
449
                        // This is a function symbol applied to a number of arguments
450
46322
                        let head = DataFunctionSymbol::new(t.get_head_symbol().name());
451

            
452
66564
                        for arg in t.arguments() {
453
66564
                            args.push(arg.protect());
454
66564
                        }
455

            
456
46322
                        Ok(Yield::Construct(head.into()))
457
                    }
458
84285
                },
459
46322
                |_tp, input, args| {
460
46322
                    let arity = args.clone().count();
461
46322
                    Ok(DataApplication::with_iter(&input, arity, args).into())
462
46322
                },
463
            )
464
17721
            .unwrap()
465
17721
            .into()
466
17721
    })
467
17721
}
468

            
469
#[cfg(test)]
470
mod tests {
471
    use super::*;
472

            
473
    use merc_aterm::ATerm;
474

            
475
    #[test]
476
1
    fn test_print() {
477
1
        let _ = merc_utilities::test_logger();
478

            
479
1
        let a = DataFunctionSymbol::new("a");
480
1
        assert_eq!("a", format!("{}", a));
481

            
482
        // Check printing of data applications.
483
1
        let f = DataFunctionSymbol::new("f");
484
1
        let appl = DataApplication::with_args(&f, &[a]);
485
1
        assert_eq!("f(a)", format!("{}", appl));
486
1
    }
487

            
488
    #[test]
489
1
    fn test_recognizers() {
490
1
        let a = DataFunctionSymbol::new("a");
491
1
        let f = DataFunctionSymbol::new("f");
492
1
        let appl = DataApplication::with_args(&f, &[a]);
493

            
494
1
        let term: ATerm = appl.into();
495
1
        assert!(is_data_application(&term));
496
1
    }
497

            
498
    #[test]
499
1
    fn test_data_arguments() {
500
1
        let a = DataFunctionSymbol::new("a");
501
1
        let f = DataFunctionSymbol::new("f");
502
1
        let appl = DataApplication::with_args(&f, &[a]);
503

            
504
1
        assert_eq!(appl.data_arguments().count(), 1);
505

            
506
1
        let data_expr: DataExpression = appl.clone().into();
507

            
508
1
        assert_eq!(data_expr.data_arguments().count(), 1);
509
1
    }
510

            
511
    #[test]
512
1
    fn test_to_data_expression() {
513
1
        let expression = DataExpression::from_string("s(s(a, b), c)").unwrap();
514

            
515
1
        assert_eq!(expression.data_arg(0).data_function_symbol().name(), "s");
516
1
        assert_eq!(expression.data_arg(0).data_arg(0).data_function_symbol().name(), "a");
517
1
    }
518
}