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 trait ThinBodyCtx: CallLowerCtx {
58 fn find_fn_def<'a>(&'a self, name: &str) -> Option<&'a FnDef>;
59}
60
61pub fn classify_body_expr_plan<'a>(expr: &'a Expr, ctx: &impl CallLowerCtx) -> BodyExprPlan<'a> {
62 if let Some(leaf) = classify_leaf_op(expr, ctx) {
63 return BodyExprPlan::Leaf(leaf);
64 }
65
66 if let Some(plan) = classify_forward_call_plan(expr, ctx) {
67 return BodyExprPlan::ForwardCall(plan);
68 }
69
70 if let Expr::FnCall(fn_expr, args) = expr {
71 let target = classify_call_plan(&fn_expr.node, ctx);
72 if !matches!(target, CallPlan::Dynamic) {
73 return BodyExprPlan::Call { target, args };
74 }
75 }
76
77 BodyExprPlan::Expr(expr)
78}
79
80pub fn classify_body_plan<'a>(body: &'a FnBody, ctx: &impl CallLowerCtx) -> Option<BodyPlan<'a>> {
81 let stmts = body.stmts();
82 let (tail_stmt, prefix) = stmts.split_last()?;
83
84 let Stmt::Expr(tail_expr) = tail_stmt else {
85 return None;
86 };
87
88 if prefix.is_empty() {
89 return Some(BodyPlan::SingleExpr(classify_body_expr_plan(
90 &tail_expr.node,
91 ctx,
92 )));
93 }
94
95 let mut bindings = Vec::with_capacity(prefix.len());
96 for stmt in prefix {
97 let Stmt::Binding(name, _type_ann, expr) = stmt else {
98 return None;
99 };
100 bindings.push(BodyBindingPlan {
101 name,
102 expr: classify_body_expr_plan(&expr.node, ctx),
103 });
104 }
105
106 Some(BodyPlan::Block {
107 stmts,
108 bindings,
109 tail: classify_body_expr_plan(&tail_expr.node, ctx),
110 })
111}
112
113pub fn classify_thin_body_plan<'a>(
114 name: &str,
115 ctx: &'a impl ThinBodyCtx,
116) -> Option<ThinBodyPlan<'a>> {
117 let fd = ctx.find_fn_def(name)?;
118 classify_thin_fn_def(fd, ctx)
119}
120
121pub fn classify_thin_fn_def<'a>(
122 fd: &'a FnDef,
123 ctx: &impl CallLowerCtx,
124) -> Option<ThinBodyPlan<'a>> {
125 let body = classify_body_plan(&fd.body, ctx)?;
126 Some(ThinBodyPlan {
127 params: &fd.params,
128 kind: classify_thin_kind(&body, ctx)?,
129 body,
130 })
131}
132
133fn classify_thin_kind(plan: &BodyPlan<'_>, ctx: &impl CallLowerCtx) -> Option<ThinKind> {
134 match plan {
135 BodyPlan::SingleExpr(expr) => classify_thin_expr_kind(expr, ctx),
136 BodyPlan::Block { bindings, tail, .. } => {
137 if bindings
138 .iter()
139 .all(|binding| body_expr_is_thin_binding(&binding.expr))
140 {
141 classify_thin_expr_kind(tail, ctx)
142 } else {
143 None
144 }
145 }
146 }
147}
148
149fn classify_thin_expr_kind(plan: &BodyExprPlan<'_>, ctx: &impl CallLowerCtx) -> Option<ThinKind> {
150 match plan {
151 BodyExprPlan::Leaf(_) => Some(ThinKind::Leaf),
152 BodyExprPlan::Call { .. } => Some(ThinKind::Direct),
153 BodyExprPlan::ForwardCall(_) => Some(ThinKind::Forward),
154 BodyExprPlan::Expr(expr) => match *expr {
155 Expr::Match { arms, .. } if classify_match_dispatch_plan(arms, ctx).is_some() => {
156 Some(ThinKind::Dispatch)
157 }
158 Expr::TailCall(_) => Some(ThinKind::Tail),
159 _ => None,
160 },
161 }
162}
163
164fn body_expr_is_thin_binding(plan: &BodyExprPlan<'_>) -> bool {
165 match plan {
166 BodyExprPlan::Leaf(_) | BodyExprPlan::Call { .. } | BodyExprPlan::ForwardCall(_) => true,
167 BodyExprPlan::Expr(expr) => match expr {
168 Expr::Literal(_) | Expr::Ident(_) | Expr::Constructor(_, _) => true,
169 Expr::BinOp(_, left, right) => {
171 is_simple_operand(&left.node) && is_simple_operand(&right.node)
172 }
173 Expr::FnCall(_, args) => args.iter().all(|a| is_simple_operand(&a.node)),
175 _ => false,
176 },
177 }
178}
179
180fn is_simple_operand(expr: &Expr) -> bool {
181 matches!(expr, Expr::Literal(_) | Expr::Ident(_) | Expr::Resolved(_))
182}