holey-bytes/cranelift-backend/src/x86_64.rs

// The classification code for the x86_64 ABI is taken from the clay language
// https://github.com/jckarter/clay/blob/db0bd2702ab0b6e48965cd85f8859bbd5f60e48e/compiler/externals.cpp

use {crate::AbiMeta, hblang::ty};

pub fn build_systemv_signature(
    sig: hblang::ty::Sig,
    types: &hblang::ty::Types,
    signature: &mut cranelift_codegen::ir::Signature,
    arg_lens: &mut Vec<AbiMeta>,
) -> bool {
    let mut alloca = Alloca::new();

    alloca.next(false, sig.ret, types, &mut signature.returns);
    let stack_ret = signature.returns.len() == 1
        && signature.returns[0].purpose == cranelift_codegen::ir::ArgumentPurpose::StructReturn;

    if stack_ret {
        signature.params.append(&mut signature.returns);
        arg_lens.push(AbiMeta { arg_count: signature.params.len(), trough_mem: true });
    } else {
        arg_lens.push(AbiMeta { arg_count: signature.returns.len(), trough_mem: false });
    }

    let mut args = sig.args.args();
    while let Some(arg) = args.next_value(types) {
        let prev = signature.params.len();
        let trough_mem = alloca.next(true, arg, types, &mut signature.params);
        arg_lens.push(AbiMeta { arg_count: signature.params.len() - prev, trough_mem });
    }

    stack_ret
}

/// Classification of "eightbyte" components.
// N.B., the order of the variants is from general to specific,
// such that `unify(a, b)` is the "smaller" of `a` and `b`.
#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Debug)]
enum Class {
    Int,
    Sse,
    SseUp,
}

#[derive(Clone, Copy, Debug)]
struct Memory;

// Currently supported vector size (AVX-512).
const LARGEST_VECTOR_SIZE: usize = 512;
const MAX_EIGHTBYTES: usize = LARGEST_VECTOR_SIZE / 64;

fn classify_arg(
    cx: &hblang::ty::Types,
    arg: hblang::ty::Id,
) -> Result<([Option<Class>; MAX_EIGHTBYTES], [u32; MAX_EIGHTBYTES]), Memory> {
    fn classify(
        cx: &hblang::ty::Types,
        layout: hblang::ty::Id,
        cls: &mut [Option<Class>],
        sizes: &mut [u32],
        off: hblang::ty::Offset,
    ) -> Result<(), Memory> {
        let size = cx.size_of(layout);
        if off & (cx.align_of(layout) - 1) != 0 {
            if size != 0 {
                return Err(Memory);
            }
            return Ok(());
        }

        let mut c = match layout.expand() {
            _ if size == 0 => return Ok(()),
            _ if layout.is_integer() || layout.is_pointer() || layout == ty::Id::BOOL => Class::Int,
            _ if layout.is_float() => Class::Sse,

            hblang::ty::Kind::Struct(s) => {
                for (f, foff) in hblang::ty::OffsetIter::new(s, cx).into_iter(cx) {
                    classify(cx, f.ty, cls, sizes, off + foff)?;
                }
                return Ok(());
            }
            hblang::ty::Kind::Tuple(tuple) => {
                for (&ty, foff) in hblang::ty::OffsetIter::new(tuple, cx).into_iter(cx) {
                    classify(cx, ty, cls, sizes, off + foff)?;
                }
                return Ok(());
            }
            hblang::ty::Kind::Enum(_) => Class::Int,
            hblang::ty::Kind::Union(union) => {
                for f in cx.union_fields(union) {
                    classify(cx, f.ty, cls, sizes, off)?;
                }
                return Ok(());
            }
            hblang::ty::Kind::Slice(slice) if let Some(len) = cx.ins.slices[slice].len() => {
                for i in 0..len as u32 {
                    classify(
                        cx,
                        cx.ins.slices[slice].elem,
                        cls,
                        sizes,
                        off + i * cx.size_of(cx.ins.slices[slice].elem),
                    )?;
                }
                return Ok(());
            }
            hblang::ty::Kind::Slice(_) => {
                classify(cx, hblang::ty::Id::UINT, cls, sizes, off)?;
                classify(cx, hblang::ty::Id::UINT, cls, sizes, off + 8)?;
                return Ok(());
            }
            hblang::ty::Kind::Opt(opt) => {
                let base = cx.ins.opts[opt].base;
                if cx.nieche_of(base).is_some() {
                    classify(cx, base, cls, sizes, off)?;
                } else {
                    classify(cx, hblang::ty::Id::BOOL, cls, sizes, off)?;
                    classify(cx, base, cls, sizes, off + cx.align_of(base))?;
                }
                return Ok(());
            }

            ty => unimplemented!("{ty:?}"),
        };

        // Fill in `cls` for scalars (Int/Sse) and vectors (Sse).
        let first = (off / 8) as usize;
        let last = ((off + size - 1) / 8) as usize;
        for (cls, sz) in cls[first..=last].iter_mut().zip(&mut sizes[first..=last]) {
            *cls = Some(cls.map_or(c, |old| old.min(c)));
            *sz = size;

            // Everything after the first Sse "eightbyte"
            // component is the upper half of a register.
            if c == Class::Sse {
                c = Class::SseUp;
            }
        }

        Ok(())
    }

    let size = cx.size_of(arg);
    let n = ((size + 7) / 8) as usize;
    if n > MAX_EIGHTBYTES {
        return Err(Memory);
    }

    let mut cls = [None; MAX_EIGHTBYTES];
    let mut sizes = [0; MAX_EIGHTBYTES];
    classify(cx, arg, &mut cls, &mut sizes, 0)?;
    if n > 2 {
        if cls[0] != Some(Class::Sse) {
            return Err(Memory);
        }
        if cls[1..n].iter().any(|&c| c != Some(Class::SseUp)) {
            return Err(Memory);
        }
    } else {
        let mut i = 0;
        while i < n {
            if cls[i] == Some(Class::SseUp) {
                cls[i] = Some(Class::Sse);
            } else if cls[i] == Some(Class::Sse) {
                i += 1;
                while i != n && cls[i] == Some(Class::SseUp) {
                    i += 1;
                }
            } else {
                i += 1;
            }
        }
    }

    Ok((cls, sizes))
}

fn reg_component(
    cls: &[Option<Class>],
    sizes: &[u32],
    i: &mut usize,
    _size: hblang::ty::Size,
) -> Option<cranelift_codegen::ir::Type> {
    if *i >= cls.len() {
        return None;
    }

    let size = sizes[*i];
    match cls[*i] {
        None => None,
        Some(Class::Int) => {
            *i += 1;
            Some(if size < 8 {
                cranelift_codegen::ir::Type::int(size as u16 * 8).unwrap()
            } else {
                cranelift_codegen::ir::types::I64
            })
        }
        Some(Class::Sse) => {
            let vec_len =
                1 + cls[*i + 1..].iter().take_while(|&&c| c == Some(Class::SseUp)).count();
            *i += vec_len;
            Some(if vec_len == 1 {
                match size {
                    4 => cranelift_codegen::ir::types::F32,
                    _ => cranelift_codegen::ir::types::F64,
                }
            } else {
                cranelift_codegen::ir::types::I64.by(vec_len as _).unwrap()
            })
        }
        Some(c) => unreachable!("reg_component: unhandled class {:?}", c),
    }
}

fn cast_target(
    cls: &[Option<Class>],
    sizes: &[u32],
    size: hblang::ty::Size,
    dest: &mut Vec<cranelift_codegen::ir::AbiParam>,
) {
    let mut i = 0;
    let lo = reg_component(cls, sizes, &mut i, size).unwrap();
    let offset = 8 * (i as u32);
    dest.push(cranelift_codegen::ir::AbiParam::new(lo));
    if size > offset {
        if let Some(hi) = reg_component(cls, sizes, &mut i, size - offset) {
            dest.push(cranelift_codegen::ir::AbiParam::new(hi));
        }
    }
    assert_eq!(reg_component(cls, sizes, &mut i, 0), None);
}

const MAX_INT_REGS: usize = 6; // RDI, RSI, RDX, RCX, R8, R9
const MAX_SSE_REGS: usize = 8; // XMM0-7

pub struct Alloca {
    int_regs: usize,
    sse_regs: usize,
}

impl Alloca {
    pub fn new() -> Self {
        Self { int_regs: MAX_INT_REGS, sse_regs: MAX_SSE_REGS }
    }

    pub fn next(
        &mut self,
        is_arg: bool,
        arg: hblang::ty::Id,
        cx: &hblang::ty::Types,
        dest: &mut Vec<cranelift_codegen::ir::AbiParam>,
    ) -> bool {
        if cx.size_of(arg) == 0 {
            return false;
        }

        let mut cls_or_mem = classify_arg(cx, arg);

        if is_arg {
            if let Ok((cls, _)) = cls_or_mem {
                let mut needed_int = 0;
                let mut needed_sse = 0;
                for c in cls {
                    match c {
                        Some(Class::Int) => needed_int += 1,
                        Some(Class::Sse) => needed_sse += 1,
                        _ => {}
                    }
                }
                match (self.int_regs.checked_sub(needed_int), self.sse_regs.checked_sub(needed_sse))
                {
                    (Some(left_int), Some(left_sse)) => {
                        self.int_regs = left_int;
                        self.sse_regs = left_sse;
                    }
                    _ => {
                        // Not enough registers for this argument, so it will be
                        // passed on the stack, but we only mark aggregates
                        // explicitly as indirect `byval` arguments, as LLVM will
                        // automatically put immediates on the stack itself.
                        if arg.is_aggregate(cx) {
                            cls_or_mem = Err(Memory);
                        }
                    }
                }
            }
        }

        match cls_or_mem {
            Err(Memory) => {
                if is_arg {
                    dest.push(cranelift_codegen::ir::AbiParam::new(
                        cranelift_codegen::ir::types::I64,
                    ));
                } else {
                    dest.push(cranelift_codegen::ir::AbiParam::special(
                        cranelift_codegen::ir::types::I64,
                        cranelift_codegen::ir::ArgumentPurpose::StructReturn,
                    ));
                }
                true
            }
            Ok((ref cls, ref sizes)) => {
                // split into sized chunks passed individually
                if arg.is_aggregate(cx) {
                    cast_target(cls, sizes, cx.size_of(arg), dest);
                } else {
                    dest.push(cranelift_codegen::ir::AbiParam::new(
                        reg_component(cls, sizes, &mut 0, cx.size_of(arg)).unwrap(),
                    ));
                }
                false
            }
        }
    }
}