1use crate::ast::processed_ast::{Program, ProgramItem};
2use crate::ast::raw_ast::{
3 Comment, Directive, DirectiveType, Immediate, Instruction, InstructionType, Label, Register,
4};
5use either::Either;
6use serde::{Deserialize, Serialize};
7
8trait FormattedDisplay {
9 fn formatted_display(&self, style: &FormatStyle) -> (Vec<String>, String, Option<String>);
11}
12
13#[derive(Debug, Default, Copy, Clone, Serialize, Deserialize)]
14#[serde(rename_all = "kebab-case")]
15pub struct FormatStyle {
16 pub indent_directive: u8, pub indent_instruction: u8, pub indent_label: u8, pub indent_min_comment_from_block: u8, pub space_block_to_comment: u8, pub space_comment_stick_to_body: u8, pub space_from_label_block: u8, pub space_from_start_end_block: u8, pub colon_after_label: bool,
25 pub fixed_body_comment_indent: bool,
26}
27
28pub struct Formatter<'a> {
29 style: &'a FormatStyle,
30 buffer: Vec<u8>,
31}
32
33impl<'a> Formatter<'a> {
34 pub fn new(style: &'a FormatStyle) -> Self {
35 Self {
36 style,
37 buffer: Vec::new(),
38 }
39 }
40
41 pub fn format(&mut self, program: Program) {
42 self.buffer.reserve(program.items().len() * 10);
43 let mut lines: Vec<(Vec<String>, String, Option<String>, usize)> = vec![];
44 for (index, line) in program.items().iter().enumerate() {
45 let (labels, body, comments) = line.formatted_display(&self.style);
46 lines.push((
47 labels,
48 body,
49 comments,
50 self.control_padding(line, program.items().get(index + 1)) + 1,
51 ));
52 }
53
54 let comment_start_column = lines.iter().map(|e| e.1.len()).max().unwrap_or(0)
55 + (self.style.indent_min_comment_from_block as usize);
56
57 for (labels, body, comment, space) in lines.into_iter() {
58 let missing_indent = comment_start_column - body.len();
59 let mut label = "".to_owned();
60 labels
61 .into_iter()
62 .map(|mut l| {
63 l.push('\n');
64 l
65 })
66 .for_each(|e| label.push_str(e.as_str()));
67 self.buffer.append(&mut label.into_bytes());
68 self.buffer.append(&mut body.into_bytes());
69 self.add_indent(
70 self.style
71 .fixed_body_comment_indent
72 .then_some(missing_indent)
73 .unwrap_or(1), );
75 match comment {
76 None => {}
77 Some(comment) => {
78 self.buffer.append(&mut comment.into_bytes());
79 }
80 }
81 self.add_newline(space);
82 }
83 }
84
85 pub fn contents(&self) -> &Vec<u8> {
86 &self.buffer
87 }
88
89 #[inline]
90 fn add_newline(&mut self, lines: usize) {
91 for _ in 0..lines {
92 self.buffer.push(b'\n');
93 }
94 }
95
96 #[inline]
97 fn add_indent(&mut self, indent: usize) {
98 for _ in 0..indent {
99 self.buffer.push(b' ');
100 }
101 }
102
103 fn control_padding(&mut self, current: &ProgramItem, next: Option<&ProgramItem>) -> usize {
104 let mut paddings = 0usize;
105
106 if self.style.space_comment_stick_to_body != 0 {
108 if current.is_comment()
109 && next.is_some()
110 && (next.unwrap().is_directive() || next.unwrap().is_instruction())
111 {
112 paddings += self.style.space_comment_stick_to_body as usize;
113 }
114 }
115
116 if self.style.space_block_to_comment != 0 {
118 if let ProgramItem::Directive(_, directive, ..) = current {
120 if matches!(directive.directive_type(), DirectiveType::ORIG(..)) {
121 } else {
123 if (current.is_directive() || current.is_instruction())
124 && next.is_some()
125 && next.unwrap().is_comment()
126 {
127 paddings += self.style.space_block_to_comment as usize;
128 }
129 }
130 } else {
131 if (current.is_directive() || current.is_instruction())
132 && next.is_some()
133 && next.unwrap().is_comment()
134 {
135 paddings += self.style.space_block_to_comment as usize;
136 }
137 }
138 }
139
140 if self.style.space_from_label_block != 0 {
142 let space: u8 = match current {
143 ProgramItem::Instruction(curr_label, ..)
144 | ProgramItem::Directive(curr_label, ..) => match next {
145 None => 0,
146 Some(next) => match next {
147 ProgramItem::Instruction(next_label, ..)
148 | ProgramItem::Directive(next_label, ..) => {
149 if curr_label.is_empty() && (!next_label.is_empty()) {
150 self.style.space_from_label_block
151 } else {
152 0
153 }
154 }
155 ProgramItem::EOL(..) | ProgramItem::Comment(..) => 0,
156 },
157 },
158 ProgramItem::EOL(..) | ProgramItem::Comment(..) => 0,
159 };
160 paddings += space as usize;
161 }
162
163 if self.style.space_from_start_end_block != 0 {
165 let space: u8 = match current {
166 ProgramItem::Directive(_, directive, ..) => {
167 if matches!(directive.directive_type(), DirectiveType::ORIG(..)) {
168 self.style.space_from_start_end_block
169 } else if next.is_some() {
170 match next.unwrap() {
171 ProgramItem::Directive(_, directive, _, _) => {
172 if matches!(directive.directive_type(), DirectiveType::END) {
173 self.style.space_from_start_end_block
174 } else {
175 0
176 }
177 }
178 _ => 0,
179 }
180 } else {
181 0
182 }
183 }
184 _ => 0,
185 };
186 paddings += space as usize;
187 }
188 paddings
189 }
190}
191
192impl FormattedDisplay for ProgramItem {
193 fn formatted_display(&self, style: &FormatStyle) -> (Vec<String>, String, Option<String>) {
194 match self {
195 ProgramItem::Comment(comment, _) => (vec![], print_comment(comment), None),
196 ProgramItem::Instruction(labels, instruction, comment, _) => {
197 let mut label_indent = "".to_owned();
198 add_indent(
199 &mut label_indent,
200 labels.is_empty().then_some(0).unwrap_or(style.indent_label),
201 );
202 let labels = labels
203 .into_iter()
204 .map(|l| format!("{label_indent}{}", print_label(style, l)))
205 .collect();
206 let mut instruction_indent = "".to_owned();
207 add_indent(&mut instruction_indent, style.indent_instruction);
208 let comment = comment.as_ref().map_or(None, |c| Some(print_comment(c)));
209 (
210 labels,
211 format!("{instruction_indent}{}", print_instruction(instruction)),
212 comment,
213 )
214 }
215 ProgramItem::Directive(labels, directive, comment, _) => {
216 let mut label_indent = "".to_owned();
217 add_indent(
218 &mut label_indent,
219 labels.is_empty().then_some(0).unwrap_or(style.indent_label),
220 );
221 let labels = labels
222 .into_iter()
223 .map(|l| format!("{label_indent}{}", print_label(style, l)))
224 .collect();
225 let mut directive_indent = "".to_owned();
226 add_indent(
227 &mut directive_indent,
228 (matches!(directive.directive_type(), DirectiveType::END)
229 || matches!(directive.directive_type(), DirectiveType::ORIG(..)))
230 .then_some(0)
231 .unwrap_or(style.indent_directive),
232 );
233 let comment = comment.as_ref().map_or(None, |c| Some(print_comment(c)));
234 (
235 labels,
236 format!("{directive_indent}{}", print_directive(directive)),
237 comment,
238 )
239 }
240 ProgramItem::EOL(labels) => {
241 let mut label_indent = "".to_owned();
242 add_indent(
243 &mut label_indent,
244 labels.is_empty().then_some(0).unwrap_or(style.indent_label),
245 );
246 let labels = labels
247 .into_iter()
248 .map(|l| format!("{label_indent}{}", print_label(style, l)))
249 .collect();
250 (labels, "".to_owned(), None)
251 }
252 }
253 }
254}
255
256fn print_instruction(instruction: &Instruction) -> String {
257 let operands: String = match instruction.instruction_type() {
258 InstructionType::Add(register1, register2, register_or_immediate)
259 | InstructionType::And(register1, register2, register_or_immediate) => {
260 format!(
261 "{}, {}, {}",
262 register1.content(),
263 register2.content(),
264 print_register_or_immediate(register_or_immediate)
265 )
266 }
267 InstructionType::Not(register1, register2) => {
268 format!("{}, {}", register1.content(), register2.content(),)
269 }
270 InstructionType::Ldr(register1, register2, immediate)
271 | InstructionType::Str(register1, register2, immediate) => {
272 format!(
273 "{}, {}, {}",
274 register1.content(),
275 register2.content(),
276 immediate.content()
277 )
278 }
279 InstructionType::Ld(register1, label_ref)
280 | InstructionType::Ldi(register1, label_ref)
281 | InstructionType::Lea(register1, label_ref)
282 | InstructionType::St(register1, label_ref)
283 | InstructionType::Sti(register1, label_ref) => {
284 format!("{}, {}", register1.content(), label_ref.content())
285 }
286 InstructionType::Br(_, label_ref) => label_ref.content().to_owned(),
287 InstructionType::Jmp(register) | InstructionType::Jsrr(register) => {
288 register.content().to_owned()
289 }
290 InstructionType::Jsr(label_ref) => label_ref.content().to_owned(),
291 InstructionType::Nop
292 | InstructionType::Ret
293 | InstructionType::Halt
294 | InstructionType::Puts
295 | InstructionType::Getc
296 | InstructionType::Out
297 | InstructionType::In => "".to_owned(),
298 InstructionType::Trap(hex_address) => hex_address.content().to_owned(),
299 };
300 if operands.is_empty() {
301 format!("{}", instruction.content())
302 } else {
303 format!("{} {}", instruction.content(), operands)
304 }
305}
306
307fn print_comment(comment: &Comment) -> String {
308 match comment.content().strip_prefix(";") {
309 None => {
310 unreachable!()
311 }
312 Some(comment) => {
313 let comment = comment.trim();
314 format!(";{comment}")
315 }
316 }
317}
318
319fn print_label(style: &FormatStyle, label: &Label) -> String {
320 if style.colon_after_label {
321 match label.content().ends_with(":") {
322 false => {
323 if style.colon_after_label {
324 format!("{}:", label.content())
325 } else {
326 label.content().into()
327 }
328 }
329 true => label.content().to_owned(),
330 }
331 } else {
332 match label.content().strip_suffix(":") {
333 None => label.content().into(),
334 Some(label) => label.to_owned(),
335 }
336 }
337}
338
339fn print_register_or_immediate(either: &Either<Register, Immediate>) -> String {
340 match either {
341 Either::Left(r) => r.content(),
342 Either::Right(im) => im.content(),
343 }
344 .to_owned()
345}
346
347fn print_directive(directive: &Directive) -> String {
348 let operands: String = match directive.directive_type() {
349 DirectiveType::ORIG(address) => address.content(),
350 DirectiveType::END => "",
351 DirectiveType::BLKW(immediate) | DirectiveType::FILL(immediate) => immediate.content(),
352 DirectiveType::STRINGZ(string) => string.content(),
353 }
354 .to_owned();
355 if operands.is_empty() {
356 format!("{}", directive.content())
357 } else {
358 format!("{} {}", directive.content(), operands)
359 }
360}
361
362fn add_indent(string: &mut String, indent: u8) {
363 for _ in 0..indent {
364 string.push(' ');
365 }
366}