Skip to main content

conjure_cp_core/ast/
reference.rs

1use crate::ast::serde::{AsId, HasId};
2use crate::representation::types::ReprGetOrInitResult;
3use crate::representation::{
4    ReferenceReprError, ReprError, ReprRule, ReprRulePtr, ReprSelectError, ReprStateStored,
5};
6use crate::{ast::DeclarationPtr, bug};
7use derivative::Derivative;
8use parking_lot::{MappedRwLockReadGuard, RwLockReadGuard};
9use serde::{Deserialize, Serialize};
10use serde_with::serde_as;
11use std::fmt::{Display, Formatter};
12use uniplate::Uniplate;
13
14use super::{
15    Atom, DeclarationKind, DomainPtr, Expression, GroundDomain, Literal, Metadata, Moo, Name,
16    categories::{Category, CategoryOf},
17    domains::HasDomain,
18};
19
20/// A reference to a declaration (variable, parameter, etc.)
21///
22/// This is a thin wrapper around [`DeclarationPtr`] with two main purposes:
23/// 1. Encapsulate the serde pragmas (e.g., serializing as IDs rather than full objects)
24/// 2. Enable type-directed traversals of references via uniplate
25#[serde_as]
26#[derive(
27    Clone, Debug, PartialEq, Eq, PartialOrd, Ord, Serialize, Deserialize, Uniplate, Derivative,
28)]
29#[derivative(Hash)]
30#[uniplate()]
31#[biplate(to=DeclarationPtr)]
32#[biplate(to=Name)]
33pub struct Reference {
34    #[serde_as(as = "AsId")]
35    pub ptr: DeclarationPtr,
36    pub repr: Option<ReprRulePtr>,
37}
38
39impl Reference {
40    pub fn new(ptr: DeclarationPtr) -> Self {
41        Reference { ptr, repr: None }
42    }
43
44    pub fn ptr(&self) -> &DeclarationPtr {
45        &self.ptr
46    }
47
48    pub fn into_ptr(self) -> DeclarationPtr {
49        self.ptr
50    }
51
52    pub fn name(&self) -> MappedRwLockReadGuard<'_, Name> {
53        self.ptr.name()
54    }
55
56    pub fn id(&self) -> crate::ast::serde::ObjId {
57        self.ptr.id()
58    }
59
60    pub fn domain(&self) -> Option<DomainPtr> {
61        self.ptr.domain()
62    }
63
64    pub fn resolved_domain(&self) -> Option<Moo<GroundDomain>> {
65        self.domain()?.resolve().ok()
66    }
67
68    /// Select the given representation for this reference, if it is currently unrepresented
69    /// and the representation exists for the underlying variable.
70    ///
71    /// # Errors
72    /// - [ReprSelectError::AlreadySelected] if a different representation is already selected for this reference
73    /// - [ReprSelectError::DoesNotExist] if the representation does not exist for this variable
74    ///
75    /// # Returns
76    /// State of the initialised representation
77    pub fn select_repr<R: ReprRule + ?Sized>(
78        &mut self,
79    ) -> Result<MappedRwLockReadGuard<'_, R::DeclLevel>, ReprSelectError> {
80        let _ = self.select_repr_via(R::STORED);
81        Ok(self.repr_state_as_unchecked::<R>())
82    }
83
84    /// Same as [Reference::select_repr], but type-erased
85    pub fn select_repr_via(
86        &mut self,
87        rule: ReprRulePtr,
88    ) -> Result<MappedRwLockReadGuard<'_, dyn ReprStateStored>, ReprSelectError> {
89        if let Some(repr) = self.repr
90            && repr != rule
91        {
92            return Err(ReprSelectError::AlreadySelected(repr));
93        }
94        if !self.ptr.reprs().has_repr(rule) {
95            return Err(ReprSelectError::DoesNotExist(self.ptr.clone(), rule.name()));
96        }
97        self.repr = Some(rule);
98        Ok(self.repr_state_unchecked())
99    }
100
101    /// Same as [Reference::select_or_init_repr], but type-erased
102    pub fn select_or_init_repr_via(
103        &mut self,
104        rule: ReprRulePtr,
105    ) -> ReprGetOrInitResult<'_, dyn ReprStateStored, ReferenceReprError> {
106        if let Some(repr) = self.repr
107            && repr != rule
108        {
109            return Err(ReprSelectError::AlreadySelected(repr).into());
110        }
111        self.update_or_init_repr_via(rule).map_err(Into::into)
112    }
113
114    /// Select the given representation for this reference, initialising it if necessary.
115    /// Will fail if a different representation is already selected.
116    ///
117    /// # Errors
118    /// - [ReprSelectError] if a different representation is already selected for this reference
119    /// - [ReprInitError] | [ReprInstantiateError] if the representation could not be initialised
120    ///
121    /// # Returns
122    ///
123    /// `(state, symbols, constraints)`
124    /// where:
125    /// - `state` is an instance of the given representation
126    /// - `symbols` are new variables created by the representation
127    /// - `constraints` are new top-level constraints created by the representation
128    pub fn select_or_init_repr<R: ReprRule + ?Sized>(
129        &mut self,
130    ) -> ReprGetOrInitResult<'_, R::DeclLevel, ReferenceReprError> {
131        let (_, symbols, constraints) = self.select_or_init_repr_via(R::STORED)?;
132        let state = self.repr_state_as_unchecked::<R>();
133        Ok((state, symbols, constraints))
134    }
135
136    /// Select the given representation for this reference, initialising it if necessary.
137    /// Will overwrite the existing selection.
138    ///
139    /// # Errors
140    /// - [ReprInitError] | [ReprInstantiateError] if the representation could not be initialised
141    ///
142    /// # Returns
143    ///
144    /// `(state, symbols, constraints)`
145    /// where:
146    /// - `state` is an instance of the given representation
147    /// - `symbols` are new variables created by the representation
148    /// - `constraints` are new top-level constraints created by the representation
149    pub fn update_or_init_repr<R: ReprRule + ?Sized>(
150        &mut self,
151    ) -> ReprGetOrInitResult<'_, R::DeclLevel, ReprError> {
152        let (_, symbols, constraints) = self.update_or_init_repr_via(R::STORED)?;
153        let state = self.repr_state_as_unchecked::<R>();
154        Ok((state, symbols, constraints))
155    }
156
157    /// Same as [Reference::update_or_init_repr], but type-erased
158    pub fn update_or_init_repr_via(
159        &mut self,
160        rule: ReprRulePtr,
161    ) -> ReprGetOrInitResult<'_, dyn ReprStateStored, ReprError> {
162        let (symbols, constraints) = rule.init_for_if_not_exists(&mut self.ptr)?;
163        self.repr = Some(rule);
164        let state = self.repr_state_unchecked();
165        Ok((state, symbols, constraints))
166    }
167
168    /// If this reference has a representation selected, return `(rule, state)`
169    /// where
170    /// - `rule` is a pointer to the representation rule
171    /// - `state` is an instance of that representation
172    pub fn get_repr(
173        &self,
174    ) -> Option<(ReprRulePtr, MappedRwLockReadGuard<'_, dyn ReprStateStored>)> {
175        let rule = self.repr?;
176        Some((rule, self.repr_state_unchecked()))
177    }
178
179    /// If this reference has this specific representation selected, get its state as a concrete type
180    pub fn get_repr_as<R: ReprRule + ?Sized>(
181        &self,
182    ) -> Option<MappedRwLockReadGuard<'_, R::DeclLevel>> {
183        if let Some(rule) = self.repr
184            && rule == R::STORED
185        {
186            return Some(self.repr_state_as_unchecked::<R>());
187        }
188        None
189    }
190
191    /// If this reference has a representation selected, return its state, otherwise crash
192    fn repr_state_unchecked(&self) -> MappedRwLockReadGuard<'_, dyn ReprStateStored> {
193        let rule = self
194            .repr
195            .unwrap_or_else(|| bug!("`{}` had no representation", self.name()));
196        RwLockReadGuard::map(self.ptr.reprs(), |reprs| {
197            reprs.get_by_rule(rule).unwrap_or_else(|| {
198                bug!(
199                    "Representation '{}' was selected for '{}' but its state was not stored!",
200                    rule.name(),
201                    self.name()
202                )
203            })
204        })
205    }
206
207    /// If this reference has this specific representation selected, return its state
208    /// as a concrete type, otherwise crash
209    fn repr_state_as_unchecked<R: ReprRule + ?Sized>(
210        &self,
211    ) -> MappedRwLockReadGuard<'_, R::DeclLevel> {
212        R::get_for(&self.ptr).unwrap_or_else(|| {
213            bug!(
214                "`{}` did not have the expected representation `{}`",
215                self.name(),
216                R::NAME
217            )
218        })
219    }
220
221    /// Returns the expression behind a value-letting reference, if this is one.
222    ///
223    /// Prefer [`Reference::with_resolved_expression`] when the expression only needs to be
224    /// inspected. Returning an owned value here necessarily clones collection-bearing
225    /// expressions such as matrix literals.
226    pub fn resolve_expression(&self) -> Option<Expression> {
227        if let Some(expr) = self.ptr().as_value_letting() {
228            return Some(expr.clone());
229        }
230
231        let generator = {
232            let kind = self.ptr.kind();
233            if let DeclarationKind::Quantified(inner) = &*kind {
234                inner.generator().cloned()
235            } else {
236                None
237            }
238        };
239
240        if let Some(generator) = generator
241            && let Some(expr) = generator.as_value_letting()
242        {
243            return Some(expr.clone());
244        }
245
246        None
247    }
248
249    /// Calls `inspect` with the expression behind a value-letting reference.
250    ///
251    /// The declaration is kept read-locked for the duration of `inspect`, allowing callers to ask
252    /// questions about large value lettings without cloning them. Quantified declarations may
253    /// delegate their value to an expression generator; those are handled here as well.
254    pub fn with_resolved_expression<T>(&self, inspect: impl FnOnce(&Expression) -> T) -> Option<T> {
255        if let Some(expr) = self.ptr().as_value_letting() {
256            return Some(inspect(&expr));
257        }
258
259        let generator = {
260            let kind = self.ptr.kind();
261            if let DeclarationKind::Quantified(inner) = &*kind {
262                inner.generator().cloned()
263            } else {
264                None
265            }
266        }?;
267        let expr = generator.as_value_letting()?;
268        Some(inspect(&expr))
269    }
270
271    /// Evaluates this reference to a literal if it resolves to a constant.
272    pub fn resolve_constant(&self) -> Option<Literal> {
273        self.with_resolved_expression(super::eval::eval_constant)
274            .flatten()
275    }
276
277    /// Resolves this reference to an atomic expression, if possible.
278    pub fn resolve_atomic(&self) -> Option<Atom> {
279        self.with_resolved_expression(|expr| match expr {
280            Expression::Atomic(_, atom) => Some(atom.clone()),
281            _ => None,
282        })
283        .flatten()
284    }
285}
286
287impl From<Reference> for Expression {
288    fn from(value: Reference) -> Self {
289        Expression::Atomic(Metadata::new(), value.into())
290    }
291}
292
293impl From<Reference> for Moo<Expression> {
294    fn from(value: Reference) -> Self {
295        Moo::new(value.into())
296    }
297}
298
299impl From<DeclarationPtr> for Reference {
300    fn from(ptr: DeclarationPtr) -> Self {
301        Reference::new(ptr)
302    }
303}
304
305impl CategoryOf for Reference {
306    fn category_of(&self) -> Category {
307        self.ptr.category_of()
308    }
309}
310
311impl HasDomain for Reference {
312    fn domain_of(&self) -> DomainPtr {
313        self.ptr
314            .domain()
315            .or_else(|| self.resolve_constant().map(|literal| literal.domain_of()))
316            .unwrap_or_else(|| {
317                bug!(
318                    "reference ({name}) should have a domain",
319                    name = self.ptr.name()
320                )
321            })
322    }
323}
324
325impl Display for Reference {
326    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
327        self.ptr.name().fmt(f)
328    }
329}