equator_macro/
lib.rs

1use proc_macro2::{Span, TokenStream};
2use quote::quote;
3use syn::*;
4
5// assert!(all(a == 0, b))
6// should expand to
7//
8// match (&(a), &(0), &(b)) {
9//     (__0, __1, __2) => {
10//         use $crate::Expr;
11//
12//         let __assert_expr = $crate::Finalize {
13//             expr: $crate::expr::AndExpr {
14//                 lhs: $crate::atomic::EqExpr {
15//                     lhs: __0,
16//                     rhs: __1,
17//                 },
18//                 rhs: *__2,
19//             },
20//             line: (),
21//             col: (),
22//             file: (),
23//         };
24//
25//         if !__assert_expr.eval_expr() {
26//             let __assert_message = $crate::DebugMessage {
27//                 result: __assert_expr.result(),
28//                 source: $crate::Finalize {
29//                     expr: $crate::expr::AndExpr {
30//                         lhs: $crate::atomic::EqExpr {
31//                             lhs: ::core::stringify!(a),
32//                             rhs: ::core::stringify!(0),
33//                         },
34//                         rhs: ::core::stringify!(b),
35//                     },
36//                     line: ::core::line!(),
37//                     col: ::core::column!(),
38//                     file: ::core::file!(),
39//                 },
40//                 vtable: $crate::vtable_for(&__assert_expr),
41//                 debug: $crate::Finalize {
42//                     expr: $crate::expr::AndExpr {
43//                         lhs: $crate::atomic::EqExpr {
44//                             lhs: __0 as *const _ as *const (),
45//                             rhs: __1 as *const _ as *const (),
46//                         },
47//                         rhs: *__2,
48//                     },
49//                     line: (),
50//                     col: (),
51//                     file: (),
52//                 },
53//             };
54//             let __marker = $crate::marker(&__assert_message);
55//             $crate::panic_failed_assert(
56//                 __marker,
57//                 __assert_message.result,
58//                 __assert_message.source,
59//                 __assert_message.vtable,
60//                 __assert_message.debug,
61//             );
62//         }
63//     }
64// }
65
66struct Operand {
67    placeholder_id: Ident,
68    expr: Expr,
69}
70
71enum AssertExpr {
72    BoolExpr(Operand),
73    EqExpr(Operand, Operand),
74    NeExpr(Operand, Operand),
75    LtExpr(Operand, Operand),
76    LeExpr(Operand, Operand),
77    GtExpr(Operand, Operand),
78    GeExpr(Operand, Operand),
79    AndExpr(Box<(AssertExpr, AssertExpr)>),
80    OrExpr(Box<(AssertExpr, AssertExpr)>),
81}
82
83struct Code {
84    assert_expr: TokenStream,
85    source: TokenStream,
86    debug: TokenStream,
87}
88
89impl AssertExpr {
90    fn code(&self, crate_name: syn::Path) -> Code {
91        match self {
92            AssertExpr::BoolExpr(Operand {
93                placeholder_id,
94                expr,
95            }) => Code {
96                assert_expr: quote! { *#placeholder_id },
97                source: quote! { ::core::stringify!(#expr) },
98                debug: quote! { *#placeholder_id },
99            },
100            AssertExpr::EqExpr(
101                Operand {
102                    placeholder_id: left_placeholder_id,
103                    expr: left_expr,
104                },
105                Operand {
106                    placeholder_id: right_placeholder_id,
107                    expr: right_expr,
108                },
109            ) => Code {
110                assert_expr: quote! {
111                    #crate_name::atomic::EqExpr {
112                        lhs: (& &#crate_name::Wrapper(#left_placeholder_id)).wrap().do_wrap(#left_placeholder_id),
113                        rhs: (& &#crate_name::Wrapper(#right_placeholder_id)).wrap().do_wrap(#right_placeholder_id),
114                    }
115                },
116                source: quote! {
117                    #crate_name::atomic::EqExpr {
118                        lhs: ::core::stringify!(#left_expr),
119                        rhs: ::core::stringify!(#right_expr),
120                    }
121                },
122                debug: quote! {
123                    #crate_name::atomic::EqExpr {
124                        lhs: (#left_placeholder_id) as *const _ as *const (),
125                        rhs: (#right_placeholder_id) as *const _ as *const (),
126                    }
127                },
128            },
129            AssertExpr::NeExpr(
130                Operand {
131                    placeholder_id: left_placeholder_id,
132                    expr: left_expr,
133                },
134                Operand {
135                    placeholder_id: right_placeholder_id,
136                    expr: right_expr,
137                },
138            ) => Code {
139                assert_expr: quote! {
140                    #crate_name::atomic::NeExpr {
141                        lhs: (& &#crate_name::Wrapper(#left_placeholder_id)).wrap().do_wrap(#left_placeholder_id),
142                        rhs: (& &#crate_name::Wrapper(#right_placeholder_id)).wrap().do_wrap(#right_placeholder_id),
143                    }
144                },
145                source: quote! {
146                    #crate_name::atomic::NeExpr {
147                        lhs: ::core::stringify!(#left_expr),
148                        rhs: ::core::stringify!(#right_expr),
149                    }
150                },
151                debug: quote! {
152                    #crate_name::atomic::NeExpr {
153                        lhs: (#left_placeholder_id) as *const _ as *const (),
154                        rhs: (#right_placeholder_id) as *const _ as *const (),
155                    }
156                },
157            },
158            AssertExpr::LtExpr(
159                Operand {
160                    placeholder_id: left_placeholder_id,
161                    expr: left_expr,
162                },
163                Operand {
164                    placeholder_id: right_placeholder_id,
165                    expr: right_expr,
166                },
167            ) => Code {
168                assert_expr: quote! {
169                    #crate_name::atomic::LtExpr {
170                        lhs: (& &#crate_name::Wrapper(#left_placeholder_id)).wrap().do_wrap(#left_placeholder_id),
171                        rhs: (& &#crate_name::Wrapper(#right_placeholder_id)).wrap().do_wrap(#right_placeholder_id),
172                    }
173                },
174                source: quote! {
175                    #crate_name::atomic::LtExpr {
176                        lhs: ::core::stringify!(#left_expr),
177                        rhs: ::core::stringify!(#right_expr),
178                    }
179                },
180                debug: quote! {
181                    #crate_name::atomic::LtExpr {
182                        lhs: (#left_placeholder_id) as *const _ as *const (),
183                        rhs: (#right_placeholder_id) as *const _ as *const (),
184                    }
185                },
186            },
187            AssertExpr::LeExpr(
188                Operand {
189                    placeholder_id: left_placeholder_id,
190                    expr: left_expr,
191                },
192                Operand {
193                    placeholder_id: right_placeholder_id,
194                    expr: right_expr,
195                },
196            ) => Code {
197                assert_expr: quote! {
198                    #crate_name::atomic::LeExpr {
199                        lhs: (& &#crate_name::Wrapper(#left_placeholder_id)).wrap().do_wrap(#left_placeholder_id),
200                        rhs: (& &#crate_name::Wrapper(#right_placeholder_id)).wrap().do_wrap(#right_placeholder_id),
201                    }
202                },
203                source: quote! {
204                    #crate_name::atomic::LeExpr {
205                        lhs: ::core::stringify!(#left_expr),
206                        rhs: ::core::stringify!(#right_expr),
207                    }
208                },
209                debug: quote! {
210                    #crate_name::atomic::LeExpr {
211                        lhs: (#left_placeholder_id) as *const _ as *const (),
212                        rhs: (#right_placeholder_id) as *const _ as *const (),
213                    }
214                },
215            },
216            AssertExpr::GtExpr(
217                Operand {
218                    placeholder_id: left_placeholder_id,
219                    expr: left_expr,
220                },
221                Operand {
222                    placeholder_id: right_placeholder_id,
223                    expr: right_expr,
224                },
225            ) => Code {
226                assert_expr: quote! {
227                    #crate_name::atomic::GtExpr {
228                        lhs: (& &#crate_name::Wrapper(#left_placeholder_id)).wrap().do_wrap(#left_placeholder_id),
229                        rhs: (& &#crate_name::Wrapper(#right_placeholder_id)).wrap().do_wrap(#right_placeholder_id),
230                    }
231                },
232                source: quote! {
233                    #crate_name::atomic::GtExpr {
234                        lhs: ::core::stringify!(#left_expr),
235                        rhs: ::core::stringify!(#right_expr),
236                    }
237                },
238                debug: quote! {
239                    #crate_name::atomic::GtExpr {
240                        lhs: (#left_placeholder_id) as *const _ as *const (),
241                        rhs: (#right_placeholder_id) as *const _ as *const (),
242                    }
243                },
244            },
245            AssertExpr::GeExpr(
246                Operand {
247                    placeholder_id: left_placeholder_id,
248                    expr: left_expr,
249                },
250                Operand {
251                    placeholder_id: right_placeholder_id,
252                    expr: right_expr,
253                },
254            ) => Code {
255                assert_expr: quote! {
256                    #crate_name::atomic::GeExpr {
257                        lhs: (& &#crate_name::Wrapper(#left_placeholder_id)).wrap().do_wrap(#left_placeholder_id),
258                        rhs: (& &#crate_name::Wrapper(#right_placeholder_id)).wrap().do_wrap(#right_placeholder_id),
259                    }
260                },
261                source: quote! {
262                    #crate_name::atomic::GeExpr {
263                        lhs: ::core::stringify!(#left_expr),
264                        rhs: ::core::stringify!(#right_expr),
265                    }
266                },
267                debug: quote! {
268                    #crate_name::atomic::GeExpr {
269                        lhs: (#left_placeholder_id) as *const _ as *const (),
270                        rhs: (#right_placeholder_id) as *const _ as *const (),
271                    }
272                },
273            },
274            AssertExpr::AndExpr(inner) => {
275                let (left, right) = &**inner;
276                let Code {
277                    assert_expr: left_assert_expr,
278                    source: left_source,
279                    debug: left_debug,
280                } = left.code(crate_name.clone());
281                let Code {
282                    assert_expr: right_assert_expr,
283                    source: right_source,
284                    debug: right_debug,
285                } = right.code(crate_name.clone());
286                Code {
287                    assert_expr: quote! {
288                        #crate_name::expr::AndExpr {
289                            lhs: (#left_assert_expr),
290                            rhs: (#right_assert_expr),
291                        }
292                    },
293                    source: quote! {
294                        #crate_name::expr::AndExpr {
295                            lhs: (#left_source),
296                            rhs: (#right_source),
297                        }
298                    },
299                    debug: quote! {
300                        #crate_name::expr::AndExpr {
301                            lhs: (#left_debug),
302                            rhs: (#right_debug),
303                        }
304                    },
305                }
306            }
307            AssertExpr::OrExpr(inner) => {
308                let (left, right) = &**inner;
309                let Code {
310                    assert_expr: left_assert_expr,
311                    source: left_source,
312                    debug: left_debug,
313                } = left.code(crate_name.clone());
314                let Code {
315                    assert_expr: right_assert_expr,
316                    source: right_source,
317                    debug: right_debug,
318                } = right.code(crate_name.clone());
319                Code {
320                    assert_expr: quote! {
321                        #crate_name::expr::OrExpr {
322                            lhs: (#left_assert_expr),
323                            rhs: (#right_assert_expr),
324                        }
325                    },
326                    source: quote! {
327                        #crate_name::expr::OrExpr {
328                            lhs: (#left_source),
329                            rhs: (#right_source),
330                        }
331                    },
332                    debug: quote! {
333                        #crate_name::expr::OrExpr {
334                            lhs: (#left_debug),
335                            rhs: (#right_debug),
336                        }
337                    },
338                }
339            }
340        }
341    }
342}
343
344fn usize_to_ident(idx: usize) -> Ident {
345    Ident::new(&format!("__{idx}"), Span::call_site())
346}
347
348fn handle_expr(
349    atomics: &mut Vec<Expr>,
350    diagnostics: &mut Vec<TokenStream>,
351    mut placeholder_id: usize,
352    expr: Expr,
353) -> (usize, AssertExpr) {
354    match expr {
355        Expr::Call(expr)
356            if match &*expr.func {
357                Expr::Path(path) => {
358                    path.path
359                        .get_ident()
360                        .map(|ident| ident.to_string() == "all")
361                        == Some(true)
362                }
363                _ => false,
364            } =>
365        {
366            let mut args = expr.args.into_iter().collect::<Vec<_>>();
367            if args.is_empty() {
368                let expr = Expr::Lit(ExprLit {
369                    attrs: Vec::new(),
370                    lit: Lit::Bool(LitBool {
371                        value: true,
372                        span: Span::call_site(),
373                    }),
374                });
375                atomics.push(expr.clone());
376                (
377                    placeholder_id + 1,
378                    AssertExpr::BoolExpr(Operand {
379                        placeholder_id: usize_to_ident(placeholder_id),
380                        expr,
381                    }),
382                )
383            } else {
384                let mut assert_expr;
385                let mut arg_expr;
386                (placeholder_id, assert_expr) =
387                    handle_expr(atomics, diagnostics, placeholder_id, args.pop().unwrap());
388                while let Some(arg) = args.pop() {
389                    (placeholder_id, arg_expr) =
390                        handle_expr(atomics, diagnostics, placeholder_id, arg);
391                    assert_expr = AssertExpr::AndExpr(Box::new((arg_expr, assert_expr)));
392                }
393                (placeholder_id, assert_expr)
394            }
395        }
396        Expr::Call(expr)
397            if match &*expr.func {
398                Expr::Path(path) => {
399                    path.path
400                        .get_ident()
401                        .map(|ident| ident.to_string() == "any")
402                        == Some(true)
403                }
404                _ => false,
405            } =>
406        {
407            let mut args = expr.args.into_iter().collect::<Vec<_>>();
408            if args.is_empty() {
409                let expr = Expr::Lit(ExprLit {
410                    attrs: Vec::new(),
411                    lit: Lit::Bool(LitBool {
412                        value: false,
413                        span: Span::call_site(),
414                    }),
415                });
416                atomics.push(expr.clone());
417                (
418                    placeholder_id + 1,
419                    AssertExpr::BoolExpr(Operand {
420                        placeholder_id: usize_to_ident(placeholder_id),
421                        expr,
422                    }),
423                )
424            } else {
425                let mut assert_expr;
426                let mut arg_expr;
427                (placeholder_id, assert_expr) =
428                    handle_expr(atomics, diagnostics, placeholder_id, args.pop().unwrap());
429                while let Some(arg) = args.pop() {
430                    (placeholder_id, arg_expr) =
431                        handle_expr(atomics, diagnostics, placeholder_id, arg);
432                    assert_expr = AssertExpr::OrExpr(Box::new((arg_expr, assert_expr)));
433                }
434                (placeholder_id, assert_expr)
435            }
436        }
437        Expr::Binary(ExprBinary {
438            left,
439            right,
440            op: BinOp::Eq(_),
441            ..
442        }) => (
443            {
444                let lhs = usize_to_ident(placeholder_id);
445                let rhs = usize_to_ident(placeholder_id + 1);
446                diagnostics.push(quote! { #lhs == #rhs });
447                atomics.push((*left).clone());
448                atomics.push((*right).clone());
449                placeholder_id + 2
450            },
451            AssertExpr::EqExpr(
452                Operand {
453                    placeholder_id: usize_to_ident(placeholder_id),
454                    expr: *left,
455                },
456                Operand {
457                    placeholder_id: usize_to_ident(placeholder_id + 1),
458                    expr: *right,
459                },
460            ),
461        ),
462
463        Expr::Binary(ExprBinary {
464            left,
465            right,
466            op: BinOp::Ne(_),
467            ..
468        }) => (
469            {
470                let lhs = usize_to_ident(placeholder_id);
471                let rhs = usize_to_ident(placeholder_id + 1);
472                diagnostics.push(quote! { #lhs != #rhs });
473                atomics.push((*left).clone());
474                atomics.push((*right).clone());
475                placeholder_id + 2
476            },
477            AssertExpr::NeExpr(
478                Operand {
479                    placeholder_id: usize_to_ident(placeholder_id),
480                    expr: *left,
481                },
482                Operand {
483                    placeholder_id: usize_to_ident(placeholder_id + 1),
484                    expr: *right,
485                },
486            ),
487        ),
488        Expr::Binary(ExprBinary {
489            left,
490            right,
491            op: BinOp::Lt(_),
492            ..
493        }) => (
494            {
495                let lhs = usize_to_ident(placeholder_id);
496                let rhs = usize_to_ident(placeholder_id + 1);
497                diagnostics.push(quote! { #lhs < #rhs });
498                atomics.push((*left).clone());
499                atomics.push((*right).clone());
500                placeholder_id + 2
501            },
502            AssertExpr::LtExpr(
503                Operand {
504                    placeholder_id: usize_to_ident(placeholder_id),
505                    expr: *left,
506                },
507                Operand {
508                    placeholder_id: usize_to_ident(placeholder_id + 1),
509                    expr: *right,
510                },
511            ),
512        ),
513        Expr::Binary(ExprBinary {
514            left,
515            right,
516            op: BinOp::Le(_),
517            ..
518        }) => (
519            {
520                let lhs = usize_to_ident(placeholder_id);
521                let rhs = usize_to_ident(placeholder_id + 1);
522                diagnostics.push(quote! { #lhs <= #rhs });
523                atomics.push((*left).clone());
524                atomics.push((*right).clone());
525                placeholder_id + 2
526            },
527            AssertExpr::LeExpr(
528                Operand {
529                    placeholder_id: usize_to_ident(placeholder_id),
530                    expr: *left,
531                },
532                Operand {
533                    placeholder_id: usize_to_ident(placeholder_id + 1),
534                    expr: *right,
535                },
536            ),
537        ),
538        Expr::Binary(ExprBinary {
539            left,
540            right,
541            op: BinOp::Gt(_),
542            ..
543        }) => (
544            {
545                let lhs = usize_to_ident(placeholder_id);
546                let rhs = usize_to_ident(placeholder_id + 1);
547                diagnostics.push(quote! { #lhs > #rhs });
548                atomics.push((*left).clone());
549                atomics.push((*right).clone());
550                placeholder_id + 2
551            },
552            AssertExpr::GtExpr(
553                Operand {
554                    placeholder_id: usize_to_ident(placeholder_id),
555                    expr: *left,
556                },
557                Operand {
558                    placeholder_id: usize_to_ident(placeholder_id + 1),
559                    expr: *right,
560                },
561            ),
562        ),
563        Expr::Binary(ExprBinary {
564            left,
565            right,
566            op: BinOp::Ge(_),
567            ..
568        }) => (
569            {
570                let lhs = usize_to_ident(placeholder_id);
571                let rhs = usize_to_ident(placeholder_id + 1);
572                diagnostics.push(quote! { #lhs >= #rhs });
573                atomics.push((*left).clone());
574                atomics.push((*right).clone());
575                placeholder_id + 2
576            },
577            AssertExpr::GeExpr(
578                Operand {
579                    placeholder_id: usize_to_ident(placeholder_id),
580                    expr: *left,
581                },
582                Operand {
583                    placeholder_id: usize_to_ident(placeholder_id + 1),
584                    expr: *right,
585                },
586            ),
587        ),
588        expr => (
589            {
590                let val = usize_to_ident(placeholder_id);
591                diagnostics.push(quote! { *#val });
592                atomics.push(expr.clone());
593                placeholder_id + 1
594            },
595            AssertExpr::BoolExpr(Operand {
596                placeholder_id: usize_to_ident(placeholder_id),
597                expr,
598            }),
599        ),
600    }
601}
602
603type FormatArgs = syn::punctuated::Punctuated<syn::Expr, syn::token::Comma>;
604
605struct Args {
606    crate_name: syn::Path,
607    expr: syn::Expr,
608    format_args: Option<FormatArgs>,
609}
610
611impl syn::parse::Parse for Args {
612    fn parse(input: syn::parse::ParseStream) -> syn::Result<Self> {
613        let crate_name = input.parse()?;
614        let _comma: syn::token::Comma = input.parse()?;
615        let expr = input.parse()?;
616        let format_args = if input.is_empty() {
617            FormatArgs::new()
618        } else {
619            input.parse::<syn::token::Comma>()?;
620            FormatArgs::parse_terminated(input)?
621        };
622
623        let format_args = Some(format_args).filter(|x| !x.is_empty());
624        Ok(Self {
625            crate_name,
626            expr,
627            format_args,
628        })
629    }
630}
631
632#[proc_macro]
633pub fn assert(item: proc_macro::TokenStream) -> proc_macro::TokenStream {
634    let item: TokenStream = item.into();
635    let Ok(input) = parse2::<Args>(item) else {
636        return quote! {
637            ::core::compile_error!("invalid arguments");
638        }
639        .into();
640    };
641
642    let crate_name = input.crate_name;
643    let args = input
644        .format_args
645        .map(|punc| punc.into_iter().collect())
646        .unwrap_or(Vec::new());
647    let body = input.expr;
648
649    let mut atomics = Vec::new();
650    let mut diagnostics = Vec::new();
651    let assert_expr = handle_expr(&mut atomics, &mut diagnostics, 0, body.clone()).1;
652    let atomics = atomics;
653    let placeholders = atomics
654        .iter()
655        .enumerate()
656        .map(|(idx, _)| Ident::new(&format!("__{idx}"), Span::call_site()));
657
658    let Code {
659        assert_expr,
660        source,
661        debug,
662    } = assert_expr.code(crate_name.clone());
663
664    let outer_block = if args.is_empty() {
665        quote! {
666            match (#(&(#atomics),)*) {
667                (#(#placeholders,)*) => {
668                    if false {
669                        #(let _: bool = #diagnostics;)*
670                    }
671                    use #crate_name::Expr;
672                    use #crate_name::TryDebugWrap;
673
674                    let __assert_expr = #crate_name::Finalize {
675                        expr: #assert_expr,
676                        line: (),
677                        col: (),
678                        file: (),
679                    };
680
681                    if !(&&&__assert_expr).eval_expr() {
682                        let __assert_message = #crate_name::DebugMessage {
683                            result: (&&&__assert_expr).result(),
684                            source: &#crate_name::Finalize {
685                                expr: #source,
686                                line: ::core::line!(),
687                                col: ::core::column!(),
688                                file: ::core::file!(),
689                            },
690                            vtable: #crate_name::vtable_for(&__assert_expr),
691                            debug: &#crate_name::Finalize {
692                                expr: #debug,
693                                line: (),
694                                col: (),
695                                file: (),
696                            },
697                            message: ::core::format_args!(""),
698                        };
699                        let __marker = #crate_name::marker(&__assert_message);
700                        #crate_name::panic_failed_assert(
701                            __marker,
702                            __assert_message.result,
703                            __assert_message.source,
704                            __assert_message.vtable,
705                            __assert_message.debug,
706                            );
707                    }
708                }
709            }
710        }
711    } else {
712        quote! {
713            match (#(&(#atomics),)* ::core::format_args!(#(#args,)*)) {
714                (#(#placeholders,)* __message) => {
715                    if false {
716                        #(let _: bool = #diagnostics;)*
717                    }
718
719                    use #crate_name::Expr;
720                    use #crate_name::TryDebugWrap;
721
722                    let __assert_expr = #crate_name::Finalize {
723                        expr: #assert_expr,
724                        line: (),
725                        col: (),
726                        file: (),
727                    };
728
729                    if !(&&&__assert_expr).eval_expr() {
730                        let __assert_message = #crate_name::DebugMessage {
731                            result: (&&&__assert_expr).result(),
732                            source: &#crate_name::Finalize {
733                                expr: #source,
734                                line: ::core::line!(),
735                                col: ::core::column!(),
736                                file: ::core::file!(),
737                            },
738                            vtable: #crate_name::vtable_for(&__assert_expr),
739                            debug: &#crate_name::Finalize {
740                                expr: #debug,
741                                line: (),
742                                col: (),
743                                file: (),
744                            },
745                            message: __message,
746                        };
747                        let __marker = #crate_name::marker(&__assert_message);
748                        #crate_name::panic_failed_assert_with_message(
749                            __marker,
750                            __assert_message.message,
751                            __assert_message.result,
752                            __assert_message.source,
753                            __assert_message.vtable,
754                            __assert_message.debug,
755                        );
756                    }
757                }
758            }
759        }
760    };
761
762    outer_block.into()
763}