Support for custom attributes

4 files changed, 153 insertions, 7 deletions

diff --git a/src/ast/ast.h b/src/ast/ast.h
index 967f27f..b272cae 100644
--- a/src/ast/ast.h
+++ b/src/ast/ast.h
@@ -157,6 +157,14 @@ typedef enum
 } NodeType;
 
 // ** AST Node Structure **
+typedef struct Attribute
+{
+    char *name;
+    char **args;
+    int arg_count;
+    struct Attribute *next;
+} Attribute;
+
 struct ASTNode
 {
     NodeType type;
@@ -212,6 +220,8 @@ struct ASTNode
             int cuda_global; // @global -> __global__
             int cuda_device; // @device -> __device__
             int cuda_host;   // @host -> __host__
+
+            Attribute *attributes; // Custom attributes
         } func;
 
         struct
@@ -419,11 +429,12 @@ struct ASTNode
             int generic_param_count; // Number of generic parameters
             char *parent;
             int is_union;
-            int is_packed;       // @packed attribute.
-            int align;           // @align(N) attribute, 0 = default.
-            int is_incomplete;   // Forward declaration (prototype)
-            int is_export;       // @export attribute
-            char **used_structs; // Names of structs used/mixed-in
+            int is_packed;         // @packed attribute.
+            int align;             // @align(N) attribute, 0 = default.
+            int is_incomplete;     // Forward declaration (prototype)
+            int is_export;         // @export attribute
+            Attribute *attributes; // Custom attributes
+            char **used_structs;   // Names of structs used/mixed-in
             int used_struct_count;
         } strct;
 
diff --git a/src/codegen/codegen_decl.c b/src/codegen/codegen_decl.c
index eb53911..5fb9f54 100644
--- a/src/codegen/codegen_decl.c
+++ b/src/codegen/codegen_decl.c
@@ -388,6 +388,38 @@ void emit_struct_defs(ParserContext *ctx, ASTNode *node, FILE *out)
             {
                 fprintf(out, " __attribute__((visibility(\"default\")))");
             }
+
+            if (node->strct.attributes)
+            {
+                fprintf(out, " __attribute__((");
+                Attribute *custom = node->strct.attributes;
+                int first = 1;
+                while (custom)
+                {
+                    if (!first)
+                    {
+                        fprintf(out, ", ");
+                    }
+                    fprintf(out, "%s", custom->name);
+                    if (custom->arg_count > 0)
+                    {
+                        fprintf(out, "(");
+                        for (int i = 0; i < custom->arg_count; i++)
+                        {
+                            if (i > 0)
+                            {
+                                fprintf(out, ", ");
+                            }
+                            fprintf(out, "%s", custom->args[i]);
+                        }
+                        fprintf(out, ")");
+                    }
+                    first = 0;
+                    custom = custom->next;
+                }
+                fprintf(out, "))");
+            }
+
             fprintf(out, ";\n\n");
         }
         else if (node->type == NODE_ENUM)
diff --git a/src/codegen/codegen_stmt.c b/src/codegen/codegen_stmt.c
index c823b13..bd0c816 100644
--- a/src/codegen/codegen_stmt.c
+++ b/src/codegen/codegen_stmt.c
@@ -684,9 +684,66 @@ void codegen_node_single(ParserContext *ctx, ASTNode *node, FILE *out)
                     }
                     fprintf(out, "section(\"%s\")", node->func.section);
                 }
+
+                Attribute *custom = node->func.attributes;
+                while (custom)
+                {
+                    if (!first)
+                    {
+                        fprintf(out, ", ");
+                    }
+                    fprintf(out, "%s", custom->name);
+                    if (custom->arg_count > 0)
+                    {
+                        fprintf(out, "(");
+                        for (int i = 0; i < custom->arg_count; i++)
+                        {
+                            if (i > 0)
+                            {
+                                fprintf(out, ", ");
+                            }
+                            fprintf(out, "%s", custom->args[i]);
+                        }
+                        fprintf(out, ")");
+                    }
+                    first = 0;
+                    custom = custom->next;
+                }
+
 #undef EMIT_ATTR
                 fprintf(out, ")) ");
             }
+            else if (node->func.attributes)
+            {
+                // Handle case where specific attributes are missing but custom ones exist
+                fprintf(out, "__attribute__((");
+                int first = 1;
+                Attribute *custom = node->func.attributes;
+                while (custom)
+                {
+                    if (!first)
+                    {
+                        fprintf(out, ", ");
+                    }
+                    fprintf(out, "%s", custom->name);
+                    if (custom->arg_count > 0)
+                    {
+                        fprintf(out, "(");
+                        for (int i = 0; i < custom->arg_count; i++)
+                        {
+                            if (i > 0)
+                            {
+                                fprintf(out, ", ");
+                            }
+                            fprintf(out, "%s", custom->args[i]);
+                        }
+                        fprintf(out, ")");
+                    }
+                    first = 0;
+                    custom = custom->next;
+                }
+                fprintf(out, ")) ");
+            }
         }
 
         if (node->func.is_inline)
diff --git a/src/parser/parser_core.c b/src/parser/parser_core.c
index ac578a1..e9a418e 100644
--- a/src/parser/parser_core.c
+++ b/src/parser/parser_core.c
@@ -72,11 +72,13 @@ ASTNode *parse_program_nodes(ParserContext *ctx, Lexer *l)
         char *derived_traits[32];
         int derived_count = 0;
 
+        Attribute *current_custom_attributes = NULL;
+
         while (t.type == TOK_AT)
         {
             lexer_next(l);
             Token attr = lexer_next(l);
-            if (attr.type != TOK_IDENT && attr.type != TOK_COMPTIME)
+            if (attr.type != TOK_IDENT && attr.type != TOK_COMPTIME && attr.type != TOK_ALIAS)
             {
                 zpanic_at(attr, "Expected attribute name after @");
             }
@@ -250,7 +252,49 @@ ASTNode *parse_program_nodes(ParserContext *ctx, Lexer *l)
                 }
                 else
                 {
-                    zwarn_at(attr, "Unknown attribute: %.*s", attr.len, attr.start);
+                    Attribute *new_attr = xmalloc(sizeof(Attribute));
+                    new_attr->name = token_strdup(attr);
+                    new_attr->args = NULL;
+                    new_attr->arg_count = 0;
+                    new_attr->next = current_custom_attributes; // Prepend
+                    current_custom_attributes = new_attr;
+
+                    if (lexer_peek(l).type == TOK_LPAREN)
+                    {
+                        lexer_next(l); // eat (
+                        while (1)
+                        {
+                            Token t = lexer_next(l);
+                            new_attr->args =
+                                realloc(new_attr->args, sizeof(char *) * (new_attr->arg_count + 1));
+
+                            if (t.type == TOK_STRING)
+                            {
+                                new_attr->args[new_attr->arg_count++] = token_strdup(t);
+                            }
+                            else
+                            {
+                                new_attr->args[new_attr->arg_count++] = token_strdup(t);
+                            }
+
+                            if (lexer_peek(l).type == TOK_COMMA)
+                            {
+                                lexer_next(l);
+                            }
+                            else if (lexer_peek(l).type == TOK_RPAREN)
+                            {
+                                break;
+                            }
+                            else
+                            {
+                                zpanic_at(lexer_peek(l), "Expected , or ) in attribute args");
+                            }
+                        }
+                        if (lexer_next(l).type != TOK_RPAREN)
+                        {
+                            zpanic_at(lexer_peek(l), "Expected )");
+                        }
+                    }
                 }
             }
 
@@ -500,6 +544,7 @@ ASTNode *parse_program_nodes(ParserContext *ctx, Lexer *l)
             s->func.cuda_global = attr_cuda_global;
             s->func.cuda_device = attr_cuda_device;
             s->func.cuda_host = attr_cuda_host;
+            s->func.attributes = current_custom_attributes;
 
             if (attr_deprecated && s->func.name)
             {
@@ -519,6 +564,7 @@ ASTNode *parse_program_nodes(ParserContext *ctx, Lexer *l)
         if (s && s->type == NODE_STRUCT)
         {
             s->strct.is_export = attr_export;
+            s->strct.attributes = current_custom_attributes;
             s->strct.is_packed = attr_packed || s->strct.is_packed;
             if (attr_align)
             {