Skip to main content

ptx_parser/pretty_print/
common.rs

1use super::{TreeDisplay, TreeFormatter};
2/// TreeDisplay implementations for common types (src/type/common.rs)
3use crate::r#type::common::*;
4
5impl TreeDisplay for CodeLinkage {
6    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
7        let (variant, span) = match self {
8            CodeLinkage::Visible { span } => ("Visible", span),
9            CodeLinkage::Extern { span } => ("Extern", span),
10            CodeLinkage::Weak { span } => ("Weak", span),
11        };
12        f.root(&format!(
13            "CodeLinkage::{} [{}]",
14            variant,
15            f.format_raw(*span, source)
16        ))
17    }
18}
19
20impl TreeDisplay for DataLinkage {
21    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
22        let (variant, span) = match self {
23            DataLinkage::Visible { span } => ("Visible", span),
24            DataLinkage::Extern { span } => ("Extern", span),
25            DataLinkage::Weak { span } => ("Weak", span),
26            DataLinkage::Common { span } => ("Common", span),
27        };
28        f.root(&format!(
29            "DataLinkage::{} [{}]",
30            variant,
31            f.format_raw(*span, source)
32        ))
33    }
34}
35
36impl TreeDisplay for AttributeDirective {
37    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
38        match self {
39            AttributeDirective::Unified { uuid1, uuid2, span } => {
40                f.root(&format!(
41                    "AttributeDirective::Unified [{}]",
42                    f.format_raw(*span, source)
43                ))?;
44                f.field(false, "uuid1", &uuid1.to_string())?;
45                f.field(true, "uuid2", &uuid2.to_string())
46            }
47            AttributeDirective::Managed { span } => f.root(&format!(
48                "AttributeDirective::Managed [{}]",
49                f.format_raw(*span, source)
50            )),
51        }
52    }
53}
54
55impl TreeDisplay for DataType {
56    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
57        let (variant, span) = match self {
58            DataType::U8 { span } => ("U8", span),
59            DataType::U16 { span } => ("U16", span),
60            DataType::U32 { span } => ("U32", span),
61            DataType::U64 { span } => ("U64", span),
62            DataType::S8 { span } => ("S8", span),
63            DataType::S16 { span } => ("S16", span),
64            DataType::S32 { span } => ("S32", span),
65            DataType::S64 { span } => ("S64", span),
66            DataType::F16 { span } => ("F16", span),
67            DataType::F16x2 { span } => ("F16x2", span),
68            DataType::F32 { span } => ("F32", span),
69            DataType::F64 { span } => ("F64", span),
70            DataType::B8 { span } => ("B8", span),
71            DataType::B16 { span } => ("B16", span),
72            DataType::B32 { span } => ("B32", span),
73            DataType::B64 { span } => ("B64", span),
74            DataType::B128 { span } => ("B128", span),
75            DataType::Pred { span } => ("Pred", span),
76            DataType::TexRef { span } => ("TexRef", span),
77            DataType::SamplerRef { span } => ("SamplerRef", span),
78            DataType::SurfRef { span } => ("SurfRef", span),
79        };
80        f.root(&format!(
81            "DataType::{} [{}]",
82            variant,
83            f.format_raw(*span, source)
84        ))
85    }
86}
87
88impl TreeDisplay for Sign {
89    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
90        let (variant, span) = match self {
91            Sign::Negative { span } => ("Negative", span),
92            Sign::Positive { span } => ("Positive", span),
93        };
94        f.root(&format!(
95            "Sign::{} [{}]",
96            variant,
97            f.format_raw(*span, source)
98        ))
99    }
100}
101
102impl TreeDisplay for Axis {
103    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
104        let (variant, span) = match self {
105            Axis::None { span } => ("None", span),
106            Axis::X { span } => ("X", span),
107            Axis::Y { span } => ("Y", span),
108            Axis::Z { span } => ("Z", span),
109        };
110        f.root(&format!(
111            "Axis::{} [{}]",
112            variant,
113            f.format_raw(*span, source)
114        ))
115    }
116}
117
118impl TreeDisplay for SpecialRegister {
119    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
120        match self {
121            SpecialRegister::AggrSmemSize { span } => f.root(&format!(
122                "SpecialRegister::AggrSmemSize [{}]",
123                f.format_raw(*span, source)
124            )),
125            SpecialRegister::DynamicSmemSize { span } => f.root(&format!(
126                "SpecialRegister::DynamicSmemSize [{}]",
127                f.format_raw(*span, source)
128            )),
129            SpecialRegister::LanemaskGt { span } => f.root(&format!(
130                "SpecialRegister::LanemaskGt [{}]",
131                f.format_raw(*span, source)
132            )),
133            SpecialRegister::ReservedSmemOffsetBegin { span } => f.root(&format!(
134                "SpecialRegister::ReservedSmemOffsetBegin [{}]",
135                f.format_raw(*span, source)
136            )),
137            SpecialRegister::Clock { span } => f.root(&format!(
138                "SpecialRegister::Clock [{}]",
139                f.format_raw(*span, source)
140            )),
141            SpecialRegister::Envreg { index, span } => {
142                f.root(&format!(
143                    "SpecialRegister::Envreg [{}]",
144                    f.format_raw(*span, source)
145                ))?;
146                f.field(true, "index", &index.to_string())
147            }
148            SpecialRegister::LanemaskLe { span } => f.root(&format!(
149                "SpecialRegister::LanemaskLe [{}]",
150                f.format_raw(*span, source)
151            )),
152            SpecialRegister::ReservedSmemOffsetCap { span } => f.root(&format!(
153                "SpecialRegister::ReservedSmemOffsetCap [{}]",
154                f.format_raw(*span, source)
155            )),
156            SpecialRegister::Clock64 { span } => f.root(&format!(
157                "SpecialRegister::Clock64 [{}]",
158                f.format_raw(*span, source)
159            )),
160            SpecialRegister::Globaltimer { span } => f.root(&format!(
161                "SpecialRegister::Globaltimer [{}]",
162                f.format_raw(*span, source)
163            )),
164            SpecialRegister::LanemaskLt { span } => f.root(&format!(
165                "SpecialRegister::LanemaskLt [{}]",
166                f.format_raw(*span, source)
167            )),
168            SpecialRegister::ReservedSmemOffsetEnd { span } => f.root(&format!(
169                "SpecialRegister::ReservedSmemOffsetEnd [{}]",
170                f.format_raw(*span, source)
171            )),
172            SpecialRegister::ClusterCtaid { axis, span } => {
173                f.root(&format!(
174                    "SpecialRegister::ClusterCtaid [{}]",
175                    f.format_raw(*span, source)
176                ))?;
177                f.field_with_child(true, "axis", axis, source)
178            }
179            SpecialRegister::GlobaltimerHi { span } => f.root(&format!(
180                "SpecialRegister::GlobaltimerHi [{}]",
181                f.format_raw(*span, source)
182            )),
183            SpecialRegister::Nclusterid { span } => f.root(&format!(
184                "SpecialRegister::Nclusterid [{}]",
185                f.format_raw(*span, source)
186            )),
187            SpecialRegister::Smid { span } => f.root(&format!(
188                "SpecialRegister::Smid [{}]",
189                f.format_raw(*span, source)
190            )),
191            SpecialRegister::ClusterCtarank { axis, span } => {
192                f.root(&format!(
193                    "SpecialRegister::ClusterCtarank [{}]",
194                    f.format_raw(*span, source)
195                ))?;
196                f.field_with_child(true, "axis", axis, source)
197            }
198            SpecialRegister::GlobaltimerLo { span } => f.root(&format!(
199                "SpecialRegister::GlobaltimerLo [{}]",
200                f.format_raw(*span, source)
201            )),
202            SpecialRegister::Nctaid { axis, span } => {
203                f.root(&format!(
204                    "SpecialRegister::Nctaid [{}]",
205                    f.format_raw(*span, source)
206                ))?;
207                f.field_with_child(true, "axis", axis, source)
208            }
209            SpecialRegister::Tid { axis, span } => {
210                f.root(&format!(
211                    "SpecialRegister::Tid [{}]",
212                    f.format_raw(*span, source)
213                ))?;
214                f.field_with_child(true, "axis", axis, source)
215            }
216            SpecialRegister::ClusterNctaid { axis, span } => {
217                f.root(&format!(
218                    "SpecialRegister::ClusterNctaid [{}]",
219                    f.format_raw(*span, source)
220                ))?;
221                f.field_with_child(true, "axis", axis, source)
222            }
223            SpecialRegister::Gridid { span } => f.root(&format!(
224                "SpecialRegister::Gridid [{}]",
225                f.format_raw(*span, source)
226            )),
227            SpecialRegister::Nsmid { span } => f.root(&format!(
228                "SpecialRegister::Nsmid [{}]",
229                f.format_raw(*span, source)
230            )),
231            SpecialRegister::TotalSmemSize { span } => f.root(&format!(
232                "SpecialRegister::TotalSmemSize [{}]",
233                f.format_raw(*span, source)
234            )),
235            SpecialRegister::ClusterNctarank { axis, span } => {
236                f.root(&format!(
237                    "SpecialRegister::ClusterNctarank [{}]",
238                    f.format_raw(*span, source)
239                ))?;
240                f.field_with_child(true, "axis", axis, source)
241            }
242            SpecialRegister::IsExplicitCluster { span } => f.root(&format!(
243                "SpecialRegister::IsExplicitCluster [{}]",
244                f.format_raw(*span, source)
245            )),
246            SpecialRegister::Ntid { axis, span } => {
247                f.root(&format!(
248                    "SpecialRegister::Ntid [{}]",
249                    f.format_raw(*span, source)
250                ))?;
251                f.field_with_child(true, "axis", axis, source)
252            }
253            SpecialRegister::Warpid { span } => f.root(&format!(
254                "SpecialRegister::Warpid [{}]",
255                f.format_raw(*span, source)
256            )),
257            SpecialRegister::Clusterid { span } => f.root(&format!(
258                "SpecialRegister::Clusterid [{}]",
259                f.format_raw(*span, source)
260            )),
261            SpecialRegister::Laneid { span } => f.root(&format!(
262                "SpecialRegister::Laneid [{}]",
263                f.format_raw(*span, source)
264            )),
265            SpecialRegister::Nwarpid { span } => f.root(&format!(
266                "SpecialRegister::Nwarpid [{}]",
267                f.format_raw(*span, source)
268            )),
269            SpecialRegister::WARPSZ { span } => f.root(&format!(
270                "SpecialRegister::WARPSZ [{}]",
271                f.format_raw(*span, source)
272            )),
273            SpecialRegister::Ctaid { axis, span } => {
274                f.root(&format!(
275                    "SpecialRegister::Ctaid [{}]",
276                    f.format_raw(*span, source)
277                ))?;
278                f.field_with_child(true, "axis", axis, source)
279            }
280            SpecialRegister::LanemaskEq { span } => f.root(&format!(
281                "SpecialRegister::LanemaskEq [{}]",
282                f.format_raw(*span, source)
283            )),
284            SpecialRegister::Pm { index, span } => {
285                f.root(&format!(
286                    "SpecialRegister::Pm [{}]",
287                    f.format_raw(*span, source)
288                ))?;
289                f.field(true, "index", &index.to_string())
290            }
291            SpecialRegister::Pm64 { index, span } => {
292                f.root(&format!(
293                    "SpecialRegister::Pm64 [{}]",
294                    f.format_raw(*span, source)
295                ))?;
296                f.field(true, "index", &index.to_string())
297            }
298            SpecialRegister::CurrentGraphExec { span } => f.root(&format!(
299                "SpecialRegister::CurrentGraphExec [{}]",
300                f.format_raw(*span, source)
301            )),
302            SpecialRegister::LanemaskGe { span } => f.root(&format!(
303                "SpecialRegister::LanemaskGe [{}]",
304                f.format_raw(*span, source)
305            )),
306            SpecialRegister::ReservedSmemOffset { index, span } => {
307                f.root(&format!(
308                    "SpecialRegister::ReservedSmemOffset [{}]",
309                    f.format_raw(*span, source)
310                ))?;
311                f.field(true, "index", &index.to_string())
312            }
313        }
314    }
315}
316
317impl TreeDisplay for Label {
318    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
319        f.root(&format!("Label [{}]", f.format_raw(self.span, source)))?;
320        f.field(true, "val", &self.val)
321    }
322}
323
324impl TreeDisplay for Predicate {
325    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
326        f.root(&format!("Predicate [{}]", f.format_raw(self.span, source)))?;
327        f.field(false, "negated", &self.negated.to_string())?;
328        f.field_with_child(true, "operand", &self.operand, source)
329    }
330}
331
332impl TreeDisplay for Instruction {
333    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
334        let pred_str = match &self.predicate {
335            Some(p) => format!(
336                "@{}{} ",
337                if p.negated { "!" } else { "" },
338                format!("{:?}", p.operand)
339            ),
340            None => String::new(),
341        };
342        f.root(&format!(
343            "Instruction [{}] {}{:?}",
344            f.format_raw(self.span, source),
345            pred_str,
346            self.inst
347        ))
348    }
349}
350
351impl TreeDisplay for Operand {
352    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
353        match self {
354            Operand::Register { operand, span } => {
355                f.root(&format!(
356                    "Operand::Register [{}]",
357                    f.format_raw(*span, source)
358                ))?;
359                f.field_with_child(true, "operand", operand, source)
360            }
361            Operand::Immediate { operand, span } => {
362                f.root(&format!(
363                    "Operand::Immediate [{}]",
364                    f.format_raw(*span, source)
365                ))?;
366                f.field_with_child(true, "operand", operand, source)
367            }
368            Operand::Symbol { name, span } => {
369                f.root(&format!(
370                    "Operand::Symbol [{}]",
371                    f.format_raw(*span, source)
372                ))?;
373                f.field(true, "name", name)
374            }
375            Operand::SymbolOffset {
376                symbol,
377                offset,
378                span,
379            } => {
380                f.root(&format!(
381                    "Operand::SymbolOffset [{}]",
382                    f.format_raw(*span, source)
383                ))?;
384                f.field(false, "symbol", symbol)?;
385                f.field_with_child(true, "offset", offset, source)
386            }
387        }
388    }
389}
390
391impl TreeDisplay for RegisterOperand {
392    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
393        f.root(&format!(
394            "RegisterOperand [{}]",
395            f.format_raw(self.span, source)
396        ))?;
397        f.field(false, "name", &self.name)?;
398        match &self.component {
399            Some(c) => f.field(true, "component", &format!("Some({})", c)),
400            None => f.field(true, "component", "None"),
401        }
402    }
403}
404
405impl TreeDisplay for Immediate {
406    fn tree_display(&self, f: &mut TreeFormatter, source: &str) -> std::fmt::Result {
407        f.root(&format!("Immediate [{}]", f.format_raw(self.span, source)))?;
408        f.field(true, "value", &self.value)
409    }
410}