1use super::{
2 BoolExpr, CallArg, CustomFieldAccess, CustomFunctionExpr, CustomListExpr, FloatExpr, IntExpr,
3 PanicExpr, StringExpr, TupleExpr,
4};
5use crate::plan::module::constant::MaterializedConstantCustomConstruction;
6use crate::plan::{
7 ConstantCustomReference, CustomConstructor, CustomConstructorRefinement, CustomLocal,
8 CustomLocalId, CustomType, CustomValueShape, FunctionInstantiation, Step, ValueShape,
9};
10use ecow::EcoString;
11use num_bigint::BigInt;
12
13#[derive(Debug, Clone, PartialEq)]
14pub struct CustomExpr {
15 shape: CustomValueShape,
16 kind: CustomExprKind,
17}
18
19#[derive(Debug, Clone, PartialEq)]
20pub(crate) struct CustomBoolCaseBranches {
21 shape: CustomValueShape,
22 true_: CustomExprKind,
23 false_: CustomExprKind,
24}
25
26#[derive(Debug, Clone, PartialEq)]
27pub(crate) struct CustomCaseBranches<Pattern> {
28 shape: CustomValueShape,
29 clauses: Vec<(Pattern, CustomExprKind)>,
30 fallback: CustomExprKind,
31}
32
33#[derive(Debug, Clone, PartialEq)]
34pub(crate) struct CustomLocalExpr {
35 local: CustomLocal,
36 value: CustomExpr,
37}
38
39#[cfg(test)]
40#[derive(Debug, Clone, PartialEq, Eq)]
41pub(crate) struct CustomArgumentCountMismatch {
42 pub(crate) expected: usize,
43 pub(crate) actual: usize,
44}
45
46#[derive(Debug, Clone, PartialEq)]
47pub(crate) struct CustomConstruction {
48 constructor: CustomConstructor,
49 fields: Box<[super::Expr]>,
50}
51
52#[derive(Debug, Clone, PartialEq)]
53pub(crate) struct CustomCallArguments {
54 values: Box<[CallArg]>,
55}
56
57#[derive(Debug, Clone, PartialEq)]
58pub(crate) struct CustomFunctionCall {
59 function: Box<CustomFunctionExpr>,
60 arguments: CustomCallArguments,
61 site: crate::plan::HostCallSite,
62}
63
64#[derive(Debug, Clone, PartialEq)]
65pub(crate) enum CustomExprKind {
66 Constructor(CustomConstruction),
67 Constant(ConstantCustomReference),
68 LocalGet {
69 local: CustomLocal,
70 name: EcoString,
71 },
72 Call {
73 function: FunctionInstantiation,
74 args: Vec<CallArg>,
75 site: crate::plan::HostCallSite,
76 },
77 FunctionCall(CustomFunctionCall),
78 TupleIndex {
79 tuple: Box<TupleExpr>,
80 index: usize,
81 },
82 CustomField(CustomFieldAccess),
83 ListIndex {
84 list: Box<CustomListExpr>,
85 index: usize,
86 },
87 Panic(PanicExpr),
88 BoolCase {
89 subject: Box<BoolExpr>,
90 true_: Box<CustomExprKind>,
91 false_: Box<CustomExprKind>,
92 },
93 IntCase {
94 subject: Box<IntExpr>,
95 clauses: Vec<(BigInt, CustomExprKind)>,
96 fallback: Box<CustomExprKind>,
97 },
98 StringCase {
99 subject: Box<StringExpr>,
100 clauses: Vec<(EcoString, CustomExprKind)>,
101 fallback: Box<CustomExprKind>,
102 },
103 FloatCase {
104 subject: Box<FloatExpr>,
105 clauses: Vec<(f64, CustomExprKind)>,
106 fallback: Box<CustomExprKind>,
107 },
108 Block {
109 steps: Vec<Step>,
110 return_: Box<CustomExprKind>,
111 },
112}
113
114pub(crate) fn custom_constructor_expr(constructor: CustomConstructor) -> super::Expr {
115 if constructor.fields().is_empty() {
116 let shape = custom_constructor_shape(&constructor);
117 super::Expr::custom(CustomExpr::new(
118 shape,
119 CustomExprKind::Constructor(CustomConstruction {
120 constructor,
121 fields: Vec::new().into_boxed_slice(),
122 }),
123 ))
124 } else {
125 super::Expr::function(super::FunctionExpr::custom(
126 CustomFunctionExpr::constructor(constructor),
127 ))
128 }
129}
130
131impl CustomExpr {
132 pub(in crate::plan::module) fn constant(reference: ConstantCustomReference) -> Self {
133 let shape = reference.shape().clone();
134 Self::new(shape, CustomExprKind::Constant(reference))
135 }
136
137 #[cfg(test)]
138 pub(crate) fn try_constructor(
139 constructor: CustomConstructor,
140 fields: Vec<super::Expr>,
141 ) -> Result<Self, CustomArgumentCountMismatch> {
142 let shape = custom_constructor_shape(&constructor);
143 CustomConstruction::try_new(constructor, fields)
144 .map(|construction| Self::new(shape, CustomExprKind::Constructor(construction)))
145 }
146
147 pub(crate) fn from_construction(
148 shape: CustomValueShape,
149 construction: CustomConstruction,
150 ) -> Self {
151 Self::new(shape, CustomExprKind::Constructor(construction))
152 }
153
154 pub(crate) fn local_get(local: CustomLocal, name: EcoString) -> Self {
155 Self::new(
156 local.shape().clone(),
157 CustomExprKind::LocalGet { local, name },
158 )
159 }
160
161 #[cfg(test)]
162 pub(crate) fn call(
163 function: FunctionInstantiation,
164 args: Vec<CallArg>,
165 shape: CustomValueShape,
166 ) -> Self {
167 Self::call_at(function, args, shape, crate::plan::HostCallSite::unknown())
168 }
169
170 pub(crate) fn call_at(
171 function: FunctionInstantiation,
172 args: Vec<CallArg>,
173 shape: CustomValueShape,
174 site: crate::plan::HostCallSite,
175 ) -> Self {
176 Self::new(
177 shape,
178 CustomExprKind::Call {
179 function,
180 args,
181 site,
182 },
183 )
184 }
185
186 #[cfg(test)]
187 pub(crate) fn function_call(function: CustomFunctionExpr, args: Vec<CallArg>) -> Self {
188 Self::function_call_at(function, args, crate::plan::HostCallSite::unknown())
189 }
190
191 pub(crate) fn function_call_at(
192 function: CustomFunctionExpr,
193 args: Vec<CallArg>,
194 site: crate::plan::HostCallSite,
195 ) -> Self {
196 let shape = function.custom_function_type().return_().clone();
197 let call = CustomFunctionCall::new(function, args, site);
198 Self::new(shape, CustomExprKind::FunctionCall(call))
199 }
200
201 pub(crate) fn tuple_index_shape(
202 tuple: TupleExpr,
203 index: usize,
204 shape: CustomValueShape,
205 ) -> Self {
206 Self::new(
207 shape,
208 CustomExprKind::TupleIndex {
209 tuple: Box::new(tuple),
210 index,
211 },
212 )
213 }
214
215 pub(crate) fn custom_field_shape(access: CustomFieldAccess, shape: CustomValueShape) -> Self {
216 Self::new(shape, CustomExprKind::CustomField(access))
217 }
218
219 pub(crate) fn list_index_shape(
220 list: CustomListExpr,
221 index: usize,
222 shape: CustomValueShape,
223 ) -> Self {
224 Self::new(
225 shape,
226 CustomExprKind::ListIndex {
227 list: Box::new(list),
228 index,
229 },
230 )
231 }
232
233 pub(crate) fn panic_shape(panic: PanicExpr, shape: CustomValueShape) -> Self {
234 Self::new(shape, CustomExprKind::Panic(panic))
235 }
236
237 pub(crate) fn bool_case(subject: BoolExpr, branches: CustomBoolCaseBranches) -> Self {
238 let (shape, true_, false_) = branches.into_parts();
239 Self::new(
240 shape,
241 CustomExprKind::BoolCase {
242 subject: Box::new(subject),
243 true_: Box::new(true_),
244 false_: Box::new(false_),
245 },
246 )
247 }
248
249 pub(crate) fn int_case(subject: IntExpr, branches: CustomCaseBranches<BigInt>) -> Self {
250 let (shape, clauses, fallback) = branches.into_parts();
251 Self::new(
252 shape,
253 CustomExprKind::IntCase {
254 subject: Box::new(subject),
255 clauses,
256 fallback: Box::new(fallback),
257 },
258 )
259 }
260
261 pub(crate) fn string_case(
262 subject: StringExpr,
263 branches: CustomCaseBranches<EcoString>,
264 ) -> Self {
265 let (shape, clauses, fallback) = branches.into_parts();
266 Self::new(
267 shape,
268 CustomExprKind::StringCase {
269 subject: Box::new(subject),
270 clauses,
271 fallback: Box::new(fallback),
272 },
273 )
274 }
275
276 pub(crate) fn float_case(subject: FloatExpr, branches: CustomCaseBranches<f64>) -> Self {
277 let (shape, clauses, fallback) = branches.into_parts();
278 Self::new(
279 shape,
280 CustomExprKind::FloatCase {
281 subject: Box::new(subject),
282 clauses,
283 fallback: Box::new(fallback),
284 },
285 )
286 }
287
288 pub(crate) fn block(steps: Vec<Step>, return_: Self) -> Self {
289 let (shape, return_) = return_.into_parts();
290 Self::new(
291 shape,
292 CustomExprKind::Block {
293 steps,
294 return_: Box::new(return_),
295 },
296 )
297 }
298
299 pub fn type_(&self) -> &CustomType {
300 self.shape.type_()
301 }
302 pub(crate) fn shape(&self) -> &CustomValueShape {
303 &self.shape
304 }
305
306 pub(super) fn with_shape(mut self, shape: CustomValueShape) -> Self {
307 self.shape = shape;
308 self
309 }
310 pub(crate) fn kind(&self) -> &CustomExprKind {
311 &self.kind
312 }
313 pub(crate) fn into_parts(self) -> (CustomValueShape, CustomExprKind) {
314 (self.shape, self.kind)
315 }
316
317 fn new(shape: CustomValueShape, kind: CustomExprKind) -> Self {
318 Self { shape, kind }
319 }
320}
321
322impl CustomConstruction {
323 pub(in crate::plan::module) fn from_constant(
324 construction: MaterializedConstantCustomConstruction,
325 ) -> Self {
326 let (constructor, fields) = construction.into_parts();
327 Self {
328 constructor,
329 fields,
330 }
331 }
332}
333
334impl CustomLocalExpr {
335 pub(crate) fn from_value(local: CustomLocalId, value: CustomExpr) -> Self {
336 let local = CustomLocal::from_shape(local, value.shape().clone());
337 Self { local, value }
338 }
339
340 pub(crate) fn local(&self) -> &CustomLocal {
341 &self.local
342 }
343
344 pub(crate) fn value(&self) -> &CustomExpr {
345 &self.value
346 }
347}
348
349impl CustomBoolCaseBranches {
350 pub(crate) fn from_resolved_shape(
351 shape: CustomValueShape,
352 true_: CustomExpr,
353 false_: CustomExpr,
354 ) -> Self {
355 Self {
356 shape,
357 true_: true_.kind,
358 false_: false_.kind,
359 }
360 }
361
362 fn into_parts(self) -> (CustomValueShape, CustomExprKind, CustomExprKind) {
363 (self.shape, self.true_, self.false_)
364 }
365}
366
367impl<Pattern> CustomCaseBranches<Pattern> {
368 pub(crate) fn from_resolved_shape(
369 shape: CustomValueShape,
370 clauses: Vec<(Pattern, CustomExpr)>,
371 fallback: CustomExpr,
372 ) -> Self {
373 Self {
374 shape,
375 clauses: clauses
376 .into_iter()
377 .map(|(pattern, branch)| (pattern, branch.kind))
378 .collect(),
379 fallback: fallback.kind,
380 }
381 }
382
383 fn into_parts(
384 self,
385 ) -> (
386 CustomValueShape,
387 Vec<(Pattern, CustomExprKind)>,
388 CustomExprKind,
389 ) {
390 (self.shape, self.clauses, self.fallback)
391 }
392}
393
394pub(super) fn custom_constructor_shape(constructor: &CustomConstructor) -> CustomValueShape {
395 CustomValueShape::new(
396 constructor.type_().type_name().clone(),
397 constructor
398 .type_()
399 .arguments()
400 .iter()
401 .cloned()
402 .map(ValueShape::from_value_type)
403 .collect(),
404 CustomConstructorRefinement::Exact(constructor.index()),
405 )
406}
407
408impl CustomConstruction {
409 pub(crate) fn from_validated(constructor: CustomConstructor, fields: Vec<super::Expr>) -> Self {
410 Self {
411 constructor,
412 fields: fields.into_boxed_slice(),
413 }
414 }
415
416 #[cfg(test)]
417 pub(crate) fn try_new(
418 constructor: CustomConstructor,
419 fields: Vec<super::Expr>,
420 ) -> Result<Self, CustomArgumentCountMismatch> {
421 if constructor.fields().len() != fields.len() {
422 return Err(CustomArgumentCountMismatch {
423 expected: constructor.fields().len(),
424 actual: fields.len(),
425 });
426 }
427
428 Ok(Self::from_validated(constructor, fields))
429 }
430
431 pub(crate) fn constructor(&self) -> &CustomConstructor {
432 &self.constructor
433 }
434
435 pub(crate) fn fields(&self) -> &[super::Expr] {
436 &self.fields
437 }
438}
439
440impl CustomFunctionCall {
441 pub(crate) fn new(
442 function: CustomFunctionExpr,
443 arguments: Vec<CallArg>,
444 site: crate::plan::HostCallSite,
445 ) -> Self {
446 Self {
447 function: Box::new(function),
448 arguments: CustomCallArguments {
449 values: arguments.into_boxed_slice(),
450 },
451 site,
452 }
453 }
454
455 pub(crate) fn function(&self) -> &CustomFunctionExpr {
456 &self.function
457 }
458
459 pub(crate) fn arguments(&self) -> &[CallArg] {
460 &self.arguments.values
461 }
462
463 pub(crate) fn site(&self) -> &crate::plan::HostCallSite {
464 &self.site
465 }
466}
467
468#[cfg(test)]
469mod tests {
470 use super::{CustomArgumentCountMismatch, CustomExpr, CustomExprKind};
471 use crate::plan::{
472 CallArg, CustomConstructor, CustomConstructorField, CustomFunctionExpr, CustomType,
473 CustomTypeName, Expr, IntExpr, ValueType,
474 };
475
476 #[test]
477 fn custom_construction_owns_an_exact_field_pack() {
478 let type_ = CustomType::new(
479 CustomTypeName::new("geam".into(), "main".into(), "Boxed".into()),
480 Vec::new(),
481 );
482 let constructor = CustomConstructor::new(
483 type_.clone(),
484 "Boxed".into(),
485 0,
486 vec![CustomConstructorField::new(None, ValueType::Int)],
487 );
488 let field = Expr::int(IntExpr::value(1.into()));
489
490 assert_eq!(
491 CustomExpr::try_constructor(constructor.clone(), Vec::new()),
492 Err(CustomArgumentCountMismatch {
493 expected: 1,
494 actual: 0,
495 }),
496 );
497 assert_eq!(
498 CustomExpr::try_constructor(constructor.clone(), vec![field.clone(), field.clone()]),
499 Err(CustomArgumentCountMismatch {
500 expected: 1,
501 actual: 2,
502 }),
503 );
504
505 let expression = CustomExpr::try_constructor(constructor.clone(), vec![field.clone()])
506 .expect("exact custom construction should be valid");
507 assert_eq!(expression.type_(), &type_);
508 assert_eq!(
509 expression.kind(),
510 &CustomExprKind::Constructor(super::CustomConstruction::from_validated(
511 constructor,
512 vec![field]
513 ),),
514 );
515 }
516
517 #[test]
518 fn zero_field_custom_construction_preserves_an_empty_pack() {
519 let type_ = CustomType::new(
520 CustomTypeName::new("geam".into(), "main".into(), "Empty".into()),
521 Vec::new(),
522 );
523 let constructor = CustomConstructor::new(type_.clone(), "Empty".into(), 0, Vec::new());
524 let expression = validated_custom_expr(constructor.clone(), Vec::new());
525
526 assert_eq!(expression.type_(), &type_);
527 assert_eq!(
528 expression.kind(),
529 &CustomExprKind::Constructor(super::CustomConstruction::from_validated(
530 constructor,
531 Vec::new()
532 ),),
533 );
534 }
535
536 #[test]
537 fn custom_function_call_owns_its_validated_argument_pack() {
538 let type_ = CustomType::new(
539 CustomTypeName::new("geam".into(), "main".into(), "Boxed".into()),
540 Vec::new(),
541 );
542 let constructor = CustomConstructor::new(
543 type_.clone(),
544 "Boxed".into(),
545 0,
546 vec![CustomConstructorField::new(None, ValueType::Int)],
547 );
548 let function = CustomFunctionExpr::constructor(constructor);
549 let argument = CallArg::new(crate::plan::Expr::int(IntExpr::value(1.into())));
550
551 let expression = CustomExpr::function_call(function.clone(), vec![argument.clone()]);
552 assert_eq!(expression.type_(), &type_);
553 assert_eq!(
554 expression.kind(),
555 &CustomExprKind::FunctionCall(super::CustomFunctionCall::new(
556 function,
557 vec![argument],
558 crate::plan::HostCallSite::unknown(),
559 ),),
560 );
561 }
562
563 #[test]
564 fn same_result_children_store_only_custom_bodies() {
565 let type_ = CustomType::new(
566 CustomTypeName::new("geam".into(), "main".into(), "Boxed".into()),
567 Vec::new(),
568 );
569 let constructor = CustomConstructor::new(type_.clone(), "Boxed".into(), 0, Vec::new());
570 let branch = validated_custom_expr(constructor.clone(), Vec::new());
571 let shape = branch.shape().clone();
572 let fallback = validated_custom_expr(constructor.clone(), Vec::new());
573
574 let expression = CustomExpr::block(
575 Vec::new(),
576 CustomExpr::bool_case(
577 crate::plan::BoolExpr::value(true),
578 super::CustomBoolCaseBranches::from_resolved_shape(shape.clone(), branch, fallback),
579 ),
580 );
581
582 assert_eq!(
583 expression.into_parts(),
584 (
585 shape,
586 CustomExprKind::Block {
587 steps: Vec::new(),
588 return_: Box::new(CustomExprKind::BoolCase {
589 subject: Box::new(crate::plan::BoolExpr::value(true)),
590 true_: Box::new(CustomExprKind::Constructor(
591 super::CustomConstruction::from_validated(
592 constructor.clone(),
593 Vec::new(),
594 ),
595 )),
596 false_: Box::new(CustomExprKind::Constructor(
597 super::CustomConstruction::from_validated(constructor, Vec::new()),
598 )),
599 }),
600 },
601 ),
602 );
603 }
604
605 fn validated_custom_expr(constructor: CustomConstructor, fields: Vec<Expr>) -> CustomExpr {
606 let shape = super::custom_constructor_shape(&constructor);
607 let construction = super::CustomConstruction::from_validated(constructor, fields);
608 CustomExpr::from_construction(shape, construction)
609 }
610}