Add addCompilationUnitElement() to PatioCodeGenContext

1 files changed, 12 insertions, 3 deletions

diff --git a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
index 1e5072e8..fdfb8b0b 100644
--- a/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
+++ b/src/main/java/org/shedlang/compiler/patio/PatioCodeGenContext.java
@@ -66,10 +66,12 @@ public record PatioCodeGenContext(
             return toResult(fullName);
         }
 
-        public CompilationUnitContext addDataElement(String name, PatioDataLayout layout, Source source) {
+        public CompilationUnitContext addCompilationUnitElement(
+            PatioCompilationUnitElement element
+        ) {
             return new CompilationUnitContext(
                 this.typeInfo,
-                this.elementsByName.plus(name, new PatioDataElement(name, layout, source))
+                this.elementsByName.plus(element.name(), element)
             );
         }
 
@@ -229,8 +231,15 @@ public record PatioCodeGenContext(
     }
 
     public PatioCodeGenContext addDataElement(String name, PatioDataLayout layout, Source source) {
+        var dataElement = new PatioDataElement(name, layout, source);
+        return addCompilationUnitElement(dataElement);
+    }
+
+    public PatioCodeGenContext addCompilationUnitElement(
+        PatioCompilationUnitElement element
+    ) {
         return withCompilationUnitContext(
-            this.compilationUnitContext.addDataElement(name, layout, source)
+            this.compilationUnitContext.addCompilationUnitElement(element)
         );
     }