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
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153

// Copyright (c) The Diem Core Contributors
// SPDX-License-Identifier: Apache-2.0

// Copyright 2021 Conflux Foundation. All rights reserved.
// Conflux is free software and distributed under GNU General Public License.
// See http://www.gnu.org/licenses/

//! This module provides traits that define the behavior of a schema and its
//! associated key and value types, along with helpers to define a new schema
//! with ease.
use crate::ColumnFamilyName;
use anyhow::Result;
use std::fmt::Debug;

/// Macro for defining a SchemaDB schema.
///
/// `define_schema!` allows a schema to be defined in the following syntax:
/// ```
/// use anyhow::Result;
/// use schemadb::{
///     define_schema,
///     schema::{KeyCodec, SeekKeyCodec, ValueCodec},
/// };
///
/// // Define key type and value type for a schema with derived traits (Clone, Debug, Eq, PartialEq)
/// #[derive(Clone, Debug, Eq, PartialEq)]
/// pub struct Key;
/// #[derive(Clone, Debug, Eq, PartialEq)]
/// pub struct Value;
///
/// // Implement KeyCodec/ValueCodec traits for key and value types
/// impl KeyCodec<ExampleSchema> for Key {
///     fn encode_key(&self) -> Result<Vec<u8>> {
///         Ok(vec![])
///     }
///
///     fn decode_key(data: &[u8]) -> Result<Self> {
///         Ok(Key)
///     }
/// }
///
/// impl ValueCodec<ExampleSchema> for Value {
///     fn encode_value(&self) -> Result<Vec<u8>> {
///         Ok(vec![])
///     }
///
///     fn decode_value(data: &[u8]) -> Result<Self> {
///         Ok(Value)
///     }
/// }
///
/// // And finally define a schema type and associate it with key and value types, as well as the
/// // column family name, by generating code that implements the `Schema` trait for the type.
/// define_schema!(ExampleSchema, Key, Value, "exmaple_cf_name");
///
/// // SeekKeyCodec is automatically implemented for KeyCodec,
/// // so you can seek an iterator with the Key type:
/// // iter.seek(&Key);
///
/// // Or if seek-by-prefix is desired, you can implement your own SeekKey
/// #[derive(Clone, Eq, PartialEq, Debug)]
/// pub struct PrefixSeekKey;
///
/// impl SeekKeyCodec<ExampleSchema> for PrefixSeekKey {
///     fn encode_seek_key(&self) -> Result<Vec<u8>> {
///         Ok(vec![])
///     }
/// }
/// // and seek like this:
/// // iter.seek(&PrefixSeekKey);
/// ```
#[macro_export]
macro_rules! define_schema {
    ($schema_type:ident, $key_type:ty, $value_type:ty, $cf_name:expr) => {
        pub(crate) struct $schema_type;

        impl $crate::schema::Schema for $schema_type {
            type Key = $key_type;
            type Value = $value_type;

            const COLUMN_FAMILY_NAME: $crate::ColumnFamilyName = $cf_name;
        }
    };
}

/// This trait defines a type that can serve as a [`Schema::Key`].
pub trait KeyCodec<S: Schema + ?Sized>: Sized + PartialEq + Debug {
    /// Converts `self` to bytes to be stored in DB.
    fn encode_key(&self) -> Result<Vec<u8>>;
    /// Converts bytes fetched from DB to `Self`.
    fn decode_key(data: &[u8]) -> Result<Self>;
}

/// This trait defines a type that can serve as a [`Schema::Value`].
pub trait ValueCodec<S: Schema + ?Sized>: Sized + PartialEq + Debug {
    /// Converts `self` to bytes to be stored in DB.
    fn encode_value(&self) -> Result<Vec<u8>>;
    /// Converts bytes fetched from DB to `Self`.
    fn decode_value(data: &[u8]) -> Result<Self>;
}

/// This defines a type that can be used to seek a
/// [`SchemaIterator`](crate::SchemaIterator), via interfaces like
/// [`seek`](crate::SchemaIterator::seek).
pub trait SeekKeyCodec<S: Schema + ?Sized>: Sized {
    /// Converts `self` to bytes which is used to seek the underlying raw
    /// iterator.
    fn encode_seek_key(&self) -> Result<Vec<u8>>;
}

/// All keys can automatically be used as seek keys.
impl<S, K> SeekKeyCodec<S> for K
where
    S: Schema,
    K: KeyCodec<S>,
{
    /// Delegates to [`KeyCodec::encode_key`].
    fn encode_seek_key(&self) -> Result<Vec<u8>> {
        <K as KeyCodec<S>>::encode_key(&self)
    }
}

/// This trait defines a schema: an association of a column family name, the key
/// type and the value type.
pub trait Schema {
    /// The column family name associated with this struct.
    /// Note: all schemas within the same SchemaDB must have distinct column
    /// family names.
    const COLUMN_FAMILY_NAME: ColumnFamilyName;

    /// Type of the key.
    type Key: KeyCodec<Self>;
    /// Type of the value.
    type Value: ValueCodec<Self>;
}

/// Helper used in tests to assert a (key, value) pair for a certain [`Schema`]
/// is able to convert to bytes and convert back.
pub fn assert_encode_decode<S: Schema>(key: &S::Key, value: &S::Value) {
    {
        let encoded = key.encode_key().expect("Encoding key should work.");
        let decoded =
            S::Key::decode_key(&encoded).expect("Decoding key should work.");
        assert_eq!(*key, decoded);
    }
    {
        let encoded =
            value.encode_value().expect("Encoding value should work.");
        let decoded = S::Value::decode_value(&encoded)
            .expect("Decoding value should work.");
        assert_eq!(*value, decoded);
    }
}