use super::graph::{TranscriptRevisionBody, TranscriptRewriteCommit};
use crate::session::{
TranscriptRewriteSelection, canonicalize_message_images_for_digest, sha256_json_digest,
transcript_messages_digest,
};
use crate::types::Message;
use std::collections::BTreeMap;
pub(super) fn heal_legacy_revision_strings(
revisions: &mut [TranscriptRevisionBody],
commits: &mut [TranscriptRewriteCommit],
head: Option<&mut String>,
) -> Result<(), serde_json::Error> {
let mut remap: BTreeMap<String, String> = BTreeMap::new();
for body in revisions.iter() {
let content = transcript_messages_digest(&body.messages)?;
if body.revision == content {
continue;
}
if body.revision == legacy_transcript_messages_digest(&body.messages)? {
remap.insert(body.revision.clone(), content);
}
}
if remap.is_empty() {
return Ok(());
}
for body in revisions.iter_mut() {
if let Some(current) = remap.get(&body.revision) {
body.revision = current.clone();
}
if let Some(parent) = body.parent_revision.as_ref()
&& let Some(current) = remap.get(parent)
{
body.parent_revision = Some(current.clone());
}
}
for commit in commits.iter_mut() {
if let Some(current) = remap.get(&commit.parent_revision) {
commit.parent_revision = current.clone();
}
if let Some(current) = remap.get(&commit.revision) {
commit.revision = current.clone();
}
heal_legacy_commit_span_digests(commit, revisions)?;
}
if let Some(head) = head
&& let Some(current) = remap.get(head.as_str())
{
*head = current.clone();
}
Ok(())
}
fn heal_legacy_commit_span_digests(
commit: &mut TranscriptRewriteCommit,
revisions: &[TranscriptRevisionBody],
) -> Result<(), serde_json::Error> {
let Some(parent_body) = revisions
.iter()
.find(|body| body.revision == commit.parent_revision)
else {
return Ok(());
};
let Some(revision_body) = revisions
.iter()
.find(|body| body.revision == commit.revision)
else {
return Ok(());
};
let (start, end) = commit.selection.bounds();
if start > end || end > parent_body.messages.len() {
return Ok(());
}
let removed_len = end - start;
let Some(retained_len) = commit.messages_before.checked_sub(removed_len) else {
return Ok(());
};
let Some(replacement_len) = commit.messages_after.checked_sub(retained_len) else {
return Ok(());
};
let Some(replacement_end) = start.checked_add(replacement_len) else {
return Ok(());
};
if replacement_end > revision_body.messages.len() {
return Ok(());
}
let original_span = &parent_body.messages[start..end];
if commit.original_span_digest == legacy_transcript_messages_digest(original_span)? {
commit.original_span_digest = transcript_messages_digest(original_span)?;
}
let replacement_span = &revision_body.messages[start..replacement_end];
if commit.replacement_digest == legacy_transcript_messages_digest(replacement_span)? {
commit.replacement_digest = transcript_messages_digest(replacement_span)?;
}
Ok(())
}
pub(super) fn heal_legacy_compaction_rewrite_semantics(
commits: &mut [TranscriptRewriteCommit],
revisions: &[TranscriptRevisionBody],
) {
for commit in commits {
if !commit.selection.is_legacy_untyped() {
continue;
}
let (start, end) = commit.selection.bounds();
if start != 0
|| end != commit.messages_before
|| commit.messages_after >= commit.messages_before
{
continue;
}
let Some(parent) = revisions
.iter()
.find(|body| body.revision == commit.parent_revision)
else {
continue;
};
let Some(revision) = revisions
.iter()
.find(|body| body.revision == commit.revision)
else {
continue;
};
if parent.messages.len() != commit.messages_before
|| revision.messages.len() != commit.messages_after
{
continue;
}
let summary_count = revision
.messages
.iter()
.filter(|message| {
matches!(message, Message::User(user) if user.transcript_role.is_compaction_summary())
})
.count();
if summary_count == 1 {
commit.selection = TranscriptRewriteSelection::migrated_legacy_compaction(start, end);
}
}
}
pub(crate) fn legacy_transcript_messages_digest(
messages: &[Message],
) -> Result<String, serde_json::Error> {
sha256_json_digest(&canonicalize_message_images_for_digest(messages))
}