ruma_lean/

lib.rs

// Copyright 2026 Shane Jaroch
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

#![no_std]

extern crate alloc;

use alloc::collections::{BTreeMap, BinaryHeap};
use alloc::string::String;
use alloc::vec::Vec;
use core::cmp::Ordering;
use serde::{Deserialize, Serialize};

#[cfg(feature = "std")]
extern crate std;

#[cfg(feature = "std")]
pub use std::collections::HashMap;

#[cfg(not(feature = "std"))]
pub use hashbrown::HashMap;

/// The version of the Matrix State Resolution algorithm to use.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Serialize, Deserialize)]
pub enum StateResVersion {
    V1,
    V2,
    V2_1,
}

/// A lightweight Matrix Event representation for Lean-equivalent resolution.
#[derive(Debug, Clone, Eq, PartialEq, Serialize, Deserialize)]
pub struct LeanEvent {
    pub event_id: String,
    pub power_level: i64,
    pub origin_server_ts: u64,
    pub prev_events: Vec<String>,
    pub depth: u64, // Required for V1
}

/// The core tie-breaking logic from Ruma Lean (StateRes.lean).
/// Matches Lean model: power_level (desc) -> origin_server_ts (asc) -> event_id (asc)
impl Ord for LeanEvent {
    fn cmp(&self, other: &Self) -> Ordering {
        // Higher power level comes FIRST (is "smaller" in terms of order)
        match other.power_level.cmp(&self.power_level) {
            Ordering::Equal => {
                // Earlier timestamp comes FIRST
                match self.origin_server_ts.cmp(&other.origin_server_ts) {
                    Ordering::Equal => {
                        // Lexicographically smaller ID comes FIRST
                        self.event_id.cmp(&other.event_id)
                    }
                    ord => ord,
                }
            }
            ord => ord,
        }
    }
}

impl PartialOrd for LeanEvent {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

/// A wrapper to ensure BinaryHeap pops the "smallest" (best) event first.
#[derive(Debug, Clone, Copy, Eq, PartialEq)]
struct SortPriority<'a> {
    event: &'a LeanEvent,
    version: StateResVersion,
}

impl<'a> Ord for SortPriority<'a> {
    fn cmp(&self, other: &Self) -> Ordering {
        match self.version {
            StateResVersion::V1 => {
                // V1 tie-breaking: depth (asc) -> event_id (asc)
                // Inverted for Max-Heap
                match other.event.depth.cmp(&self.event.depth) {
                    Ordering::Equal => other.event.event_id.cmp(&self.event.event_id),
                    ord => ord,
                }
            }
            StateResVersion::V2 | StateResVersion::V2_1 => {
                // V2 tie-breaking: power_level (desc) -> origin_server_ts (asc) -> event_id (asc)
                // Priority popping (best first)
                match self.event.power_level.cmp(&other.event.power_level) {
                    Ordering::Equal => {
                        match other
                            .event
                            .origin_server_ts
                            .cmp(&self.event.origin_server_ts)
                        {
                            Ordering::Equal => other.event.event_id.cmp(&self.event.event_id),
                            ord => ord,
                        }
                    }
                    ord => ord,
                }
            }
        }
    }
}

impl<'a> PartialOrd for SortPriority<'a> {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

/// A simplified implementation of Kahn's Topological Sort.
pub fn lean_kahn_sort(
    events: &HashMap<String, LeanEvent>,
    version: StateResVersion,
) -> Vec<String> {
    let mut in_degree: HashMap<String, usize> = HashMap::new();
    let mut adjacency: HashMap<String, Vec<String>> = HashMap::new();

    for (id, event) in events {
        in_degree.entry(id.clone()).or_insert(0);
        for prev in &event.prev_events {
            if events.contains_key(prev) {
                adjacency.entry(prev.clone()).or_default().push(id.clone());
                *in_degree.entry(id.clone()).or_insert(0) += 1;
            }
        }
    }

    let mut queue: BinaryHeap<SortPriority> = BinaryHeap::new();
    for (id, &degree) in &in_degree {
        if degree == 0 {
            if let Some(event) = events.get(id) {
                queue.push(SortPriority { event, version });
            }
        }
    }

    let mut result = Vec::new();
    while let Some(priority) = queue.pop() {
        let event = priority.event;
        result.push(event.event_id.clone());
        if let Some(neighbors) = adjacency.get(&event.event_id) {
            for next_id in neighbors {
                let degree = in_degree.get_mut(next_id).unwrap();
                *degree -= 1;
                if *degree == 0 {
                    queue.push(SortPriority {
                        event: events.get(next_id).unwrap(),
                        version,
                    });
                }
            }
        }
    }
    result
}

pub fn resolve_lean(
    unconflicted_state: BTreeMap<(String, String), String>,
    conflicted_events: HashMap<String, LeanEvent>,
    version: StateResVersion,
) -> BTreeMap<(String, String), String> {
    let resolved = unconflicted_state;
    let _sorted_ids = lean_kahn_sort(&conflicted_events, version);
    resolved
}

#[cfg(feature = "zkvm")]
pub fn verify_signature(_public_key: &[u8; 32], _message: &[u8], _signature: &[u8; 64]) {
    // Verifiable signature check for ZKVM environment
}

#[cfg(all(feature = "std", not(feature = "zkvm")))]
pub fn verify_signature(public_key: &[u8; 32], message: &[u8], signature: &[u8; 64]) {
    use ed25519_consensus::{Signature, VerificationKey};
    let vk = VerificationKey::try_from(*public_key).expect("Invalid public key");
    let sig = Signature::from(*signature);
    vk.verify(&sig, message)
        .expect("Signature verification failed");
}

#[cfg(all(not(feature = "std"), not(feature = "zkvm")))]
pub fn verify_signature(_public_key: &[u8; 32], _message: &[u8], _signature: &[u8; 64]) {
    // No-op for other configurations
}

#[cfg(test)]
mod tests {
    use super::*;
    use alloc::string::ToString;
    use alloc::vec;

    #[cfg(not(feature = "std"))]
    use hashbrown::HashMap;
    #[cfg(feature = "std")]
    use std::collections::HashMap;

    #[test]
    fn test_v1_resolution_happy_path() {
        let mut events = HashMap::new();
        events.insert(
            "A".into(),
            LeanEvent {
                event_id: "A".into(),
                power_level: 0,
                origin_server_ts: 100,
                prev_events: vec![],
                depth: 1,
            },
        );
        events.insert(
            "B".into(),
            LeanEvent {
                event_id: "B".into(),
                power_level: 0,
                origin_server_ts: 50,
                prev_events: vec![],
                depth: 2,
            },
        );
        let sorted = lean_kahn_sort(&events, StateResVersion::V1);
        assert_eq!(sorted, vec!["A", "B"]);
    }

    #[test]
    fn test_v1_tie_break_by_id() {
        let mut events = HashMap::new();
        events.insert(
            "B".into(),
            LeanEvent {
                event_id: "B".into(),
                power_level: 0,
                origin_server_ts: 100,
                prev_events: vec![],
                depth: 1,
            },
        );
        events.insert(
            "A".into(),
            LeanEvent {
                event_id: "A".into(),
                power_level: 0,
                origin_server_ts: 100,
                prev_events: vec![],
                depth: 1,
            },
        );
        let sorted = lean_kahn_sort(&events, StateResVersion::V1);
        assert_eq!(sorted, vec!["A", "B"]);
    }

    #[test]
    fn test_v2_resolution_happy_path() {
        let mut events = HashMap::new();
        events.insert(
            "A".into(),
            LeanEvent {
                event_id: "A".into(),
                power_level: 100,
                origin_server_ts: 100,
                prev_events: vec![],
                depth: 10,
            },
        );
        events.insert(
            "B".into(),
            LeanEvent {
                event_id: "B".into(),
                power_level: 50,
                origin_server_ts: 10,
                prev_events: vec![],
                depth: 1,
            },
        );
        let sorted = lean_kahn_sort(&events, StateResVersion::V2);
        assert_eq!(sorted, vec!["A", "B"]);
    }

    #[test]
    fn test_v2_deep_tie_break() {
        let mut events = HashMap::new();
        events.insert(
            "B".into(),
            LeanEvent {
                event_id: "B".into(),
                power_level: 100,
                origin_server_ts: 10,
                prev_events: vec![],
                depth: 1,
            },
        );
        events.insert(
            "A".into(),
            LeanEvent {
                event_id: "A".into(),
                power_level: 100,
                origin_server_ts: 10,
                prev_events: vec![],
                depth: 1,
            },
        );
        let sorted = lean_kahn_sort(&events, StateResVersion::V2);
        assert_eq!(sorted, vec!["A", "B"]);
    }

    #[test]
    fn test_v1_v2_v2_1_comparison_determinism() {
        let mut events = HashMap::new();
        events.insert(
            "A".into(),
            LeanEvent {
                event_id: "A".into(),
                power_level: 10,
                origin_server_ts: 10,
                prev_events: vec![],
                depth: 1,
            },
        );
        events.insert(
            "B".into(),
            LeanEvent {
                event_id: "B".into(),
                power_level: 100,
                origin_server_ts: 100,
                prev_events: vec![],
                depth: 10,
            },
        );
        let sorted_v1 = lean_kahn_sort(&events, StateResVersion::V1);
        let sorted_v2 = lean_kahn_sort(&events, StateResVersion::V2);
        let sorted_v2_1 = lean_kahn_sort(&events, StateResVersion::V2_1);
        assert_eq!(sorted_v1, vec!["A", "B"]);
        assert_eq!(sorted_v2, vec!["B", "A"]);
        assert_eq!(sorted_v2_1, vec!["B", "A"]);
    }

    #[test]
    fn test_unhappy_path_cycle_detection() {
        let mut events = HashMap::new();
        events.insert(
            "A".into(),
            LeanEvent {
                event_id: "A".into(),
                power_level: 100,
                origin_server_ts: 100,
                prev_events: vec!["B".into()],
                depth: 1,
            },
        );
        events.insert(
            "B".into(),
            LeanEvent {
                event_id: "B".into(),
                power_level: 100,
                origin_server_ts: 100,
                prev_events: vec!["A".into()],
                depth: 1,
            },
        );
        let sorted = lean_kahn_sort(&events, StateResVersion::V2);
        assert!(sorted.is_empty());
    }

    #[test]
    fn test_signature_verification_failure() {
        #[cfg(all(feature = "std", not(feature = "zkvm")))]
        {
            let pk = [
                215, 90, 152, 1, 130, 177, 10, 183, 213, 75, 254, 211, 201, 100, 7, 58, 14, 225,
                114, 243, 218, 166, 35, 37, 175, 2, 26, 104, 247, 7, 81, 26,
            ];
            let sig = [0u8; 64];
            let msg = b"test";
            let result = std::panic::catch_unwind(|| {
                verify_signature(&pk, msg, &sig);
            });
            assert!(result.is_err());
        }
    }

    #[test]
    fn test_serialization_roundtrip() {
        let event = LeanEvent {
            event_id: "$abc".into(),
            power_level: 100,
            origin_server_ts: 12345,
            prev_events: vec![],
            depth: 5,
        };
        let serialized = serde_json::to_string(&event).unwrap();
        let deserialized: LeanEvent = serde_json::from_str(&serialized).unwrap();
        assert_eq!(event, deserialized);
    }

    #[test]
    fn test_partial_ord_implementations() {
        let e1 = LeanEvent {
            event_id: "a".into(),
            power_level: 100,
            origin_server_ts: 10,
            prev_events: vec![],
            depth: 1,
        };
        let e2 = LeanEvent {
            event_id: "b".into(),
            power_level: 100,
            origin_server_ts: 10,
            prev_events: vec![],
            depth: 1,
        };
        assert!(e1.partial_cmp(&e2).is_some());

        let p1 = SortPriority {
            event: &e1,
            version: StateResVersion::V2,
        };
        let p2 = SortPriority {
            event: &e2,
            version: StateResVersion::V2,
        };
        assert!(p1.partial_cmp(&p2).is_some());
    }

    #[test]
    fn test_trait_coverage() {
        let v = StateResVersion::V2;
        assert_eq!(v, StateResVersion::V2);
        let _ = alloc::format!("{:?}", v);

        let e = LeanEvent {
            event_id: "a".into(),
            power_level: 100,
            origin_server_ts: 10,
            prev_events: vec![],
            depth: 1,
        };
        let _ = e.clone();
        let _ = alloc::format!("{:?}", e);
    }

    #[test]
    fn test_complex_dag_sort() {
        let mut events = HashMap::new();
        events.insert(
            "1".into(),
            LeanEvent {
                event_id: "1".into(),
                power_level: 100,
                origin_server_ts: 10,
                prev_events: vec![],
                depth: 1,
            },
        );
        events.insert(
            "2".into(),
            LeanEvent {
                event_id: "2".into(),
                power_level: 50,
                origin_server_ts: 20,
                prev_events: vec!["1".into()],
                depth: 2,
            },
        );
        events.insert(
            "3".into(),
            LeanEvent {
                event_id: "3".into(),
                power_level: 50,
                origin_server_ts: 15,
                prev_events: vec!["1".into()],
                depth: 2,
            },
        );
        events.insert(
            "4".into(),
            LeanEvent {
                event_id: "4".into(),
                power_level: 10,
                origin_server_ts: 30,
                prev_events: vec!["2".into(), "3".into()],
                depth: 3,
            },
        );
        let sorted = lean_kahn_sort(&events, StateResVersion::V2);
        assert_eq!(sorted, vec!["1", "3", "2", "4"]);
    }

    #[test]
    fn test_kahn_missing_parents() {
        let mut events = HashMap::new();
        events.insert(
            "A".into(),
            LeanEvent {
                event_id: "A".into(),
                power_level: 100,
                origin_server_ts: 10,
                prev_events: vec!["MISSING".into()],
                depth: 1,
            },
        );
        let sorted = lean_kahn_sort(&events, StateResVersion::V2);
        assert_eq!(sorted, vec!["A"]);
    }

    #[test]
    fn test_resolve_lean_functionality() {
        let mut unconflicted = BTreeMap::new();
        unconflicted.insert(("type".into(), "key".into()), "id".into());
        let conflicted = HashMap::new();
        let resolved = resolve_lean(unconflicted.clone(), conflicted, StateResVersion::V2);
        assert_eq!(resolved, unconflicted);
    }

    fn run_batch_test(
        version: StateResVersion,
        rows: &[(&str, i64, u64, u64, &[&str])],
        expected: &[&str],
    ) {
        let mut events = HashMap::new();
        for r in rows {
            events.insert(
                r.0.to_string(),
                LeanEvent {
                    event_id: r.0.to_string(),
                    power_level: r.1,
                    origin_server_ts: r.2,
                    depth: r.3,
                    prev_events: r.4.iter().map(|s| s.to_string()).collect(),
                },
            );
        }
        let result = lean_kahn_sort(&events, version);
        assert_eq!(
            result,
            expected.iter().map(|s| s.to_string()).collect::<Vec<_>>()
        );
    }

    #[test]
    fn test_resolution_batch() {
        run_batch_test(
            StateResVersion::V2,
            &[("Alice", 100, 500, 1, &[]), ("Bob", 50, 100, 1, &[])],
            &["Alice", "Bob"],
        );
        run_batch_test(
            StateResVersion::V1,
            &[("Deep", 100, 100, 10, &[]), ("Shallow", 10, 100, 1, &[])],
            &["Shallow", "Deep"],
        );
    }

    #[test]
    fn test_native_resolution_bootstrap_parity() {
        let mut events = HashMap::new();
        events.insert(
            "1".into(),
            LeanEvent {
                event_id: "1".into(),
                power_level: 100,
                origin_server_ts: 10,
                prev_events: vec![],
                depth: 1,
            },
        );
        events.insert(
            "2".into(),
            LeanEvent {
                event_id: "2".into(),
                power_level: 0,
                origin_server_ts: 20,
                prev_events: vec!["1".into()],
                depth: 2,
            },
        );
        let sorted = lean_kahn_sort(&events, StateResVersion::V2);
        let mut resolved_state = BTreeMap::new();
        for id in sorted {
            let ev = events.get(&id).unwrap();
            let key = ("m.room.member".to_string(), "@user:example.com".to_string());
            resolved_state.insert(key, ev.event_id.clone());
        }
        assert_eq!(
            resolved_state.get(&("m.room.member".to_string(), "@user:example.com".to_string())),
            Some(&"2".to_string())
        );
    }

    #[test]
    fn test_enum_coverage() {
        let v = StateResVersion::V2;
        let v2 = v;
        assert_eq!(v, v2);
        let debug_str = alloc::format!("{:?}", v);
        assert!(debug_str.contains("V2"));
    }

    #[test]
    fn test_event_traits_coverage() {
        let e = LeanEvent {
            event_id: "a".into(),
            power_level: 100,
            origin_server_ts: 10,
            prev_events: vec![],
            depth: 1,
        };
        let e2 = e.clone();
        assert_eq!(e, e2);
        let debug_str = alloc::format!("{:?}", e);
        assert!(debug_str.contains("event_id"));
    }

    #[test]
    fn test_sort_priority_traits() {
        let e = LeanEvent {
            event_id: "a".into(),
            power_level: 100,
            origin_server_ts: 10,
            prev_events: vec![],
            depth: 1,
        };
        let p = SortPriority {
            event: &e,
            version: StateResVersion::V2,
        };
        let p2 = p;
        assert_eq!(p, p2);
        let debug_str = alloc::format!("{:?}", p);
        assert!(debug_str.contains("version"));
    }

    #[test]
    fn test_v1_equal_depth_tie_break() {
        let mut events = HashMap::new();
        events.insert(
            "B".into(),
            LeanEvent {
                event_id: "B".into(),
                power_level: 0,
                origin_server_ts: 10,
                prev_events: vec![],
                depth: 1,
            },
        );
        events.insert(
            "A".into(),
            LeanEvent {
                event_id: "A".into(),
                power_level: 0,
                origin_server_ts: 10,
                prev_events: vec![],
                depth: 1,
            },
        );
        let sorted = lean_kahn_sort(&events, StateResVersion::V1);
        assert_eq!(sorted, vec!["A", "B"]);
    }

    #[test]
    fn test_kahn_no_neighbors() {
        let mut events = HashMap::new();
        events.insert(
            "1".into(),
            LeanEvent {
                event_id: "1".into(),
                power_level: 100,
                origin_server_ts: 10,
                prev_events: vec![],
                depth: 1,
            },
        );
        let sorted = lean_kahn_sort(&events, StateResVersion::V2);
        assert_eq!(sorted, vec!["1"]);
    }

    #[test]
    fn test_v2_1_full_coverage() {
        let mut events = HashMap::new();
        events.insert(
            "A".into(),
            LeanEvent {
                event_id: "A".into(),
                power_level: 100,
                origin_server_ts: 10,
                prev_events: vec![],
                depth: 1,
            },
        );
        let sorted = lean_kahn_sort(&events, StateResVersion::V2_1);
        assert_eq!(sorted, vec!["A"]);
    }

    #[test]
    fn test_total_order_properties() {
        let e1 = LeanEvent {
            event_id: "a".into(),
            power_level: 100,
            origin_server_ts: 10,
            prev_events: vec![],
            depth: 1,
        };
        let e2 = LeanEvent {
            event_id: "b".into(),
            power_level: 100,
            origin_server_ts: 10,
            prev_events: vec![],
            depth: 1,
        };
        let e3 = LeanEvent {
            event_id: "c".into(),
            power_level: 50,
            origin_server_ts: 10,
            prev_events: vec![],
            depth: 1,
        };
        assert_eq!(e1.cmp(&e1), Ordering::Equal);
        assert!(e1 <= e1);
        assert!(e1 <= e2 || e2 <= e1);
        if e1 <= e2 && e2 <= e3 {
            assert!(e1 <= e3);
        }
        let e1_copy = e1.clone();
        if e1 <= e1_copy && e1_copy <= e1 {
            assert_eq!(e1, e1_copy);
        }
    }

    #[test]
    fn test_coverage_booster_all_branches() {
        let e_base = LeanEvent {
            event_id: "m".into(),
            power_level: 50,
            origin_server_ts: 50,
            prev_events: vec![],
            depth: 50,
        };
        let p_base = SortPriority {
            event: &e_base,
            version: StateResVersion::V2,
        };
        let e_high_power = LeanEvent {
            power_level: 100,
            ..e_base.clone()
        };
        let p_high_power = SortPriority {
            event: &e_high_power,
            version: StateResVersion::V2,
        };
        assert_eq!(p_base.cmp(&p_high_power), Ordering::Less);
        let e_early_ts = LeanEvent {
            origin_server_ts: 10,
            ..e_base.clone()
        };
        let p_early_ts = SortPriority {
            event: &e_early_ts,
            version: StateResVersion::V2,
        };
        assert_eq!(p_base.cmp(&p_early_ts), Ordering::Less);
        let e_early_id = LeanEvent {
            event_id: "a".into(),
            ..e_base.clone()
        };
        let p_early_id = SortPriority {
            event: &e_early_id,
            version: StateResVersion::V2,
        };
        assert_eq!(p_base.cmp(&p_early_id), Ordering::Less);
        let p_v1_base = SortPriority {
            event: &e_base,
            version: StateResVersion::V1,
        };
        let e_shallow = LeanEvent {
            depth: 1,
            ..e_base.clone()
        };
        let p_shallow = SortPriority {
            event: &e_shallow,
            version: StateResVersion::V1,
        };
        assert_eq!(p_v1_base.cmp(&p_shallow), Ordering::Less);
        let p_v1_early_id = SortPriority {
            event: &e_early_id,
            version: StateResVersion::V1,
        };
        assert_eq!(p_v1_base.cmp(&p_v1_early_id), Ordering::Less);
        assert_eq!(p_v1_base.cmp(&p_v1_base), Ordering::Equal);
    }
}