Yeast: Fix handling of captures with multiple results

Author: asgerf

shared/yeast/src/captures.rs

@@ -63,16 +63,20 @@ impl Captures {
     }
 
     /// Apply a fallible function to every captured id (across all keys),
-    /// replacing each id with the result. Stops and returns the error on
-    /// the first failure.
+    /// replacing each id with the results. A function returning an empty
+    /// vector removes the capture; returning multiple ids splices them
+    /// into the capture's value list (suitable for `*`/`+` captures).
+    /// Stops and returns the error on the first failure.
     pub fn try_map_all_captures<E>(
         &mut self,
-        mut f: impl FnMut(Id) -> Result<Id, E>,
+        mut f: impl FnMut(Id) -> Result<Vec<Id>, E>,
     ) -> Result<(), E> {
         for ids in self.captures.values_mut() {
-            for id in ids {
-                *id = f(*id)?;
+            let mut new_ids = Vec::with_capacity(ids.len());
+            for &id in ids.iter() {
+                new_ids.extend(f(id)?);
             }
+            *ids = new_ids;
         }
         Ok(())
     }

shared/yeast/src/lib.rs

@@ -563,6 +563,15 @@ impl Node {
             NodeContent::DynamicString(s) => Some(s.to_string()),
         }
     }
+
+    /// Read the child ids stored under a given field, or an empty slice if
+    /// no such field is present on this node.
+    pub fn field_children(&self, field_id: FieldId) -> &[Id] {
+        self.fields
+            .get(&field_id)
+            .map(|v| v.as_slice())
+            .unwrap_or(&[])
+    }
 }
 
 /// The contents of a node is either a range in the original source file,
@@ -836,17 +845,9 @@ fn apply_one_shot_rules_inner(
                 // pattern root): re-analyzing it would match the same rule
                 // again indefinitely.
                 if captured_id == id {
-                    return Ok(captured_id);
-                }
-                let result =
-                    apply_one_shot_rules_inner(index, ast, captured_id, fresh, rewrite_depth + 1)?;
-                if result.len() != 1 {
-                    return Err(format!(
-                        "OneShot: recursion on captured node produced {} results, expected exactly 1",
-                        result.len()
-                    ));
+                    return Ok(vec![captured_id]);
                 }
-                Ok(result[0])
+                apply_one_shot_rules_inner(index, ast, captured_id, fresh, rewrite_depth + 1)
             })?;
             return Ok(rule.run_transform(ast, captures, id, fresh));
         }