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::ATermIntRef;
42
    use merc_aterm::ATermStringRef;
43
    use merc_utilities::MercError;
44

            
45
    use super::*;
46

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

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

            
83
        /// Returns the data sub-expressions of a data expression.
84
        ///     - function symbol                  f -> []
85
        ///     - variable                         x -> []
86
        ///     - machine number                   n -> []
87
        ///     - application       f(t_0, ..., t_n) -> [t_0, ..., t_n]
88
        ///
89
        /// Panics for binders and where clauses, which have structured sub-terms that do not
90
        /// map cleanly to a flat argument list.
91
        #[merc_ignore]
92
1594807
        pub fn data_arguments(&self) -> impl ExactSizeIterator<Item = DataExpressionRef<'_>> + use<'_> {
93
1594807
            let skip = data_argument_skip_count(&self.term)
94
1594807
                .unwrap_or_else(|| panic!("data_arguments is not defined for binders and where clauses: {self}"));
95
1678900
            self.term.arguments().skip(skip).map(|t| t.into())
96
1594807
        }
97

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

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

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

            
119
4
            self.term.arg(index + 1).into()
120
4
        }
121

            
122
        /// Returns the sort of a data expression.
123
        ///
124
        /// Only defined for function symbols and variables. Panics for applications (the result
125
        /// sort requires traversing the SortArrow chain), machine numbers, binders, and where
126
        /// clauses.
127
1
        pub fn data_sort(&self) -> SortExpression {
128
1
            if is_data_function_symbol(&self.term) {
129
                DataFunctionSymbolRef::from(self.term.copy()).sort().protect()
130
1
            } else if is_data_variable(&self.term) {
131
1
                DataVariableRef::from(self.term.copy()).sort().protect()
132
            } else {
133
                panic!("data_sort not implemented for {self}");
134
            }
135
1
        }
136
    }
137

            
138
    impl fmt::Display for DataExpression {
139
39125
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
140
39125
            if is_data_function_symbol(&self.term) {
141
14386
                write!(f, "{}", DataFunctionSymbolRef::from(self.term.copy()))
142
24739
            } else if is_data_application(&self.term) {
143
19866
                write!(f, "{}", DataApplicationRef::from(self.term.copy()))
144
4873
            } else if is_data_variable(&self.term) {
145
315
                write!(f, "{}", DataVariableRef::from(self.term.copy()))
146
4558
            } else if is_data_machine_number(&self.term) {
147
4558
                write!(f, "{}", MachineNumberRef::from(self.term.copy()))
148
            } else {
149
                write!(f, "{}", self.term)
150
            }
151
39125
        }
152
    }
153

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

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

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

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

            
184
        /// Returns the internal operation id (a unique number) for the data::function_symbol.
185
        ///
186
        /// This is the term's pool index, which is only a stable identifier for indexed `OpId`
187
        /// symbols; it is not meaningful for the `OpIdNoIndex` variant.
188
267555960
        pub fn operation_id(&self) -> usize {
189
267555960
            self.term.index()
190
267555960
        }
191
    }
192

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

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

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

            
213
50748
                DataVariable {
214
50748
                    term: ATerm::with_args(ds.data_variable.deref(), args).protect(),
215
50748
                }
216
50748
            })
217
50748
        }
218

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

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

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

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

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

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

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

            
262
182026
                let args = iter::once(head.copy()).chain(arguments.iter().map(|t| t.copy()));
263
125825
                let term = ATerm::with_iter(&symbol, args);
264

            
265
125825
                DataApplication { term }
266
125825
            })
267
125825
        }
268

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

            
282
12836231
                let term = ATerm::with_iter_head(&symbol, head, arguments);
283

            
284
12836231
                DataApplication { term }
285
12836231
            })
286
12836231
        }
287

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

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

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

            
307
            self.term.arg(index + 1).into()
308
        }
309
    }
310

            
311
    impl fmt::Display for DataApplication {
312
19867
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
313
19867
            write!(f, "{}", self.data_function_symbol())?;
314

            
315
19867
            let mut first = true;
316
29359
            for arg in self.data_arguments() {
317
29359
                if !first {
318
9492
                    write!(f, ", ")?;
319
                } else {
320
19867
                    write!(f, "(")?;
321
                }
322

            
323
29359
                write!(f, "{}", DataExpressionRef::from(arg.copy()))?;
324
29359
                first = false;
325
            }
326

            
327
19867
            if !first {
328
19867
                write!(f, ")")?;
329
            }
330

            
331
19867
            Ok(())
332
19867
        }
333
    }
334

            
335
    #[merc_term(is_data_machine_number)]
336
    struct MachineNumber {
337
        pub term: ATerm,
338
    }
339

            
340
    impl MachineNumber {
341
        /// Obtain the underlying value of a machine number.
342
        ///
343
        /// Assumes the term is an integer term, which is guaranteed by the constructor
344
        /// and [`is_data_machine_number`]. The cast reinterprets the stored `i64` bit
345
        /// pattern as `u64`, recovering values in `[0, 2^64-1]`.
346
4558
        pub(crate) fn value(&self) -> u64 {
347
4558
            Into::<ATermIntRef<'_>>::into(self.term.copy()).value() as u64
348
4558
        }
349
    }
350

            
351
    impl fmt::Display for MachineNumber {
352
4558
        fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
353
4558
            write!(f, "{}", self.value())
354
4558
        }
355
    }
356

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

            
365
    #[merc_ignore]
366
    impl From<DataApplication> for DataExpression {
367
89884033
        fn from(value: DataApplication) -> Self {
368
89884033
            value.term.into()
369
89884033
        }
370
    }
371

            
372
    #[merc_ignore]
373
    impl From<DataVariable> for DataExpression {
374
6301
        fn from(value: DataVariable) -> Self {
375
6301
            value.term.into()
376
6301
        }
377
    }
378

            
379
    #[merc_ignore]
380
    impl From<DataExpression> for DataFunctionSymbol {
381
52731
        fn from(value: DataExpression) -> Self {
382
52731
            value.term.into()
383
52731
        }
384
    }
385

            
386
    #[merc_ignore]
387
    impl From<DataExpression> for DataVariable {
388
606312
        fn from(value: DataExpression) -> Self {
389
606312
            value.term.into()
390
606312
        }
391
    }
392

            
393
    #[merc_ignore]
394
    impl<'a> From<DataExpressionRef<'a>> for DataVariableRef<'a> {
395
762048
        fn from(value: DataExpressionRef<'a>) -> Self {
396
762048
            value.term.into()
397
762048
        }
398
    }
399
}
400

            
401
pub use inner::*;
402

            
403
/// Returns the number of leading `ATerm` arguments that are *not* data sub-expressions and
404
/// must therefore be skipped by `data_arguments`, or `None` for binders/where clauses which
405
/// have no flat argument list.
406
///
407
///   - application `f(t_0, ..., t_n)` -> skip 1 (the head function symbol)
408
///   - function symbol / variable     -> skip 2 (name and sort)
409
///   - machine number                 -> skip 0 (int terms carry no `ATerm` children)
410
54488977
fn data_argument_skip_count<'a, 'b, T: Term<'a, 'b>>(term: &'b T) -> Option<usize> {
411
54488977
    if is_data_application(term) {
412
42055610
        Some(1)
413
12433367
    } else if is_data_function_symbol(term) || is_data_variable(term) {
414
12433366
        Some(2)
415
1
    } else if is_data_machine_number(term) {
416
1
        Some(0)
417
    } else {
418
        None
419
    }
420
54488977
}
421

            
422
impl<'a> DataExpressionRef<'a> {
423
52894170
    pub fn data_arguments(&self) -> impl ExactSizeIterator<Item = DataExpressionRef<'a>> + use<'a> {
424
52894170
        let skip = data_argument_skip_count(&self.term)
425
52894170
            .unwrap_or_else(|| panic!("data_arguments is not defined for binders and where clauses: {self}"));
426
52894170
        self.term.arguments().skip(skip).map(|t| t.into())
427
52894170
    }
428

            
429
    /// Returns the ith argument of a data application.
430
343529089
    pub fn data_arg(&self, index: usize) -> DataExpressionRef<'a> {
431
343529089
        debug_assert!(is_data_application(self), "Term {self:?} is not a data application");
432
343529089
        debug_assert!(
433
343529089
            index + 1 < self.get_head_symbol().arity(),
434
            "data_arg({index}) is not defined for term {self:?}"
435
        );
436

            
437
343529089
        self.term.arg(index + 1).into()
438
343529089
    }
439
}
440

            
441
/// Converts an [ATerm] to an untyped data expression.
442
49017
pub fn to_untyped_data_expression(t: ATerm, variables: Option<&AHashSet<String>>) -> DataExpression {
443
49017
    let mut builder = TermBuilder::<ATerm, ATerm>::new();
444
49017
    THREAD_TERM_POOL.with(|tp| {
445
49017
        builder
446
49017
            .evaluate(
447
49017
                tp,
448
49017
                t,
449
224584
                |_tp, args, t| {
450
224584
                    let name = t.get_head_symbol().name();
451
224584
                    if t.get_head_symbol().arity() == 0 {
452
102381
                        if variables.is_some_and(|v| v.contains(name)) {
453
                            // Convert a constant identifier, for example 'x', into an untyped variable.
454
40363
                            Ok(Yield::Term(DataVariable::new(name).into()))
455
                        } else {
456
62018
                            Ok(Yield::Term(DataFunctionSymbol::new(name).into()))
457
                        }
458
                    } else {
459
                        // This is a function symbol applied to a number of arguments. Variables are
460
                        // only recognised in nullary position, so an applied identifier keeps its
461
                        // arguments instead of being silently collapsed to a variable.
462
122203
                        let head = DataFunctionSymbol::new(name);
463

            
464
175567
                        for arg in t.arguments() {
465
175567
                            args.push(arg.protect());
466
175567
                        }
467

            
468
122203
                        Ok(Yield::Construct(head.into()))
469
                    }
470
224584
                },
471
122203
                |_tp, input, args| {
472
122203
                    let args: Vec<ATerm> = args.cloned().collect();
473
122203
                    Ok(DataApplication::with_args(&input, &args).into())
474
122203
                },
475
            )
476
49017
            .unwrap()
477
49017
            .into()
478
49017
    })
479
49017
}
480

            
481
#[cfg(test)]
482
mod tests {
483
    use ahash::AHashSet;
484
    use merc_aterm::ATerm;
485
    use merc_aterm::ATermInt;
486

            
487
    use crate::is_data_application;
488
    use crate::is_data_machine_number;
489
    use crate::is_data_variable;
490

            
491
    use super::DataApplication;
492
    use super::DataExpression;
493
    use super::DataFunctionSymbol;
494
    use super::DataVariable;
495

            
496
    #[test]
497
1
    fn test_print() {
498
1
        merc_utilities::test_logger();
499

            
500
1
        let a = DataFunctionSymbol::new("a");
501
1
        assert_eq!("a", format!("{}", a));
502

            
503
        // Check printing of data applications.
504
1
        let f = DataFunctionSymbol::new("f");
505
1
        let appl = DataApplication::with_args(&f, &[a]);
506
1
        assert_eq!("f(a)", format!("{}", appl));
507
1
    }
508

            
509
    #[test]
510
1
    fn test_recognizers() {
511
1
        let a = DataFunctionSymbol::new("a");
512
1
        let f = DataFunctionSymbol::new("f");
513
1
        let appl = DataApplication::with_args(&f, &[a]);
514

            
515
1
        let term: ATerm = appl.into();
516
1
        assert!(is_data_application(&term));
517
1
    }
518

            
519
    #[test]
520
1
    fn test_data_arguments() {
521
1
        let a = DataFunctionSymbol::new("a");
522
1
        let f = DataFunctionSymbol::new("f");
523
1
        let appl = DataApplication::with_args(&f, &[a]);
524

            
525
1
        assert_eq!(appl.data_arguments().count(), 1);
526

            
527
1
        let data_expr: DataExpression = appl.clone().into();
528

            
529
1
        assert_eq!(data_expr.data_arguments().count(), 1);
530
1
    }
531

            
532
    #[test]
533
1
    fn test_to_data_expression() {
534
1
        let expression = DataExpression::from_string("s(s(a, b), c)").unwrap();
535

            
536
1
        assert_eq!(expression.data_arg(0).data_function_symbol().name(), "s");
537
1
        assert_eq!(expression.data_arg(0).data_arg(0).data_function_symbol().name(), "a");
538
1
    }
539

            
540
    #[test]
541
1
    fn test_machine_number() {
542
1
        let term: ATerm = ATermInt::new(42).into();
543
1
        assert!(is_data_machine_number(&term));
544

            
545
1
        let expr: DataExpression = term.into();
546
1
        assert_eq!(format!("{expr}"), "42");
547
        // Machine numbers have no data sub-expressions.
548
1
        assert_eq!(expr.data_arguments().count(), 0);
549
1
    }
550

            
551
    #[test]
552
1
    fn test_variable_sort() {
553
1
        let var = DataVariable::new("x");
554
1
        assert_eq!(var.name(), "x");
555

            
556
1
        let expr: DataExpression = var.into();
557
1
        assert!(is_data_variable(&expr));
558
1
        assert_eq!(expr.data_sort().name(), "@no_value@");
559
1
        assert_eq!(expr.data_arguments().count(), 0);
560
1
    }
561

            
562
    #[test]
563
1
    fn test_from_string_untyped_variable() {
564
1
        let vars = AHashSet::from_iter(["x".to_string()]);
565
1
        let expr = DataExpression::from_string_untyped("f(x, a)", &vars).unwrap();
566

            
567
        // 'x' is recognised as a variable, 'a' stays a function symbol.
568
1
        assert!(is_data_variable(&expr.data_arg(0)));
569
1
        assert_eq!(expr.data_arg(1).data_function_symbol().name(), "a");
570
1
    }
571

            
572
    #[test]
573
1
    fn test_from_string_untyped_applied_identifier_keeps_args() {
574
        // 'x' is in the variable set but appears applied; it must stay an application rather than
575
        // collapsing to a variable and silently dropping its argument.
576
1
        let vars = AHashSet::from_iter(["x".to_string()]);
577
1
        let expr = DataExpression::from_string_untyped("x(a)", &vars).unwrap();
578

            
579
1
        assert!(is_data_application(&expr));
580
1
        assert_eq!(expr.data_function_symbol().name(), "x");
581
1
        assert_eq!(expr.data_arguments().count(), 1);
582
1
    }
583
}