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

// Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.
// SPDX-License-Identifier: Apache-2.0

//! This module provides `Query` and `QueryMut` traits, which are used for querying
//! and executing functions on several different providers. This includes
//! [`Subscriber::ConnectionContext`](crate::event::Subscriber::ConnectionContext)
//! on a Subscriber and the [`Sender`](crate::datagram::Sender) and
//! [`Receiver`](crate::datagram::Receiver) types on a datagram [`Endpoint`](crate::datagram::Endpoint).

use core::marker::PhantomData;

pub trait Query {
    fn execute(&mut self, context: &dyn core::any::Any) -> ControlFlow;
}

pub trait QueryMut {
    fn execute_mut(&mut self, context: &mut dyn core::any::Any) -> ControlFlow;
}

/// Used to tell a query whether it should exit early or go on as usual.
#[derive(Debug, Clone, Copy, PartialEq)]
pub enum ControlFlow {
    Continue,
    Break,
}

impl ControlFlow {
    #[inline]
    #[must_use]
    pub fn and_then(self, f: impl FnOnce() -> Self) -> Self {
        match self {
            Self::Continue => f(),
            Self::Break => Self::Break,
        }
    }
}

/// A type that implements Query and QueryMut traits and only executes once.
///
/// This will execute and short-circuit on the first match.
pub struct Once<F, Context, Outcome> {
    query: Option<F>,
    result: Option<Outcome>,
    context: PhantomData<Context>,
}

impl<F, Context, Outcome> From<Once<F, Context, Outcome>> for Result<Outcome, Error> {
    #[inline]
    fn from(query: Once<F, Context, Outcome>) -> Self {
        query.result.ok_or(Error::ContextTypeMismatch)
    }
}

impl<F, ConnectionContext, Outcome> Once<F, ConnectionContext, Outcome>
where
    F: FnOnce(&ConnectionContext) -> Outcome,
    ConnectionContext: 'static,
{
    #[inline]
    pub fn new(query: F) -> Self {
        Self {
            query: Some(query),
            result: None,
            context: PhantomData,
        }
    }
}

impl<F, ConnectionContext, Outcome> Once<F, ConnectionContext, Outcome>
where
    F: FnOnce(&mut ConnectionContext) -> Outcome,
    ConnectionContext: 'static,
{
    #[inline]
    pub fn new_mut(query: F) -> Self {
        Self {
            query: Some(query),
            result: None,
            context: PhantomData,
        }
    }
}

impl<F, Context, Outcome> Query for Once<F, Context, Outcome>
where
    F: FnOnce(&Context) -> Outcome,
    Context: 'static,
{
    fn execute(&mut self, context: &dyn core::any::Any) -> ControlFlow {
        match context.downcast_ref::<Context>() {
            Some(context) => {
                let query = self.query.take().expect("can only match once");
                self.result = Some(query(context));
                ControlFlow::Break
            }
            None => ControlFlow::Continue,
        }
    }
}

impl<F, Context, Outcome> QueryMut for Once<F, Context, Outcome>
where
    F: FnOnce(&mut Context) -> Outcome,
    Context: 'static,
{
    fn execute_mut(&mut self, context: &mut dyn core::any::Any) -> ControlFlow {
        match context.downcast_mut::<Context>() {
            Some(context) => {
                let query = self.query.take().expect("can only match once");
                self.result = Some(query(context));
                ControlFlow::Break
            }
            None => ControlFlow::Continue,
        }
    }
}

#[non_exhaustive]
#[derive(Debug, Clone)]
/// Reason for the failed query.
pub enum Error {
    /// The connection lock is poisoned and the connection unusable.
    ConnectionLockPoisoned,

    /// The expected query type failed to match any of the configured types.
    ContextTypeMismatch,
}

impl core::fmt::Display for Error {
    fn fmt(&self, f: &mut core::fmt::Formatter) -> core::fmt::Result {
        write!(f, "{self:?}")
    }
}

#[cfg(feature = "std")]
impl std::error::Error for Error {}