diff --git a/cranelift/isle/isle/src/codegen.rs b/cranelift/isle/isle/src/codegen.rs index 21f35b4f21ee..9840bed6b0ac 100644 --- a/cranelift/isle/isle/src/codegen.rs +++ b/cranelift/isle/isle/src/codegen.rs @@ -72,10 +72,25 @@ struct BodyContext<'a, W> { is_bound: StableSet, term_name: &'a str, emit_logging: bool, + + // Extra fields for iterator-returning terms. + // These fields are used to generate optimized Rust code for iterator-returning terms. + /// The number of match splits that have been generated. + /// This is used to generate unique names for the match splits. + match_split: usize, + + /// The action to take when the iterator overflows. + iter_overflow_action: &'static str, } impl<'a, W: Write> BodyContext<'a, W> { - fn new(out: &'a mut W, ruleset: &'a RuleSet, term_name: &'a str, emit_logging: bool) -> Self { + fn new( + out: &'a mut W, + ruleset: &'a RuleSet, + term_name: &'a str, + emit_logging: bool, + iter_overflow_action: &'static str, + ) -> Self { Self { out, ruleset, @@ -84,6 +99,8 @@ impl<'a, W: Write> BodyContext<'a, W> { is_bound: Default::default(), term_name, emit_logging, + match_split: Default::default(), + iter_overflow_action, } } @@ -426,7 +443,17 @@ impl Length for ContextIterWrapper {{ let termdata = &self.termenv.terms[termid.index()]; let term_name = &self.typeenv.syms[termdata.name.index()]; - let mut ctx = BodyContext::new(code, ruleset, term_name, options.emit_logging); + + // Split a match if the term returns an iterator. + let mut ctx = BodyContext::new( + code, + ruleset, + term_name, + options.emit_logging, + "return;", // At top level, we just return. + ); + + // Generate the function signature. writeln!(ctx.out)?; writeln!( ctx.out, @@ -470,6 +497,7 @@ impl Length for ContextIterWrapper {{ ReturnKind::Option => write!(ctx.out, "Option<{ret}>")?, ReturnKind::Plain => write!(ctx.out, "{ret}")?, }; + // Generating the function signature is done. let last_expr = if let Some(EvalStep { check: ControlFlow::Return { .. }, @@ -530,6 +558,21 @@ impl Length for ContextIterWrapper {{ Nested::Cases(block.steps.iter()) } + fn block_weight(block: &Block) -> usize { + fn cf_weight(cf: &ControlFlow) -> usize { + match cf { + ControlFlow::Match { arms, .. } => { + arms.iter().map(|a| Codegen::block_weight(&a.body)).sum() + } + ControlFlow::Equal { body, .. } => Codegen::block_weight(body), + ControlFlow::Loop { body, .. } => Codegen::block_weight(body), + ControlFlow::Return { .. } => 0, + } + } + + block.steps.iter().map(|s| 1 + cf_weight(&s.check)).sum() + } + fn emit_block( &self, ctx: &mut BodyContext, @@ -538,8 +581,19 @@ impl Length for ContextIterWrapper {{ last_expr: &str, scope: StableSet, ) -> std::fmt::Result { - let mut stack = Vec::new(); ctx.begin_block()?; + self.emit_block_contents(ctx, block, ret_kind, last_expr, scope) + } + + fn emit_block_contents( + &self, + ctx: &mut BodyContext, + block: &Block, + ret_kind: ReturnKind, + last_expr: &str, + scope: StableSet, + ) -> std::fmt::Result { + let mut stack = Vec::new(); stack.push((Self::validate_block(ret_kind, block), last_expr, scope)); while let Some((mut nested, last_line, scope)) = stack.pop() { @@ -706,8 +760,8 @@ impl Length for ContextIterWrapper {{ writeln!(ctx.out, "));")?; writeln!( ctx.out, - "{}if returns.len() >= MAX_ISLE_RETURNS {{ return; }}", - ctx.indent + "{}if returns.len() >= MAX_ISLE_RETURNS {{ {} }}", + ctx.indent, ctx.iter_overflow_action )?; } } @@ -729,7 +783,39 @@ impl Length for ContextIterWrapper {{ self.emit_constraint(ctx, source, arm)?; write!(ctx.out, " =>")?; ctx.begin_block()?; - stack.push((Self::validate_block(ret_kind, &arm.body), "", scope)); + + // Compile-time optimization: huge function bodies (often from very large match arms + // of constructor bodies)cause rustc to spend a lot of time in analysis passes. + // Wrap such bodies in a local closure to move the bulk of the work into a separate body + // without needing to know the types of captured locals. + const MATCH_ARM_BODY_CLOSURE_THRESHOLD: usize = 256; + if ret_kind == ReturnKind::Iterator + && Codegen::block_weight(&arm.body) > MATCH_ARM_BODY_CLOSURE_THRESHOLD + { + let closure_id = ctx.match_split; + ctx.match_split += 1; + + write!(ctx.out, "{}if (|| -> bool", &ctx.indent)?; + ctx.begin_block()?; + + let old_overflow_action = ctx.iter_overflow_action; + ctx.iter_overflow_action = "return true;"; + let closure_scope = ctx.enter_scope(); + self.emit_block_contents(ctx, &arm.body, ret_kind, "false", closure_scope)?; + ctx.iter_overflow_action = old_overflow_action; + + // Close `if (|| -> bool { ... })()` and stop the outer function on + // iterator-overflow. + writeln!( + ctx.out, + "{})() {{ {} }} // __isle_arm_{}", + &ctx.indent, ctx.iter_overflow_action, closure_id + )?; + + ctx.end_block("", scope)?; + } else { + stack.push((Self::validate_block(ret_kind, &arm.body), "", scope)); + } } } }