neyuki 0.1.1

A small Lua-like scripting language implemented in Rust
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93

// Linear scan CFG-aware register allocator.
// Maps unbounded virtual IrVars to physical VM registers (0..255) using live intervals.

use std::collections::{BTreeSet, HashMap};

use crate::compiler::ir::block::ControlFlowGraph;
use crate::compiler::ir::liveness::LivenessInfo;
use crate::compiler::ir::types::IrVar;

#[derive(Debug, Clone)]
pub struct RegisterAllocation {
    pub mapping: HashMap<IrVar, u8>,
    pub max_registers: u8,
}

impl RegisterAllocation {
    pub fn get(&self, var: IrVar) -> Option<u8> {
        self.mapping.get(&var).copied()
    }
}

pub fn allocate_registers(
    cfg: &ControlFlowGraph,
    num_params: u8,
) -> Result<RegisterAllocation, String> {
    let liveness = LivenessInfo::compute(cfg);
    let intervals_map = liveness.compute_live_intervals(cfg);

    let mut intervals: Vec<_> = intervals_map.into_values().collect();

    // A closure holds the register of each variable it captures, so that
    // register must not be handed to anything else once the variable's last
    // read goes by: the closure would see whatever overwrote it.
    let captured = crate::compiler::ir::opt::captured_vars(cfg);
    for interval in &mut intervals {
        if captured.contains(&interval.var) {
            interval.end = usize::MAX;
        }
    }

    // Sort by start position
    intervals.sort_by_key(|intv| (intv.start, intv.end));

    let mut mapping: HashMap<IrVar, u8> = HashMap::new();

    // 1. Pre-allocate parameter registers: 0..num_params
    for i in 0..num_params {
        mapping.insert(IrVar(i as u32), i);
    }

    // 2. Linear scan register allocation
    let mut free_registers: BTreeSet<u8> = (num_params..255).collect();
    // Active intervals: (end_point, register, IrVar)
    let mut active: Vec<(usize, u8, IrVar)> = Vec::new();
    let mut max_reg = num_params;

    for interval in intervals {
        // Skip already mapped parameters
        if mapping.contains_key(&interval.var) {
            continue;
        }

        // Expire old intervals
        active.retain(|&(end, reg, _var)| {
            if end < interval.start {
                free_registers.insert(reg);
                false
            } else {
                true
            }
        });

        // Allocate the lowest available register
        if let Some(&reg) = free_registers.iter().next() {
            free_registers.remove(&reg);
            mapping.insert(interval.var, reg);
            active.push((interval.end, reg, interval.var));
            if reg.saturating_add(1) > max_reg {
                max_reg = reg.saturating_add(1);
            }
        } else {
            return Err(format!(
                "register allocation failure: register limit (255) exceeded for variable {:?}",
                interval.var
            ));
        }
    }

    Ok(RegisterAllocation {
        mapping,
        max_registers: max_reg,
    })
}