Preserve source comments and layout through model token edits

crates/oa-core/src/language.rs

@@ -933,6 +933,143 @@ pub(crate) fn parse_unit(
     .unit(entry, true)
     .map_err(|e| Box::new(diagnostic(source, e)))
 }
+
+// Editing positions are derived from the existing, validated parser AST. They
+// do not define another grammar or change the shared V2 migration AST.
+pub(crate) struct EditField {
+    pub(crate) span: Span,
+    pub(crate) value: EditValue,
+}
+pub(crate) struct EditValue {
+    pub(crate) span: Span,
+    pub(crate) kind: EditKind,
+}
+pub(crate) enum EditKind {
+    Atom,
+    Reference(String),
+    Record(BTreeMap<String, EditField>),
+    List(Vec<EditValue>),
+}
+pub(crate) struct EditDeclaration {
+    pub(crate) header: BTreeMap<&'static str, Span>,
+    pub(crate) fields: EditValue,
+}
+
+pub(crate) fn edit_header(source: &str) -> Result<BTreeMap<&'static str, Span>, Box<Diagnostic>> {
+    let tokens = strict_tokens(source).map_err(|e| Box::new(diagnostic(source, e)))?;
+    let named = tokens
+        .first()
+        .is_none_or(|token| token.text(source) != "node");
+    let roles: &[&str] = if named {
+        &["kind", "alias", "name", "id_keyword", "id"]
+    } else {
+        &["kind", "alias", "id_keyword", "id"]
+    };
+    if tokens.len() <= roles.len() {
+        return Err(Box::new(diagnostic(
+            source,
+            Error {
+                code: "invalid_source",
+                offset: 0,
+                message: "Expected a complete stabilized declaration header.".into(),
+            },
+        )));
+    }
+    Ok(roles
+        .iter()
+        .zip(tokens)
+        .map(|(role, token)| (*role, token.span))
+        .collect())
+}
+
+pub(crate) fn edit_declaration(source: &str) -> Result<EditDeclaration, Box<Diagnostic>> {
+    let unit = parse_unit(source, false, ParseLimits::UNLIMITED)?;
+    if unit.declarations.len() != 1 || !unit.includes.is_empty() {
+        return Err(Box::new(diagnostic(
+            source,
+            Error {
+                code: "invalid_source",
+                offset: 0,
+                message: "Expected one declaration for a source edit.".into(),
+            },
+        )));
+    }
+    let tokens = strict_tokens(source).map_err(|e| Box::new(diagnostic(source, e)))?;
+    let mut ends: Vec<_> = tokens.iter().map(|token| token.span.end).collect();
+    let mut stack = Vec::new();
+    for (index, token) in tokens.iter().enumerate() {
+        match token.text(source) {
+            "{" | "[" => stack.push(index),
+            "}" | "]" => ends[stack.pop().expect("parser checked delimiters")] = token.span.end,
+            _ => {}
+        }
+    }
+    struct Positions<'a> {
+        tokens: &'a [Lexeme],
+        ends: Vec<usize>,
+    }
+    impl Positions<'_> {
+        fn token(&self, start: usize) -> usize {
+            self.tokens
+                .partition_point(|token| token.span.start < start)
+        }
+        fn fields(&self, fields: &BTreeMap<String, Ast>) -> BTreeMap<String, EditField> {
+            fields
+                .iter()
+                .map(|(key, ast)| {
+                    let value = self.value(ast);
+                    let start = self.tokens[self.token(value.span.start) - 1].span.start;
+                    let end = self.tokens[self.token(value.span.end)].span.end;
+                    (
+                        key.clone(),
+                        EditField {
+                            span: Span { start, end },
+                            value,
+                        },
+                    )
+                })
+                .collect()
+        }
+        fn value(&self, ast: &Ast) -> EditValue {
+            let index = self.token(ast.offset());
+            let span = Span {
+                start: ast.offset(),
+                end: match ast {
+                    Ast::Reference(..) => self.tokens[index + 1].span.end,
+                    _ => self.ends[index],
+                },
+            };
+            let kind = match ast {
+                Ast::Object(fields, _) => EditKind::Record(self.fields(fields)),
+                Ast::Array(items, _) => {
+                    EditKind::List(items.iter().map(|item| self.value(item)).collect())
+                }
+                Ast::Reference(alias, _) => EditKind::Reference(alias.clone()),
+                _ => EditKind::Atom,
+            };
+            EditValue { span, kind }
+        }
+    }
+    let positions = Positions {
+        tokens: &tokens,
+        ends,
+    };
+    let body = tokens
+        .iter()
+        .position(|token| token.text(source) == "{")
+        .expect("declaration body");
+    let fields = EditValue {
+        span: Span {
+            start: tokens[body].span.start,
+            end: positions.ends[body],
+        },
+        kind: EditKind::Record(positions.fields(&unit.declarations[0].fields)),
+    };
+    Ok(EditDeclaration {
+        header: edit_header(source)?,
+        fields,
+    })
+}
 struct Parser<'a> {
     source: &'a str,
     tokens: Vec<Lexeme>,