rv32_asm/
assembler.rs

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}