1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
use core::ops::ControlFlow;

use rancor::{Fallible, Source};

use crate::{
    alloc::collections::BTreeSet,
    collections::btree_set::{ArchivedBTreeSet, BTreeSetResolver},
    ser::{Allocator, Writer},
    Archive, Deserialize, Place, Serialize,
};

impl<K: Archive + Ord> Archive for BTreeSet<K>
where
    K::Archived: Ord,
{
    type Archived = ArchivedBTreeSet<K::Archived>;
    type Resolver = BTreeSetResolver;

    fn resolve(&self, resolver: Self::Resolver, out: Place<Self::Archived>) {
        ArchivedBTreeSet::<K::Archived>::resolve_from_len(
            self.len(),
            resolver,
            out,
        );
    }
}

impl<K, S> Serialize<S> for BTreeSet<K>
where
    K: Serialize<S> + Ord,
    K::Archived: Ord,
    S: Fallible + Allocator + Writer + ?Sized,
    S::Error: Source,
{
    fn serialize(
        &self,
        serializer: &mut S,
    ) -> Result<Self::Resolver, S::Error> {
        Self::Archived::serialize_from_ordered_iter::<_, K, _>(
            self.iter(),
            serializer,
        )
    }
}

impl<K, D> Deserialize<BTreeSet<K>, D> for ArchivedBTreeSet<K::Archived>
where
    K: Archive + Ord,
    K::Archived: Deserialize<K, D> + Ord,
    D: Fallible + ?Sized,
{
    fn deserialize(
        &self,
        deserializer: &mut D,
    ) -> Result<BTreeSet<K>, D::Error> {
        let mut result = BTreeSet::new();
        let r = self.visit(|ak| {
            let k = match ak.deserialize(deserializer) {
                Ok(k) => k,
                Err(e) => return ControlFlow::Break(e),
            };
            result.insert(k);
            ControlFlow::Continue(())
        });
        match r {
            Some(e) => Err(e),
            None => Ok(result),
        }
    }
}

impl<K, AK: PartialEq<K>> PartialEq<BTreeSet<K>> for ArchivedBTreeSet<AK> {
    fn eq(&self, other: &BTreeSet<K>) -> bool {
        if self.len() != other.len() {
            false
        } else {
            let mut iter = other.iter();
            self.visit(|ak| {
                if let Some(k) = iter.next() {
                    if ak.eq(k) {
                        return ControlFlow::Continue(());
                    }
                }
                ControlFlow::Break(())
            })
            .is_none()
        }
    }
}

#[cfg(test)]
mod tests {
    use crate::{
        alloc::{collections::BTreeSet, string::ToString},
        api::test::roundtrip,
    };

    #[test]
    fn roundtrip_btree_set() {
        let mut value = BTreeSet::new();
        value.insert("foo".to_string());
        value.insert("bar".to_string());
        value.insert("baz".to_string());
        value.insert("bat".to_string());

        roundtrip(&value);
    }

    #[test]
    fn roundtrip_btree_set_zst() {
        let mut value = BTreeSet::new();
        value.insert(());

        roundtrip(&value);
    }
}