1use crate::ast::{Expr, FnBody, FnDef, Spanned, Stmt};
2
3use super::{
4 CallLowerCtx, CallPlan, ForwardCallPlan, LeafOp, classify_call_plan,
5 classify_forward_call_plan, classify_leaf_op, classify_match_dispatch_plan,
6};
7
8#[derive(Debug, Clone, PartialEq)]
15pub enum BodyExprPlan<'a> {
16 Expr(&'a Expr),
17 Leaf(LeafOp<'a>),
18 Call {
19 target: CallPlan,
20 args: &'a [Spanned<Expr>],
21 },
22 ForwardCall(ForwardCallPlan),
23}
24
25#[derive(Debug, Clone, PartialEq)]
26pub struct BodyBindingPlan<'a> {
27 pub name: &'a str,
28 pub expr: BodyExprPlan<'a>,
29}
30
31#[derive(Debug, Clone, PartialEq)]
32pub enum BodyPlan<'a> {
33 SingleExpr(BodyExprPlan<'a>),
34 Block {
35 stmts: &'a [Stmt],
36 bindings: Vec<BodyBindingPlan<'a>>,
37 tail: BodyExprPlan<'a>,
38 },
39}
40
41#[derive(Debug, Clone, Copy, PartialEq, Eq)]
42pub enum ThinKind {
43 Leaf,
44 Direct,
45 Forward,
46 Dispatch,
47 Tail,
48}
49
50#[derive(Debug, Clone, PartialEq)]
51pub struct ThinBodyPlan<'a> {
52 pub params: &'a [(String, String)],
53 pub body: BodyPlan<'a>,
54 pub kind: ThinKind,
55}
56
57pub fn thin_kind_is_parent_thin_candidate(kind: ThinKind) -> bool {
58 matches!(
59 kind,
60 ThinKind::Leaf | ThinKind::Direct | ThinKind::Forward | ThinKind::Dispatch
61 )
62}
63
64pub fn thin_body_plan_is_parent_thin_candidate(plan: &ThinBodyPlan<'_>) -> bool {
65 thin_kind_is_parent_thin_candidate(plan.kind) && matches!(plan.body, BodyPlan::SingleExpr(_))
66}
67
68pub trait ThinBodyCtx: CallLowerCtx {
69 fn find_fn_def<'a>(&'a self, name: &str) -> Option<&'a FnDef>;
70}
71
72pub fn classify_body_expr_plan<'a>(expr: &'a Expr, ctx: &impl CallLowerCtx) -> BodyExprPlan<'a> {
73 if let Some(leaf) = classify_leaf_op(expr, ctx) {
74 return BodyExprPlan::Leaf(leaf);
75 }
76
77 if let Some(plan) = classify_forward_call_plan(expr, ctx) {
78 return BodyExprPlan::ForwardCall(plan);
79 }
80
81 if let Expr::FnCall(fn_expr, args) = expr {
82 let target = classify_call_plan(&fn_expr.node, ctx);
83 if !matches!(target, CallPlan::Dynamic) {
84 return BodyExprPlan::Call { target, args };
85 }
86 }
87
88 BodyExprPlan::Expr(expr)
89}
90
91pub fn classify_body_plan<'a>(body: &'a FnBody, ctx: &impl CallLowerCtx) -> Option<BodyPlan<'a>> {
92 let stmts = body.stmts();
93 let (tail_stmt, prefix) = stmts.split_last()?;
94
95 let Stmt::Expr(tail_expr) = tail_stmt else {
96 return None;
97 };
98
99 if prefix.is_empty() {
100 return Some(BodyPlan::SingleExpr(classify_body_expr_plan(
101 &tail_expr.node,
102 ctx,
103 )));
104 }
105
106 let mut bindings = Vec::with_capacity(prefix.len());
107 for stmt in prefix {
108 let Stmt::Binding(name, _type_ann, expr) = stmt else {
109 return None;
110 };
111 bindings.push(BodyBindingPlan {
112 name,
113 expr: classify_body_expr_plan(&expr.node, ctx),
114 });
115 }
116
117 Some(BodyPlan::Block {
118 stmts,
119 bindings,
120 tail: classify_body_expr_plan(&tail_expr.node, ctx),
121 })
122}
123
124pub fn classify_thin_body_plan<'a>(
125 name: &str,
126 ctx: &'a impl ThinBodyCtx,
127) -> Option<ThinBodyPlan<'a>> {
128 let fd = ctx.find_fn_def(name)?;
129 classify_thin_fn_def(fd, ctx)
130}
131
132pub fn classify_thin_fn_def<'a>(
133 fd: &'a FnDef,
134 ctx: &impl CallLowerCtx,
135) -> Option<ThinBodyPlan<'a>> {
136 let body = classify_body_plan(&fd.body, ctx)?;
137 Some(ThinBodyPlan {
138 params: &fd.params,
139 kind: classify_thin_kind(&body, ctx)?,
140 body,
141 })
142}
143
144fn classify_thin_kind(plan: &BodyPlan<'_>, ctx: &impl CallLowerCtx) -> Option<ThinKind> {
145 match plan {
146 BodyPlan::SingleExpr(expr) => classify_thin_expr_kind(expr, ctx),
147 BodyPlan::Block { bindings, tail, .. } => {
148 if bindings
149 .iter()
150 .all(|binding| body_expr_is_thin_binding(&binding.expr))
151 {
152 classify_thin_expr_kind(tail, ctx)
153 } else {
154 None
155 }
156 }
157 }
158}
159
160fn classify_thin_expr_kind(plan: &BodyExprPlan<'_>, ctx: &impl CallLowerCtx) -> Option<ThinKind> {
161 match plan {
162 BodyExprPlan::Leaf(_) => Some(ThinKind::Leaf),
163 BodyExprPlan::Call { .. } => Some(ThinKind::Direct),
164 BodyExprPlan::ForwardCall(_) => Some(ThinKind::Forward),
165 BodyExprPlan::Expr(expr) => match *expr {
166 Expr::Match { arms, .. } if classify_match_dispatch_plan(arms, ctx).is_some() => {
167 Some(ThinKind::Dispatch)
168 }
169 Expr::TailCall(_) => Some(ThinKind::Tail),
170 _ => None,
171 },
172 }
173}
174
175fn body_expr_is_thin_binding(plan: &BodyExprPlan<'_>) -> bool {
176 match plan {
177 BodyExprPlan::Leaf(_) | BodyExprPlan::Call { .. } | BodyExprPlan::ForwardCall(_) => true,
178 BodyExprPlan::Expr(expr) => match expr {
179 Expr::Literal(_) | Expr::Ident(_) | Expr::Constructor(_, _) => true,
180 Expr::BinOp(_, left, right) => {
182 is_simple_operand(&left.node) && is_simple_operand(&right.node)
183 }
184 Expr::FnCall(_, args) => args.iter().all(|a| is_simple_operand(&a.node)),
186 _ => false,
187 },
188 }
189}
190
191fn is_simple_operand(expr: &Expr) -> bool {
192 matches!(
193 expr,
194 Expr::Literal(_) | Expr::Ident(_) | Expr::Resolved { .. }
195 )
196}