loro_internal/

encoding.rs

pub(crate) mod arena;
pub(crate) mod fast_snapshot;
pub(crate) mod json_schema;
mod outdated_encode_reordered;
mod shallow_snapshot;
pub(crate) mod value;
pub(crate) mod value_register;
pub(crate) use outdated_encode_reordered::{
    decode_op, encode_op, get_op_prop, EncodedDeleteStartId, IterableEncodedDeleteStartId,
};
use outdated_encode_reordered::{import_changes_to_oplog, ImportChangesResult};
pub(crate) use value::OwnedValue;

use crate::change::Change;
use crate::version::{Frontiers, VersionRange};
use crate::LoroDoc;
use crate::{oplog::OpLog, LoroError, VersionVector};
use loro_common::{HasIdSpan, IdSpan, InternalString, LoroEncodeError, LoroResult, ID};
use num_traits::{FromPrimitive, ToPrimitive};
use std::borrow::Cow;

/// The mode of the export.
///
/// Loro CRDT internally consists of two parts: document history and current document state.
/// The export modes offer various options to meet different requirements.
///
/// - CRDT property: Documents maintain consistent states when they receive the same set of updates.
/// - In real-time collaboration, peers typically only need to synchronize updates
///   (operations/history) to achieve consistency.
///
/// ## Update Export
///
/// - Exports only the history part, containing multiple operations.
/// - Suitable for real-time collaboration scenarios where peers only need to synchronize updates.
///
/// ## Snapshot Export
///
/// ### Default Snapshot
///
/// - Includes complete history and current full state.
///
/// ### Shallow Snapshot
///
/// - Contains the complete current state.
/// - Retains partial history starting from a specified version.
///
/// ### State-only Snapshot
///
/// - Exports the state of the target version.
/// - Includes a minimal set of operation history.
#[non_exhaustive]
#[derive(Debug, Clone)]
pub enum ExportMode<'a> {
    /// It contains the full history and the current state of the document.
    Snapshot,
    /// It contains the history since the `from` version vector.
    Updates { from: Cow<'a, VersionVector> },
    /// This mode exports the history in the specified range.
    UpdatesInRange { spans: Cow<'a, [IdSpan]> },
    /// The shallow snapshot only contains the history since the target frontiers
    ShallowSnapshot(Cow<'a, Frontiers>),
    /// The state only snapshot exports the state of the target version
    /// with a minimal set of history (a few ops).
    ///
    /// It's a shallow snapshot with depth=1 at the target version.
    /// If the target version is None, it will use the latest version as the target version.
    StateOnly(Option<Cow<'a, Frontiers>>),
    /// The snapshot at the specified frontiers. It contains the full history
    /// till the target frontiers and the state at the target frontiers.
    SnapshotAt { version: Cow<'a, Frontiers> },
}

impl<'a> ExportMode<'a> {
    /// It contains the full history and the current state of the document.
    pub fn snapshot() -> Self {
        ExportMode::Snapshot
    }

    /// It contains the history since the `from` version vector.
    pub fn updates(from: &'a VersionVector) -> Self {
        ExportMode::Updates {
            from: Cow::Borrowed(from),
        }
    }

    /// It contains the history since the `from` version vector.
    pub fn updates_owned(from: VersionVector) -> Self {
        ExportMode::Updates {
            from: Cow::Owned(from),
        }
    }

    /// It contains all the history of the document.
    pub fn all_updates() -> Self {
        ExportMode::Updates {
            from: Cow::Owned(Default::default()),
        }
    }

    /// This mode exports the history in the specified range.
    pub fn updates_in_range(spans: impl Into<Cow<'a, [IdSpan]>>) -> Self {
        ExportMode::UpdatesInRange {
            spans: spans.into(),
        }
    }

    /// The shallow snapshot only contains the history since the target frontiers.
    pub fn shallow_snapshot(frontiers: &'a Frontiers) -> Self {
        ExportMode::ShallowSnapshot(Cow::Borrowed(frontiers))
    }

    /// The shallow snapshot only contains the history since the target frontiers.
    pub fn shallow_snapshot_owned(frontiers: Frontiers) -> Self {
        ExportMode::ShallowSnapshot(Cow::Owned(frontiers))
    }

    /// The shallow snapshot only contains the history since the target frontiers.
    pub fn shallow_snapshot_since(id: ID) -> Self {
        let frontiers = Frontiers::from_id(id);
        ExportMode::ShallowSnapshot(Cow::Owned(frontiers))
    }

    /// The state only snapshot exports the state of the target version
    /// with a minimal set of history (a few ops).
    ///
    /// It's a shallow snapshot with depth=1 at the target version.
    /// If the target version is None, it will use the latest version as the target version.
    pub fn state_only(frontiers: Option<&'a Frontiers>) -> Self {
        ExportMode::StateOnly(frontiers.map(Cow::Borrowed))
    }

    /// The snapshot at the specified frontiers. It contains the full history
    /// till the target frontiers and the state at the target frontiers.
    pub fn snapshot_at(frontiers: &'a Frontiers) -> Self {
        ExportMode::SnapshotAt {
            version: Cow::Borrowed(frontiers),
        }
    }

    /// This mode exports the history within the specified version vector.
    pub fn updates_till(vv: &VersionVector) -> ExportMode<'static> {
        let mut spans = Vec::with_capacity(vv.len());
        for (peer, counter) in vv.iter() {
            if *counter > 0 {
                spans.push(IdSpan::new(*peer, 0, *counter));
            }
        }

        ExportMode::UpdatesInRange {
            spans: Cow::Owned(spans),
        }
    }
}

const MAGIC_BYTES: [u8; 4] = *b"loro";

#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum EncodeMode {
    // This is a config option, it won't be used in encoding.
    Auto = 255,
    OutdatedRle = 1,
    OutdatedSnapshot = 2,
    FastSnapshot = 3,
    FastUpdates = 4,
}

impl num_traits::FromPrimitive for EncodeMode {
    #[allow(trivial_numeric_casts)]
    #[inline]
    fn from_i64(n: i64) -> Option<Self> {
        match n {
            n if n == EncodeMode::Auto as i64 => Some(EncodeMode::Auto),
            n if n == EncodeMode::OutdatedRle as i64 => Some(EncodeMode::OutdatedRle),
            n if n == EncodeMode::OutdatedSnapshot as i64 => Some(EncodeMode::OutdatedSnapshot),
            n if n == EncodeMode::FastSnapshot as i64 => Some(EncodeMode::FastSnapshot),
            n if n == EncodeMode::FastUpdates as i64 => Some(EncodeMode::FastUpdates),
            _ => None,
        }
    }
    #[inline]
    fn from_u64(n: u64) -> Option<Self> {
        Self::from_i64(n as i64)
    }
}

impl num_traits::ToPrimitive for EncodeMode {
    #[inline]
    #[allow(trivial_numeric_casts)]
    fn to_i64(&self) -> Option<i64> {
        Some(match *self {
            EncodeMode::Auto => EncodeMode::Auto as i64,
            EncodeMode::OutdatedRle => EncodeMode::OutdatedRle as i64,
            EncodeMode::OutdatedSnapshot => EncodeMode::OutdatedSnapshot as i64,
            EncodeMode::FastSnapshot => EncodeMode::FastSnapshot as i64,
            EncodeMode::FastUpdates => EncodeMode::FastUpdates as i64,
        })
    }
    #[inline]
    fn to_u64(&self) -> Option<u64> {
        self.to_i64().map(|x| x as u64)
    }
}

impl EncodeMode {
    pub fn to_bytes(self) -> [u8; 2] {
        let value = self.to_u16().unwrap();
        value.to_be_bytes()
    }

    pub fn is_snapshot(self) -> bool {
        matches!(
            self,
            EncodeMode::OutdatedSnapshot | EncodeMode::FastSnapshot
        )
    }
}

impl TryFrom<[u8; 2]> for EncodeMode {
    type Error = LoroError;

    fn try_from(value: [u8; 2]) -> Result<Self, Self::Error> {
        let value = u16::from_be_bytes(value);
        Self::from_u16(value).ok_or(LoroError::IncompatibleFutureEncodingError(value as usize))
    }
}

#[derive(Debug, Clone, PartialEq, Eq, Default)]
pub struct ImportStatus {
    pub success: VersionRange,
    pub pending: Option<VersionRange>,
}

pub(crate) fn decode_oplog(
    oplog: &mut OpLog,
    parsed: ParsedHeaderAndBody,
) -> Result<ImportStatus, LoroError> {
    let changes = decode_oplog_changes(oplog, parsed)?;
    let result = apply_decoded_changes_to_oplog(oplog, changes);
    if result.has_deps_before_shallow_root {
        return Err(LoroError::ImportUpdatesThatDependsOnOutdatedVersion);
    }

    Ok(result.status)
}

pub(crate) fn decode_oplog_changes(
    oplog: &mut OpLog,
    parsed: ParsedHeaderAndBody,
) -> Result<Vec<Change>, LoroError> {
    let ParsedHeaderAndBody { mode, body, .. } = parsed;
    match mode {
        EncodeMode::OutdatedRle | EncodeMode::OutdatedSnapshot => {
            Err(LoroError::ImportUnsupportedEncodingMode)
        }
        EncodeMode::FastSnapshot => fast_snapshot::decode_oplog(oplog, body),
        EncodeMode::FastUpdates => fast_snapshot::decode_updates(oplog, body.to_vec().into()),
        EncodeMode::Auto => unreachable!(),
    }
}

pub(crate) struct ApplyDecodedChangesResult {
    pub status: ImportStatus,
    pub has_deps_before_shallow_root: bool,
}

pub(crate) fn apply_decoded_changes_to_oplog(
    oplog: &mut OpLog,
    changes: Vec<Change>,
) -> ApplyDecodedChangesResult {
    let ImportChangesResult {
        mut imported,
        latest_ids,
        pending_changes,
        changes_that_have_deps_before_shallow_root,
    } = import_changes_to_oplog(changes, oplog);

    let mut pending = VersionRange::default();
    pending_changes.iter().for_each(|c| {
        pending.extends_to_include_id_span(c.id_span());
    });
    // TODO: PERF: should we use hashmap to filter latest_ids with the same peer first?
    oplog.try_apply_pending(latest_ids, Some(&mut imported));
    oplog
        .import_unknown_lamport_pending_changes(pending_changes)
        .expect("importing unknown-lamport pending changes is infallible");
    ApplyDecodedChangesResult {
        status: ImportStatus {
            success: imported,
            pending: (!pending.is_empty()).then_some(pending),
        },
        has_deps_before_shallow_root: !changes_that_have_deps_before_shallow_root.is_empty(),
    }
}

pub(crate) struct ParsedHeaderAndBody<'a> {
    pub checksum: [u8; 16],
    pub checksum_body: &'a [u8],
    pub mode: EncodeMode,
    pub body: &'a [u8],
}

const XXH_SEED: u32 = u32::from_le_bytes(*b"LORO");
impl ParsedHeaderAndBody<'_> {
    /// Return if the checksum is correct.
    fn check_checksum(&self) -> LoroResult<()> {
        match self.mode {
            EncodeMode::OutdatedRle | EncodeMode::OutdatedSnapshot => {
                if md5::compute(self.checksum_body).0 != self.checksum {
                    return Err(LoroError::DecodeChecksumMismatchError);
                }
            }
            EncodeMode::FastSnapshot | EncodeMode::FastUpdates => {
                let mut expected_bytes = [0; 4];
                expected_bytes.copy_from_slice(&self.checksum[12..16]);
                let expected = u32::from_le_bytes(expected_bytes);
                if xxhash_rust::xxh32::xxh32(self.checksum_body, XXH_SEED) != expected {
                    return Err(LoroError::DecodeChecksumMismatchError);
                }
            }
            EncodeMode::Auto => {
                return Err(LoroError::DecodeError(
                    "Invalid import mode `Auto` in encoded blob"
                        .to_string()
                        .into_boxed_str(),
                ));
            }
        }

        Ok(())
    }
}

const MIN_HEADER_SIZE: usize = 22;
pub(crate) fn parse_header_and_body(
    bytes: &[u8],
    check_checksum: bool,
) -> Result<ParsedHeaderAndBody<'_>, LoroError> {
    let reader = &bytes;
    if bytes.len() < MIN_HEADER_SIZE {
        return Err(LoroError::DecodeError("Invalid import data".into()));
    }

    let (magic_bytes, reader) = reader.split_at(4);
    if magic_bytes != MAGIC_BYTES {
        return Err(LoroError::DecodeError("Invalid magic bytes".into()));
    }

    let (checksum, reader) = reader.split_at(16);
    let checksum_body = reader;
    let (mode_bytes, reader) = reader.split_at(2);
    let mode: EncodeMode = [mode_bytes[0], mode_bytes[1]].try_into()?;
    if mode == EncodeMode::Auto {
        return Err(LoroError::DecodeError(
            "Invalid import mode `Auto` in encoded blob"
                .to_string()
                .into_boxed_str(),
        ));
    }
    let mut checksum_arr = [0; 16];
    checksum_arr.copy_from_slice(checksum);

    let ans = ParsedHeaderAndBody {
        mode,
        checksum_body,
        checksum: checksum_arr,
        body: reader,
    };

    if check_checksum {
        ans.check_checksum()?;
    }
    Ok(ans)
}

pub(crate) fn export_fast_snapshot(doc: &LoroDoc) -> Result<Vec<u8>, LoroEncodeError> {
    let snapshot = fast_snapshot::encode_snapshot_inner(doc)?;
    let expected_len = snapshot
        .encoded_len()
        .and_then(|len| MIN_HEADER_SIZE.checked_add(len))
        .ok_or_else(|| LoroEncodeError::internal("snapshot length overflow"))?;
    let encoded = encode_with_capacity(EncodeMode::FastSnapshot, expected_len, &mut |ans| {
        fast_snapshot::_encode_snapshot(&snapshot, ans);
        Ok(())
    })?;
    debug_assert_eq!(encoded.len(), expected_len);
    Ok(encoded)
}

pub(crate) fn export_snapshot_at(
    doc: &LoroDoc,
    frontiers: &Frontiers,
) -> Result<Vec<u8>, LoroEncodeError> {
    check_target_version_reachable(doc, frontiers)?;
    encode_with(EncodeMode::FastSnapshot, &mut |ans| {
        shallow_snapshot::encode_snapshot_at(doc, frontiers, ans)
    })
}

pub(crate) fn export_fast_updates(doc: &LoroDoc, vv: &VersionVector) -> Vec<u8> {
    encode_with(EncodeMode::FastUpdates, &mut |ans| {
        fast_snapshot::encode_updates(doc, vv, ans);
        Ok(())
    })
    .unwrap()
}

pub(crate) fn export_fast_updates_in_range(oplog: &OpLog, spans: &[IdSpan]) -> Vec<u8> {
    encode_with(EncodeMode::FastUpdates, &mut |ans| {
        fast_snapshot::encode_updates_in_range(oplog, spans, ans);
        Ok(())
    })
    .unwrap()
}

pub(crate) fn export_shallow_snapshot(
    doc: &LoroDoc,
    f: &Frontiers,
) -> Result<Vec<u8>, LoroEncodeError> {
    check_target_version_reachable(doc, f)?;
    encode_with(EncodeMode::FastSnapshot, &mut |ans| {
        shallow_snapshot::export_shallow_snapshot(doc, f, ans)?;
        Ok(())
    })
}

fn check_target_version_reachable(doc: &LoroDoc, f: &Frontiers) -> Result<(), LoroEncodeError> {
    let oplog = doc.oplog.lock();
    if !oplog.dag.can_export_shallow_snapshot_on(f) {
        return Err(LoroEncodeError::FrontiersNotFound(format!("{f:?}")));
    }

    Ok(())
}

pub(crate) fn export_state_only_snapshot(
    doc: &LoroDoc,
    f: &Frontiers,
) -> Result<Vec<u8>, LoroEncodeError> {
    check_target_version_reachable(doc, f)?;
    encode_with(EncodeMode::FastSnapshot, &mut |ans| {
        shallow_snapshot::export_state_only_snapshot(doc, f, ans)?;
        Ok(())
    })
}

fn encode_with(
    mode: EncodeMode,
    f: &mut dyn FnMut(&mut Vec<u8>) -> Result<(), LoroEncodeError>,
) -> Result<Vec<u8>, LoroEncodeError> {
    encode_with_capacity(mode, MIN_HEADER_SIZE, f)
}

fn encode_with_capacity(
    mode: EncodeMode,
    capacity: usize,
    f: &mut dyn FnMut(&mut Vec<u8>) -> Result<(), LoroEncodeError>,
) -> Result<Vec<u8>, LoroEncodeError> {
    // HEADER
    let mut ans = Vec::with_capacity(capacity);
    ans.extend(MAGIC_BYTES);
    let checksum = [0; 16];
    ans.extend(checksum);
    ans.extend(mode.to_bytes());

    // BODY
    f(&mut ans)?;

    // CHECKSUM in HEADER
    let checksum_body = &ans[20..];
    let checksum = xxhash_rust::xxh32::xxh32(checksum_body, XXH_SEED);
    ans[16..20].copy_from_slice(&checksum.to_le_bytes());
    Ok(ans)
}

pub(crate) fn decode_snapshot(
    doc: &LoroDoc,
    mode: EncodeMode,
    body: &[u8],
    origin: InternalString,
) -> Result<ImportStatus, LoroError> {
    match mode {
        EncodeMode::OutdatedSnapshot => {
            return Err(LoroError::ImportUnsupportedEncodingMode);
        }
        EncodeMode::FastSnapshot => {
            fast_snapshot::decode_snapshot(doc, body.to_vec().into(), origin)?
        }
        _ => {
            return Err(LoroError::DecodeError(
                format!("Invalid snapshot encoding mode: {mode:?}").into_boxed_str(),
            ));
        }
    };
    Ok(ImportStatus {
        success: VersionRange::from_vv(&doc.oplog_vv()),
        pending: None,
    })
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord)]
pub enum EncodedBlobMode {
    Snapshot,
    OutdatedSnapshot,
    ShallowSnapshot,
    OutdatedRle,
    Updates,
}

impl std::fmt::Display for EncodedBlobMode {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_str(match self {
            EncodedBlobMode::OutdatedRle => "outdated-update",
            EncodedBlobMode::OutdatedSnapshot => "outdated-snapshot",
            EncodedBlobMode::Snapshot => "snapshot",
            EncodedBlobMode::ShallowSnapshot => "shallow-snapshot",
            EncodedBlobMode::Updates => "update",
        })
    }
}

impl EncodedBlobMode {
    pub fn is_snapshot(&self) -> bool {
        matches!(
            self,
            EncodedBlobMode::Snapshot
                | EncodedBlobMode::ShallowSnapshot
                | EncodedBlobMode::OutdatedSnapshot
        )
    }
}

#[derive(Debug, Clone)]
pub struct ImportBlobMetadata {
    /// The partial start version vector.
    ///
    /// Import blob includes all the ops from `partial_start_vv` to `partial_end_vv`.
    /// However, it does not constitute a complete version vector, as it only contains counters
    /// from peers included within the import blob.
    pub partial_start_vv: VersionVector,
    /// The partial end version vector.
    ///
    /// Import blob includes all the ops from `partial_start_vv` to `partial_end_vv`.
    /// However, it does not constitute a complete version vector, as it only contains counters
    /// from peers included within the import blob.
    pub partial_end_vv: VersionVector,
    pub start_timestamp: i64,
    pub start_frontiers: Frontiers,
    pub end_timestamp: i64,
    pub change_num: u32,
    pub mode: EncodedBlobMode,
}

impl LoroDoc {
    /// Decodes the metadata for an imported blob from the provided bytes.
    pub fn decode_import_blob_meta(
        blob: &[u8],
        check_checksum: bool,
    ) -> LoroResult<ImportBlobMetadata> {
        let parsed = parse_header_and_body(blob, check_checksum)?;
        match parsed.mode {
            EncodeMode::Auto => unreachable!(),
            EncodeMode::OutdatedRle | EncodeMode::OutdatedSnapshot => {
                Err(LoroError::ImportUnsupportedEncodingMode)
            }
            EncodeMode::FastSnapshot => fast_snapshot::decode_snapshot_blob_meta(parsed),
            EncodeMode::FastUpdates => fast_snapshot::decode_updates_blob_meta(parsed),
        }
    }
}

#[cfg(test)]
mod test {
    use super::*;
    use loro_common::{loro_value, ContainerID, ContainerType, LoroValue, ID};

    #[test]
    fn fast_snapshot_envelope_has_exact_length_and_valid_checksum() {
        let doc = LoroDoc::new_auto_commit();
        doc.get_map("root").insert("key", "value").unwrap();
        let encoded = doc.export(ExportMode::Snapshot).unwrap();
        assert_eq!(&encoded[..4], &MAGIC_BYTES);
        assert_eq!(&encoded[20..22], &EncodeMode::FastSnapshot.to_bytes());
        let parsed = parse_header_and_body(&encoded, true).unwrap();
        assert_eq!(parsed.mode, EncodeMode::FastSnapshot);
        assert_eq!(encoded.len(), MIN_HEADER_SIZE + parsed.body.len());

        let mut corrupted = encoded;
        *corrupted.last_mut().unwrap() ^= 1;
        assert!(matches!(
            parse_header_and_body(&corrupted, true),
            Err(LoroError::DecodeChecksumMismatchError)
        ));
    }

    #[test]
    fn test_value_encode_size() {
        fn assert_size(value: LoroValue, max_size: usize) {
            let size = postcard::to_allocvec(&value).unwrap().len();
            assert!(
                size <= max_size,
                "value: {:?}, size: {}, max_size: {}",
                value,
                size,
                max_size
            );
        }

        assert_size(LoroValue::Null, 1);
        assert_size(LoroValue::I64(1), 2);
        assert_size(LoroValue::Double(1.), 9);
        assert_size(LoroValue::Bool(true), 2);
        assert_size(LoroValue::String("123".to_string().into()), 5);
        assert_size(LoroValue::Binary(vec![1, 2, 3].into()), 5);
        assert_size(
            loro_value!({
                "a": 1,
                "b": 2,
            }),
            10,
        );
        assert_size(loro_value!([1, 2, 3]), 8);
        assert_size(
            LoroValue::Container(ContainerID::new_normal(ID::new(1, 1), ContainerType::Map)),
            5,
        );
        assert_size(
            LoroValue::Container(ContainerID::new_root("a", ContainerType::Map)),
            5,
        );
    }
}