1use sim_kernel::Expr;
10
11use crate::build::sym;
12
13#[derive(Clone, Debug, PartialEq)]
15pub enum Segment {
16 Key(Expr),
18 Index(usize),
20}
21
22#[derive(Clone, Debug, Default, PartialEq)]
24pub struct Path(pub Vec<Segment>);
25
26#[derive(Clone, Debug, PartialEq, Eq)]
28pub enum PathError {
29 NotAList,
31 BadSegment,
33 NotAMap,
35 NotASequence,
37 MissingKey,
39 IndexOutOfBounds,
41 RemoveIndexUnsupported,
43 EmptyRemove,
45}
46
47impl Path {
48 pub fn new() -> Self {
50 Self(Vec::new())
51 }
52
53 pub fn key(mut self, key: Expr) -> Self {
55 self.0.push(Segment::Key(key));
56 self
57 }
58
59 pub fn index(mut self, index: usize) -> Self {
61 self.0.push(Segment::Index(index));
62 self
63 }
64
65 pub fn to_expr(&self) -> Expr {
67 Expr::List(
68 self.0
69 .iter()
70 .map(|segment| match segment {
71 Segment::Key(key) => Expr::Vector(vec![sym("k"), key.clone()]),
72 Segment::Index(index) => {
73 Expr::Vector(vec![sym("i"), Expr::String(index.to_string())])
74 }
75 })
76 .collect(),
77 )
78 }
79
80 pub fn from_expr(expr: &Expr) -> Result<Self, PathError> {
82 let Expr::List(segments) = expr else {
83 return Err(PathError::NotAList);
84 };
85 let mut out = Vec::with_capacity(segments.len());
86 for segment in segments {
87 out.push(parse_segment(segment)?);
88 }
89 Ok(Path(out))
90 }
91}
92
93fn parse_segment(segment: &Expr) -> Result<Segment, PathError> {
94 let Expr::Vector(parts) = segment else {
95 return Err(PathError::BadSegment);
96 };
97 match parts.as_slice() {
98 [Expr::Symbol(tag), key] if &*tag.name == "k" => Ok(Segment::Key(key.clone())),
99 [Expr::Symbol(tag), Expr::String(index)] if &*tag.name == "i" => index
100 .parse::<usize>()
101 .map(Segment::Index)
102 .map_err(|_| PathError::BadSegment),
103 _ => Err(PathError::BadSegment),
104 }
105}
106
107fn seq_items(value: &Expr) -> Option<&Vec<Expr>> {
108 match value {
109 Expr::List(items) | Expr::Vector(items) | Expr::Set(items) => Some(items),
110 _ => None,
111 }
112}
113
114fn rewrap(original: &Expr, items: Vec<Expr>) -> Expr {
115 match original {
116 Expr::Vector(_) => Expr::Vector(items),
117 Expr::Set(_) => Expr::Set(items),
118 _ => Expr::List(items),
119 }
120}
121
122pub fn get<'a>(root: &'a Expr, path: &Path) -> Option<&'a Expr> {
124 let mut current = root;
125 for segment in &path.0 {
126 current = match segment {
127 Segment::Key(key) => {
128 let Expr::Map(entries) = current else {
129 return None;
130 };
131 entries
132 .iter()
133 .find_map(|(entry_key, value)| (entry_key == key).then_some(value))?
134 }
135 Segment::Index(index) => seq_items(current)?.get(*index)?,
136 };
137 }
138 Some(current)
139}
140
141pub fn set_at(root: &Expr, path: &Path, value: Expr) -> Result<Expr, PathError> {
144 set_rec(root, &path.0, value)
145}
146
147fn set_rec(node: &Expr, segments: &[Segment], value: Expr) -> Result<Expr, PathError> {
148 let Some((first, rest)) = segments.split_first() else {
149 return Ok(value);
150 };
151 match first {
152 Segment::Key(key) => {
153 let Expr::Map(entries) = node else {
154 return Err(PathError::NotAMap);
155 };
156 let mut entries = entries.clone();
157 match entries.iter_mut().find(|(entry_key, _)| entry_key == key) {
158 Some(slot) => slot.1 = set_rec(&slot.1.clone(), rest, value)?,
159 None if rest.is_empty() => entries.push((key.clone(), value)),
160 None => return Err(PathError::MissingKey),
161 }
162 Ok(Expr::Map(entries))
163 }
164 Segment::Index(index) => {
165 let mut items = seq_items(node).ok_or(PathError::NotASequence)?.clone();
166 let slot = items.get(*index).ok_or(PathError::IndexOutOfBounds)?;
167 let replaced = set_rec(&slot.clone(), rest, value)?;
168 items[*index] = replaced;
169 Ok(rewrap(node, items))
170 }
171 }
172}
173
174pub fn remove_at(root: &Expr, path: &Path) -> Result<Expr, PathError> {
177 let Some((last, parents)) = path.0.split_last() else {
178 return Err(PathError::EmptyRemove);
179 };
180 let Segment::Key(key) = last else {
181 return Err(PathError::RemoveIndexUnsupported);
182 };
183 remove_rec(root, parents, key)
184}
185
186fn remove_rec(node: &Expr, parents: &[Segment], key: &Expr) -> Result<Expr, PathError> {
187 let Some((first, rest)) = parents.split_first() else {
188 let Expr::Map(entries) = node else {
189 return Err(PathError::NotAMap);
190 };
191 return Ok(Expr::Map(
192 entries
193 .iter()
194 .filter(|(entry_key, _)| entry_key != key)
195 .cloned()
196 .collect(),
197 ));
198 };
199 match first {
200 Segment::Key(parent_key) => {
201 let Expr::Map(entries) = node else {
202 return Err(PathError::NotAMap);
203 };
204 let mut entries = entries.clone();
205 let slot = entries
206 .iter_mut()
207 .find(|(entry_key, _)| entry_key == parent_key)
208 .ok_or(PathError::MissingKey)?;
209 slot.1 = remove_rec(&slot.1.clone(), rest, key)?;
210 Ok(Expr::Map(entries))
211 }
212 Segment::Index(index) => {
213 let mut items = seq_items(node).ok_or(PathError::NotASequence)?.clone();
214 let slot = items.get(*index).ok_or(PathError::IndexOutOfBounds)?;
215 items[*index] = remove_rec(&slot.clone(), rest, key)?;
216 Ok(rewrap(node, items))
217 }
218 }
219}