Skip to content
35 changes: 35 additions & 0 deletions src/irx/builders/llvmliteir.py
Original file line number Diff line number Diff line change
Expand Up @@ -343,6 +343,8 @@ class LLVMLiteIRVisitor(BuilderVisitor):
type: RuntimeFeatureState
const_vars:
type: set[str]
struct_types:
type: dict[str, ir.IdentifiedStructType]
_fast_math_enabled:
type: bool
target:
Expand Down Expand Up @@ -377,6 +379,9 @@ def __init__(
self.const_vars: set[str] = set()
self.function_protos: dict[str, astx.FunctionPrototype] = {}
self.result_stack: list[ir.Value | ir.Function] = []

self.struct_types: dict[str, ir.IdentifiedStructType] = {}

self._fast_math_enabled: bool = False

self.initialize()
Expand Down Expand Up @@ -3612,6 +3617,36 @@ def visit(self, node: astx.VariableDeclaration) -> None:
self.const_vars.add(node.name)
self.named_values[node.name] = alloca

@dispatch # type: ignore[no-redef]
def visit(self, node: astx.StructDefStmt) -> None:
"""
title: Struct Definition Codegen
summary: Translate ASTx StructDefStmt to LLVM-IR.
parameters:
node:
type: astx.StructDefStmt
"""

# Create identified struct type
struct_type = self._llvm.module.context.get_identified_type(node.name)

# Prevent duplicate struct definitions
if not struct_type.is_opaque:
raise ValueError(f"Struct '{node.name}' already defined.")

# Convert AST field types → LLVM types
field_types = []
for attr in node.attributes:
type_str = attr.type_.__class__.__name__.lower()
field_type = self._llvm.get_data_type(type_str)
field_types.append(field_type)

# Set struct body
struct_type.set_body(*field_types)

# Register struct
self.struct_types[node.name] = struct_type

@dispatch # type: ignore[no-redef]
def visit(self, node: astx.LiteralInt16) -> None:
"""
Expand Down
147 changes: 147 additions & 0 deletions tests/test_struct_definition.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,147 @@
"""
title: Test Struct Definition
summary: Verify StructDefStmt generates an LLVM identified struct type.
"""

import astx
import pytest

from irx.builders.base import Builder
from irx.builders.llvmliteir import LLVMLiteIR


@pytest.mark.parametrize("builder_class", [LLVMLiteIR])
def test_struct_definition(builder_class: type[Builder]) -> None:
"""
title: Struct definition code generation
summary: Ensure StructDefStmt translates to an LLVM struct type.
parameters:
builder_class:
type: type[Builder]
"""

builder = builder_class()
module = builder.module()

# Define struct: Point { x: int32, y: int32 }
struct_def = astx.StructDefStmt(
name="Point",
attributes=[
astx.VariableDeclaration(name="x", type_=astx.Int32()),
astx.VariableDeclaration(name="y", type_=astx.Int32()),
],
)

# Define main() -> int32
main_proto = astx.FunctionPrototype(
name="main",
args=astx.Arguments(),
return_type=astx.Int32(),
)

main_block = astx.Block()
main_block.append(struct_def)
main_block.append(astx.FunctionReturn(astx.LiteralInt32(0)))

main_fn = astx.FunctionDef(
prototype=main_proto,
body=main_block,
)

module.block.append(main_fn)

ir_text = builder.translate(module)
assert '%"Point" = type {i32, i32}' in ir_text


@pytest.mark.parametrize("builder_class", [LLVMLiteIR])
def test_struct_definition_single_field(builder_class: type[Builder]) -> None:
"""
title: Single field struct definition
summary: Ensure StructDefStmt works with a single attribute.
parameters:
builder_class:
type: type[Builder]
"""

builder = builder_class()
module = builder.module()

struct_def = astx.StructDefStmt(
name="Value",
attributes=[
astx.VariableDeclaration(name="x", type_=astx.Int32()),
],
)

main_proto = astx.FunctionPrototype(
name="main",
args=astx.Arguments(),
return_type=astx.Int32(),
)

main_block = astx.Block()
main_block.append(struct_def)
main_block.append(astx.FunctionReturn(astx.LiteralInt32(0)))

main_fn = astx.FunctionDef(
prototype=main_proto,
body=main_block,
)

module.block.append(main_fn)

ir_text = builder.translate(module)
assert '%"Value" = type {i32}' in ir_text


@pytest.mark.parametrize("builder_class", [LLVMLiteIR])
def test_struct_definition_duplicate_name(
builder_class: type[Builder],
) -> None:
"""
title: Duplicate struct name raises error
summary: >-
Ensure defining a struct with the same name twice raises ValueError.
parameters:
builder_class:
type: type[Builder]
"""

builder = builder_class()
module = builder.module()

struct_a = astx.StructDefStmt(
name="Duplicate",
attributes=[
astx.VariableDeclaration(name="x", type_=astx.Int32()),
],
)

struct_b = astx.StructDefStmt(
name="Duplicate",
attributes=[
astx.VariableDeclaration(name="y", type_=astx.Int32()),
],
)

main_proto = astx.FunctionPrototype(
name="main",
args=astx.Arguments(),
return_type=astx.Int32(),
)

main_block = astx.Block()
main_block.append(struct_a)
main_block.append(struct_b)
main_block.append(astx.FunctionReturn(astx.LiteralInt32(0)))

main_fn = astx.FunctionDef(
prototype=main_proto,
body=main_block,
)

module.block.append(main_fn)

with pytest.raises(ValueError, match="already defined"):
builder.translate(module)
Loading