//! Computing stack layout.

use crate::ir::stackslot::{StackOffset, StackSize, StackSlotKind};
use crate::ir::{StackLayoutInfo, StackSlots};
use crate::result::{CodegenError, CodegenResult};
use core::cmp::{max, min};

/// Compute the stack frame layout.
///
/// Determine the total size of this stack frame and assign offsets to all `Spill` and `Explicit`
/// stack slots.
///
/// The total frame size will be a multiple of `alignment` which must be a power of two, unless the
/// function doesn't perform any call.
///
/// Returns the total stack frame size which is also saved in `frame.frame_size`.
///
/// If the stack frame is too big, returns an `ImplLimitExceeded` error.
pub fn layout_stack(
    frame: &mut StackSlots,
    is_leaf: bool,
    alignment: StackSize,
) -> CodegenResult<StackSize> {
    // Each object and the whole stack frame must fit in 2 GB such that any relative offset within
    // the frame fits in a `StackOffset`.
    let max_size = StackOffset::max_value() as StackSize;
    debug_assert!(alignment.is_power_of_two() && alignment <= max_size);

    // We assume a stack that grows toward lower addresses as implemented by modern ISAs. The
    // stack layout from high to low addresses will be:
    //
    // 1. incoming arguments.
    // 2. spills + explicits + struct returns.
    // 3. outgoing arguments.
    //
    // The incoming arguments can have both positive and negative offsets. A negative offset
    // incoming arguments is usually the x86 return address pushed by the call instruction, but
    // it can also be fixed stack slots pushed by an externally generated prologue.
    //
    // Both incoming and outgoing argument slots have fixed offsets that are treated as
    // reserved zones by the layout algorithm.
    //
    // If a function only has incoming arguments and does not perform any calls, then it doesn't
    // require the stack to be aligned.

    let mut incoming_min = 0;
    let mut incoming_max = 0;
    let mut outgoing_max = 0;
    let mut min_align = alignment;
    let mut must_align = !is_leaf;

    for slot in frame.values() {
        if slot.size > max_size {
            return Err(CodegenError::ImplLimitExceeded);
        }

        match slot.kind {
            StackSlotKind::IncomingArg => {
                incoming_min = min(incoming_min, slot.offset.unwrap());
                incoming_max = max(incoming_max, slot.offset.unwrap() + slot.size as i32);
            }
            StackSlotKind::OutgoingArg => {
                let offset = slot
                    .offset
                    .unwrap()
                    .checked_add(slot.size as StackOffset)
                    .ok_or(CodegenError::ImplLimitExceeded)?;
                outgoing_max = max(outgoing_max, offset);
                must_align = true;
            }
            StackSlotKind::StructReturnSlot
            | StackSlotKind::SpillSlot
            | StackSlotKind::ExplicitSlot
            | StackSlotKind::EmergencySlot => {
                // Determine the smallest alignment of any explicit or spill slot.
                min_align = slot.alignment(min_align);
                must_align = true;
            }
        }
    }

    // Lay out spill slots, struct return slots, and explicit slots below the
    // incoming arguments. The offset is negative, growing downwards. Start with
    // the smallest alignments for better packing.
    let mut offset = incoming_min;
    debug_assert!(min_align.is_power_of_two());
    while min_align <= alignment {
        for slot in frame.values_mut() {
            // Pick out explicit and spill slots with exact alignment `min_align`.
            match slot.kind {
                StackSlotKind::SpillSlot
                | StackSlotKind::StructReturnSlot
                | StackSlotKind::ExplicitSlot
                | StackSlotKind::EmergencySlot => {
                    if slot.alignment(alignment) != min_align {
                        continue;
                    }
                }
                StackSlotKind::IncomingArg | StackSlotKind::OutgoingArg => continue,
            }

            offset = offset
                .checked_sub(slot.size as StackOffset)
                .ok_or(CodegenError::ImplLimitExceeded)?;

            // Aligning the negative offset can never cause overflow. We're only clearing bits.
            offset &= -(min_align as StackOffset);
            slot.offset = Some(offset);
        }

        // Move on to the next higher alignment.
        min_align *= 2;
    }

    // Finally, make room for the outgoing arguments.
    offset = offset
        .checked_sub(outgoing_max)
        .ok_or(CodegenError::ImplLimitExceeded)?;

    if must_align {
        offset &= -(alignment as StackOffset);
    }

    // Set the computed layout information for the frame
    let frame_size = (offset as StackSize).wrapping_neg();
    let inbound_args_size = incoming_max as u32;
    frame.layout_info = Some(StackLayoutInfo {
        frame_size,
        inbound_args_size,
    });

    Ok(frame_size)
}

#[cfg(test)]
mod tests {
    use super::layout_stack;
    use crate::ir::stackslot::StackOffset;
    use crate::ir::types;
    use crate::ir::{StackSlotData, StackSlotKind, StackSlots};
    use crate::result::CodegenError;

    #[test]
    fn layout() {
        let sss = &mut StackSlots::new();

        // For all these test cases, assume it will call.
        let is_leaf = false;

        // An empty layout should have 0-sized stack frame.
        assert_eq!(layout_stack(sss, is_leaf, 1), Ok(0));
        assert_eq!(layout_stack(sss, is_leaf, 16), Ok(0));

        // Same for incoming arguments with non-negative offsets.
        let in0 = sss.make_incoming_arg(8, 0);
        let in1 = sss.make_incoming_arg(8, 8);

        assert_eq!(layout_stack(sss, is_leaf, 1), Ok(0));
        assert_eq!(layout_stack(sss, is_leaf, 16), Ok(0));
        assert_eq!(sss[in0].offset, Some(0));
        assert_eq!(sss[in1].offset, Some(8));

        // Add some spill slots.
        let ss0 = sss.make_spill_slot(types::I64);
        let ss1 = sss.make_spill_slot(types::I32);

        assert_eq!(layout_stack(sss, is_leaf, 1), Ok(12));
        assert_eq!(sss[in0].offset, Some(0));
        assert_eq!(sss[in1].offset, Some(8));
        assert_eq!(sss[ss0].offset, Some(-8));
        assert_eq!(sss[ss1].offset, Some(-12));

        assert_eq!(layout_stack(sss, is_leaf, 16), Ok(16));
        assert_eq!(sss[in0].offset, Some(0));
        assert_eq!(sss[in1].offset, Some(8));
        assert_eq!(sss[ss0].offset, Some(-16));
        assert_eq!(sss[ss1].offset, Some(-4));

        // An incoming argument with negative offset counts towards the total frame size, but it
        // should still pack nicely with the spill slots.
        let in2 = sss.make_incoming_arg(4, -4);

        assert_eq!(layout_stack(sss, is_leaf, 1), Ok(16));
        assert_eq!(sss[in0].offset, Some(0));
        assert_eq!(sss[in1].offset, Some(8));
        assert_eq!(sss[in2].offset, Some(-4));
        assert_eq!(sss[ss0].offset, Some(-12));
        assert_eq!(sss[ss1].offset, Some(-16));

        assert_eq!(layout_stack(sss, is_leaf, 16), Ok(16));
        assert_eq!(sss[in0].offset, Some(0));
        assert_eq!(sss[in1].offset, Some(8));
        assert_eq!(sss[in2].offset, Some(-4));
        assert_eq!(sss[ss0].offset, Some(-16));
        assert_eq!(sss[ss1].offset, Some(-8));

        // Finally, make sure there is room for the outgoing args.
        let out0 = sss.get_outgoing_arg(4, 0);

        assert_eq!(layout_stack(sss, is_leaf, 1), Ok(20));
        assert_eq!(sss[in0].offset, Some(0));
        assert_eq!(sss[in1].offset, Some(8));
        assert_eq!(sss[in2].offset, Some(-4));
        assert_eq!(sss[ss0].offset, Some(-12));
        assert_eq!(sss[ss1].offset, Some(-16));
        assert_eq!(sss[out0].offset, Some(0));

        assert_eq!(layout_stack(sss, is_leaf, 16), Ok(32));
        assert_eq!(sss[in0].offset, Some(0));
        assert_eq!(sss[in1].offset, Some(8));
        assert_eq!(sss[in2].offset, Some(-4));
        assert_eq!(sss[ss0].offset, Some(-16));
        assert_eq!(sss[ss1].offset, Some(-8));
        assert_eq!(sss[out0].offset, Some(0));

        // Also test that an unsupported offset is rejected.
        sss.get_outgoing_arg(1, StackOffset::max_value() - 1);
        assert_eq!(
            layout_stack(sss, is_leaf, 1),
            Err(CodegenError::ImplLimitExceeded)
        );
    }

    #[test]
    fn slot_kinds() {
        let sss = &mut StackSlots::new();

        // Add some slots of various kinds.
        let ss0 = sss.make_spill_slot(types::I32);
        let ss1 = sss.push(StackSlotData::new(
            StackSlotKind::ExplicitSlot,
            types::I32.bytes(),
        ));
        let ss2 = sss.get_emergency_slot(types::I32, &[]);

        assert_eq!(layout_stack(sss, true, 1), Ok(12));
        assert_eq!(sss[ss0].offset, Some(-4));
        assert_eq!(sss[ss1].offset, Some(-8));
        assert_eq!(sss[ss2].offset, Some(-12));
    }
}