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>,

crates/oa-core/src/source_project.rs

@@ -719,8 +719,9 @@ pub fn format_bundle(project: &Project) -> SourceBundle {
 }
 
 /// Apply a model change to a compiled document workspace. Existing declaration
-/// ownership and untouched bytes remain stable; only changed declarations are
-/// canonicalized. Alias edits also rewrite declarations that reference them.
+/// ownership and surviving syntax trivia remain stable. Changed values and
+/// headers receive token edits; canonical formatting is only used for new syntax.
+/// Alias edits also rewrite declarations that reference them.
 /// No files are deleted, even when their last declaration is removed.
 pub fn update_bundle(
     current: &CompiledSourceProject,
@@ -761,34 +762,133 @@ pub fn update_bundle(
     let old = &current.decoded.project;
     let before: BTreeMap<_, _> = language::format_declarations(old).into_iter().collect();
     let after: BTreeMap<_, _> = language::format_declarations(project).into_iter().collect();
+    let old_aliases: BTreeMap<_, _> = old
+        .identities()
+        .into_iter()
+        .map(|(id, alias)| (alias.to_owned(), id))
+        .collect();
+    let new_aliases: BTreeMap<_, _> = project
+        .identities()
+        .into_iter()
+        .map(|(id, alias)| (alias.to_owned(), id))
+        .collect();
     let mut bundle = current.bundle.clone();
     let mut edits: Edits = BTreeMap::new();
     for (id, old_text) in &before {
         let location = &current.index.declarations[id];
         if after.get(id) != Some(old_text) {
-            let text = after.get(id).map_or_else(String::new, |text| {
-                replacement(text, &bundle.files[&location.path], location.span.start)
-            });
+            let source = &bundle.files[&location.path][location.span.start..location.span.end];
+            let mut local = Vec::new();
+            if let Some(new_text) = after.get(id) {
+                let parse = |text| {
+                    language::edit_declaration(text).map_err(|error| {
+                        vec![problem(
+                            &bundle,
+                            &location.path,
+                            location.span,
+                            "source_edit",
+                            error.message,
+                        )]
+                    })
+                };
+                let actual = parse(source)?;
+                let previous = parse(old_text)?;
+                let next = parse(new_text)?;
+                let update = SyntaxUpdate {
+                    source,
+                    before: old_text,
+                    after: new_text,
+                    old_aliases: &old_aliases,
+                    new_aliases: &new_aliases,
+                };
+                update.header(&actual.header, &previous.header, &next.header, &mut local);
+                update.value(
+                    &actual.fields,
+                    Some(&previous.fields),
+                    &next.fields,
+                    &mut local,
+                );
+            } else {
+                erase_syntax(
+                    source,
+                    Span {
+                        start: 0,
+                        end: source.len(),
+                    },
+                    &mut local,
+                );
+            }
             edits
                 .entry(location.path.clone())
                 .or_default()
-                .push((location.span, text));
+                .extend(local.into_iter().map(|(span, text)| {
+                    (
+                        Span {
+                            start: span.start + location.span.start,
+                            end: span.end + location.span.start,
+                        },
+                        text,
+                    )
+                }));
         }
     }
     if old.id != project.id || old.alias != project.alias || old.name != project.name {
         let location = &current.index.project_header;
+        let source = &bundle.files[&location.path][location.span.start..location.span.end];
+        let before = project_header(old);
+        let after = project_header(project);
+        let parse = |text| {
+            language::edit_header(text).map_err(|error| {
+                vec![problem(
+                    &bundle,
+                    &location.path,
+                    location.span,
+                    "source_edit",
+                    error.message,
+                )]
+            })
+        };
+        let update = SyntaxUpdate {
+            source,
+            before: &before,
+            after: &after,
+            old_aliases: &old_aliases,
+            new_aliases: &new_aliases,
+        };
+        let mut local = Vec::new();
+        update.header(
+            &parse(source)?,
+            &parse(&before)?,
+            &parse(&after)?,
+            &mut local,
+        );
         edits
             .entry(location.path.clone())
             .or_default()
-            .push((location.span, project_header(project)));
+            .extend(local.into_iter().map(|(span, text)| {
+                (
+                    Span {
+                        start: span.start + location.span.start,
+                        end: span.end + location.span.start,
+                    },
+                    text,
+                )
+            }));
     }
     if old.description != project.description {
         let field = format!("description {};", quote(&project.description));
         if let Some(location) = &current.index.project_description {
-            edits
-                .entry(location.path.clone())
-                .or_default()
-                .push((location.span, field));
+            let value =
+                significant(&bundle.files[&location.path][location.span.start..location.span.end])
+                    [1]
+                .span;
+            edits.entry(location.path.clone()).or_default().push((
+                Span {
+                    start: location.span.start + value.start,
+                    end: location.span.start + value.end,
+                },
+                quote(&project.description),
+            ));
         } else {
             let location = &current.index.project_body;
             edits.entry(location.path.clone()).or_default().push((
@@ -909,19 +1009,436 @@ fn unindent(text: &str) -> String {
         .map(|line| line.strip_prefix("    ").unwrap_or(line))
         .collect()
 }
-fn replacement(text: &str, source: &str, start: usize) -> String {
-    let prefix = source[..start].rsplit('\n').next().unwrap_or("");
-    let indent = if prefix.chars().all(char::is_whitespace) {
-        prefix
-    } else {
-        ""
-    };
-    let text = unindent(text);
-    let text = text.trim_end().replace('\n', &format!("\n{indent}"));
+fn significant(source: &str) -> Vec<language::TokenSpan> {
+    language::tokenize(source)
+        .into_iter()
+        .filter(|token| {
+            !matches!(
+                token.class,
+                language::TokenClass::Whitespace | language::TokenClass::Comment
+            )
+        })
+        .collect()
+}
+type LocalEdits = Vec<(Span, String)>;
+fn erase_syntax(source: &str, span: Span, edits: &mut LocalEdits) {
+    edits.extend(
+        significant(&source[span.start..span.end])
+            .into_iter()
+            .map(|token| {
+                (
+                    Span {
+                        start: span.start + token.span.start,
+                        end: span.start + token.span.end,
+                    },
+                    String::new(),
+                )
+            }),
+    );
+}
+fn trivia(source: &str) -> String {
+    language::tokenize(source)
+        .into_iter()
+        .filter(|token| {
+            matches!(
+                token.class,
+                language::TokenClass::Whitespace | language::TokenClass::Comment
+            )
+        })
+        .map(|token| &source[token.span.start..token.span.end])
+        .collect()
+}
+fn changed_text(source: &str, span: Span, mut edits: LocalEdits) -> String {
+    let mut result = source[span.start..span.end].to_owned();
+    edits.sort_by_key(|(span, _)| (span.start, span.end));
+    for (edit, text) in edits.into_iter().rev() {
+        result.replace_range(edit.start - span.start..edit.end - span.start, &text);
+    }
+    result
+}
+fn newline(source: &str) -> &'static str {
     if source.contains("\r\n") {
-        text.replace('\n', "\r\n")
+        "\r\n"
     } else {
-        text
+        "\n"
+    }
+}
+fn indent_at(source: &str, at: usize) -> &str {
+    let line = source[..at].rsplit('\n').next().unwrap_or("");
+    &line[..line.len() - line.trim_start_matches([' ', '\t']).len()]
+}
+fn inserted_value(text: &str, indent: &str, source: &str) -> String {
+    let lines: Vec<_> = text.lines().collect();
+    let base = lines
+        .iter()
+        .skip(1)
+        .filter(|line| !line.trim().is_empty())
+        .map(|line| line.len() - line.trim_start_matches(' ').len())
+        .min()
+        .unwrap_or(0);
+    let mut output = lines.first().copied().unwrap_or("").to_owned();
+    for line in lines.iter().skip(1) {
+        output.push_str(newline(source));
+        output.push_str(indent);
+        output.push_str(line.get(base..).unwrap_or(line));
+    }
+    output
+}
+
+struct SyntaxUpdate<'a> {
+    source: &'a str,
+    before: &'a str,
+    after: &'a str,
+    old_aliases: &'a BTreeMap<String, Id>,
+    new_aliases: &'a BTreeMap<String, Id>,
+}
+impl SyntaxUpdate<'_> {
+    fn header(
+        &self,
+        actual: &BTreeMap<&str, Span>,
+        before: &BTreeMap<&str, Span>,
+        after: &BTreeMap<&str, Span>,
+        edits: &mut LocalEdits,
+    ) {
+        for (key, span) in actual {
+            match after.get(key) {
+                Some(next)
+                    if before.get(key).is_none_or(|old| {
+                        self.before[old.start..old.end] != self.after[next.start..next.end]
+                    }) =>
+                {
+                    edits.push((*span, self.after[next.start..next.end].to_owned()))
+                }
+                None => erase_syntax(self.source, *span, edits),
+                _ => {}
+            }
+        }
+        if let Some(name) = after.get("name")
+            && !actual.contains_key("name")
+        {
+            let at = actual["id_keyword"].start;
+            edits.push((
+                Span { start: at, end: at },
+                format!("{} ", &self.after[name.start..name.end]),
+            ));
+        }
+    }
+    fn same(&self, before: &language::EditValue, after: &language::EditValue) -> bool {
+        use language::EditKind;
+        match (&before.kind, &after.kind) {
+            (EditKind::Record(a), EditKind::Record(b)) => {
+                a.len() == b.len()
+                    && a.iter().all(|(key, value)| {
+                        b.get(key)
+                            .is_some_and(|other| self.same(&value.value, &other.value))
+                    })
+            }
+            (EditKind::List(a), EditKind::List(b)) => {
+                a.len() == b.len() && a.iter().zip(b).all(|(a, b)| self.same(a, b))
+            }
+            _ => {
+                self.before[before.span.start..before.span.end]
+                    == self.after[after.span.start..after.span.end]
+            }
+        }
+    }
+    fn value(
+        &self,
+        actual: &language::EditValue,
+        before: Option<&language::EditValue>,
+        after: &language::EditValue,
+        edits: &mut LocalEdits,
+    ) {
+        use language::EditKind;
+        if before.is_some_and(|before| self.same(before, after)) {
+            return;
+        }
+        match (&actual.kind, &after.kind) {
+            (EditKind::Record(fields), EditKind::Record(next)) => {
+                let previous = before.and_then(|old| {
+                    if let EditKind::Record(fields) = &old.kind {
+                        Some(fields)
+                    } else {
+                        None
+                    }
+                });
+                self.record(actual.span, fields, previous, next, edits);
+            }
+            (EditKind::List(items), EditKind::List(next)) => {
+                let previous = before.and_then(|old| {
+                    if let EditKind::List(items) = &old.kind {
+                        Some(items.as_slice())
+                    } else {
+                        None
+                    }
+                });
+                self.list(actual.span, items, previous, next, edits);
+            }
+            (EditKind::Reference(_), EditKind::Reference(alias)) => {
+                let token = significant(&self.source[actual.span.start..actual.span.end])
+                    .pop()
+                    .expect("reference alias");
+                edits.push((
+                    Span {
+                        start: actual.span.start + token.span.start,
+                        end: actual.span.start + token.span.end,
+                    },
+                    alias.clone(),
+                ));
+            }
+            (EditKind::Atom, EditKind::Atom) => edits.push((
+                actual.span,
+                self.after[after.span.start..after.span.end].into(),
+            )),
+            _ => {
+                erase_syntax(self.source, actual.span, edits);
+                let text = inserted_value(
+                    &self.after[after.span.start..after.span.end],
+                    indent_at(self.source, actual.span.start),
+                    self.source,
+                );
+                edits.push((
+                    Span {
+                        start: actual.span.start,
+                        end: actual.span.start,
+                    },
+                    text,
+                ));
+            }
+        }
+    }
+    fn record(
+        &self,
+        span: Span,
+        actual: &BTreeMap<String, language::EditField>,
+        before: Option<&BTreeMap<String, language::EditField>>,
+        after: &BTreeMap<String, language::EditField>,
+        edits: &mut LocalEdits,
+    ) {
+        // A typed constructor tag change retains shared fields and their trivia.
+        if actual.len() == 1
+            && after.len() == 1
+            && let Some(before) = before.filter(|fields| fields.len() == 1)
+        {
+            let (old_key, old) = before.first_key_value().unwrap();
+            let (new_key, new) = after.first_key_value().unwrap();
+            if old_key != new_key
+                && old_key.starts_with(|c: char| c.is_ascii_uppercase())
+                && new_key.starts_with(|c: char| c.is_ascii_uppercase())
+                && let Some(field) = actual.get(old_key)
+            {
+                let key =
+                    significant(&self.source[field.span.start..field.value.span.start])[0].span;
+                edits.push((
+                    Span {
+                        start: key.start + field.span.start,
+                        end: key.end + field.span.start,
+                    },
+                    new_key.clone(),
+                ));
+                self.value(&field.value, Some(&old.value), &new.value, edits);
+                return;
+            }
+        }
+        for (key, field) in actual {
+            if let Some(next) = after.get(key) {
+                self.value(
+                    &field.value,
+                    before
+                        .and_then(|fields| fields.get(key))
+                        .map(|field| &field.value),
+                    &next.value,
+                    edits,
+                );
+            } else {
+                erase_syntax(self.source, field.span, edits);
+            }
+        }
+        let added: Vec<_> = after
+            .iter()
+            .filter(|(key, field)| {
+                !actual.contains_key(*key)
+                    && before
+                        .and_then(|fields| fields.get(*key))
+                        .is_none_or(|old| !self.same(&old.value, &field.value))
+            })
+            .collect();
+        if added.is_empty() {
+            return;
+        }
+        let close = span.end - 1;
+        let multiline = self.source[span.start..span.end].contains('\n');
+        let indent = actual
+            .values()
+            .min_by_key(|field| field.span.start)
+            .map(|field| indent_at(self.source, field.span.start).to_owned())
+            .filter(|indent| !indent.is_empty())
+            .unwrap_or_else(|| format!("{}    ", indent_at(self.source, span.start)));
+        let mut text = String::new();
+        for (key, field) in added {
+            let value = inserted_value(
+                &self.after[field.value.span.start..field.value.span.end],
+                &indent,
+                self.source,
+            );
+            if multiline {
+                text.push_str(&indent);
+                text.push_str(&format!("{key} {value};{}", newline(self.source)));
+            } else {
+                text.push_str(&format!(" {key} {value}; "));
+            }
+        }
+        let prefix = self.source[..close].rsplit('\n').next().unwrap_or("");
+        let at = if multiline && prefix.chars().all(char::is_whitespace) {
+            close - prefix.len()
+        } else {
+            close
+        };
+        if multiline && at == close && !prefix.is_empty() {
+            text.insert_str(0, newline(self.source));
+        }
+        edits.push((Span { start: at, end: at }, text));
+    }
+    fn list(
+        &self,
+        span: Span,
+        actual: &[language::EditValue],
+        before: Option<&[language::EditValue]>,
+        after: &[language::EditValue],
+        edits: &mut LocalEdits,
+    ) {
+        let previous = before.unwrap_or(&[]);
+        let old_keys: Option<Vec<_>> = previous
+            .iter()
+            .map(|item| item_key(item, self.before, self.old_aliases))
+            .collect();
+        let new_keys: Option<Vec<_>> = after
+            .iter()
+            .map(|item| item_key(item, self.after, self.new_aliases))
+            .collect();
+        let keyed = old_keys.zip(new_keys).filter(|(old, new)| {
+            old.iter().collect::<BTreeSet<_>>().len() == old.len()
+                && new.iter().collect::<BTreeSet<_>>().len() == new.len()
+        });
+        let mapping: Vec<_> = if let Some((old, new)) = keyed {
+            let positions: BTreeMap<_, _> = old
+                .into_iter()
+                .enumerate()
+                .map(|(i, key)| (key, i))
+                .collect();
+            new.into_iter()
+                .map(|key| positions.get(&key).copied())
+                .collect()
+        } else {
+            (0..after.len())
+                .map(|i| (i < actual.len()).then_some(i))
+                .collect()
+        };
+        if actual.len() == after.len()
+            && mapping
+                .iter()
+                .enumerate()
+                .all(|(i, source)| *source == Some(i))
+        {
+            for (i, (current, next)) in actual.iter().zip(after).enumerate() {
+                self.value(current, previous.get(i), next, edits);
+            }
+            return;
+        }
+        let mut chunks = Vec::new();
+        let mut start = span.start + 1;
+        for (i, item) in actual.iter().enumerate() {
+            let end = if let Some(next) = actual.get(i + 1) {
+                item.span.end
+                    + significant(&self.source[item.span.end..next.span.start])[0]
+                        .span
+                        .start
+            } else {
+                span.end - 1
+            };
+            chunks.push(Span { start, end });
+            start = end + 1;
+        }
+        let multiline = self.source[span.start..span.end].contains('\n');
+        let indent = actual
+            .first()
+            .map(|item| indent_at(self.source, item.span.start).to_owned())
+            .unwrap_or_else(|| format!("{}    ", indent_at(self.source, span.start)));
+        let mut text = String::new();
+        let mut retained = BTreeSet::new();
+        for (i, next) in after.iter().enumerate() {
+            let suffix;
+            if let Some(index) = mapping[i].filter(|index| *index < actual.len()) {
+                retained.insert(index);
+                let current = &actual[index];
+                let chunk = chunks[index];
+                text.push_str(&self.source[chunk.start..current.span.start]);
+                let mut local = Vec::new();
+                self.value(current, previous.get(index), next, &mut local);
+                text.push_str(&changed_text(self.source, current.span, local));
+                suffix = &self.source[current.span.end..chunk.end];
+            } else {
+                if multiline {
+                    text.push_str(newline(self.source));
+                    text.push_str(&indent);
+                } else {
+                    text.push(' ');
+                }
+                text.push_str(&inserted_value(
+                    &self.after[next.span.start..next.span.end],
+                    &indent,
+                    self.source,
+                ));
+                suffix = "";
+            }
+            if i + 1 < after.len() {
+                text.push(',');
+            }
+            text.push_str(suffix);
+        }
+        for (i, chunk) in chunks
+            .iter()
+            .enumerate()
+            .filter(|(i, _)| !retained.contains(i))
+        {
+            let _ = i;
+            text.push_str(&trivia(&self.source[chunk.start..chunk.end]));
+        }
+        if actual.is_empty() {
+            text.push_str(&self.source[span.start + 1..span.end - 1]);
+        }
+        edits.push((
+            Span {
+                start: span.start + 1,
+                end: span.end - 1,
+            },
+            text,
+        ));
+    }
+}
+
+fn item_key(
+    value: &language::EditValue,
+    source: &str,
+    aliases: &BTreeMap<String, Id>,
+) -> Option<String> {
+    use language::EditKind;
+    match &value.kind {
+        EditKind::Reference(alias) => aliases.get(alias).map(|id| format!("ref:{id}")),
+        EditKind::Atom => Some(format!(
+            "literal:{}",
+            &source[value.span.start..value.span.end]
+        )),
+        EditKind::Record(fields) => ["field", "referenced"].into_iter().find_map(|key| {
+            fields.get(key).and_then(|field| {
+                if let EditKind::Reference(alias) = &field.value.kind {
+                    aliases.get(alias).map(|id| format!("{key}:{id}"))
+                } else {
+                    None
+                }
+            })
+        }),
+        EditKind::List(_) => None,
     }
 }