ptx_parser/unparser/
module.rs

1use super::PtxUnparser;
2use crate::{
3    lexer::{PtxToken, tokenize},
4    r#type::{function::DwarfDirective, module::*},
5    unparser::{push_decimal, push_directive, push_identifier, push_token_from_str},
6};
7
8fn numeric_token_from_str(value: &str) -> Option<PtxToken> {
9    if value.starts_with("0f")
10        || value.starts_with("0F")
11        || value.starts_with("0d")
12        || value.starts_with("0D")
13    {
14        Some(PtxToken::HexFloat(value.to_string()))
15    } else if value.starts_with("0x") || value.starts_with("0X") {
16        Some(PtxToken::HexInteger(value.to_string()))
17    } else if value.starts_with("0b") || value.starts_with("0B") {
18        Some(PtxToken::BinaryInteger(value.to_string()))
19    } else if value.len() > 1
20        && value.starts_with('0')
21        && value.chars().all(|c| c >= '0' && c <= '7')
22    {
23        Some(PtxToken::OctalInteger(value.to_string()))
24    } else if value.chars().all(|c| c.is_ascii_digit()) {
25        Some(PtxToken::DecimalInteger(value.to_string()))
26    } else if is_float_exponent(value) {
27        Some(PtxToken::FloatExponent(value.to_string()))
28    } else if is_decimal_float(value) {
29        Some(PtxToken::Float(value.to_string()))
30    } else {
31        None
32    }
33}
34
35fn is_decimal_float(value: &str) -> bool {
36    let mut has_digits = false;
37    let mut has_dot = false;
38    for ch in value.chars() {
39        match ch {
40            '0'..='9' => has_digits = true,
41            '.' if !has_dot => has_dot = true,
42            _ => return false,
43        }
44    }
45    has_digits && has_dot
46}
47
48fn is_float_exponent(value: &str) -> bool {
49    let mut chars = value.chars().peekable();
50    let mut has_digits = false;
51    while let Some(ch) = chars.peek() {
52        match ch {
53            '0'..='9' | '+' | '-' => {
54                has_digits |= ch.is_ascii_digit();
55                chars.next();
56            }
57            '.' => {
58                chars.next();
59            }
60            'e' | 'E' => {
61                chars.next();
62                if chars.peek().is_none() {
63                    return false;
64                }
65                let mut exponent_digits = false;
66                if let Some(sign) = chars.peek() {
67                    if *sign == '+' || *sign == '-' {
68                        chars.next();
69                    }
70                }
71                while let Some(next) = chars.peek() {
72                    match next {
73                        '0'..='9' => {
74                            exponent_digits = true;
75                            chars.next();
76                        }
77                        _ => return false,
78                    }
79                }
80                return has_digits && exponent_digits;
81            }
82            _ => return false,
83        }
84    }
85    false
86}
87
88fn push_token_from_repr(tokens: &mut Vec<PtxToken>, value: &str) {
89    let trimmed = value.trim();
90    match trimmed {
91        "{" => tokens.push(PtxToken::LBrace),
92        "}" => tokens.push(PtxToken::RBrace),
93        "[" => tokens.push(PtxToken::LBracket),
94        "]" => tokens.push(PtxToken::RBracket),
95        "(" => tokens.push(PtxToken::LParen),
96        ")" => tokens.push(PtxToken::RParen),
97        "," => tokens.push(PtxToken::Comma),
98        ":" => tokens.push(PtxToken::Colon),
99        ";" => tokens.push(PtxToken::Semicolon),
100        "." => tokens.push(PtxToken::Dot),
101        "*" => tokens.push(PtxToken::Star),
102        "+" => tokens.push(PtxToken::Plus),
103        "-" => tokens.push(PtxToken::Minus),
104        "/" => tokens.push(PtxToken::Slash),
105        "%" => tokens.push(PtxToken::Percent),
106        "=" => tokens.push(PtxToken::Equals),
107        "|" => tokens.push(PtxToken::Pipe),
108        "&" => tokens.push(PtxToken::Ampersand),
109        "^" => tokens.push(PtxToken::Caret),
110        "~" => tokens.push(PtxToken::Tilde),
111        "!" => tokens.push(PtxToken::Exclaim),
112        "<" => tokens.push(PtxToken::LAngle),
113        ">" => tokens.push(PtxToken::RAngle),
114        "@" => tokens.push(PtxToken::At),
115        "::" => tokens.push(PtxToken::DoubleColon),
116        _ => {
117            if let Some(rest) = trimmed.strip_prefix('.') {
118                push_directive(tokens, rest);
119            } else if trimmed.starts_with('"') && trimmed.ends_with('"') && trimmed.len() >= 2 {
120                tokens.push(PtxToken::StringLiteral(
121                    trimmed[1..trimmed.len() - 1].to_string(),
122                ));
123            } else if trimmed.starts_with('%') {
124                tokens.push(PtxToken::Register(trimmed.to_string()));
125            } else if let Some(numeric) = numeric_token_from_str(trimmed) {
126                tokens.push(numeric);
127            } else {
128                push_token_from_str(tokens, trimmed);
129            }
130        }
131    }
132}
133
134fn push_entries(tokens: &mut Vec<PtxToken>, entries: &[String]) {
135    for (index, entry) in entries.iter().enumerate() {
136        if index > 0 {
137            tokens.push(PtxToken::Comma);
138        }
139        if let Some(rest) = entry.strip_prefix('.') {
140            push_directive(tokens, rest);
141        } else if let Some(numeric) = numeric_token_from_str(entry) {
142            tokens.push(numeric);
143        } else {
144            push_identifier(tokens, entry);
145        }
146    }
147}
148
149fn push_raw_tokens(tokens: &mut Vec<PtxToken>, raw: &str) {
150    if raw.trim().is_empty() {
151        return;
152    }
153    let lexemes = tokenize(raw).expect("tokenizing stored PTX snippet");
154    tokens.extend(lexemes.into_iter().map(|(token, _)| token));
155}
156
157impl PtxUnparser for VersionDirective {
158    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
159        push_directive(tokens, "version");
160        let value = format!("{}.{:}", self.major, self.minor);
161        tokens.push(PtxToken::Float(value));
162    }
163}
164
165impl PtxUnparser for TargetDirective {
166    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
167        push_directive(tokens, "target");
168        push_entries(tokens, &self.entries);
169    }
170}
171
172impl PtxUnparser for AddressSizeDirective {
173    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
174        push_directive(tokens, "address_size");
175        push_decimal(tokens, self.size);
176    }
177}
178
179impl PtxUnparser for ModuleInfoDirectiveKind {
180    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
181        match self {
182            ModuleInfoDirectiveKind::Version { directive, .. } => directive.unparse_tokens(tokens),
183            ModuleInfoDirectiveKind::Target { directive, .. } => directive.unparse_tokens(tokens),
184            ModuleInfoDirectiveKind::AddressSize { directive, .. } => directive.unparse_tokens(tokens),
185        }
186    }
187}
188
189impl PtxUnparser for FileDirective {
190    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
191        push_directive(tokens, "file");
192        push_decimal(tokens, self.index);
193        tokens.push(PtxToken::StringLiteral(self.path.clone()));
194    }
195}
196
197impl PtxUnparser for SectionDirective {
198    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
199        push_directive(tokens, "section");
200        if let Some(rest) = self.name.strip_prefix('.') {
201            push_directive(tokens, rest);
202        } else {
203            push_identifier(tokens, &self.name);
204        }
205        for attribute in &self.attributes {
206            push_token_from_repr(tokens, attribute);
207        }
208    }
209}
210
211impl PtxUnparser for DwarfDirective {
212    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
213        // Reconstruct the dwarf directive from its fields
214        push_raw_tokens(tokens, &self.keyword);
215        for arg in &self.arguments {
216            push_raw_tokens(tokens, arg);
217        }
218        if let Some(comment) = &self.comment {
219            push_raw_tokens(tokens, comment);
220        }
221    }
222}
223
224impl PtxUnparser for ModuleDebugDirective {
225    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
226        match self {
227            ModuleDebugDirective::File { directive, .. } => directive.unparse_tokens(tokens),
228            ModuleDebugDirective::Section { directive, .. } => directive.unparse_tokens(tokens),
229            ModuleDebugDirective::Dwarf { directive, .. } => directive.unparse_tokens(tokens),
230        }
231    }
232}
233
234impl PtxUnparser for LinkingDirective {
235    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
236        self.kind.unparse_tokens(tokens);
237        push_raw_tokens(tokens, &self.prototype);
238        tokens.push(PtxToken::Semicolon);
239    }
240}
241
242impl PtxUnparser for ModuleDirective {
243    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
244        match self {
245            ModuleDirective::ModuleVariable { directive, .. } => directive.unparse_tokens(tokens),
246            ModuleDirective::FunctionKernel { directive, .. } => directive.unparse_tokens(tokens),
247            ModuleDirective::ModuleInfo { directive, .. } => directive.unparse_tokens(tokens),
248            ModuleDirective::Debug { directive, .. } => directive.unparse_tokens(tokens),
249            ModuleDirective::Linking { directive, .. } => directive.unparse_tokens(tokens),
250        }
251    }
252}
253
254impl PtxUnparser for Module {
255    fn unparse_tokens(&self, tokens: &mut Vec<PtxToken>) {
256        for directive in &self.directives {
257            directive.unparse_tokens(tokens);
258        }
259    }
260}