From a3cde8d6c25772e85d665aeecf004e3a8714b9a2 Mon Sep 17 00:00:00 2001
From: Evan Haas <evan@lagerdata.com>
Date: Wed, 4 Dec 2024 11:30:45 -0800
Subject: [PATCH] Attribute: implement alloc_align

---
 src/aro/Attribute.zig              | 15 ++++++++++++++-
 src/aro/Diagnostics/messages.def   | 14 ++++++++++++++
 test/cases/alloc_align attribute.c | 12 ++++++++++++
 3 files changed, 40 insertions(+), 1 deletion(-)
 create mode 100644 test/cases/alloc_align attribute.c

diff --git a/src/aro/Attribute.zig b/src/aro/Attribute.zig
index 30f728ed..13270b93 100644
--- a/src/aro/Attribute.zig
+++ b/src/aro/Attribute.zig
@@ -945,8 +945,21 @@ pub fn applyFunctionAttributes(p: *Parser, ty: Type, attr_buf_start: usize) !Typ
                 try ignoredAttrErr(p, tok, attr.tag, "functions that do not return pointers");
             }
         },
+        .alloc_align => {
+            if (base_ty.returnType().isPtr()) {
+                const params = base_ty.params();
+                if (attr.args.alloc_align.position == 0 or attr.args.alloc_align.position > params.len) {
+                    try p.errExtra(.attribute_param_out_of_bounds, tok, .{ .attr_arg_count = .{ .attribute = .alloc_align, .expected = 1 } });
+                } else if (!params[attr.args.alloc_align.position - 1].ty.isInt()) {
+                    try p.errTok(.alloc_align_required_int_param, tok);
+                } else {
+                    try p.attr_application_buf.append(p.gpa, attr);
+                }
+            } else {
+                try p.errTok(.alloc_align_requires_ptr_return, tok);
+            }
+        },
         .access,
-        .alloc_align,
         .alloc_size,
         .artificial,
         .assume_aligned,
diff --git a/src/aro/Diagnostics/messages.def b/src/aro/Diagnostics/messages.def
index 1ae3dba1..9050b41f 100644
--- a/src/aro/Diagnostics/messages.def
+++ b/src/aro/Diagnostics/messages.def
@@ -2566,3 +2566,17 @@ packed_member_address
     .opt = W("address-of-packed-member")
     .kind = .warning
     .extra = .str
+
+alloc_align_requires_ptr_return
+    .msg = "'alloc_align' attribute only applies to return values that are pointers"
+    .opt = W("ignored-attributes")
+    .kind = .warning
+
+attribute_param_out_of_bounds
+    .msg = "'{s}' attribute parameter {d} is out of bounds"
+    .kind = .@"error"
+    .extra = .attr_arg_count
+
+alloc_align_required_int_param
+    .msg = "'alloc_align' attribute argument may only refer to a function parameter of integer type"
+    .kind = .@"error"
diff --git a/test/cases/alloc_align attribute.c b/test/cases/alloc_align attribute.c
new file mode 100644
index 00000000..98e76c77
--- /dev/null
+++ b/test/cases/alloc_align attribute.c	
@@ -0,0 +1,12 @@
+__attribute__((alloc_align(0))) int* foo1(int x, double y);
+__attribute__((alloc_align(1))) int* foo2(int x, double y);
+__attribute__((alloc_align(2))) int* foo3(int x, double y);
+__attribute__((alloc_align(3))) int* foo4(int x, double y);
+__attribute__((alloc_align(1))) int foo5(int x, double y);
+
+#define EXPECTED_ERRORS \
+    "alloc_align attribute.c:1:16: error: 'alloc_align' attribute parameter 1 is out of bounds" \
+    "alloc_align attribute.c:3:16: error: 'alloc_align' attribute argument may only refer to a function parameter of integer type" \
+    "alloc_align attribute.c:4:16: error: 'alloc_align' attribute parameter 1 is out of bounds" \
+    "alloc_align attribute.c:5:16: warning: 'alloc_align' attribute only applies to return values that are pointers [-Wignored-attributes]" \
+