swamp_code_gen/
reg_pool.rs

1/*
2 * Copyright (c) Peter Bjorklund. All rights reserved. https://github.com/swamp/swamp
3 * Licensed under the MIT License. See LICENSE in the project root for license information.
4 */
5use std::fmt::{Debug, Display, Formatter};
6use swamp_vm_types::types::{TypedRegister, VmType};
7
8#[derive(Debug)]
9pub struct TempRegister {
10    pub(crate) register: TypedRegister,
11}
12
13impl Display for TempRegister {
14    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
15        write!(f, "{}", self.register)
16    }
17}
18
19impl TempRegister {
20    pub(crate) const fn register(&self) -> &TypedRegister {
21        &self.register
22    }
23    #[must_use]
24    pub const fn addressing(&self) -> u8 {
25        self.register.addressing()
26    }
27}
28
29#[derive(Debug)]
30pub struct RegisterInfo {
31    pub index: u8,
32}
33
34pub struct HwmTempRegisterPool {
35    start_index: u8,
36    capacity: u8,
37    num_allocated: u8,
38}
39
40impl Debug for HwmTempRegisterPool {
41    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
42        if self.num_allocated == 0 {
43            write!(f, "hwm(no protected)")
44        } else {
45            write!(
46                f,
47                "hwm({}-{})",
48                self.start_index,
49                self.start_index + self.num_allocated - 1
50            )
51        }
52    }
53}
54
55impl HwmTempRegisterPool {
56    /// # Panics
57    ///
58    #[must_use]
59    pub fn new(start: u8, count: usize) -> Self {
60        if count == 0 {
61            return Self {
62                start_index: start,
63                capacity: 0,
64                num_allocated: 0,
65            };
66        }
67        assert!(
68            count <= u8::MAX as usize + 1,
69            "Register count too large for u8 capacity"
70        );
71        assert!(
72            start.checked_add((count - 1) as u8).is_some(),
73            "Register index range would overflow u8"
74        );
75
76        Self {
77            start_index: start,
78            capacity: count as u8,
79            num_allocated: 0, // Initially, no registers are allocated
80        }
81    }
82
83    pub(crate) const fn start_index_and_number_of_allocated(&self) -> (u8, u8) {
84        (self.start_index, self.num_allocated)
85    }
86
87    /// # Panics
88    /// if out of registers
89    pub fn allocate(&mut self, ty: VmType, comment: &str) -> TempRegister {
90        assert!(
91            (self.num_allocated < self.capacity),
92            "HwmTempRegisterPool: Out of temporary registers. Requested for: '{comment}'",
93        );
94
95        let register_index = self.start_index + self.num_allocated;
96        self.num_allocated += 1;
97
98        TempRegister {
99            register: TypedRegister {
100                index: register_index,
101                ty,
102                comment: comment.to_string(),
103            },
104        }
105    }
106
107    #[must_use]
108    pub const fn save_mark(&self) -> u8 {
109        self.num_allocated
110    }
111
112    /// # Panics
113    ///
114    pub fn restore_to_mark(&mut self, mark: u8) {
115        // TODO: add a check that we aren't using marks in the pinned range
116        assert!(
117            (mark <= self.num_allocated),
118            "HwmTempRegisterPool: Invalid mark {} provided. Current allocation count is {}.",
119            mark,
120            self.num_allocated
121        );
122
123        self.num_allocated = mark;
124    }
125}
126
127#[derive(Debug)]
128pub struct RegisterPool {
129    pub start_index: u8,
130    pub end_index: u8,
131    pub current_index: u8,
132}
133
134impl RegisterPool {
135    #[must_use]
136    pub const fn new(start: u8, count: u8) -> Self {
137        Self {
138            start_index: start,
139            end_index: start + count,
140            current_index: start,
141        }
142    }
143
144    /// # Panics
145    ///
146    pub fn alloc_register(&mut self, ty: VmType, comment: &str) -> TypedRegister {
147        assert!(
148            self.current_index + 1 < self.end_index,
149            "out of registers {} {}",
150            self.current_index,
151            self.end_index
152        );
153        let allocated_register = self.current_index;
154        self.current_index += 1;
155        TypedRegister {
156            index: allocated_register,
157            ty,
158            comment: comment.to_string(),
159        }
160    }
161}