Skip to main content

cranelift_isle/
codegen.rs

1//! Generate Rust code from a series of Sequences.
2
3use crate::files::Files;
4use crate::sema::{
5    BuiltinType, ExternalSig, Fields, IntType, ReturnKind, Term, TermEnv, TermId, Type, TypeEnv,
6    TypeId,
7};
8use crate::serialize::{Block, ControlFlow, EvalStep, MatchArm};
9use crate::stablemapset::StableSet;
10use crate::trie_again::{Binding, BindingId, Constraint, RuleSet};
11use std::borrow::Cow;
12use std::fmt::Write;
13use std::slice::Iter;
14use std::sync::Arc;
15
16const DEFAULT_MATCH_ARM_BODY_CLOSURE_THRESHOLD: usize = 256;
17
18/// Options for code generation.
19#[derive(Clone, Debug, Default)]
20pub struct CodegenOptions {
21    /// Do not include the `#![allow(...)]` pragmas in the generated
22    /// source. Useful if it must be include!()'d elsewhere.
23    pub exclude_global_allow_pragmas: bool,
24
25    /// Prefixes to remove when printing file names in generated files. This
26    /// helps keep codegen deterministic.
27    pub prefixes: Vec<Prefix>,
28
29    /// Emit `log::debug!` and `log::trace!` invocations in the generated code to help
30    /// debug rule matching and execution.
31    ///
32    /// In Cranelift this is typically controlled by a cargo feature on the
33    /// crate that includes the generated code (e.g. `cranelift-codegen`).
34    pub emit_logging: bool,
35
36    /// Split large match arms into local closures when generating iterator terms.
37    ///
38    /// In Cranelift this is typically controlled by a cargo feature on the
39    /// crate that includes the generated code (e.g. `cranelift-codegen`).
40    pub split_match_arms: bool,
41
42    /// Threshold for splitting match arms into local closures.
43    ///
44    /// If `None`, a default threshold is used.
45    pub match_arm_split_threshold: Option<usize>,
46}
47
48/// A path prefix which should be replaced when printing file names.
49#[derive(Clone, Debug)]
50pub struct Prefix {
51    /// Prefix to strip
52    pub prefix: String,
53
54    /// Name replacing the stripped prefix.
55    pub name: String,
56}
57
58/// Emit Rust source code for the given type and term environments.
59pub fn codegen(
60    files: Arc<Files>,
61    typeenv: &TypeEnv,
62    termenv: &TermEnv,
63    terms: &[(TermId, RuleSet)],
64    options: &CodegenOptions,
65) -> String {
66    Codegen::compile(files, typeenv, termenv, terms).generate_rust(options)
67}
68
69#[derive(Clone, Debug)]
70struct Codegen<'a> {
71    files: Arc<Files>,
72    typeenv: &'a TypeEnv,
73    termenv: &'a TermEnv,
74    terms: &'a [(TermId, RuleSet)],
75}
76
77enum Nested<'a> {
78    Cases(Iter<'a, EvalStep>),
79    Arms(BindingId, Iter<'a, MatchArm>),
80}
81
82struct BodyContext<'a, W> {
83    out: &'a mut W,
84    ruleset: &'a RuleSet,
85    indent: String,
86    is_ref: StableSet<BindingId>,
87    is_bound: StableSet<BindingId>,
88    term_name: &'a str,
89    emit_logging: bool,
90    split_match_arms: bool,
91    match_arm_split_threshold: Option<usize>,
92
93    // Extra fields for iterator-returning terms.
94    // These fields are used to generate optimized Rust code for iterator-returning terms.
95    /// The number of match splits that have been generated.
96    /// This is used to generate unique names for the match splits.
97    match_split: usize,
98
99    /// The action to take when the iterator overflows.
100    iter_overflow_action: &'static str,
101}
102
103impl<'a, W: Write> BodyContext<'a, W> {
104    fn new(
105        out: &'a mut W,
106        ruleset: &'a RuleSet,
107        term_name: &'a str,
108        emit_logging: bool,
109        split_match_arms: bool,
110        match_arm_split_threshold: Option<usize>,
111        iter_overflow_action: &'static str,
112    ) -> Self {
113        Self {
114            out,
115            ruleset,
116            indent: Default::default(),
117            is_ref: Default::default(),
118            is_bound: Default::default(),
119            term_name,
120            emit_logging,
121            split_match_arms,
122            match_arm_split_threshold,
123            match_split: Default::default(),
124            iter_overflow_action,
125        }
126    }
127
128    fn enter_scope(&mut self) -> StableSet<BindingId> {
129        let new = self.is_bound.clone();
130        std::mem::replace(&mut self.is_bound, new)
131    }
132
133    fn begin_block(&mut self) -> std::fmt::Result {
134        self.indent.push_str("    ");
135        writeln!(self.out, " {{")
136    }
137
138    fn end_block(&mut self, last_line: &str, scope: StableSet<BindingId>) -> std::fmt::Result {
139        if !last_line.is_empty() {
140            writeln!(self.out, "{}{}", &self.indent, last_line)?;
141        }
142        self.is_bound = scope;
143        self.end_block_without_newline()?;
144        writeln!(self.out)
145    }
146
147    fn end_block_without_newline(&mut self) -> std::fmt::Result {
148        self.indent.truncate(self.indent.len() - 4);
149        write!(self.out, "{}}}", &self.indent)
150    }
151
152    fn set_ref(&mut self, binding: BindingId, is_ref: bool) {
153        if is_ref {
154            self.is_ref.insert(binding);
155        } else {
156            debug_assert!(!self.is_ref.contains(&binding));
157        }
158    }
159}
160
161impl<'a> Codegen<'a> {
162    fn compile(
163        files: Arc<Files>,
164        typeenv: &'a TypeEnv,
165        termenv: &'a TermEnv,
166        terms: &'a [(TermId, RuleSet)],
167    ) -> Codegen<'a> {
168        Codegen {
169            files,
170            typeenv,
171            termenv,
172            terms,
173        }
174    }
175
176    fn generate_rust(&self, options: &CodegenOptions) -> String {
177        let mut code = String::new();
178
179        self.generate_header(&mut code, options);
180        self.generate_ctx_trait(&mut code);
181        self.generate_internal_types(&mut code);
182        self.generate_internal_term_constructors(&mut code, options)
183            .unwrap();
184
185        code
186    }
187
188    fn generate_header(&self, code: &mut String, options: &CodegenOptions) {
189        writeln!(code, "// GENERATED BY ISLE. DO NOT EDIT!").unwrap();
190        writeln!(code, "//").unwrap();
191        writeln!(
192            code,
193            "// Generated automatically from the instruction-selection DSL code in:",
194        )
195        .unwrap();
196        for file in &self.files.file_names {
197            writeln!(code, "// - {file}").unwrap();
198        }
199
200        if !options.exclude_global_allow_pragmas {
201            writeln!(
202                code,
203                "\n#![allow(dead_code, unreachable_code, unreachable_patterns)]"
204            )
205            .unwrap();
206            writeln!(
207                code,
208                "#![allow(unused_imports, unused_variables, non_snake_case, unused_mut)]"
209            )
210            .unwrap();
211            writeln!(
212                code,
213                "#![allow(irrefutable_let_patterns, unused_assignments, non_camel_case_types)]"
214            )
215            .unwrap();
216        }
217
218        writeln!(code, "\nuse super::*;  // Pulls in all external types.").unwrap();
219        writeln!(code, "use core::marker::PhantomData;").unwrap();
220    }
221
222    fn generate_trait_sig(&self, code: &mut String, indent: &str, sig: &ExternalSig) {
223        let ret_tuple = format!(
224            "{open_paren}{rets}{close_paren}",
225            open_paren = if sig.ret_tys.len() != 1 { "(" } else { "" },
226            rets = sig
227                .ret_tys
228                .iter()
229                .map(|&ty| self.type_name(ty, /* by_ref = */ false))
230                .collect::<Vec<_>>()
231                .join(", "),
232            close_paren = if sig.ret_tys.len() != 1 { ")" } else { "" },
233        );
234
235        if sig.ret_kind == ReturnKind::Iterator {
236            writeln!(
237                code,
238                "{indent}type {name}_returns: Default + IntoContextIter<Context = Self, Output = {output}>;",
239                indent = indent,
240                name = sig.func_name,
241                output = ret_tuple,
242            )
243            .unwrap();
244        }
245
246        let ret_ty = match sig.ret_kind {
247            ReturnKind::Plain => ret_tuple,
248            ReturnKind::Option => format!("Option<{ret_tuple}>"),
249            ReturnKind::Iterator => format!("()"),
250        };
251
252        writeln!(
253            code,
254            "{indent}fn {name}(&mut self, {params}) -> {ret_ty};",
255            indent = indent,
256            name = sig.func_name,
257            params = sig
258                .param_tys
259                .iter()
260                .enumerate()
261                .map(|(i, &ty)| format!("arg{}: {}", i, self.type_name(ty, /* by_ref = */ true)))
262                .chain(if sig.ret_kind == ReturnKind::Iterator {
263                    Some(format!("returns: &mut Self::{}_returns", sig.func_name))
264                } else {
265                    None
266                })
267                .collect::<Vec<_>>()
268                .join(", "),
269            ret_ty = ret_ty,
270        )
271        .unwrap();
272    }
273
274    fn generate_ctx_trait(&self, code: &mut String) {
275        writeln!(code).unwrap();
276        writeln!(
277            code,
278            "/// Context during lowering: an implementation of this trait"
279        )
280        .unwrap();
281        writeln!(
282            code,
283            "/// must be provided with all external constructors and extractors."
284        )
285        .unwrap();
286        writeln!(
287            code,
288            "/// A mutable borrow is passed along through all lowering logic."
289        )
290        .unwrap();
291        writeln!(code, "pub trait Context {{").unwrap();
292        for term in &self.termenv.terms {
293            if term.has_external_extractor() {
294                let ext_sig = term.extractor_sig(self.typeenv).unwrap();
295                self.generate_trait_sig(code, "    ", &ext_sig);
296            }
297            if term.has_external_constructor() {
298                let ext_sig = term.constructor_sig(self.typeenv).unwrap();
299                self.generate_trait_sig(code, "    ", &ext_sig);
300            }
301        }
302        writeln!(code, "}}").unwrap();
303        writeln!(
304            code,
305            r#"
306pub trait ContextIter {{
307    type Context;
308    type Output;
309    fn next(&mut self, ctx: &mut Self::Context) -> Option<Self::Output>;
310    fn size_hint(&self) -> (usize, Option<usize>) {{ (0, None) }}
311}}
312
313pub trait IntoContextIter {{
314    type Context;
315    type Output;
316    type IntoIter: ContextIter<Context = Self::Context, Output = Self::Output>;
317    fn into_context_iter(self) -> Self::IntoIter;
318}}
319
320pub trait Length {{
321    fn len(&self) -> usize;
322}}
323
324impl<T> Length for alloc::vec::Vec<T> {{
325    fn len(&self) -> usize {{
326        alloc::vec::Vec::len(self)
327    }}
328}}
329
330pub struct ContextIterWrapper<I, C> {{
331    iter: I,
332    _ctx: core::marker::PhantomData<C>,
333}}
334impl<I: Default, C> Default for ContextIterWrapper<I, C> {{
335    fn default() -> Self {{
336        ContextIterWrapper {{
337            iter: I::default(),
338            _ctx: core::marker::PhantomData
339        }}
340    }}
341}}
342impl<I, C> core::ops::Deref for ContextIterWrapper<I, C> {{
343    type Target = I;
344    fn deref(&self) -> &I {{
345        &self.iter
346    }}
347}}
348impl<I, C> core::ops::DerefMut for ContextIterWrapper<I, C> {{
349    fn deref_mut(&mut self) -> &mut I {{
350        &mut self.iter
351    }}
352}}
353impl<I: Iterator, C: Context> From<I> for ContextIterWrapper<I, C> {{
354    fn from(iter: I) -> Self {{
355        Self {{ iter, _ctx: core::marker::PhantomData }}
356    }}
357}}
358impl<I: Iterator, C: Context> ContextIter for ContextIterWrapper<I, C> {{
359    type Context = C;
360    type Output = I::Item;
361    fn next(&mut self, _ctx: &mut Self::Context) -> Option<Self::Output> {{
362        self.iter.next()
363    }}
364    fn size_hint(&self) -> (usize, Option<usize>) {{
365        self.iter.size_hint()
366    }}
367}}
368impl<I: IntoIterator, C: Context> IntoContextIter for ContextIterWrapper<I, C> {{
369    type Context = C;
370    type Output = I::Item;
371    type IntoIter = ContextIterWrapper<I::IntoIter, C>;
372    fn into_context_iter(self) -> Self::IntoIter {{
373        ContextIterWrapper {{
374            iter: self.iter.into_iter(),
375            _ctx: core::marker::PhantomData
376        }}
377    }}
378}}
379impl<T, E: Extend<T>, C> Extend<T> for ContextIterWrapper<E, C> {{
380    fn extend<I: IntoIterator<Item = T>>(&mut self, iter: I) {{
381        self.iter.extend(iter);
382    }}
383}}
384impl<L: Length, C> Length for ContextIterWrapper<L, C> {{
385    fn len(&self) -> usize {{
386        self.iter.len()
387    }}
388}}
389           "#,
390        )
391        .unwrap();
392    }
393
394    fn generate_internal_types(&self, code: &mut String) {
395        for ty in &self.typeenv.types {
396            match ty {
397                &Type::Enum {
398                    name,
399                    is_extern,
400                    is_nodebug,
401                    ref variants,
402                    pos,
403                    ..
404                } if !is_extern => {
405                    let name = &self.typeenv.syms[name.index()];
406                    writeln!(
407                        code,
408                        "\n/// Internal type {}: defined at {}.",
409                        name,
410                        pos.pretty_print_line(&self.files)
411                    )
412                    .unwrap();
413
414                    // Generate the `derive`s.
415                    let debug_derive = if is_nodebug { "" } else { ", Debug" };
416                    if variants.iter().all(|v| v.fields == Fields::Unit) {
417                        writeln!(code, "#[derive(Copy, Clone, PartialEq, Eq{debug_derive})]")
418                            .unwrap();
419                    } else {
420                        writeln!(code, "#[derive(Clone{debug_derive})]").unwrap();
421                    }
422
423                    writeln!(code, "pub enum {name} {{").unwrap();
424                    for variant in variants {
425                        let name = &self.typeenv.syms[variant.name.index()];
426                        write!(code, "    {name}").unwrap();
427                        self.generate_fields(&mut *code, &variant.fields, "    ", false)
428                            .unwrap();
429                        writeln!(code, ",").unwrap();
430                    }
431                    writeln!(code, "}}").unwrap();
432                }
433                &Type::Struct {
434                    name,
435                    is_extern,
436                    is_nodebug,
437                    ref fields,
438                    pos,
439                    ..
440                } if !is_extern => {
441                    let name = &self.typeenv.syms[name.index()];
442                    writeln!(
443                        code,
444                        "\n/// Internal type {}: defined at {}.",
445                        name,
446                        pos.pretty_print_line(&self.files)
447                    )
448                    .unwrap();
449
450                    // Generate the `derive`s.
451                    let debug_derive = if is_nodebug { "" } else { ", Debug" };
452                    match fields {
453                        Fields::Unit => {
454                            writeln!(code, "#[derive(Copy, Clone, PartialEq, Eq{debug_derive})]")
455                                .unwrap();
456                            writeln!(code, "pub struct {name};").unwrap();
457                        }
458                        Fields::Tuple(_) | Fields::Struct(_) => {
459                            writeln!(code, "#[derive(Clone{debug_derive})]").unwrap();
460                            write!(code, "pub struct {name}").unwrap();
461                            self.generate_fields(&mut *code, &fields, "", true).unwrap();
462                            if matches!(fields, Fields::Tuple(_)) {
463                                write!(code, ";").unwrap();
464                            }
465                        }
466                    }
467                }
468                _ => {}
469            }
470        }
471    }
472
473    fn generate_fields(
474        &self,
475        code: &mut String,
476        fields: &Fields,
477        pad: &str,
478        allow_pub: bool,
479    ) -> std::fmt::Result {
480        let _pub = if allow_pub { "pub " } else { "" };
481        match fields {
482            Fields::Unit => Ok(()),
483            Fields::Struct(fields) => {
484                writeln!(code, " {{")?;
485                for field in &fields.fields {
486                    let name = &self.typeenv.syms[field.name.index()];
487                    let ty_name = self.typeenv.types[field.ty.index()].name(self.typeenv);
488                    writeln!(code, "{pad}    {_pub}{name}: {ty_name},")?;
489                }
490                write!(code, "{pad}}}")
491            }
492            Fields::Tuple(fields) => {
493                write!(code, "(")?;
494                for (i, field) in fields.fields.iter().enumerate() {
495                    let ty_name = self.typeenv.types[field.ty.index()].name(self.typeenv);
496                    write!(code, "{_pub}{ty_name}")?;
497                    if i < fields.fields.len() - 1 {
498                        write!(code, ", ")?;
499                    }
500                }
501                write!(code, ")")
502            }
503        }
504    }
505
506    fn type_name(&self, typeid: TypeId, by_ref: bool) -> String {
507        match self.typeenv.types[typeid.index()] {
508            Type::Builtin(bt) => String::from(bt.name()),
509            Type::Primitive(_, sym, _) => self.typeenv.syms[sym.index()].clone(),
510            Type::Enum { name, .. } | Type::Struct { name, .. } => {
511                let r = if by_ref { "&" } else { "" };
512                format!("{}{}", r, self.typeenv.syms[name.index()])
513            }
514        }
515    }
516
517    fn generate_internal_term_constructors(
518        &self,
519        code: &mut String,
520        options: &CodegenOptions,
521    ) -> std::fmt::Result {
522        for &(termid, ref ruleset) in self.terms.iter() {
523            let root = crate::serialize::serialize(ruleset);
524
525            let termdata = &self.termenv.terms[termid.index()];
526            let term_name = &self.typeenv.syms[termdata.name.index()];
527
528            // Split a match if the term returns an iterator.
529            let mut ctx = BodyContext::new(
530                code,
531                ruleset,
532                term_name,
533                options.emit_logging,
534                options.split_match_arms,
535                options.match_arm_split_threshold,
536                "return;", // At top level, we just return.
537            );
538
539            // Generate the function signature.
540            writeln!(ctx.out)?;
541            writeln!(
542                ctx.out,
543                "{}// Generated as internal constructor for term {}.",
544                &ctx.indent, term_name,
545            )?;
546
547            let sig = termdata.constructor_sig(self.typeenv).unwrap();
548            writeln!(
549                ctx.out,
550                "{}pub fn {}<C: Context>(",
551                &ctx.indent, sig.func_name
552            )?;
553
554            writeln!(ctx.out, "{}    ctx: &mut C,", &ctx.indent)?;
555            for (i, &ty) in sig.param_tys.iter().enumerate() {
556                let (is_ref, ty) = self.ty(ty);
557                write!(ctx.out, "{}    arg{}: ", &ctx.indent, i)?;
558                write!(ctx.out, "{}{}", if is_ref { "&" } else { "" }, ty)?;
559                if let Some(binding) = ctx.ruleset.find_binding(&Binding::Argument {
560                    index: i.try_into().unwrap(),
561                }) {
562                    ctx.set_ref(binding, is_ref);
563                }
564                writeln!(ctx.out, ",")?;
565            }
566
567            let (_, ret) = self.ty(sig.ret_tys[0]);
568
569            if let ReturnKind::Iterator = sig.ret_kind {
570                writeln!(
571                    ctx.out,
572                    "{}    returns: &mut (impl Extend<{}> + Length),",
573                    &ctx.indent, ret
574                )?;
575            }
576
577            write!(ctx.out, "{}) -> ", &ctx.indent)?;
578            match sig.ret_kind {
579                ReturnKind::Iterator => write!(ctx.out, "()")?,
580                ReturnKind::Option => write!(ctx.out, "Option<{ret}>")?,
581                ReturnKind::Plain => write!(ctx.out, "{ret}")?,
582            };
583            // Generating the function signature is done.
584
585            let last_expr = if let Some(EvalStep {
586                check: ControlFlow::Return { .. },
587                ..
588            }) = root.steps.last()
589            {
590                // If there's an outermost fallback, no need for another `return` statement.
591                String::new()
592            } else {
593                match sig.ret_kind {
594                    ReturnKind::Iterator => String::new(),
595                    ReturnKind::Option => "None".to_string(),
596                    ReturnKind::Plain => format!(
597                        "unreachable!(\"no rule matched for term {{}} at {{}}; should it be partial?\", {:?}, {:?})",
598                        term_name,
599                        termdata.decl_pos.pretty_print_line(&self.files)
600                    ),
601                }
602            };
603
604            let scope = ctx.enter_scope();
605            self.emit_block(&mut ctx, &root, sig.ret_kind, &last_expr, scope, options)?;
606        }
607        Ok(())
608    }
609
610    fn ty(&self, typeid: TypeId) -> (bool, String) {
611        let ty = &self.typeenv.types[typeid.index()];
612        let name = ty.name(self.typeenv);
613        let is_ref = match ty {
614            Type::Builtin(_) | Type::Primitive(..) => false,
615            Type::Enum { .. } | Type::Struct { .. } => true,
616        };
617        (is_ref, String::from(name))
618    }
619
620    fn validate_block(ret_kind: ReturnKind, block: &Block) -> Nested<'_> {
621        if !matches!(ret_kind, ReturnKind::Iterator) {
622            // Loops are only allowed if we're returning an iterator.
623            assert!(
624                !block
625                    .steps
626                    .iter()
627                    .any(|c| matches!(c.check, ControlFlow::Loop { .. }))
628            );
629
630            // Unless we're returning an iterator, a case which returns a result must be the last
631            // case in a block.
632            if let Some(result_pos) = block
633                .steps
634                .iter()
635                .position(|c| matches!(c.check, ControlFlow::Return { .. }))
636            {
637                assert_eq!(block.steps.len() - 1, result_pos);
638            }
639        }
640
641        Nested::Cases(block.steps.iter())
642    }
643
644    fn block_weight(block: &Block) -> usize {
645        fn cf_weight(cf: &ControlFlow) -> usize {
646            match cf {
647                ControlFlow::Match { arms, .. } => {
648                    arms.iter().map(|a| Codegen::block_weight(&a.body)).sum()
649                }
650                ControlFlow::Equal { body, .. } => Codegen::block_weight(body),
651                ControlFlow::Loop { body, .. } => Codegen::block_weight(body),
652                ControlFlow::Return { .. } => 0,
653            }
654        }
655
656        block.steps.iter().map(|s| 1 + cf_weight(&s.check)).sum()
657    }
658
659    fn emit_block<W: Write>(
660        &self,
661        ctx: &mut BodyContext<W>,
662        block: &Block,
663        ret_kind: ReturnKind,
664        last_expr: &str,
665        scope: StableSet<BindingId>,
666        _options: &CodegenOptions,
667    ) -> std::fmt::Result {
668        ctx.begin_block()?;
669        self.emit_block_contents(ctx, block, ret_kind, last_expr, scope)
670    }
671
672    fn emit_block_contents<W: Write>(
673        &self,
674        ctx: &mut BodyContext<W>,
675        block: &Block,
676        ret_kind: ReturnKind,
677        last_expr: &str,
678        scope: StableSet<BindingId>,
679    ) -> std::fmt::Result {
680        let mut stack = Vec::new();
681        stack.push((Self::validate_block(ret_kind, block), last_expr, scope));
682
683        while let Some((mut nested, last_line, scope)) = stack.pop() {
684            match &mut nested {
685                Nested::Cases(cases) => {
686                    let Some(case) = cases.next() else {
687                        ctx.end_block(last_line, scope)?;
688                        continue;
689                    };
690                    // Iterator isn't done, put it back on the stack.
691                    stack.push((nested, last_line, scope));
692
693                    for &expr in case.bind_order.iter() {
694                        let iter_return = match &ctx.ruleset.bindings[expr.index()] {
695                            Binding::Extractor { term, .. } => {
696                                let termdata = &self.termenv.terms[term.index()];
697                                let sig = termdata.extractor_sig(self.typeenv).unwrap();
698                                if sig.ret_kind == ReturnKind::Iterator {
699                                    if termdata.has_external_extractor() {
700                                        Some(format!("C::{}_returns", sig.func_name))
701                                    } else {
702                                        Some(format!("ContextIterWrapper::<ConstructorVec<_>, _>"))
703                                    }
704                                } else {
705                                    None
706                                }
707                            }
708                            Binding::Constructor { term, .. } => {
709                                let termdata = &self.termenv.terms[term.index()];
710                                let sig = termdata.constructor_sig(self.typeenv).unwrap();
711                                if sig.ret_kind == ReturnKind::Iterator {
712                                    if termdata.has_external_constructor() {
713                                        Some(format!("C::{}_returns", sig.func_name))
714                                    } else {
715                                        Some(format!("ContextIterWrapper::<ConstructorVec<_>, _>"))
716                                    }
717                                } else {
718                                    None
719                                }
720                            }
721                            _ => None,
722                        };
723                        if let Some(ty) = iter_return {
724                            writeln!(
725                                ctx.out,
726                                "{}let mut v{} = {}::default();",
727                                &ctx.indent,
728                                expr.index(),
729                                ty
730                            )?;
731                            write!(ctx.out, "{}", &ctx.indent)?;
732                        } else {
733                            write!(ctx.out, "{}let v{} = ", &ctx.indent, expr.index())?;
734                        }
735                        self.emit_expr(ctx, expr)?;
736                        writeln!(ctx.out, ";")?;
737                        ctx.is_bound.insert(expr);
738                    }
739
740                    match &case.check {
741                        // Use a shorthand notation if there's only one match arm.
742                        ControlFlow::Match { source, arms } if arms.len() == 1 => {
743                            let arm = &arms[0];
744                            let scope = ctx.enter_scope();
745                            match arm.constraint {
746                                Constraint::ConstBool { .. }
747                                | Constraint::ConstInt { .. }
748                                | Constraint::ConstPrim { .. } => {
749                                    write!(ctx.out, "{}if ", &ctx.indent)?;
750                                    self.emit_expr(ctx, *source)?;
751                                    write!(ctx.out, " == ")?;
752                                    self.emit_constraint(ctx, *source, arm)?;
753                                }
754                                Constraint::Variant { .. }
755                                | Constraint::Struct { .. }
756                                | Constraint::Some => {
757                                    write!(ctx.out, "{}if let ", &ctx.indent)?;
758                                    self.emit_constraint(ctx, *source, arm)?;
759                                    write!(ctx.out, " = ")?;
760                                    self.emit_source(ctx, *source, arm.constraint)?;
761                                }
762                            }
763                            ctx.begin_block()?;
764                            stack.push((Self::validate_block(ret_kind, &arm.body), "", scope));
765                        }
766
767                        ControlFlow::Match { source, arms } => {
768                            let scope = ctx.enter_scope();
769                            write!(ctx.out, "{}match ", &ctx.indent)?;
770                            self.emit_source(ctx, *source, arms[0].constraint)?;
771                            ctx.begin_block()?;
772
773                            // Always add a catchall arm, because we
774                            // don't do exhaustiveness checking on the
775                            // match arms.
776                            stack.push((Nested::Arms(*source, arms.iter()), "_ => {}", scope));
777                        }
778
779                        ControlFlow::Equal { a, b, body } => {
780                            let scope = ctx.enter_scope();
781                            write!(ctx.out, "{}if ", &ctx.indent)?;
782                            self.emit_expr(ctx, *a)?;
783                            write!(ctx.out, " == ")?;
784                            self.emit_expr(ctx, *b)?;
785                            ctx.begin_block()?;
786                            stack.push((Self::validate_block(ret_kind, body), "", scope));
787                        }
788
789                        ControlFlow::Loop { result, body } => {
790                            let source = match &ctx.ruleset.bindings[result.index()] {
791                                Binding::Iterator { source } => source,
792                                _ => unreachable!("Loop from a non-Iterator"),
793                            };
794                            let scope = ctx.enter_scope();
795
796                            writeln!(
797                                ctx.out,
798                                "{}let mut v{} = v{}.into_context_iter();",
799                                &ctx.indent,
800                                source.index(),
801                                source.index(),
802                            )?;
803
804                            write!(
805                                ctx.out,
806                                "{}while let Some(v{}) = v{}.next(ctx)",
807                                &ctx.indent,
808                                result.index(),
809                                source.index()
810                            )?;
811                            ctx.is_bound.insert(*result);
812                            ctx.begin_block()?;
813                            stack.push((Self::validate_block(ret_kind, body), "", scope));
814                        }
815
816                        &ControlFlow::Return { pos, result } => {
817                            writeln!(
818                                ctx.out,
819                                "{}// Rule at {}.",
820                                &ctx.indent,
821                                pos.pretty_print_line(&self.files)
822                            )?;
823                            if ctx.emit_logging {
824                                // Produce a valid Rust string literal with escapes.
825                                let pp = pos.pretty_print_line(&self.files);
826                                writeln!(
827                                    ctx.out,
828                                    "{}log::debug!(\"ISLE {{}} {{}}\", {:?}, {:?});",
829                                    &ctx.indent, ctx.term_name, pp
830                                )?;
831                            }
832                            write!(ctx.out, "{}", &ctx.indent)?;
833                            match ret_kind {
834                                ReturnKind::Plain | ReturnKind::Option => {
835                                    write!(ctx.out, "return ")?
836                                }
837                                ReturnKind::Iterator => write!(ctx.out, "returns.extend(Some(")?,
838                            }
839                            self.emit_expr(ctx, result)?;
840                            if ctx.is_ref.contains(&result) {
841                                write!(ctx.out, ".clone()")?;
842                            }
843                            match ret_kind {
844                                ReturnKind::Plain | ReturnKind::Option => writeln!(ctx.out, ";")?,
845                                ReturnKind::Iterator => {
846                                    writeln!(ctx.out, "));")?;
847                                    writeln!(
848                                        ctx.out,
849                                        "{}if returns.len() >= MAX_ISLE_RETURNS {{ {} }}",
850                                        ctx.indent, ctx.iter_overflow_action
851                                    )?;
852                                }
853                            }
854                        }
855                    }
856                }
857
858                Nested::Arms(source, arms) => {
859                    let Some(arm) = arms.next() else {
860                        ctx.end_block(last_line, scope)?;
861                        continue;
862                    };
863                    let source = *source;
864                    // Iterator isn't done, put it back on the stack.
865                    stack.push((nested, last_line, scope));
866
867                    let scope = ctx.enter_scope();
868                    write!(ctx.out, "{}", &ctx.indent)?;
869                    self.emit_constraint(ctx, source, arm)?;
870                    write!(ctx.out, " =>")?;
871                    ctx.begin_block()?;
872
873                    // Compile-time optimization: huge function bodies (often from very large match arms
874                    // of constructor bodies)cause rustc to spend a lot of time in analysis passes.
875                    // Wrap such bodies in a local closure to move the bulk of the work into a separate body
876                    // without needing to know the types of captured locals.
877                    let match_arm_body_closure_threshold = ctx
878                        .match_arm_split_threshold
879                        .unwrap_or(DEFAULT_MATCH_ARM_BODY_CLOSURE_THRESHOLD);
880                    if ctx.split_match_arms
881                        && ret_kind == ReturnKind::Iterator
882                        && Codegen::block_weight(&arm.body) > match_arm_body_closure_threshold
883                    {
884                        let closure_id = ctx.match_split;
885                        ctx.match_split += 1;
886
887                        write!(ctx.out, "{}if (|| -> bool", &ctx.indent)?;
888                        ctx.begin_block()?;
889
890                        let old_overflow_action = ctx.iter_overflow_action;
891                        ctx.iter_overflow_action = "return true;";
892                        let closure_scope = ctx.enter_scope();
893                        self.emit_block_contents(ctx, &arm.body, ret_kind, "false", closure_scope)?;
894                        ctx.iter_overflow_action = old_overflow_action;
895
896                        // Close `if (|| -> bool { ... })()` and stop the outer function on
897                        // iterator-overflow.
898                        writeln!(
899                            ctx.out,
900                            "{})() {{ {} }} // __isle_arm_{}",
901                            &ctx.indent, ctx.iter_overflow_action, closure_id
902                        )?;
903
904                        ctx.end_block("", scope)?;
905                    } else {
906                        stack.push((Self::validate_block(ret_kind, &arm.body), "", scope));
907                    }
908                }
909            }
910        }
911
912        Ok(())
913    }
914
915    fn emit_expr<W: Write>(&self, ctx: &mut BodyContext<W>, result: BindingId) -> std::fmt::Result {
916        if ctx.is_bound.contains(&result) {
917            return write!(ctx.out, "v{}", result.index());
918        }
919
920        let binding = &ctx.ruleset.bindings[result.index()];
921
922        let call = |ctx: &mut BodyContext<W>,
923                    term: TermId,
924                    parameters: &[BindingId],
925                    get_sig: fn(&Term, &TypeEnv) -> Option<ExternalSig>| {
926            let termdata = &self.termenv.terms[term.index()];
927            let sig = get_sig(termdata, self.typeenv).unwrap();
928            if let &[ret_ty] = &sig.ret_tys[..] {
929                let (is_ref, _) = self.ty(ret_ty);
930                if is_ref {
931                    ctx.set_ref(result, true);
932                    write!(ctx.out, "&")?;
933                }
934            }
935            write!(ctx.out, "{}(ctx", sig.full_name)?;
936            debug_assert_eq!(parameters.len(), sig.param_tys.len());
937            for (&parameter, &arg_ty) in parameters.iter().zip(sig.param_tys.iter()) {
938                let (is_ref, _) = self.ty(arg_ty);
939                write!(ctx.out, ", ")?;
940                let (before, after) = match (is_ref, ctx.is_ref.contains(&parameter)) {
941                    (false, true) => ("", ".clone()"),
942                    (true, false) => ("&", ""),
943                    _ => ("", ""),
944                };
945                write!(ctx.out, "{before}")?;
946                self.emit_expr(ctx, parameter)?;
947                write!(ctx.out, "{after}")?;
948            }
949            if let ReturnKind::Iterator = sig.ret_kind {
950                write!(ctx.out, ", &mut v{}", result.index())?;
951            }
952            write!(ctx.out, ")")
953        };
954
955        let extract_fields = |ctx: &mut BodyContext<W>,
956                              field_bindings: &[BindingId],
957                              fields: &Fields|
958         -> std::fmt::Result {
959            if !field_bindings.is_empty() {
960                ctx.begin_block()?;
961                for (i, value) in field_bindings.iter().enumerate() {
962                    let field_name = match fields {
963                        Fields::Unit => panic!(),
964                        Fields::Struct(fields) => {
965                            Cow::Borrowed(&self.typeenv.syms[fields.fields[i].name.index()])
966                        }
967                        Fields::Tuple(_) => Cow::Owned(format!("{i}")),
968                    };
969                    write!(ctx.out, "{}{field_name}: ", &ctx.indent)?;
970                    self.emit_expr(ctx, *value)?;
971                    if ctx.is_ref.contains(value) {
972                        write!(ctx.out, ".clone()")?;
973                    }
974                    writeln!(ctx.out, ",")?;
975                }
976                ctx.end_block_without_newline()?;
977            }
978            Ok(())
979        };
980
981        match binding {
982            &Binding::ConstBool { val, .. } => self.emit_bool(ctx, val),
983            &Binding::ConstInt { val, ty } => self.emit_int(ctx, val, ty),
984            Binding::ConstPrim { val } => write!(ctx.out, "{}", &self.typeenv.syms[val.index()]),
985            Binding::Argument { index } => write!(ctx.out, "arg{}", index.index()),
986            Binding::Extractor { term, parameter } => call(
987                ctx,
988                *term,
989                std::slice::from_ref(parameter),
990                Term::extractor_sig,
991            ),
992            Binding::Constructor {
993                term, parameters, ..
994            } => call(ctx, *term, &parameters[..], Term::constructor_sig),
995
996            Binding::MakeVariant {
997                ty,
998                variant,
999                fields,
1000            } => {
1001                let (name, variants) = match &self.typeenv.types[ty.index()] {
1002                    Type::Enum { name, variants, .. } => (name, variants),
1003                    _ => unreachable!("MakeVariant with non-enum type"),
1004                };
1005                let variant = &variants[variant.index()];
1006                write!(
1007                    ctx.out,
1008                    "{}::{}",
1009                    &self.typeenv.syms[name.index()],
1010                    &self.typeenv.syms[variant.name.index()]
1011                )?;
1012                extract_fields(ctx, fields, &variant.fields)
1013            }
1014            Binding::MakeStruct { ty, fields } => {
1015                let (name, type_fields) = match &self.typeenv.types[ty.index()] {
1016                    Type::Struct { name, fields, .. } => (name, fields),
1017                    _ => unreachable!("MakeStruct with non-struct type"),
1018                };
1019                write!(ctx.out, "{}", &self.typeenv.syms[name.index()],)?;
1020                extract_fields(ctx, fields, type_fields)
1021            }
1022
1023            &Binding::MakeSome { inner } => {
1024                write!(ctx.out, "Some(")?;
1025                self.emit_expr(ctx, inner)?;
1026                write!(ctx.out, ")")
1027            }
1028            &Binding::MatchSome { source } => {
1029                self.emit_expr(ctx, source)?;
1030                // When isle uses an implicit `convert` from A to B and the conversion declaration is `partial`,
1031                // it emits:
1032                // ```
1033                // let v1: &Option<B> = &C::a_to_b(ctx, arg0);
1034                // let v2: &B = v1?;
1035                // ```
1036                // This fails to compile since you can't `?` a `&Option<T>`, only on an `Option` itself.
1037                // So add a `.as_ref()` to convert it to `Option<&T>` before `?`.
1038                if ctx.is_ref.contains(&source) {
1039                    write!(ctx.out, ".as_ref()?")
1040                } else {
1041                    write!(ctx.out, "?")
1042                }
1043            }
1044            &Binding::MatchTuple { source, field } => {
1045                self.emit_expr(ctx, source)?;
1046                write!(ctx.out, ".{}", field.index())
1047            }
1048
1049            // These are not supposed to happen. If they do, make the generated code fail to compile
1050            // so this is easier to debug than if we panic during codegen.
1051            &Binding::MatchVariant { source, field, .. }
1052            | &Binding::ExtractStruct { source, field, .. } => {
1053                self.emit_expr(ctx, source)?;
1054                write!(ctx.out, ".{} /*FIXME*/", field.index())
1055            }
1056            &Binding::Iterator { source } => {
1057                self.emit_expr(ctx, source)?;
1058                write!(ctx.out, ".next() /*FIXME*/")
1059            }
1060        }
1061    }
1062
1063    fn emit_source<W: Write>(
1064        &self,
1065        ctx: &mut BodyContext<W>,
1066        source: BindingId,
1067        constraint: Constraint,
1068    ) -> std::fmt::Result {
1069        if let Constraint::Variant { .. } | Constraint::Struct { .. } = constraint {
1070            if !ctx.is_ref.contains(&source) {
1071                write!(ctx.out, "&")?;
1072            }
1073        }
1074        self.emit_expr(ctx, source)
1075    }
1076
1077    fn emit_constraint<W: Write>(
1078        &self,
1079        ctx: &mut BodyContext<W>,
1080        source: BindingId,
1081        arm: &MatchArm,
1082    ) -> std::fmt::Result {
1083        let MatchArm {
1084            constraint,
1085            bindings,
1086            ..
1087        } = arm;
1088        for binding in bindings.iter() {
1089            if let &Some(binding) = binding {
1090                ctx.is_bound.insert(binding);
1091            }
1092        }
1093        match *constraint {
1094            Constraint::ConstBool { val, .. } => self.emit_bool(ctx, val),
1095            Constraint::ConstInt { val, ty } => self.emit_int(ctx, val, ty),
1096            Constraint::ConstPrim { val } => {
1097                write!(ctx.out, "{}", &self.typeenv.syms[val.index()])
1098            }
1099            Constraint::Variant { ty, variant, .. } => {
1100                let (name, variants) = match &self.typeenv.types[ty.index()] {
1101                    Type::Enum { name, variants, .. } => (name, variants),
1102                    _ => unreachable!("Variant constraint on non-enum type"),
1103                };
1104                let variant = &variants[variant.index()];
1105                write!(
1106                    ctx.out,
1107                    "&{}::{}",
1108                    &self.typeenv.syms[name.index()],
1109                    &self.typeenv.syms[variant.name.index()]
1110                )?;
1111                self.emit_fields(ctx, bindings, &variant.fields)?;
1112                Ok(())
1113            }
1114            Constraint::Struct { ty, .. } => {
1115                let (name, fields) = match &self.typeenv.types[ty.index()] {
1116                    Type::Struct { name, fields, .. } => (name, fields),
1117                    _ => unreachable!("Struct constraint on non-struct type"),
1118                };
1119                write!(ctx.out, "&{}", &self.typeenv.syms[name.index()],)?;
1120                self.emit_fields(ctx, bindings, &fields)?;
1121                Ok(())
1122            }
1123            Constraint::Some => {
1124                write!(ctx.out, "Some(")?;
1125                if let Some(binding) = bindings[0] {
1126                    ctx.set_ref(binding, ctx.is_ref.contains(&source));
1127                    write!(ctx.out, "v{}", binding.index())?;
1128                } else {
1129                    write!(ctx.out, "_")?;
1130                }
1131                write!(ctx.out, ")")
1132            }
1133        }
1134    }
1135
1136    fn emit_fields<W: Write>(
1137        &self,
1138        ctx: &mut BodyContext<W>,
1139        bindings: &[Option<BindingId>],
1140        fields: &Fields,
1141    ) -> std::fmt::Result {
1142        if !bindings.is_empty() {
1143            ctx.begin_block()?;
1144            let mut skipped_some = false;
1145            for (i, &binding) in bindings.iter().enumerate() {
1146                if let Some(binding) = binding {
1147                    let (field_name, field_ty) = match fields {
1148                        Fields::Unit => panic!(),
1149                        Fields::Struct(fields) => {
1150                            let field = &fields.fields[i];
1151                            let name = &self.typeenv.syms[field.name.index()];
1152                            (Cow::Borrowed(name), field.ty)
1153                        }
1154                        Fields::Tuple(fields) => (Cow::Owned(format!("{i}")), fields.fields[i].ty),
1155                    };
1156                    write!(ctx.out, "{}{field_name}: ", &ctx.indent)?;
1157                    let (is_ref, _) = self.ty(field_ty);
1158                    if is_ref {
1159                        ctx.set_ref(binding, true);
1160                        write!(ctx.out, "ref ")?;
1161                    }
1162                    writeln!(ctx.out, "v{},", binding.index())?;
1163                } else {
1164                    skipped_some = true;
1165                }
1166            }
1167            if skipped_some {
1168                writeln!(ctx.out, "{}..", &ctx.indent)?;
1169            }
1170            ctx.end_block_without_newline()?;
1171        }
1172        Ok(())
1173    }
1174
1175    fn emit_bool<W: Write>(
1176        &self,
1177        ctx: &mut BodyContext<W>,
1178        val: bool,
1179    ) -> Result<(), std::fmt::Error> {
1180        write!(ctx.out, "{val}")
1181    }
1182
1183    fn emit_int<W: Write>(
1184        &self,
1185        ctx: &mut BodyContext<W>,
1186        val: i128,
1187        ty: TypeId,
1188    ) -> Result<(), std::fmt::Error> {
1189        let ty_data = &self.typeenv.types[ty.index()];
1190        match ty_data {
1191            Type::Builtin(BuiltinType::Int(ty)) => {
1192                write!(ctx.out, "{}", rust_int_literal(*ty, val))
1193            }
1194            _ => write!(ctx.out, "{val:#x}"),
1195        }
1196    }
1197}
1198
1199fn rust_int_literal(ty: IntType, val: i128) -> String {
1200    match ty {
1201        IntType::U8 => format!("{:#x}_u8", val as u8),
1202        IntType::U16 => format!("{:#x}_u16", val as u16),
1203        IntType::U32 => format!("{:#x}_u32", val as u32),
1204        IntType::U64 => format!("{:#x}_u64", val as u64),
1205        IntType::U128 => format!("{:#x}_u128", val as u128),
1206        IntType::USize => format!("{val:#x}_usize"),
1207        IntType::I8 => format!("{}_i8", val as i8),
1208        IntType::I16 => format!("{}_i16", val as i16),
1209        IntType::I32 => format!("{}_i32", val as i32),
1210        IntType::I64 => format!("{}_i64", val as i64),
1211        IntType::I128 => format!("{val}_i128"),
1212        IntType::ISize => format!("{val}_isize"),
1213    }
1214}
1215
1216#[cfg(test)]
1217mod tests {
1218    use super::rust_int_literal;
1219    use crate::sema::IntType;
1220
1221    #[test]
1222    fn formats_wrapped_unsigned_literals() {
1223        assert_eq!(rust_int_literal(IntType::U8, -2), "0xfe_u8");
1224        assert_eq!(rust_int_literal(IntType::U64, -1), "0xffffffffffffffff_u64");
1225        assert_eq!(
1226            rust_int_literal(IntType::U128, -1),
1227            "0xffffffffffffffffffffffffffffffff_u128"
1228        );
1229    }
1230
1231    #[test]
1232    fn formats_wrapped_signed_literals() {
1233        assert_eq!(rust_int_literal(IntType::I8, 255), "-1_i8");
1234        assert_eq!(rust_int_literal(IntType::I64, -1), "-1_i64");
1235        assert_eq!(rust_int_literal(IntType::I16, 65535), "-1_i16");
1236    }
1237
1238    #[test]
1239    fn preserves_positive_unsigned_literals() {
1240        assert_eq!(rust_int_literal(IntType::U64, 5), "0x5_u64");
1241    }
1242}