Add section attributes in cocas

.gitignore

@@ -263,3 +263,6 @@ gen
 notes.txt
 
 .DS_Store
+
+# ANTLR generated files
+.antlr

cocas/assembler/ast_builder.py

@@ -2,12 +2,11 @@
 import warnings
 from base64 import b64decode
 from pathlib import Path
+from typing import Literal, overload
 
 from antlr4 import CommonTokenStream, InputStream
 
-from cocas.object_module import CodeLocation
-
-from .ast_nodes import (
+from cocas.assembler.ast_nodes import (
     AbsoluteSectionNode,
     BreakStatementNode,
     ConditionalStatementNode,
@@ -26,8 +25,9 @@
     UntilLoopNode,
     WhileLoopNode,
 )
-from .exceptions import AntlrErrorListener, AssemblerException, AssemblerExceptionTag
-from .generated import AsmLexer, AsmParser, AsmParserVisitor
+from cocas.assembler.exceptions import AntlrErrorListener, AssemblerException, AssemblerExceptionTag
+from cocas.assembler.generated import AsmLexer, AsmParser, AsmParserVisitor
+from cocas.object_module import CodeLocation
 
 
 # noinspection PyPep8Naming
@@ -96,17 +96,35 @@ def check_label_is_ext(label: LabelDeclarationNode):
             raise AssemblerException(AssemblerExceptionTag.ASM, label.location.file, label.location.line,
                                      "Only external labels are allowed at the top of a file")
 
+    def visitSection_attr(self, ctx: AsmParser.Section_attrContext | None) -> str | None:
+        if ctx is None:
+            return None
+
+        return ctx.WORD().getText()
+
+    def visitSection_attrs(self, ctx: AsmParser.Section_attrsContext | None) -> list[str]:
+        if ctx is None:
+            return []
+
+        return [
+            attribute
+            for sa in ctx.section_attr()
+            if (attribute := self.visitSection_attr(sa)) is not None
+        ]
+
     def visitAbsoluteSection(self, ctx: AsmParser.AbsoluteSectionContext) -> AbsoluteSectionNode:
         header = ctx.asect_header()
         lines = self.visitSection_body(ctx.section_body())
+        attributes = self.visitSection_attrs(header.section_attrs())
         address = self.visitNumber(header.number())
-        return AbsoluteSectionNode(lines, address)
+        return AbsoluteSectionNode(lines, attributes, address)
 
     def visitRelocatableSection(self, ctx: AsmParser.RelocatableSectionContext) -> RelocatableSectionNode:
         header = ctx.rsect_header()
         lines = self.visitSection_body(ctx.section_body())
+        attributes = self.visitSection_attrs(header.section_attrs())
         name = header.name().getText()
-        return RelocatableSectionNode(lines, name)
+        return RelocatableSectionNode(lines, attributes, name)
 
     def visitTemplateSection(self, ctx: AsmParser.TemplateSectionContext) -> TemplateSectionNode:
         header = ctx.tplate_header()
@@ -185,7 +203,11 @@ def visitUntil_loop(self, ctx: AsmParser.Until_loopContext):
         mnemonic = ctx.branch_mnemonic()
         return UntilLoopNode(lines, mnemonic.getText(), self._ctx_location(mnemonic))
 
-    def visitCode_block(self, ctx: AsmParser.Code_blockContext, return_locations=False):
+    @overload
+    def visitCode_block(self, ctx: AsmParser.Code_blockContext, return_locations: Literal[True]) -> tuple[list, list]: ...
+    @overload
+    def visitCode_block(self, ctx: AsmParser.Code_blockContext, return_locations: Literal[False] = False) -> list: ...
+    def visitCode_block(self, ctx: AsmParser.Code_blockContext, return_locations: bool = False):
         if ctx.children is None:
             if return_locations:
                 return [], []
@@ -313,7 +335,7 @@ def visitArguments(self, ctx: AsmParser.ArgumentsContext):
         return [self.visitArgument(i) for i in ctx.children if isinstance(i, AsmParser.ArgumentContext)]
 
 
-def build_ast(input_stream: InputStream, filepath: Path):
+def build_ast(input_stream: InputStream, filepath: Path) -> Node:
     str_path = filepath.absolute().as_posix()
     lexer = AsmLexer(input_stream)
     lexer.removeErrorListeners()

cocas/assembler/ast_nodes.py

@@ -102,16 +102,21 @@ class ContinueStatementNode(ExportLocationNode):
 
 @dataclass
 class SectionNode(Node):
-    lines: list
+    lines: list[Node]
+
+
+@dataclass
+class AttributeMixin:
+    attributes: list[str]
 
 
 @dataclass
-class AbsoluteSectionNode(SectionNode):
+class AbsoluteSectionNode(AttributeMixin, SectionNode):
     address: int
 
 
 @dataclass
-class RelocatableSectionNode(SectionNode):
+class RelocatableSectionNode(AttributeMixin, SectionNode):
     name: str
 
 

cocas/assembler/code_block.py

@@ -176,11 +176,21 @@ def __init__(self, sn: SectionNode, target_instructions: TargetInstructions):
             address = 0
         else:
             raise Exception('Section is neither abs nor rel, can it happen? It was elif instead of else here')
+
+        self.attributes = sn.attributes
         super().__init__(address, sn.lines, target_instructions)
 
     def to_object_section_record(self, labels: dict[str, int], templates: dict[str, dict[str, int]]):
         entries = dict(p for p in self.labels.items() if p[0] in self.ents)
-        out = ObjectSectionRecord(self.name, self.address, bytearray(), entries, [], self.code_locations)
+        out = ObjectSectionRecord(
+            self.name,
+            self.address,
+            bytearray(),
+            entries,
+            [],
+            self.code_locations,
+            attributes=self.attributes,
+        )
         for seg in self.segments:
             seg.fill(out, self, labels, templates)
         return out