vibesql/catalog/
function.rs

1//! Function signature definitions.
2
3use crate::types::SqlType;
4
5/// Function signature for built-in and user-defined functions.
6#[derive(Debug, Clone, PartialEq)]
7pub struct FunctionSignature {
8    /// Function name (uppercase).
9    pub name: String,
10    /// Parameter types (None means any type is accepted).
11    pub parameters: Vec<FunctionParameter>,
12    /// Return type.
13    pub return_type: SqlType,
14    /// Whether this is an aggregate function.
15    pub is_aggregate: bool,
16    /// Whether this is a window function.
17    pub is_window: bool,
18    /// Whether the function is deterministic.
19    pub is_deterministic: bool,
20    /// Minimum number of arguments.
21    pub min_args: usize,
22    /// Maximum number of arguments (None = unlimited).
23    pub max_args: Option<usize>,
24}
25
26/// A function parameter.
27#[derive(Debug, Clone, PartialEq)]
28pub struct FunctionParameter {
29    /// Parameter name (optional).
30    pub name: Option<String>,
31    /// Expected type (None = any type).
32    pub data_type: Option<SqlType>,
33    /// Whether the parameter is optional.
34    pub optional: bool,
35    /// Whether this parameter accepts variable arguments.
36    pub variadic: bool,
37}
38
39impl FunctionSignature {
40    /// Create a scalar function signature.
41    pub fn scalar(name: impl Into<String>, return_type: SqlType) -> Self {
42        Self {
43            name: name.into().to_uppercase(),
44            parameters: Vec::new(),
45            return_type,
46            is_aggregate: false,
47            is_window: false,
48            is_deterministic: true,
49            min_args: 0,
50            max_args: None,
51        }
52    }
53
54    /// Create an aggregate function signature.
55    pub fn aggregate(name: impl Into<String>, return_type: SqlType) -> Self {
56        Self {
57            name: name.into().to_uppercase(),
58            parameters: Vec::new(),
59            return_type,
60            is_aggregate: true,
61            is_window: false,
62            is_deterministic: true,
63            min_args: 0,
64            max_args: None,
65        }
66    }
67
68    /// Create a window function signature.
69    pub fn window(name: impl Into<String>, return_type: SqlType) -> Self {
70        Self {
71            name: name.into().to_uppercase(),
72            parameters: Vec::new(),
73            return_type,
74            is_aggregate: false,
75            is_window: true,
76            is_deterministic: true,
77            min_args: 0,
78            max_args: None,
79        }
80    }
81
82    /// Set minimum number of arguments.
83    pub fn with_min_args(mut self, min: usize) -> Self {
84        self.min_args = min;
85        self
86    }
87
88    /// Set maximum number of arguments.
89    pub fn with_max_args(mut self, max: usize) -> Self {
90        self.max_args = Some(max);
91        self
92    }
93
94    /// Set exact number of arguments.
95    pub fn with_args(mut self, count: usize) -> Self {
96        self.min_args = count;
97        self.max_args = Some(count);
98        self
99    }
100
101    /// Add a parameter.
102    pub fn with_param(mut self, param: FunctionParameter) -> Self {
103        self.parameters.push(param);
104        self
105    }
106
107    /// Mark as non-deterministic.
108    pub fn non_deterministic(mut self) -> Self {
109        self.is_deterministic = false;
110        self
111    }
112
113    /// Check if the given number of arguments is valid.
114    pub fn accepts_arg_count(&self, count: usize) -> bool {
115        if count < self.min_args {
116            return false;
117        }
118        match self.max_args {
119            Some(max) => count <= max,
120            None => true,
121        }
122    }
123
124    /// Check if the function can be used in a window context.
125    pub fn can_be_window(&self) -> bool {
126        self.is_window || self.is_aggregate
127    }
128}
129
130impl FunctionParameter {
131    /// Create a new parameter.
132    pub fn new(name: impl Into<String>, data_type: SqlType) -> Self {
133        Self {
134            name: Some(name.into()),
135            data_type: Some(data_type),
136            optional: false,
137            variadic: false,
138        }
139    }
140
141    /// Create a parameter that accepts any type.
142    pub fn any(name: impl Into<String>) -> Self {
143        Self {
144            name: Some(name.into()),
145            data_type: None,
146            optional: false,
147            variadic: false,
148        }
149    }
150
151    /// Create an unnamed parameter.
152    pub fn unnamed(data_type: SqlType) -> Self {
153        Self {
154            name: None,
155            data_type: Some(data_type),
156            optional: false,
157            variadic: false,
158        }
159    }
160
161    /// Mark as optional.
162    pub fn optional(mut self) -> Self {
163        self.optional = true;
164        self
165    }
166
167    /// Mark as variadic (accepts multiple values).
168    pub fn variadic(mut self) -> Self {
169        self.variadic = true;
170        self
171    }
172}
173
174#[cfg(test)]
175mod tests {
176    use super::*;
177
178    #[test]
179    fn test_function_signature() {
180        let count = FunctionSignature::aggregate("COUNT", SqlType::Int64)
181            .with_min_args(1)
182            .with_max_args(1);
183
184        assert!(count.is_aggregate);
185        assert!(!count.is_window);
186        assert!(count.accepts_arg_count(1));
187        assert!(!count.accepts_arg_count(0));
188        assert!(!count.accepts_arg_count(2));
189    }
190
191    #[test]
192    fn test_window_function() {
193        let row_number = FunctionSignature::window("ROW_NUMBER", SqlType::Int64).with_args(0);
194
195        assert!(row_number.is_window);
196        assert!(row_number.can_be_window());
197        assert!(row_number.accepts_arg_count(0));
198    }
199}