1use anyhow::{anyhow, Result};
2use std::collections::HashMap;
3
4use super::{ArgCount, FunctionCategory, FunctionSignature, SqlFunction};
5use crate::data::datatable::DataValue;
6
7pub struct AvogadroFunction;
9
10impl SqlFunction for AvogadroFunction {
11 fn signature(&self) -> FunctionSignature {
12 FunctionSignature {
13 name: "AVOGADRO",
14 category: FunctionCategory::Chemical,
15 arg_count: ArgCount::Fixed(0),
16 description: "Returns Avogadro's number (6.022 × 10^23)",
17 returns: "FLOAT",
18 examples: vec![
19 "SELECT AVOGADRO()",
20 "SELECT molecules / AVOGADRO() AS moles",
21 ],
22 }
23 }
24
25 fn evaluate(&self, args: &[DataValue]) -> Result<DataValue> {
26 self.validate_args(args)?;
27 Ok(DataValue::Float(6.022140857e23))
28 }
29}
30
31pub struct AtomicMassFunction;
33
34impl AtomicMassFunction {
35 fn get_atomic_mass(element: &str) -> Option<f64> {
36 let masses: HashMap<&str, f64> = [
37 ("H", 1.008),
39 ("HYDROGEN", 1.008),
40 ("HE", 4.003),
41 ("HELIUM", 4.003),
42 ("LI", 6.941),
43 ("LITHIUM", 6.941),
44 ("BE", 9.012),
45 ("BERYLLIUM", 9.012),
46 ("B", 10.81),
47 ("BORON", 10.81),
48 ("C", 12.01),
49 ("CARBON", 12.01),
50 ("N", 14.01),
51 ("NITROGEN", 14.01),
52 ("O", 16.00),
53 ("OXYGEN", 16.00),
54 ("F", 19.00),
55 ("FLUORINE", 19.00),
56 ("NE", 20.18),
57 ("NEON", 20.18),
58 ("NA", 22.99),
59 ("SODIUM", 22.99),
60 ("MG", 24.31),
61 ("MAGNESIUM", 24.31),
62 ("AL", 26.98),
63 ("ALUMINUM", 26.98),
64 ("ALUMINIUM", 26.98),
65 ("SI", 28.09),
66 ("SILICON", 28.09),
67 ("P", 30.97),
68 ("PHOSPHORUS", 30.97),
69 ("S", 32.07),
70 ("SULFUR", 32.07),
71 ("SULPHUR", 32.07),
72 ("CL", 35.45),
73 ("CHLORINE", 35.45),
74 ("AR", 39.95),
75 ("ARGON", 39.95),
76 ("K", 39.10),
77 ("POTASSIUM", 39.10),
78 ("CA", 40.08),
79 ("CALCIUM", 40.08),
80 ("FE", 55.85),
82 ("IRON", 55.85),
83 ("CU", 63.55),
84 ("COPPER", 63.55),
85 ("ZN", 65.39),
86 ("ZINC", 65.39),
87 ("AG", 107.87),
88 ("SILVER", 107.87),
89 ("AU", 196.97),
90 ("GOLD", 196.97),
91 ("HG", 200.59),
92 ("MERCURY", 200.59),
93 ("PB", 207.2),
94 ("LEAD", 207.2),
95 ("U", 238.03),
96 ("URANIUM", 238.03),
97 ]
98 .iter()
99 .cloned()
100 .collect();
101
102 masses.get(element.to_uppercase().as_str()).copied()
103 }
104}
105
106impl SqlFunction for AtomicMassFunction {
107 fn signature(&self) -> FunctionSignature {
108 FunctionSignature {
109 name: "ATOMIC_MASS",
110 category: FunctionCategory::Chemical,
111 arg_count: ArgCount::Fixed(1),
112 description: "Returns the atomic mass of an element in amu",
113 returns: "FLOAT",
114 examples: vec![
115 "SELECT ATOMIC_MASS('H')",
116 "SELECT ATOMIC_MASS('Carbon')",
117 "SELECT ATOMIC_MASS('Au') AS gold_mass",
118 ],
119 }
120 }
121
122 fn evaluate(&self, args: &[DataValue]) -> Result<DataValue> {
123 self.validate_args(args)?;
124
125 match &args[0] {
126 DataValue::String(element) => match Self::get_atomic_mass(element) {
127 Some(mass) => Ok(DataValue::Float(mass)),
128 None => Err(anyhow!("Unknown element: {}", element)),
129 },
130 DataValue::InternedString(element) => match Self::get_atomic_mass(element) {
131 Some(mass) => Ok(DataValue::Float(mass)),
132 None => Err(anyhow!("Unknown element: {}", element)),
133 },
134 _ => Err(anyhow!("ATOMIC_MASS() requires a string argument")),
135 }
136 }
137}
138
139pub struct AtomicNumberFunction;
141
142impl AtomicNumberFunction {
143 fn get_atomic_number(element: &str) -> Option<i64> {
144 let numbers: HashMap<&str, i64> = [
145 ("H", 1),
146 ("HYDROGEN", 1),
147 ("HE", 2),
148 ("HELIUM", 2),
149 ("LI", 3),
150 ("LITHIUM", 3),
151 ("BE", 4),
152 ("BERYLLIUM", 4),
153 ("B", 5),
154 ("BORON", 5),
155 ("C", 6),
156 ("CARBON", 6),
157 ("N", 7),
158 ("NITROGEN", 7),
159 ("O", 8),
160 ("OXYGEN", 8),
161 ("F", 9),
162 ("FLUORINE", 9),
163 ("NE", 10),
164 ("NEON", 10),
165 ("NA", 11),
166 ("SODIUM", 11),
167 ("MG", 12),
168 ("MAGNESIUM", 12),
169 ("AL", 13),
170 ("ALUMINUM", 13),
171 ("ALUMINIUM", 13),
172 ("SI", 14),
173 ("SILICON", 14),
174 ("P", 15),
175 ("PHOSPHORUS", 15),
176 ("S", 16),
177 ("SULFUR", 16),
178 ("SULPHUR", 16),
179 ("CL", 17),
180 ("CHLORINE", 17),
181 ("AR", 18),
182 ("ARGON", 18),
183 ("K", 19),
184 ("POTASSIUM", 19),
185 ("CA", 20),
186 ("CALCIUM", 20),
187 ("FE", 26),
189 ("IRON", 26),
190 ("CU", 29),
191 ("COPPER", 29),
192 ("ZN", 30),
193 ("ZINC", 30),
194 ("AG", 47),
195 ("SILVER", 47),
196 ("AU", 79),
197 ("GOLD", 79),
198 ("HG", 80),
199 ("MERCURY", 80),
200 ("PB", 82),
201 ("LEAD", 82),
202 ("U", 92),
203 ("URANIUM", 92),
204 ]
205 .iter()
206 .cloned()
207 .collect();
208
209 numbers.get(element.to_uppercase().as_str()).copied()
210 }
211}
212
213impl SqlFunction for AtomicNumberFunction {
214 fn signature(&self) -> FunctionSignature {
215 FunctionSignature {
216 name: "ATOMIC_NUMBER",
217 category: FunctionCategory::Chemical,
218 arg_count: ArgCount::Fixed(1),
219 description: "Returns the atomic number of an element",
220 returns: "INTEGER",
221 examples: vec![
222 "SELECT ATOMIC_NUMBER('H')",
223 "SELECT ATOMIC_NUMBER('Carbon')",
224 "SELECT ATOMIC_NUMBER('Au') AS gold_number",
225 ],
226 }
227 }
228
229 fn evaluate(&self, args: &[DataValue]) -> Result<DataValue> {
230 self.validate_args(args)?;
231
232 match &args[0] {
233 DataValue::String(element) => match Self::get_atomic_number(element) {
234 Some(number) => Ok(DataValue::Integer(number)),
235 None => Err(anyhow!("Unknown element: {}", element)),
236 },
237 DataValue::InternedString(element) => match Self::get_atomic_number(element) {
238 Some(number) => Ok(DataValue::Integer(number)),
239 None => Err(anyhow!("Unknown element: {}", element)),
240 },
241 _ => Err(anyhow!("ATOMIC_NUMBER() requires a string argument")),
242 }
243 }
244}
245
246#[cfg(test)]
247mod tests {
248 use super::*;
249
250 #[test]
251 fn test_avogadro() {
252 let func = AvogadroFunction;
253 let result = func.evaluate(&[]).unwrap();
254 match result {
255 DataValue::Float(val) => assert!((val - 6.022140857e23).abs() < 1e20),
256 _ => panic!("Expected Float"),
257 }
258 }
259
260 #[test]
261 fn test_atomic_mass_hydrogen() {
262 let func = AtomicMassFunction;
263 let result = func
264 .evaluate(&[DataValue::String("H".to_string())])
265 .unwrap();
266 match result {
267 DataValue::Float(val) => assert!((val - 1.008).abs() < 0.001),
268 _ => panic!("Expected Float"),
269 }
270 }
271
272 #[test]
273 fn test_atomic_mass_carbon() {
274 let func = AtomicMassFunction;
275 let result = func
276 .evaluate(&[DataValue::String("Carbon".to_string())])
277 .unwrap();
278 match result {
279 DataValue::Float(val) => assert!((val - 12.01).abs() < 0.01),
280 _ => panic!("Expected Float"),
281 }
282 }
283
284 #[test]
285 fn test_atomic_mass_gold() {
286 let func = AtomicMassFunction;
287 let result = func
288 .evaluate(&[DataValue::String("Au".to_string())])
289 .unwrap();
290 match result {
291 DataValue::Float(val) => assert!((val - 196.97).abs() < 0.01),
292 _ => panic!("Expected Float"),
293 }
294 }
295
296 #[test]
297 fn test_atomic_mass_unknown_element() {
298 let func = AtomicMassFunction;
299 let result = func.evaluate(&[DataValue::String("Xyz".to_string())]);
300 assert!(result.is_err());
301 }
302
303 #[test]
304 fn test_atomic_number_carbon() {
305 let func = AtomicNumberFunction;
306 let result = func
307 .evaluate(&[DataValue::String("C".to_string())])
308 .unwrap();
309 match result {
310 DataValue::Integer(val) => assert_eq!(val, 6),
311 _ => panic!("Expected Integer"),
312 }
313 }
314}