1use crate::canonical::{CExpr, Stage};
10use serde::{Deserialize, Serialize};
11
12#[derive(Debug, Clone, Serialize, Deserialize, PartialEq)]
13#[serde(tag = "op", rename_all = "snake_case")]
14pub enum Patch {
15 Replace {
19 target: String,
20 with: CExpr,
21 },
22 Delete {
26 target: String,
27 },
28 WrapWith {
32 target: String,
33 wrapper: CExpr,
34 },
35}
36
37#[derive(Debug, Clone, thiserror::Error, Serialize, Deserialize)]
38#[serde(tag = "kind", rename_all = "snake_case")]
39pub enum PatchError {
40 #[error("unknown node id `{at}`")]
41 UnknownNode { at: String },
42 #[error("cannot patch non-expression at `{at}` ({reason})")]
43 NonExprTarget { at: String, reason: String },
44 #[error("Delete target is not in a list-shaped parent at `{at}`")]
45 DeleteNotInList { at: String },
46 #[error("WrapWith fragment must contain exactly one `_HOLE_` var")]
47 WrapWithMissingHole,
48 #[error("malformed NodeId `{0}`")]
49 BadNodeId(String),
50}
51
52pub fn apply_patch(stage: &Stage, patch: &Patch) -> Result<Stage, PatchError> {
55 let mut out = stage.clone();
56 let target = match patch {
57 Patch::Replace { target, .. } | Patch::Delete { target } | Patch::WrapWith { target, .. } => target,
58 };
59 let path = parse_node_id(target)?;
60
61 let body = match &mut out {
62 Stage::FnDecl(fd) => &mut fd.body,
63 Stage::TypeDecl(_) | Stage::Import(_) => {
64 return Err(PatchError::NonExprTarget {
65 at: target.clone(),
66 reason: "only fn-decl bodies are patchable".into(),
67 });
68 }
69 };
70
71 if path.is_empty() {
78 return Err(PatchError::NonExprTarget {
80 at: target.clone(),
81 reason: "cannot replace the stage root".into(),
82 });
83 }
84 let body_path = &path[1..];
87
88 apply_inside_expr(body, body_path, patch, target)?;
89 Ok(out)
90}
91
92fn parse_node_id(id: &str) -> Result<Vec<usize>, PatchError> {
94 let s = id.strip_prefix("n_").ok_or_else(|| PatchError::BadNodeId(id.into()))?;
95 let mut parts = s.split('.');
96 let head = parts.next().ok_or_else(|| PatchError::BadNodeId(id.into()))?;
97 if head != "0" {
98 return Err(PatchError::BadNodeId(id.into()));
99 }
100 let mut out = Vec::new();
101 for p in parts {
102 out.push(p.parse::<usize>().map_err(|_| PatchError::BadNodeId(id.into()))?);
103 }
104 Ok(out)
105}
106
107fn apply_inside_expr(
108 e: &mut CExpr,
109 path: &[usize],
110 patch: &Patch,
111 target_id: &str,
112) -> Result<(), PatchError> {
113 if path.is_empty() {
114 match patch {
116 Patch::Replace { with, .. } => {
117 *e = with.clone();
118 Ok(())
119 }
120 Patch::Delete { .. } => Err(PatchError::DeleteNotInList { at: target_id.into() }),
121 Patch::WrapWith { wrapper, .. } => {
122 let mut wrapped = wrapper.clone();
123 if !substitute_hole(&mut wrapped, e) {
124 return Err(PatchError::WrapWithMissingHole);
125 }
126 *e = wrapped;
127 Ok(())
128 }
129 }
130 } else {
131 let i = path[0];
133 let rest = &path[1..];
134 descend(e, i, rest, patch, target_id)
135 }
136}
137
138fn descend(
139 e: &mut CExpr,
140 i: usize,
141 rest: &[usize],
142 patch: &Patch,
143 target_id: &str,
144) -> Result<(), PatchError> {
145 let unknown = || PatchError::UnknownNode { at: target_id.into() };
146 match e {
147 CExpr::Call { callee, args } => {
148 if i == 0 { return apply_inside_expr(callee, rest, patch, target_id); }
149 let idx = i - 1;
150 args.get_mut(idx).ok_or_else(unknown)
151 .and_then(|c| apply_inside_expr(c, rest, patch, target_id))
152 }
153 CExpr::Let { value, body, .. } => match i {
154 0 => apply_inside_expr(value, rest, patch, target_id),
155 1 => apply_inside_expr(body, rest, patch, target_id),
156 _ => Err(unknown()),
157 },
158 CExpr::Match { scrutinee, arms } => {
159 if i == 0 { return apply_inside_expr(scrutinee, rest, patch, target_id); }
160 let arm_pos = i - 1;
163 let arm_idx = arm_pos / 2;
164 let is_pat = arm_pos.is_multiple_of(2);
165 let arm = arms.get_mut(arm_idx).ok_or_else(unknown)?;
166 if is_pat {
167 Err(PatchError::NonExprTarget {
168 at: target_id.into(),
169 reason: "patches on patterns are not supported (use Replace on the arm body or WrapWith)".into(),
170 })
171 } else {
172 apply_inside_expr(&mut arm.body, rest, patch, target_id)
173 }
174 }
175 CExpr::Block { statements, result } => {
176 if i < statements.len() {
178 if rest.is_empty() {
179 match patch {
181 Patch::Delete { .. } => {
182 statements.remove(i);
183 Ok(())
184 }
185 Patch::Replace { with, .. } => {
186 statements[i] = with.clone();
187 Ok(())
188 }
189 Patch::WrapWith { wrapper, .. } => {
190 let mut wrapped = wrapper.clone();
191 let original = std::mem::replace(&mut statements[i], CExpr::Literal { value: crate::canonical::CLit::Unit });
192 if !substitute_hole(&mut wrapped, &original) {
193 statements[i] = original; return Err(PatchError::WrapWithMissingHole);
195 }
196 statements[i] = wrapped;
197 Ok(())
198 }
199 }
200 } else {
201 apply_inside_expr(&mut statements[i], rest, patch, target_id)
202 }
203 } else if i == statements.len() {
204 if matches!(patch, Patch::Delete { .. }) {
206 Err(PatchError::DeleteNotInList { at: target_id.into() })
207 } else {
208 apply_inside_expr(result, rest, patch, target_id)
209 }
210 } else {
211 Err(unknown())
212 }
213 }
214 CExpr::Constructor { args, .. } => {
215 args.get_mut(i).ok_or_else(unknown)
216 .and_then(|c| apply_inside_expr(c, rest, patch, target_id))
217 }
218 CExpr::RecordLit { fields } => {
219 fields.get_mut(i).ok_or_else(unknown)
220 .and_then(|f| apply_inside_expr(&mut f.value, rest, patch, target_id))
221 }
222 CExpr::TupleLit { items } | CExpr::ListLit { items } => {
223 if rest.is_empty() && matches!(patch, Patch::Delete { .. }) {
224 if i >= items.len() { return Err(unknown()); }
225 items.remove(i);
226 return Ok(());
227 }
228 items.get_mut(i).ok_or_else(unknown)
229 .and_then(|c| apply_inside_expr(c, rest, patch, target_id))
230 }
231 CExpr::FieldAccess { value, .. } => {
232 if i == 0 { apply_inside_expr(value, rest, patch, target_id) } else { Err(unknown()) }
233 }
234 CExpr::Lambda { body, .. } => {
235 if i == 0 { apply_inside_expr(body, rest, patch, target_id) } else { Err(unknown()) }
236 }
237 CExpr::BinOp { lhs, rhs, .. } => match i {
238 0 => apply_inside_expr(lhs, rest, patch, target_id),
239 1 => apply_inside_expr(rhs, rest, patch, target_id),
240 _ => Err(unknown()),
241 },
242 CExpr::UnaryOp { expr, .. } => {
243 if i == 0 { apply_inside_expr(expr, rest, patch, target_id) } else { Err(unknown()) }
244 }
245 CExpr::Return { value } => {
246 if i == 0 { apply_inside_expr(value, rest, patch, target_id) } else { Err(unknown()) }
247 }
248 CExpr::Literal { .. } | CExpr::Var { .. } => Err(unknown()),
249 }
250}
251
252fn substitute_hole(node: &mut CExpr, target: &CExpr) -> bool {
255 let mut count = 0;
256 walk_substitute(node, target, &mut count);
257 count == 1
258}
259
260fn walk_substitute(e: &mut CExpr, target: &CExpr, count: &mut u32) {
261 if let CExpr::Var { name } = e {
262 if name == "_HOLE_" {
263 *e = target.clone();
264 *count += 1;
265 return;
266 }
267 }
268 match e {
269 CExpr::Literal { .. } | CExpr::Var { .. } => {}
270 CExpr::Call { callee, args } => {
271 walk_substitute(callee, target, count);
272 for a in args { walk_substitute(a, target, count); }
273 }
274 CExpr::Let { value, body, .. } => {
275 walk_substitute(value, target, count);
276 walk_substitute(body, target, count);
277 }
278 CExpr::Match { scrutinee, arms } => {
279 walk_substitute(scrutinee, target, count);
280 for a in arms { walk_substitute(&mut a.body, target, count); }
281 }
282 CExpr::Block { statements, result } => {
283 for s in statements { walk_substitute(s, target, count); }
284 walk_substitute(result, target, count);
285 }
286 CExpr::Constructor { args, .. } => for a in args { walk_substitute(a, target, count); },
287 CExpr::RecordLit { fields } => for f in fields { walk_substitute(&mut f.value, target, count); },
288 CExpr::TupleLit { items } | CExpr::ListLit { items } => {
289 for i in items { walk_substitute(i, target, count); }
290 }
291 CExpr::FieldAccess { value, .. } => walk_substitute(value, target, count),
292 CExpr::Lambda { body, .. } => walk_substitute(body, target, count),
293 CExpr::BinOp { lhs, rhs, .. } => {
294 walk_substitute(lhs, target, count);
295 walk_substitute(rhs, target, count);
296 }
297 CExpr::UnaryOp { expr, .. } => walk_substitute(expr, target, count),
298 CExpr::Return { value } => walk_substitute(value, target, count),
299 }
300}