Skip to main content

ptx_parser/pretty_print/
function.rs

1// TreeDisplay implementations for function types (src/type/function.rs)
2
3use super::{TreeDisplay, TreeFormatter};
4use crate::r#type::function::*;
5
6impl TreeDisplay for AliasFunctionDirective {
7    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
8        f.root(&format!(
9            "AliasFunctionDirective [{}]",
10            f.format_raw(self.span, source)
11        ))?;
12        f.field(false, "alias", &self.alias.val)?;
13        f.field(true, "target", &self.target.val)
14    }
15}
16
17impl TreeDisplay for FuncFunctionDirective {
18    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
19        f.root(&format!(
20            "FuncFunctionDirective [{}]",
21            f.format_raw(self.span, source)
22        ))?;
23        f.field_vec(false, "attributes", &self.attributes, source)?;
24        f.field_option(false, "return_param", &self.return_param, source)?;
25        f.field(false, "name", &self.name.val)?;
26        f.field_vec(false, "params", &self.params, source)?;
27        f.field_vec(false, "directives", &self.directives, source)?;
28        f.field_option(true, "body", &self.body, source)
29    }
30}
31
32impl TreeDisplay for EntryFunctionDirective {
33    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
34        f.root(&format!(
35            "EntryFunctionDirective [{}]",
36            f.format_raw(self.span, source)
37        ))?;
38        f.field(false, "name", &self.name.val)?;
39        f.field_vec(false, "params", &self.params, source)?;
40        f.field_vec(false, "directives", &self.directives, source)?;
41        f.field_option(true, "body", &self.body, source)
42    }
43}
44
45impl TreeDisplay for FuncFunctionHeaderDirective {
46    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
47        match self {
48            FuncFunctionHeaderDirective::NoReturn { span } => f.root(&format!(
49                "FuncFunctionHeaderDirective::NoReturn [{}]",
50                f.format_raw(*span, source)
51            )),
52            FuncFunctionHeaderDirective::Pragma { args, span } => {
53                f.root(&format!(
54                    "FuncFunctionHeaderDirective::Pragma [{}]",
55                    f.format_raw(*span, source)
56                ))?;
57                f.field_vec(true, "args", args, source)
58            }
59            FuncFunctionHeaderDirective::AbiPreserve { value, span } => {
60                f.root(&format!(
61                    "FuncFunctionHeaderDirective::AbiPreserve [{}]",
62                    f.format_raw(*span, source)
63                ))?;
64                f.field(true, "value", &value.to_string())
65            }
66            FuncFunctionHeaderDirective::AbiPreserveControl { value, span } => {
67                f.root(&format!(
68                    "FuncFunctionHeaderDirective::AbiPreserveControl [{}]",
69                    f.format_raw(*span, source)
70                ))?;
71                f.field(true, "value", &value.to_string())
72            }
73        }
74    }
75}
76
77impl TreeDisplay for EntryFunctionHeaderDirective {
78    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
79        match self {
80            EntryFunctionHeaderDirective::MaxNReg { value, span } => {
81                f.root(&format!(
82                    "EntryFunctionHeaderDirective::MaxNReg [{}]",
83                    f.format_raw(*span, source)
84                ))?;
85                f.field(true, "value", &value.to_string())
86            }
87            EntryFunctionHeaderDirective::MaxNTid { dim, span } => {
88                f.root(&format!(
89                    "EntryFunctionHeaderDirective::MaxNTid [{}]",
90                    f.format_raw(*span, source)
91                ))?;
92                f.field_with_child(true, "dim", dim, source)
93            }
94            EntryFunctionHeaderDirective::ReqNTid { dim, span } => {
95                f.root(&format!(
96                    "EntryFunctionHeaderDirective::ReqNTid [{}]",
97                    f.format_raw(*span, source)
98                ))?;
99                f.field_with_child(true, "dim", dim, source)
100            }
101            EntryFunctionHeaderDirective::MinNCtaPerSm { value, span } => {
102                f.root(&format!(
103                    "EntryFunctionHeaderDirective::MinNCtaPerSm [{}]",
104                    f.format_raw(*span, source)
105                ))?;
106                f.field(true, "value", &value.to_string())
107            }
108            EntryFunctionHeaderDirective::MaxNCtaPerSm { value, span } => {
109                f.root(&format!(
110                    "EntryFunctionHeaderDirective::MaxNCtaPerSm [{}]",
111                    f.format_raw(*span, source)
112                ))?;
113                f.field(true, "value", &value.to_string())
114            }
115            EntryFunctionHeaderDirective::Pragma { args, span } => {
116                f.root(&format!(
117                    "EntryFunctionHeaderDirective::Pragma [{}]",
118                    f.format_raw(*span, source)
119                ))?;
120                f.field_vec(true, "args", args, source)
121            }
122            EntryFunctionHeaderDirective::ReqNctaPerCluster { dim, span } => {
123                f.root(&format!(
124                    "EntryFunctionHeaderDirective::ReqNctaPerCluster [{}]",
125                    f.format_raw(*span, source)
126                ))?;
127                f.field_with_child(true, "dim", dim, source)
128            }
129            EntryFunctionHeaderDirective::ExplicitCluster { span } => f.root(&format!(
130                "EntryFunctionHeaderDirective::ExplicitCluster [{}]",
131                f.format_raw(*span, source)
132            )),
133            EntryFunctionHeaderDirective::MaxClusterRank { value, span } => {
134                f.root(&format!(
135                    "EntryFunctionHeaderDirective::MaxClusterRank [{}]",
136                    f.format_raw(*span, source)
137                ))?;
138                f.field(true, "value", &value.to_string())
139            }
140            EntryFunctionHeaderDirective::BlocksAreClusters { span } => f.root(&format!(
141                "EntryFunctionHeaderDirective::BlocksAreClusters [{}]",
142                f.format_raw(*span, source)
143            )),
144        }
145    }
146}
147
148impl TreeDisplay for FunctionBody {
149    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
150        f.root(&format!(
151            "FunctionBody [{}]",
152            f.format_raw(self.span, source)
153        ))?;
154        f.field_vec(true, "statements", &self.statements, source)
155    }
156}
157
158impl TreeDisplay for FunctionStatement {
159    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
160        match self {
161            FunctionStatement::Label { label, span } => {
162                f.root(&format!(
163                    "FunctionStatement::Label [{}]",
164                    f.format_raw(*span, source)
165                ))?;
166                f.field_with_child(true, "label", label, source)
167            }
168            FunctionStatement::Directive { directive, span } => {
169                f.root(&format!(
170                    "FunctionStatement::Directive [{}]",
171                    f.format_raw(*span, source)
172                ))?;
173                f.field_with_child(true, "directive", directive, source)
174            }
175            FunctionStatement::Instruction {
176                instruction,
177                span: _,
178            } => {
179                // Display instruction inline since it contains all the details
180                instruction.tree_display(f, source)
181            }
182            FunctionStatement::Block { statements, span } => {
183                f.root(&format!(
184                    "FunctionStatement::Block [{}]",
185                    f.format_raw(*span, source)
186                ))?;
187                f.field_vec(true, "statements", statements, source)
188            }
189        }
190    }
191}
192
193impl TreeDisplay for RegisterDirective {
194    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
195        f.root(&format!(
196            "RegisterDirective [{}]",
197            f.format_raw(self.span, source)
198        ))?;
199        f.field(false, "vector", &format!("{:?}", self.vector))?;
200        f.field_with_child(false, "ty", &self.ty, source)?;
201        f.field_vec(true, "registers", &self.registers, source)
202    }
203}
204
205impl TreeDisplay for RegisterTarget {
206    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
207        f.root(&format!(
208            "RegisterTarget [{}]",
209            f.format_raw(self.span, source)
210        ))?;
211        f.field(false, "name", &self.name.val)?;
212        match self.range {
213            Some(r) => f.field(true, "range", &format!("Some({})", r))?,
214            None => f.field(true, "range", "None")?,
215        };
216        Ok(())
217    }
218}
219
220impl TreeDisplay for StatementDirective {
221    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
222        match self {
223            StatementDirective::Loc { directive, span } => {
224                f.root(&format!(
225                    "StatementDirective::Loc [{}]",
226                    f.format_raw(*span, source)
227                ))?;
228                f.field_with_child(true, "directive", directive, source)
229            }
230            StatementDirective::Pragma { directive, span } => {
231                f.root(&format!(
232                    "StatementDirective::Pragma [{}]",
233                    f.format_raw(*span, source)
234                ))?;
235                f.field_with_child(true, "directive", directive, source)
236            }
237            StatementDirective::Section { directive, span } => {
238                f.root(&format!(
239                    "StatementDirective::Section [{}]",
240                    f.format_raw(*span, source)
241                ))?;
242                f.field_with_child(true, "directive", directive, source)
243            }
244            StatementDirective::Reg { directive, span } => {
245                f.root(&format!(
246                    "StatementDirective::Reg [{}]",
247                    f.format_raw(*span, source)
248                ))?;
249                f.field_with_child(true, "directive", directive, source)
250            }
251            StatementDirective::Local { directive, span } => {
252                f.root(&format!(
253                    "StatementDirective::Local [{}]",
254                    f.format_raw(*span, source)
255                ))?;
256                f.field_with_child(true, "directive", directive, source)
257            }
258            StatementDirective::Param { directive, span } => {
259                f.root(&format!(
260                    "StatementDirective::Param [{}]",
261                    f.format_raw(*span, source)
262                ))?;
263                f.field_with_child(true, "directive", directive, source)
264            }
265            StatementDirective::Shared { directive, span } => {
266                f.root(&format!(
267                    "StatementDirective::Shared [{}]",
268                    f.format_raw(*span, source)
269                ))?;
270                f.field_with_child(true, "directive", directive, source)
271            }
272            StatementDirective::Dwarf { directive, span } => {
273                f.root(&format!(
274                    "StatementDirective::Dwarf [{}]",
275                    f.format_raw(*span, source)
276                ))?;
277                f.field_with_child(true, "directive", directive, source)
278            }
279            StatementDirective::BranchTargets { directive, span } => {
280                f.root(&format!(
281                    "StatementDirective::BranchTargets [{}]",
282                    f.format_raw(*span, source)
283                ))?;
284                f.field_with_child(true, "directive", directive, source)
285            }
286            StatementDirective::CallTargets { directive, span } => {
287                f.root(&format!(
288                    "StatementDirective::CallTargets [{}]",
289                    f.format_raw(*span, source)
290                ))?;
291                f.field_with_child(true, "directive", directive, source)
292            }
293            StatementDirective::CallPrototype { directive, span } => {
294                f.root(&format!(
295                    "StatementDirective::CallPrototype [{}]",
296                    f.format_raw(*span, source)
297                ))?;
298                f.field_with_child(true, "directive", directive, source)
299            }
300        }
301    }
302}
303
304impl TreeDisplay for DwarfDirective {
305    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
306        f.root(&format!(
307            "DwarfDirective [{}]",
308            f.format_raw(self.span, source)
309        ))?;
310        f.field(true, "kind", &format!("{:?}", self.kind))
311    }
312}
313
314impl TreeDisplay for SectionDirective {
315    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
316        f.root(&format!(
317            "SectionDirective [{}]",
318            f.format_raw(self.span, source)
319        ))?;
320        f.field(false, "name", &format!("\"{}\"", self.name))?;
321        f.field(
322            true,
323            "entries",
324            &format!("<{} entries>", self.entries.len()),
325        )
326    }
327}
328
329impl TreeDisplay for LocationDirective {
330    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
331        f.root(&format!(
332            "LocationDirective [{}]",
333            f.format_raw(self.span, source)
334        ))?;
335        f.field(false, "file_index", &self.file_index.to_string())?;
336        f.field(false, "line", &self.line.to_string())?;
337        f.field(false, "column", &self.column.to_string())?;
338        f.field_option(true, "function", &self.function, source)
339    }
340}
341
342impl TreeDisplay for LocationFunctionInfo {
343    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
344        f.root(&format!(
345            "LocationFunctionInfo [{}]",
346            f.format_raw(self.span, source)
347        ))?;
348        f.field(false, "label", &self.label.val)?;
349        match self.label_offset {
350            Some(offset) => f.field(false, "label_offset", &format!("Some({offset})"))?,
351            None => f.field(false, "label_offset", "None")?,
352        }
353        f.field_with_child(true, "inlined_at", &self.inlined_at, source)
354    }
355}
356
357impl TreeDisplay for LocationInlinedAt {
358    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
359        f.root(&format!(
360            "LocationInlinedAt [{}]",
361            f.format_raw(self.span, source)
362        ))?;
363        f.field(false, "file_index", &self.file_index.to_string())?;
364        f.field(false, "line", &self.line.to_string())?;
365        f.field(true, "column", &self.column.to_string())
366    }
367}
368
369impl TreeDisplay for PragmaDirective {
370    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
371        f.root(&format!(
372            "PragmaDirective [{}]",
373            f.format_raw(self.span, source)
374        ))?;
375        f.field(true, "kind", &format!("{:?}", self.kind))
376    }
377}
378
379impl TreeDisplay for BranchTargetsDirective {
380    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
381        f.root(&format!(
382            "BranchTargetsDirective [{}]",
383            f.format_raw(self.span, source)
384        ))?;
385        f.field_vec(true, "labels", &self.labels, source)
386    }
387}
388
389impl TreeDisplay for CallTargetsDirective {
390    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
391        f.root(&format!(
392            "CallTargetsDirective [{}]",
393            f.format_raw(self.span, source)
394        ))?;
395        let targets: Vec<String> = self.targets.iter().map(|t| t.val.clone()).collect();
396        f.field_vec(true, "targets", &targets, source)
397    }
398}
399
400impl TreeDisplay for CallPrototypeDirective {
401    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
402        f.root(&format!(
403            "CallPrototypeDirective [{}]",
404            f.format_raw(self.span, source)
405        ))?;
406        f.field_option(false, "return_param", &self.return_param, source)?;
407        f.field_vec(false, "params", &self.params, source)?;
408        f.field(false, "noreturn", &self.noreturn.to_string())?;
409        match self.abi_preserve {
410            Some(v) => f.field(false, "abi_preserve", &format!("Some({})", v))?,
411            None => f.field(false, "abi_preserve", "None")?,
412        }
413        match self.abi_preserve_control {
414            Some(v) => f.field(true, "abi_preserve_control", &format!("Some({})", v)),
415            None => f.field(true, "abi_preserve_control", "None"),
416        }
417    }
418}
419
420impl TreeDisplay for FunctionDim {
421    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
422        match self {
423            FunctionDim::X { x, span } => {
424                f.root(&format!("FunctionDim::X [{}]", f.format_raw(*span, source)))?;
425                f.field(true, "x", &x.to_string())
426            }
427            FunctionDim::XY { x, y, span } => {
428                f.root(&format!(
429                    "FunctionDim::XY [{}]",
430                    f.format_raw(*span, source)
431                ))?;
432                f.field(false, "x", &x.to_string())?;
433                f.field(true, "y", &y.to_string())
434            }
435            FunctionDim::XYZ { x, y, z, span } => {
436                f.root(&format!(
437                    "FunctionDim::XYZ [{}]",
438                    f.format_raw(*span, source)
439                ))?;
440                f.field(false, "x", &x.to_string())?;
441                f.field(false, "y", &y.to_string())?;
442                f.field(true, "z", &z.to_string())
443            }
444        }
445    }
446}