toasty_core/stmt/
assignments.rs

1use crate::stmt::Statement;
2
3use super::{Expr, Projection};
4
5use indexmap::{Equivalent, IndexMap};
6use std::{hash::Hash, ops};
7
8#[derive(Clone, Debug, PartialEq)]
9pub struct Assignments {
10    /// Map from UpdateTarget field projection to assignment for that field. The
11    /// projection may reference an application-level model field or a lowered
12    /// table column. Supports both single-step (e.g., [0]) and multi-step
13    /// projections (e.g., [0, 1] for nested fields).
14    assignments: IndexMap<Projection, Assignment>,
15}
16
17#[derive(Debug, Clone, PartialEq)]
18pub struct Assignment {
19    /// Assignment operation
20    pub op: AssignmentOp,
21
22    /// Expression use for assignment
23    pub expr: Expr,
24}
25
26#[derive(Debug, Clone, Copy, PartialEq)]
27pub enum AssignmentOp {
28    /// Set a field, replacing the current value.
29    Set,
30
31    /// Insert one or more values into a set
32    Insert,
33
34    /// Remove one or more values from a set.
35    Remove,
36}
37
38impl Statement {
39    pub fn assignments(&self) -> Option<&Assignments> {
40        match self {
41            Statement::Update(update) => Some(&update.assignments),
42            _ => None,
43        }
44    }
45}
46
47impl Assignments {
48    pub fn with_capacity(capacity: usize) -> Self {
49        Self {
50            assignments: IndexMap::with_capacity(capacity),
51        }
52    }
53
54    pub fn is_empty(&self) -> bool {
55        self.assignments.is_empty()
56    }
57
58    pub fn len(&self) -> usize {
59        self.assignments.len()
60    }
61
62    pub fn contains<Q>(&self, key: &Q) -> bool
63    where
64        Q: ?Sized + Hash + Equivalent<Projection>,
65    {
66        self.assignments.contains_key(key)
67    }
68
69    pub fn get<Q>(&self, key: &Q) -> Option<&Assignment>
70    where
71        Q: ?Sized + Hash + Equivalent<Projection>,
72    {
73        self.assignments.get(key)
74    }
75
76    pub fn get_mut<Q>(&mut self, key: &Q) -> Option<&mut Assignment>
77    where
78        Q: ?Sized + Hash + Equivalent<Projection>,
79    {
80        self.assignments.get_mut(key)
81    }
82
83    pub fn set<Q>(&mut self, key: Q, expr: impl Into<Expr>)
84    where
85        Q: Into<Projection>,
86    {
87        let key = key.into();
88        self.assignments.insert(
89            key,
90            Assignment {
91                op: AssignmentOp::Set,
92                expr: expr.into(),
93            },
94        );
95    }
96
97    pub fn unset<Q>(&mut self, key: &Q)
98    where
99        Q: ?Sized + Hash + Equivalent<Projection>,
100    {
101        self.assignments.swap_remove(key);
102    }
103
104    /// Insert a value into a set. The expression should evaluate to a single
105    /// value that is inserted into the set.
106    pub fn insert<Q>(&mut self, key: Q, expr: impl Into<Expr>)
107    where
108        Q: Into<Projection>,
109    {
110        use indexmap::map::Entry;
111
112        let key = key.into();
113        match self.assignments.entry(key) {
114            Entry::Occupied(_) => {
115                todo!()
116            }
117            Entry::Vacant(entry) => {
118                entry.insert(Assignment {
119                    op: AssignmentOp::Insert,
120                    expr: expr.into(),
121                });
122            }
123        }
124    }
125
126    pub fn remove<Q>(&mut self, key: Q, expr: impl Into<Expr>)
127    where
128        Q: Into<Projection>,
129    {
130        use indexmap::map::Entry;
131
132        let key = key.into();
133        match self.assignments.entry(key) {
134            Entry::Occupied(_) => {
135                todo!()
136            }
137            Entry::Vacant(entry) => {
138                entry.insert(Assignment {
139                    op: AssignmentOp::Remove,
140                    expr: expr.into(),
141                });
142            }
143        }
144    }
145
146    pub fn take<Q>(&mut self, key: &Q) -> Option<Assignment>
147    where
148        Q: ?Sized + Hash + Equivalent<Projection>,
149    {
150        self.assignments.swap_remove(key)
151    }
152
153    pub fn keys(&self) -> impl Iterator<Item = &Projection> + '_ {
154        self.assignments.keys()
155    }
156
157    pub fn exprs(&self) -> impl Iterator<Item = &Expr> + '_ {
158        self.assignments.values().map(|assignment| &assignment.expr)
159    }
160
161    pub fn iter(&self) -> impl Iterator<Item = (&Projection, &Assignment)> + '_ {
162        self.assignments.iter()
163    }
164
165    pub fn iter_mut(&mut self) -> impl Iterator<Item = (&Projection, &mut Assignment)> + '_ {
166        self.assignments.iter_mut()
167    }
168}
169
170impl IntoIterator for Assignments {
171    type Item = (Projection, Assignment);
172    type IntoIter = indexmap::map::IntoIter<Projection, Assignment>;
173
174    fn into_iter(self) -> Self::IntoIter {
175        self.assignments.into_iter()
176    }
177}
178
179impl<'a> IntoIterator for &'a Assignments {
180    type Item = (&'a Projection, &'a Assignment);
181    type IntoIter = indexmap::map::Iter<'a, Projection, Assignment>;
182
183    fn into_iter(self) -> Self::IntoIter {
184        self.assignments.iter()
185    }
186}
187
188impl Default for Assignments {
189    fn default() -> Self {
190        Self {
191            assignments: IndexMap::new(),
192        }
193    }
194}
195
196impl<Q> ops::Index<Q> for Assignments
197where
198    Q: Hash + Equivalent<Projection>,
199{
200    type Output = Assignment;
201
202    fn index(&self, index: Q) -> &Self::Output {
203        match self.assignments.get(&index) {
204            Some(ret) => ret,
205            None => panic!("no assignment for projection"),
206        }
207    }
208}
209
210impl<Q> ops::IndexMut<Q> for Assignments
211where
212    Q: Hash + Equivalent<Projection>,
213{
214    fn index_mut(&mut self, index: Q) -> &mut Self::Output {
215        match self.assignments.get_mut(&index) {
216            Some(ret) => ret,
217            None => panic!("no assignment for projection"),
218        }
219    }
220}
221
222impl AssignmentOp {
223    pub fn is_set(self) -> bool {
224        matches!(self, Self::Set)
225    }
226
227    pub fn is_remove(self) -> bool {
228        matches!(self, Self::Remove)
229    }
230}
toasty_core/stmt/assignments.rs

toasty_core/stmt/
assignments.rs
