1
use crate::ast::serde::{AsId, HasId};
2
use crate::{ast::DeclarationPtr, bug};
3
use derivative::Derivative;
4
use parking_lot::MappedRwLockReadGuard;
5
use serde::{Deserialize, Serialize};
6
use serde_with::serde_as;
7
use std::fmt::{Display, Formatter};
8
use uniplate::Uniplate;
9

            
10
use super::{
11
    Atom, DeclarationKind, DomainPtr, Expression, GroundDomain, Literal, Metadata, Moo, Name,
12
    categories::{Category, CategoryOf},
13
    domains::HasDomain,
14
};
15

            
16
/// A reference to a declaration (variable, parameter, etc.)
17
///
18
/// This is a thin wrapper around [`DeclarationPtr`] with two main purposes:
19
/// 1. Encapsulate the serde pragmas (e.g., serializing as IDs rather than full objects)
20
/// 2. Enable type-directed traversals of references via uniplate
21
#[serde_as]
22
#[derive(
23
    Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, Uniplate, Derivative,
24
)]
25
#[derivative(Hash)]
26
#[uniplate()]
27
#[biplate(to=DeclarationPtr)]
28
#[biplate(to=Name)]
29
pub struct Reference {
30
    #[serde_as(as = "AsId")]
31
    pub ptr: DeclarationPtr,
32
}
33

            
34
impl Reference {
35
4122219
    pub fn new(ptr: DeclarationPtr) -> Self {
36
4122219
        Reference { ptr }
37
4122219
    }
38

            
39
44973824
    pub fn ptr(&self) -> &DeclarationPtr {
40
44973824
        &self.ptr
41
44973824
    }
42

            
43
1040
    pub fn into_ptr(self) -> DeclarationPtr {
44
1040
        self.ptr
45
1040
    }
46

            
47
280813256
    pub fn name(&self) -> MappedRwLockReadGuard<'_, Name> {
48
280813256
        self.ptr.name()
49
280813256
    }
50

            
51
22524
    pub fn id(&self) -> crate::ast::serde::ObjId {
52
22524
        self.ptr.id()
53
22524
    }
54

            
55
8960184
    pub fn domain(&self) -> Option<DomainPtr> {
56
8960184
        self.ptr.domain()
57
8960184
    }
58

            
59
8880964
    pub fn resolved_domain(&self) -> Option<Moo<GroundDomain>> {
60
8880964
        self.domain()?.resolve().ok()
61
8880964
    }
62

            
63
    /// Returns the expression behind a value-letting reference, if this is one.
64
42048756
    pub fn resolve_expression(&self) -> Option<Expression> {
65
42048756
        if let Some(expr) = self.ptr().as_value_letting() {
66
2428500
            return Some(expr.clone());
67
39620256
        }
68

            
69
39620256
        let generator = {
70
39620256
            let kind = self.ptr.kind();
71
39620256
            if let DeclarationKind::Quantified(inner) = &*kind {
72
636112
                inner.generator().cloned()
73
            } else {
74
38984144
                None
75
            }
76
        };
77

            
78
39620256
        if let Some(generator) = generator
79
            && let Some(expr) = generator.as_value_letting()
80
        {
81
            return Some(expr.clone());
82
39620256
        }
83

            
84
39620256
        None
85
42048756
    }
86

            
87
    /// Evaluates this reference to a literal if it resolves to a constant.
88
23336888
    pub fn resolve_constant(&self) -> Option<Literal> {
89
23336888
        self.resolve_expression()
90
23336888
            .and_then(|expr| super::eval::eval_constant(&expr))
91
23336888
    }
92

            
93
    /// Resolves this reference to an atomic expression, if possible.
94
262528
    pub fn resolve_atomic(&self) -> Option<Atom> {
95
262528
        self.resolve_expression().and_then(|expr| match expr {
96
            Expression::Atomic(_, atom) => Some(atom),
97
            _ => None,
98
        })
99
262528
    }
100
}
101

            
102
impl From<Reference> for Expression {
103
532
    fn from(value: Reference) -> Self {
104
532
        Expression::Atomic(Metadata::new(), value.into())
105
532
    }
106
}
107

            
108
impl From<DeclarationPtr> for Reference {
109
    fn from(ptr: DeclarationPtr) -> Self {
110
        Reference::new(ptr)
111
    }
112
}
113

            
114
impl CategoryOf for Reference {
115
2838652
    fn category_of(&self) -> Category {
116
2838652
        self.ptr.category_of()
117
2838652
    }
118
}
119

            
120
impl HasDomain for Reference {
121
5085638
    fn domain_of(&self) -> DomainPtr {
122
5085638
        self.ptr
123
5085638
            .domain()
124
5085638
            .or_else(|| self.resolve_constant().map(|literal| literal.domain_of()))
125
5085638
            .unwrap_or_else(|| {
126
                bug!(
127
                    "reference ({name}) should have a domain",
128
                    name = self.ptr.name()
129
                )
130
            })
131
5085638
    }
132
}
133

            
134
impl Display for Reference {
135
79837124
    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
136
79837124
        self.ptr.name().fmt(f)
137
79837124
    }
138
}