1use std::fs::File;
2use std::io::Write;
3use std::path::Path;
4
5use anyhow::Result;
6use colored::*;
7use rustc_hash::FxHashMap;
8
9use crate::error::*;
10use crate::instruction::*;
11
12struct Context {
13 label_addrs: FxHashMap<String, usize>,
14}
15
16impl Context {
17 fn new() -> Context {
18 Context {
19 label_addrs: FxHashMap::default(),
20 }
21 }
22
23 fn get_addr_by_label(&self, name: &str) -> Result<usize> {
24 Ok(self
25 .label_addrs
26 .get(name)
27 .cloned()
28 .ok_or_else(|| Error::LabelNotDefined(name.to_string()))?)
29 }
30}
31
32type Code = u32;
33type Codes = Vec<Code>;
34
35pub fn dump_instructions(insts: &[InstructionWithLabel]) {
36 println!("{}", "RISC-V Instructions:".red());
37 for (addr, InstructionWithLabel { inst, labels, ir }) in insts.iter().enumerate() {
38 for label in labels {
39 let addr = format!("; 0x{:x}", addr * 4);
40 println!("{}: {}", &label.name, addr.dimmed());
41 }
42
43 if let Some(ir) = ir {
44 let ir = format!(";{}", ir);
45 println!(" {}", ir.dimmed());
46 }
47 println!(" {}", inst);
48 }
49 println!();
50}
51
52fn replace_label(imm: Immediate, ctx: &Context) -> Result<Immediate> {
53 match imm {
54 Immediate::Value(_) => Ok(imm),
55 Immediate::Label(label) => Ok(Immediate::new(ctx.get_addr_by_label(&label.name)? as i32)),
56 }
57}
58
59fn replace_redaddr_label(rel_addr: RelAddress, addr: usize, ctx: &Context) -> Result<RelAddress> {
60 match rel_addr {
61 RelAddress::Immediate(_) => Ok(rel_addr),
62 RelAddress::Label(label) => Ok(RelAddress::Immediate(Immediate::Value(
63 ctx.get_addr_by_label(&label.name)? as i32 - addr as i32,
64 ))),
65 }
66}
67
68fn replace_labels(inst: InstructionWithLabel, ctx: &Context) -> Result<InstructionWithLabel> {
69 use Instruction::*;
70
71 let InstructionWithLabel { inst, labels, ir } = inst;
72 let replaced = match inst {
73 R(_) => inst,
74 I(IInstruction { op, imm, rs1, rd }) => I(IInstruction {
75 op,
76 imm: replace_label(imm, ctx)?,
77 rs1,
78 rd,
79 }),
80 S(SInstruction { op, imm, rs1, rs2 }) => S(SInstruction {
81 op,
82 imm: replace_label(imm, ctx)?,
83 rs1,
84 rs2,
85 }),
86 J(_) => inst,
87 U(UInstruction { op, imm, rd }) => U(UInstruction {
88 op,
89 imm: replace_label(imm, ctx)?,
90 rd,
91 }),
92 SB(_) => inst,
93 };
94 Ok(InstructionWithLabel::new(replaced, labels, ir))
95}
96
97fn replace_reladdr_labels(
98 inst: InstructionWithLabel,
99 addr: usize,
100 ctx: &Context,
101) -> Result<InstructionWithLabel> {
102 use Instruction::*;
103
104 let InstructionWithLabel { inst, labels, ir } = inst;
105 let replaced = match inst {
106 J(JInstruction { op, imm, rd }) => J(JInstruction {
107 op,
108 imm: replace_redaddr_label(imm, addr, ctx)?,
109 rd,
110 }),
111 SB(SBInstruction { op, imm, rs1, rs2 }) => SB(SBInstruction {
112 op,
113 imm: replace_redaddr_label(imm, addr, ctx)?,
114 rs1,
115 rs2,
116 }),
117 _ => inst,
118 };
119 Ok(InstructionWithLabel::new(replaced, labels, ir))
120}
121
122pub fn assemble<P>(instructions: Vec<InstructionWithLabel>, dump_to: Option<P>) -> Result<Codes>
123where
124 P: AsRef<Path>,
125{
126 use Instruction::*;
127
128 let mut ctx = Context::new();
129
130 for (idx, inst) in instructions.iter().enumerate() {
131 for label in &inst.labels {
132 ctx.label_addrs.insert(label.name.clone(), idx * 4);
133 }
134 }
135
136 let mut insts = Vec::with_capacity(instructions.len());
137 for (addr, inst) in instructions.into_iter().enumerate() {
138 let inst = replace_labels(inst, &ctx)?;
139 insts.push(replace_reladdr_labels(inst, addr * 4, &ctx)?);
140 }
141
142 if let Some(dump_to) = dump_to {
143 let mut asm = File::create(dump_to)?;
144 for (addr, InstructionWithLabel { inst, ir, labels }) in insts.iter().enumerate() {
145 for label in labels {
146 writeln!(asm, "{}: # 0x{:x}", &label.name, addr * 4)?;
147 }
148
149 if let Some(ir) = ir {
150 writeln!(asm, " #{}", ir)?;
151 }
152 let a = match inst {
153 R(ri) => ri.generate_asm(),
154 I(ii) => ii.generate_asm(),
155 S(si) => si.generate_asm(),
156 J(ji) => ji.generate_asm(),
157 U(ui) => ui.generate_asm(),
158 SB(sbi) => sbi.generate_asm(),
159 };
160 writeln!(asm, " {}", a)?;
161 }
162 }
163
164 let result = insts
165 .into_iter()
166 .map(|InstructionWithLabel { inst, .. }| match inst {
167 R(ri) => ri.generate_code(),
168 I(ii) => ii.generate_code(),
169 S(si) => si.generate_code(),
170 J(ji) => ji.generate_code(),
171 U(ui) => ui.generate_code(),
172 SB(sbi) => sbi.generate_code(),
173 })
174 .collect();
175
176 Ok(result)
177}