chomsky_uir/
regalloc.rs

1use crate::IKunTree;
2use std::collections::HashMap;
3
4#[derive(Debug, Clone, PartialEq, Eq, Hash)]
5pub enum Register {
6    Physical(String),
7    Stack(usize),
8}
9
10pub struct RegisterAllocation {
11    pub mapping: HashMap<String, Register>,
12}
13
14pub struct LinearScanAllocator {
15    registers: Vec<String>,
16}
17
18impl LinearScanAllocator {
19    pub fn new(registers: Vec<String>) -> Self {
20        Self { registers }
21    }
22
23    /// 执行简单的寄存器分配
24    pub fn allocate(&self, tree: &IKunTree) -> RegisterAllocation {
25        let mut mapping = HashMap::new();
26        let mut available = self.registers.clone();
27        let mut stack_offset = 0;
28
29        // 简单的线性扫描演示实现
30        self.collect_and_alloc(tree, &mut mapping, &mut available, &mut stack_offset);
31
32        RegisterAllocation { mapping }
33    }
34
35    fn collect_and_alloc(
36        &self,
37        tree: &IKunTree,
38        mapping: &mut HashMap<String, Register>,
39        available: &mut Vec<String>,
40        stack_offset: &mut usize,
41    ) {
42        match tree {
43            IKunTree::Symbol(s) => {
44                if !mapping.contains_key(s) {
45                    if let Some(reg) = available.pop() {
46                        mapping.insert(s.clone(), Register::Physical(reg));
47                    } else {
48                        mapping.insert(s.clone(), Register::Stack(*stack_offset));
49                        *stack_offset += 8;
50                    }
51                }
52            }
53            IKunTree::Apply(f, args) => {
54                self.collect_and_alloc(f, mapping, available, stack_offset);
55                for arg in args {
56                    self.collect_and_alloc(arg, mapping, available, stack_offset);
57                }
58            }
59            IKunTree::Map(f, x) | IKunTree::Filter(f, x) => {
60                self.collect_and_alloc(f, mapping, available, stack_offset);
61                self.collect_and_alloc(x, mapping, available, stack_offset);
62            }
63            IKunTree::Reduce(f, init, x) => {
64                self.collect_and_alloc(f, mapping, available, stack_offset);
65                self.collect_and_alloc(init, mapping, available, stack_offset);
66                self.collect_and_alloc(x, mapping, available, stack_offset);
67            }
68            IKunTree::Seq(ids) => {
69                for id in ids {
70                    self.collect_and_alloc(id, mapping, available, stack_offset);
71                }
72            }
73            IKunTree::StateUpdate(target, val) => {
74                self.collect_and_alloc(target, mapping, available, stack_offset);
75                self.collect_and_alloc(val, mapping, available, stack_offset);
76            }
77            IKunTree::Lambda(_, body) => {
78                self.collect_and_alloc(body, mapping, available, stack_offset);
79            }
80            IKunTree::Extension(_, args) => {
81                for arg in args {
82                    self.collect_and_alloc(arg, mapping, available, stack_offset);
83                }
84            }
85            IKunTree::Choice(cond, then_b, else_b) => {
86                self.collect_and_alloc(cond, mapping, available, stack_offset);
87                self.collect_and_alloc(then_b, mapping, available, stack_offset);
88                self.collect_and_alloc(else_b, mapping, available, stack_offset);
89            }
90            _ => {}
91        }
92    }
93}
chomsky_uir/regalloc.rs

chomsky_uir/
regalloc.rs
