Skip to main content

vicis_core/ir/function/
print.rs

1use super::{
2    super::module::name::Name,
3    super::types::Types,
4    super::value::{InlineAsm, Value},
5    basic_block::BasicBlockId,
6    data::Data,
7    instruction::{
8        Alloca, Cast, GetElementPtr, ICmp, Instruction, InstructionId, IntBinary, Load, Opcode,
9        Operand, Phi, Store,
10    },
11    Function,
12};
13use crate::ir::{
14    function::instruction::{
15        Br, Call, CondBr, ExtractValue, InsertValue, Invoke, LandingPad, Resume, Ret,
16    },
17    types::Type,
18};
19use rustc_hash::FxHashMap;
20use std::fmt;
21
22pub type Index = usize;
23pub type Indexes = FxHashMap<Ids, Name>;
24
25pub struct FunctionAsmPrinter<'a, 'b: 'a> {
26    fmt: &'a mut fmt::Formatter<'b>,
27    indexes: Indexes,
28    cur_index: Index,
29}
30
31#[derive(Debug, Copy, Clone, Hash, PartialEq, Eq)]
32pub enum Ids {
33    Block(BasicBlockId),
34    Inst(InstructionId),
35    Arg(usize),
36}
37
38impl<'a, 'b: 'a> FunctionAsmPrinter<'a, 'b> {
39    pub fn new(fmt: &'a mut fmt::Formatter<'b>) -> Self {
40        Self {
41            fmt,
42            indexes: FxHashMap::default(),
43            cur_index: 0,
44        }
45    }
46
47    pub fn print(&mut self, f: &Function) -> fmt::Result {
48        if f.is_prototype() {
49            write!(self.fmt, "declare ")?
50        } else {
51            write!(self.fmt, "define ")?
52        }
53
54        write!(self.fmt, "{:?} ", f.linkage)?;
55        write!(self.fmt, "{:?} ", f.preemption_specifier)?;
56        write!(self.fmt, "{:?} ", f.visibility)?;
57        for attr in &f.ret_attrs {
58            write!(self.fmt, "{} ", attr.to_string(&f.types))?
59        }
60        write!(self.fmt, "{} ", f.types.to_string(f.result_ty))?;
61        write!(self.fmt, "@{}(", f.name)?;
62
63        for (i, param) in f.params.iter().enumerate() {
64            write!(self.fmt, "{} ", f.types.to_string(param.ty))?;
65            for attr in &param.attrs {
66                write!(self.fmt, "{} ", attr.to_string(&f.types))?;
67            }
68            match param.name.to_string() {
69                Some(name) => {
70                    write!(self.fmt, "%{}", name)?;
71                    self.indexes.insert(Ids::Arg(i), Name::Name(name.clone()));
72                }
73                None => {
74                    let name = self.new_name_for_arg(i);
75                    write!(self.fmt, "%{:?}", name)?
76                }
77            }
78            write!(
79                self.fmt,
80                "{}",
81                if i == f.params.len() - 1 { "" } else { ", " }
82            )?;
83        }
84
85        if f.is_var_arg {
86            write!(self.fmt, ", ...")?;
87        }
88
89        write!(self.fmt, ") ")?;
90
91        if let Some(unnamed_addr) = f.unnamed_addr {
92            write!(self.fmt, "{:?} ", unnamed_addr)?
93        }
94
95        for attr in &f.func_attrs {
96            write!(self.fmt, "{:?} ", attr)?
97        }
98
99        if let Some((ty, func)) = &f.personality {
100            write!(
101                self.fmt,
102                "personality {} {} ",
103                f.types.to_string(*ty),
104                func.to_string(&f.types)
105            )?
106        }
107
108        if f.is_prototype() {
109            return writeln!(self.fmt);
110        }
111
112        writeln!(self.fmt, "{{")?;
113
114        for block_id in f.layout.block_iter() {
115            if let Some(name) = &f.data.block_ref(block_id).name {
116                match name.to_string() {
117                    Some(name) => {
118                        self.indexes
119                            .insert(Ids::Block(block_id), Name::Name(name.clone()));
120                    }
121                    None => {
122                        self.new_name_for_block(block_id);
123                    }
124                }
125            } else {
126                self.new_name_for_block(block_id);
127            }
128
129            for inst_id in f.layout.inst_iter(block_id) {
130                let inst = f.data.inst_ref(inst_id);
131                if matches!(
132                    inst.opcode,
133                    Opcode::Store | Opcode::Br | Opcode::CondBr | Opcode::Ret | Opcode::Resume
134                ) || (inst
135                    .operand
136                    .call_result_ty()
137                    .as_ref()
138                    .map_or(false, Type::is_void))
139                {
140                    continue;
141                }
142                if let Some(name) = &inst.dest {
143                    match name {
144                        Name::Name(name) => {
145                            self.indexes
146                                .insert(Ids::Inst(inst_id), Name::Name(name.clone()));
147                        }
148                        Name::Number(_) => {
149                            self.new_name_for_inst(inst_id);
150                        }
151                    }
152                } else {
153                    self.new_name_for_inst(inst_id);
154                }
155            }
156        }
157
158        for block_id in f.layout.block_iter() {
159            writeln!(
160                self.fmt,
161                "{:?}:",
162                self.indexes.get(&Ids::Block(block_id)).unwrap()
163            )?;
164
165            for inst_id in f.layout.inst_iter(block_id) {
166                let inst = f.data.inst_ref(inst_id);
167                write!(self.fmt, "    ")?;
168                self.print_inst(inst, &f.types, &f.data)?;
169                writeln!(self.fmt)?;
170            }
171        }
172
173        writeln!(self.fmt, "}}")
174    }
175
176    fn print_inst(&mut self, inst: &Instruction, types: &Types, data: &Data) -> fmt::Result {
177        let dest = self
178            .indexes
179            .get(&Ids::Inst(inst.id.unwrap()))
180            .unwrap_or(&Name::Number(usize::MAX));
181
182        match &inst.operand {
183            Operand::Alloca(Alloca {
184                tys,
185                num_elements,
186                align,
187            }) => {
188                write!(
189                    self.fmt,
190                    "%{:?} = alloca {}, {} {}{}",
191                    dest,
192                    types.to_string(tys[0]),
193                    types.to_string(tys[1]),
194                    num_elements.to_string(types),
195                    if *align > 0 {
196                        format!(", align {}", align)
197                    } else {
198                        "".to_string()
199                    }
200                )
201            }
202            Operand::Phi(Phi { ty, args, blocks }) => {
203                write!(
204                    self.fmt,
205                    "%{:?} = phi {} {}",
206                    dest,
207                    types.to_string(*ty),
208                    args.iter()
209                        .zip(blocks.iter())
210                        .fold("".to_string(), |acc, (arg, &block)| {
211                            format!(
212                                "{}[{}, %{:?}], ",
213                                acc,
214                                self.value_to_string(data.value_ref(*arg), types),
215                                self.indexes[&Ids::Block(block)]
216                            )
217                        })
218                        .trim_end_matches(", ")
219                )
220            }
221            Operand::Load(Load { tys, addr, align }) => {
222                write!(
223                    self.fmt,
224                    "%{:?} = load {}, {} {}{}",
225                    dest,
226                    types.to_string(tys[0]),
227                    types.to_string(tys[1]),
228                    self.value_to_string(data.value_ref(*addr), types),
229                    if *align == 0 {
230                        "".to_string()
231                    } else {
232                        format!(", align {}", align)
233                    }
234                )
235            }
236            Operand::Store(Store { tys, args, align }) => {
237                write!(
238                    self.fmt,
239                    "store {} {}, {} {}{}",
240                    types.to_string(tys[0]),
241                    self.value_to_string(data.value_ref(args[0]), types),
242                    types.to_string(tys[1]),
243                    self.value_to_string(data.value_ref(args[1]), types),
244                    if *align == 0 {
245                        "".to_string()
246                    } else {
247                        format!(", align {}", align)
248                    }
249                )
250            }
251            Operand::InsertValue(InsertValue { tys, args }) => {
252                write!(
253                    self.fmt,
254                    "%{:?} = insertvalue {} {}, {} {}, {}",
255                    dest,
256                    types.to_string(tys[0]),
257                    self.value_to_string(data.value_ref(args[0]), types),
258                    types.to_string(tys[1]),
259                    self.value_to_string(data.value_ref(args[1]), types),
260                    args[2..]
261                        .iter()
262                        .fold("".to_string(), |acc, &arg| {
263                            format!(
264                                "{}{}, ",
265                                acc,
266                                self.value_to_string(data.value_ref(arg), types)
267                            )
268                        })
269                        .trim_end_matches(", ")
270                )
271            }
272            Operand::ExtractValue(ExtractValue { ty, args }) => {
273                write!(
274                    self.fmt,
275                    "%{:?} = extractvalue {} {}, {}",
276                    dest,
277                    types.to_string(*ty),
278                    self.value_to_string(data.value_ref(args[0]), types),
279                    args[1..]
280                        .iter()
281                        .fold("".to_string(), |acc, &arg| {
282                            format!(
283                                "{}{}, ",
284                                acc,
285                                self.value_to_string(data.value_ref(arg), types)
286                            )
287                        })
288                        .trim_end_matches(", ")
289                )
290            }
291            Operand::IntBinary(IntBinary {
292                ty,
293                nuw,
294                nsw,
295                exact,
296                args,
297            }) => {
298                write!(
299                    self.fmt,
300                    "%{:?} = {:?}{}{}{} {} {}, {}",
301                    dest,
302                    inst.opcode,
303                    if *nuw { " nuw" } else { "" },
304                    if *nsw { " nsw" } else { "" },
305                    if *exact { " exact" } else { "" },
306                    types.to_string(*ty),
307                    self.value_to_string(data.value_ref(args[0]), types),
308                    self.value_to_string(data.value_ref(args[1]), types),
309                )
310            }
311            Operand::ICmp(ICmp { ty, args, cond }) => {
312                write!(
313                    self.fmt,
314                    "%{:?} = icmp {:?} {} {}, {}",
315                    dest,
316                    cond,
317                    types.to_string(*ty),
318                    self.value_to_string(data.value_ref(args[0]), types),
319                    self.value_to_string(data.value_ref(args[1]), types)
320                )
321            }
322            Operand::Cast(Cast { tys, arg }) => {
323                write!(
324                    self.fmt,
325                    "%{:?} = {:?} {} {} to {}",
326                    dest,
327                    inst.opcode,
328                    types.to_string(tys[0]),
329                    self.value_to_string(data.value_ref(*arg), types),
330                    types.to_string(tys[1]),
331                )
332            }
333            Operand::GetElementPtr(GetElementPtr {
334                inbounds,
335                tys,
336                args,
337            }) => {
338                write!(
339                    self.fmt,
340                    "%{:?} = getelementptr {}{}, {}",
341                    dest,
342                    if *inbounds { "inbounds " } else { "" },
343                    types.to_string(tys[0]),
344                    tys[1..]
345                        .iter()
346                        .zip(args.iter())
347                        .fold("".to_string(), |acc, (ty, arg)| {
348                            format!(
349                                "{}{} {}, ",
350                                acc,
351                                types.to_string(*ty),
352                                self.value_to_string(data.value_ref(*arg), types),
353                            )
354                        })
355                        .trim_end_matches(", ")
356                )
357            }
358            Operand::Call(Call {
359                tys,
360                args,
361                param_attrs,
362                ret_attrs,
363                func_attrs,
364                ..
365            }) => {
366                write!(
367                    self.fmt,
368                    "{}call {}{} {}({}) {}",
369                    if tys[0].is_void() {
370                        "".to_string()
371                    } else {
372                        format!("%{:?} = ", dest)
373                    },
374                    ret_attrs.iter().fold("".to_string(), |acc, attr| format!(
375                        "{}{} ",
376                        acc,
377                        attr.to_string(types)
378                    )),
379                    types.to_string(tys[0]),
380                    self.value_to_string(data.value_ref(args[0]), types),
381                    tys[1..]
382                        .iter()
383                        .zip(args[1..].iter())
384                        .zip(param_attrs.iter())
385                        .into_iter()
386                        .fold("".to_string(), |acc, ((&ty, &arg), attrs)| {
387                            format!(
388                                "{}{} {}{}, ",
389                                acc,
390                                types.to_string(ty),
391                                attrs.iter().fold("".to_string(), |acc, attr| {
392                                    format!("{}{} ", acc, attr.to_string(types))
393                                }),
394                                self.value_to_string(data.value_ref(arg), types),
395                            )
396                        })
397                        .trim_end_matches(", "),
398                    func_attrs
399                        .iter()
400                        .fold("".to_string(), |acc, attr| format!("{}{:?} ", acc, attr))
401                )
402            }
403            Operand::Invoke(Invoke {
404                tys,
405                args,
406                param_attrs,
407                ret_attrs,
408                func_attrs,
409                blocks,
410            }) => {
411                write!(
412                    self.fmt,
413                    "{}invoke {}{} {}({}) {}to label %{:?} unwind label %{:?}",
414                    if tys[0].is_void() {
415                        "".to_string()
416                    } else {
417                        format!("%{:?} = ", dest)
418                    },
419                    ret_attrs.iter().fold("".to_string(), |acc, attr| format!(
420                        "{}{} ",
421                        acc,
422                        attr.to_string(types)
423                    )),
424                    types.to_string(tys[0]),
425                    self.value_to_string(data.value_ref(args[0]), types),
426                    tys[1..]
427                        .iter()
428                        .zip(args[1..].iter())
429                        .zip(param_attrs.iter())
430                        .into_iter()
431                        .fold("".to_string(), |acc, ((&ty, &arg), attrs)| {
432                            format!(
433                                "{}{} {}{}, ",
434                                acc,
435                                types.to_string(ty),
436                                attrs.iter().fold("".to_string(), |acc, attr| {
437                                    format!("{}{} ", acc, attr.to_string(types))
438                                }),
439                                self.value_to_string(data.value_ref(arg), types),
440                            )
441                        })
442                        .trim_end_matches(", "),
443                    func_attrs
444                        .iter()
445                        .fold("".to_string(), |acc, attr| format!("{}{:?} ", acc, attr)),
446                    self.indexes[&Ids::Block(blocks[0])],
447                    self.indexes[&Ids::Block(blocks[1])],
448                )
449            }
450            Operand::LandingPad(LandingPad { ty }) => {
451                write!(
452                    self.fmt,
453                    "{}landingpad {} cleanup",
454                    if ty.is_void() {
455                        "".to_string()
456                    } else {
457                        format!("%{:?} = ", dest)
458                    },
459                    types.to_string(*ty),
460                )
461            }
462            Operand::Resume(Resume { ty, arg }) => {
463                write!(
464                    self.fmt,
465                    "resume {} {}",
466                    types.to_string(*ty),
467                    self.value_to_string(data.value_ref(*arg), types),
468                )
469            }
470            Operand::Br(Br { block }) => {
471                write!(
472                    self.fmt,
473                    "br label %{:?}",
474                    self.indexes[&Ids::Block(*block)]
475                )
476            }
477            Operand::CondBr(CondBr { arg, blocks }) => {
478                write!(
479                    self.fmt,
480                    "br i1 {}, label %{:?}, label %{:?}",
481                    self.value_to_string(data.value_ref(*arg), types),
482                    self.indexes[&Ids::Block(blocks[0])],
483                    self.indexes[&Ids::Block(blocks[1])],
484                )
485            }
486            Operand::Ret(Ret { val: None, .. }) => write!(self.fmt, "ret void"),
487            Operand::Ret(Ret { val: Some(val), ty }) => {
488                write!(
489                    self.fmt,
490                    "ret {} {}",
491                    types.to_string(*ty),
492                    self.value_to_string(data.value_ref(*val), types),
493                )
494            }
495            Operand::Unreachable => {
496                write!(self.fmt, "unreachable")
497            }
498            Operand::Invalid => panic!(),
499        }?;
500
501        for (kind, meta) in &inst.metadata {
502            write!(self.fmt, ", !{} {:?}", kind, meta)?;
503        }
504
505        Ok(())
506    }
507
508    fn value_to_string(&self, val: &Value, types: &Types) -> String {
509        match val {
510            Value::Constant(c) => c.to_string(types),
511            Value::Instruction(id) => {
512                format!("%{:?}", self.indexes[&Ids::Inst(*id)])
513            }
514            Value::Argument(n) => format!("%{:?}", self.indexes[&Ids::Arg(*n)]),
515            Value::InlineAsm(InlineAsm {
516                body,
517                constraints,
518                sideeffect,
519            }) => {
520                format!(
521                    "asm {}\"{}\", \"{}\"",
522                    if *sideeffect { "sideeffect " } else { "" },
523                    constraints,
524                    body
525                )
526            }
527        }
528    }
529
530    fn new_name_for_block(&mut self, id: BasicBlockId) -> Name {
531        let idx = self.cur_index;
532        self.cur_index += 1;
533        self.indexes.insert(Ids::Block(id), Name::Number(idx));
534        Name::Number(idx)
535    }
536
537    pub fn new_name_for_inst(&mut self, id: InstructionId) -> Name {
538        let idx = self.cur_index;
539        self.cur_index += 1;
540        self.indexes.insert(Ids::Inst(id), Name::Number(idx));
541        Name::Number(idx)
542    }
543
544    pub fn new_name_for_arg(&mut self, arg: usize) -> Name {
545        let idx = self.cur_index;
546        self.cur_index += 1;
547        self.indexes.insert(Ids::Arg(arg), Name::Number(idx));
548        Name::Number(idx)
549    }
550}