Use more precise source for primary key field

3 files changed, 19 insertions, 8 deletions

diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/ast/untyped/UntypedTableDefinitionNode.java b/src/main/java/org/zwobble/hobgoblin/compiler/ast/untyped/UntypedTableDefinitionNode.java
index ebff463..072d12a 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/ast/untyped/UntypedTableDefinitionNode.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/ast/untyped/UntypedTableDefinitionNode.java
@@ -8,18 +8,18 @@ package org.zwobble.hobgoblin.compiler.ast.untyped;
 public record UntypedTableDefinitionNode(
     String name,
     org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTypeLevelExpressionNode rowType,
-    String primaryKey,
+    org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTypeLevelReferenceNode primaryKey,
     org.zwobble.hobgoblin.compiler.ast.DocComment docComment,
     org.zwobble.hobgoblin.compiler.sources.Source source
 ) implements org.zwobble.hobgoblin.compiler.ast.untyped.UntypedDatabaseSchemaStatementNode {
     public static org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTableDefinitionNode.Builder arbitrary() {
-        return new org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTableDefinitionNode.Builder("", org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTypeLevelReferenceNode.arbitrary().build(), "", org.zwobble.hobgoblin.compiler.ast.HobgoblinNativeAst.arbitraryDocComment(), org.zwobble.hobgoblin.compiler.sources.HobgoblinNativeSources.arbitrarySource());
+        return new org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTableDefinitionNode.Builder("", org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTypeLevelReferenceNode.arbitrary().build(), org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTypeLevelReferenceNode.arbitrary().build(), org.zwobble.hobgoblin.compiler.ast.HobgoblinNativeAst.arbitraryDocComment(), org.zwobble.hobgoblin.compiler.sources.HobgoblinNativeSources.arbitrarySource());
     }
 
     public record Builder(
         String name,
         org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTypeLevelExpressionNode rowType,
-        String primaryKey,
+        org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTypeLevelReferenceNode primaryKey,
         org.zwobble.hobgoblin.compiler.ast.DocComment docComment,
         org.zwobble.hobgoblin.compiler.sources.Source source
     ) implements org.zwobble.hobgoblin.compiler.ast.untyped.UntypedDatabaseSchemaStatementNode.Builder {
@@ -39,10 +39,14 @@ public record UntypedTableDefinitionNode(
             return new org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTableDefinitionNode.Builder(name, rowType.build(), primaryKey, docComment, source);
         }
 
-        public org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTableDefinitionNode.Builder withPrimaryKey(String primaryKey) {
+        public org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTableDefinitionNode.Builder withPrimaryKey(org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTypeLevelReferenceNode primaryKey) {
             return new org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTableDefinitionNode.Builder(name, rowType, primaryKey, docComment, source);
         }
 
+        public org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTableDefinitionNode.Builder withPrimaryKey(org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTypeLevelReferenceNode.Builder primaryKey) {
+            return new org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTableDefinitionNode.Builder(name, rowType, primaryKey.build(), docComment, source);
+        }
+
         public org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTableDefinitionNode.Builder withDocComment(org.zwobble.hobgoblin.compiler.ast.DocComment docComment) {
             return new org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTableDefinitionNode.Builder(name, rowType, primaryKey, docComment, source);
         }
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/parser/Parser.java b/src/main/java/org/zwobble/hobgoblin/compiler/parser/Parser.java
index 104eabb..508ab70 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/parser/Parser.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/parser/Parser.java
@@ -409,7 +409,7 @@ public class Parser {
         tokens.skip(TokenType.SYMBOL_SEMICOLON);
 
         tokens.skip(TokenType.IDENTIFIER, "primaryKey");
-        var primaryKey = parseIdentifier(tokens);
+        var primaryKey = parseTypeLevelReference(tokens);
         tokens.skip(TokenType.SYMBOL_SEMICOLON);
 
         tokens.skip(TokenType.SYMBOL_BRACE_CLOSE);
@@ -493,6 +493,10 @@ public class Parser {
     }
 
     private static UntypedTypeLevelReferenceNode parsePrimaryTypeLevelExpression(TokenIterator tokens) {
+        return parseTypeLevelReference(tokens);
+    }
+
+    private static UntypedTypeLevelReferenceNode parseTypeLevelReference(TokenIterator tokens) {
         var start = tokens.startPosition();
 
         var name = parseIdentifier(tokens);
diff --git a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeChecker.java b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeChecker.java
index da83e6f..d7fb222 100644
--- a/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeChecker.java
+++ b/src/main/java/org/zwobble/hobgoblin/compiler/typechecker/TypeChecker.java
@@ -545,11 +545,14 @@ public class TypeChecker {
 
         var fields = context.fieldsOf(rowType).orElse(List.of());
         var primaryKeyField = fields.stream()
-            .filter(field -> field.name().equals(untyped.primaryKey()))
+            .filter(field -> field.name().equals(untyped.primaryKey().name()))
             .findFirst();
         if (primaryKeyField.isEmpty()) {
-            // TODO: more precise source.
-            throw new NoSuchFieldError(rowType, untyped.primaryKey(), untyped.source());
+            throw new NoSuchFieldError(
+                rowType,
+                untyped.primaryKey().name(),
+                untyped.primaryKey().source()
+            );
         }
 
         var table = new Table(databaseSchema, untyped.name());