use crate::entity::{Iter, IterMut, Keys, PrimaryMap};
use crate::ir::{StackSlot, Type};
use crate::packed_option::PackedOption;
use alloc::vec::Vec;
use core::cmp;
use core::fmt;
use core::ops::{Index, IndexMut};
use core::slice;
use core::str::FromStr;
#[cfg(feature = "enable-serde")]
use serde::{Deserialize, Serialize};
pub type StackSize = u32;
pub type StackOffset = i32;
const MIN_SPILL_SLOT_SIZE: StackSize = 4;
fn spill_size(ty: Type) -> StackSize {
cmp::max(MIN_SPILL_SLOT_SIZE, ty.bytes())
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
pub enum StackSlotKind {
SpillSlot,
ExplicitSlot,
IncomingArg,
OutgoingArg,
StructReturnSlot,
EmergencySlot,
}
impl FromStr for StackSlotKind {
type Err = ();
fn from_str(s: &str) -> Result<Self, ()> {
use self::StackSlotKind::*;
match s {
"explicit_slot" => Ok(ExplicitSlot),
"spill_slot" => Ok(SpillSlot),
"incoming_arg" => Ok(IncomingArg),
"outgoing_arg" => Ok(OutgoingArg),
"sret_slot" => Ok(StructReturnSlot),
"emergency_slot" => Ok(EmergencySlot),
_ => Err(()),
}
}
}
impl fmt::Display for StackSlotKind {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
use self::StackSlotKind::*;
f.write_str(match *self {
ExplicitSlot => "explicit_slot",
SpillSlot => "spill_slot",
IncomingArg => "incoming_arg",
OutgoingArg => "outgoing_arg",
StructReturnSlot => "sret_slot",
EmergencySlot => "emergency_slot",
})
}
}
#[derive(Clone, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
pub struct StackSlotData {
pub kind: StackSlotKind,
pub size: StackSize,
pub offset: Option<StackOffset>,
}
impl StackSlotData {
pub fn new(kind: StackSlotKind, size: StackSize) -> Self {
Self {
kind,
size,
offset: None,
}
}
pub fn alignment(&self, max_align: StackSize) -> StackSize {
debug_assert!(max_align.is_power_of_two());
let x = self.size | max_align;
x & x.wrapping_neg()
}
}
impl fmt::Display for StackSlotData {
fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
write!(f, "{} {}", self.kind, self.size)?;
if let Some(offset) = self.offset {
write!(f, ", offset {}", offset)?;
}
Ok(())
}
}
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
pub struct StackLayoutInfo {
pub frame_size: StackSize,
pub inbound_args_size: StackSize,
}
#[derive(Clone, Debug, PartialEq, Eq, Default)]
#[cfg_attr(feature = "enable-serde", derive(Serialize, Deserialize))]
pub struct StackSlots {
slots: PrimaryMap<StackSlot, StackSlotData>,
outgoing: Vec<StackSlot>,
emergency: Vec<StackSlot>,
pub layout_info: Option<StackLayoutInfo>,
}
impl StackSlots {
pub fn new() -> Self {
StackSlots::default()
}
pub fn clear(&mut self) {
self.slots.clear();
self.outgoing.clear();
self.emergency.clear();
self.layout_info = None;
}
pub fn push(&mut self, data: StackSlotData) -> StackSlot {
self.slots.push(data)
}
pub fn is_valid(&self, ss: StackSlot) -> bool {
self.slots.is_valid(ss)
}
pub fn iter(&self) -> Iter<StackSlot, StackSlotData> {
self.slots.iter()
}
pub fn iter_mut(&mut self) -> IterMut<StackSlot, StackSlotData> {
self.slots.iter_mut()
}
pub fn values(&self) -> slice::Iter<StackSlotData> {
self.slots.values()
}
pub fn values_mut(&mut self) -> slice::IterMut<StackSlotData> {
self.slots.values_mut()
}
pub fn keys(&self) -> Keys<StackSlot> {
self.slots.keys()
}
pub fn next_key(&self) -> StackSlot {
self.slots.next_key()
}
}
impl Index<StackSlot> for StackSlots {
type Output = StackSlotData;
fn index(&self, ss: StackSlot) -> &StackSlotData {
&self.slots[ss]
}
}
impl IndexMut<StackSlot> for StackSlots {
fn index_mut(&mut self, ss: StackSlot) -> &mut StackSlotData {
&mut self.slots[ss]
}
}
impl StackSlots {
pub fn make_spill_slot(&mut self, ty: Type) -> StackSlot {
self.push(StackSlotData::new(StackSlotKind::SpillSlot, spill_size(ty)))
}
pub fn make_incoming_arg(&mut self, size: u32, offset: StackOffset) -> StackSlot {
let mut data = StackSlotData::new(StackSlotKind::IncomingArg, size);
debug_assert!(offset <= StackOffset::max_value() - data.size as StackOffset);
data.offset = Some(offset);
self.push(data)
}
pub fn get_outgoing_arg(&mut self, size: u32, offset: StackOffset) -> StackSlot {
let inspos = match self.outgoing.binary_search_by_key(&(offset, size), |&ss| {
(self[ss].offset.unwrap(), self[ss].size)
}) {
Ok(idx) => return self.outgoing[idx],
Err(idx) => idx,
};
let mut data = StackSlotData::new(StackSlotKind::OutgoingArg, size);
debug_assert!(offset <= StackOffset::max_value() - size as StackOffset);
data.offset = Some(offset);
let ss = self.slots.push(data);
self.outgoing.insert(inspos, ss);
ss
}
pub fn get_emergency_slot(
&mut self,
ty: Type,
in_use: &[PackedOption<StackSlot>],
) -> StackSlot {
let size = spill_size(ty);
if let Some(&ss) = self
.emergency
.iter()
.filter(|&&ss| self[ss].size >= size && !in_use.contains(&ss.into()))
.min_by_key(|&&ss| self[ss].size)
{
return ss;
}
if let Some(&ss) = self
.emergency
.iter()
.filter(|&&ss| !in_use.contains(&ss.into()))
.max_by_key(|&&ss| self[ss].size)
{
self.slots[ss].size = size;
return ss;
}
let data = StackSlotData::new(StackSlotKind::EmergencySlot, size);
let ss = self.slots.push(data);
self.emergency.push(ss);
ss
}
}
#[cfg(test)]
mod tests {
use super::*;
use crate::ir::types;
use crate::ir::Function;
use alloc::string::ToString;
#[test]
fn stack_slot() {
let mut func = Function::new();
let ss0 = func.create_stack_slot(StackSlotData::new(StackSlotKind::IncomingArg, 4));
let ss1 = func.create_stack_slot(StackSlotData::new(StackSlotKind::SpillSlot, 8));
assert_eq!(ss0.to_string(), "ss0");
assert_eq!(ss1.to_string(), "ss1");
assert_eq!(func.stack_slots[ss0].size, 4);
assert_eq!(func.stack_slots[ss1].size, 8);
assert_eq!(func.stack_slots[ss0].to_string(), "incoming_arg 4");
assert_eq!(func.stack_slots[ss1].to_string(), "spill_slot 8");
}
#[test]
fn outgoing() {
let mut sss = StackSlots::new();
let ss0 = sss.get_outgoing_arg(4, 8);
let ss1 = sss.get_outgoing_arg(4, 4);
let ss2 = sss.get_outgoing_arg(8, 8);
assert_eq!(sss[ss0].offset, Some(8));
assert_eq!(sss[ss0].size, 4);
assert_eq!(sss[ss1].offset, Some(4));
assert_eq!(sss[ss1].size, 4);
assert_eq!(sss[ss2].offset, Some(8));
assert_eq!(sss[ss2].size, 8);
assert_eq!(sss.get_outgoing_arg(4, 8), ss0);
assert_eq!(sss.get_outgoing_arg(4, 4), ss1);
assert_eq!(sss.get_outgoing_arg(8, 8), ss2);
}
#[test]
fn alignment() {
let slot = StackSlotData::new(StackSlotKind::SpillSlot, 8);
assert_eq!(slot.alignment(4), 4);
assert_eq!(slot.alignment(8), 8);
assert_eq!(slot.alignment(16), 8);
let slot2 = StackSlotData::new(StackSlotKind::ExplicitSlot, 24);
assert_eq!(slot2.alignment(4), 4);
assert_eq!(slot2.alignment(8), 8);
assert_eq!(slot2.alignment(16), 8);
assert_eq!(slot2.alignment(32), 8);
}
#[test]
fn emergency() {
let mut sss = StackSlots::new();
let ss0 = sss.get_emergency_slot(types::I32, &[]);
assert_eq!(sss[ss0].size, 4);
assert_eq!(sss.get_emergency_slot(types::I8, &[]), ss0);
assert_eq!(sss[ss0].size, 4);
assert_eq!(sss.get_emergency_slot(types::F32, &[]), ss0);
assert_eq!(sss[ss0].size, 4);
assert_eq!(sss.get_emergency_slot(types::F64, &[]), ss0);
assert_eq!(sss[ss0].size, 8);
let ss1 = sss.get_emergency_slot(types::I32, &[None.into(), ss0.into()]);
assert_eq!(sss[ss0].size, 8);
assert_eq!(sss[ss1].size, 4);
assert_eq!(sss.get_emergency_slot(types::F32, &[]), ss1);
assert_eq!(sss.get_emergency_slot(types::F64, &[]), ss0);
}
}