Use more precise source for primary key field

7 files changed, 24 insertions, 16 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());
diff --git a/src/test/java/org/zwobble/hobgoblin/compiler/ast/untyped/UntypedTableDefinitionNodeMatcher.java b/src/test/java/org/zwobble/hobgoblin/compiler/ast/untyped/UntypedTableDefinitionNodeMatcher.java
index 6f8bdfc..d2b55d2 100644
--- a/src/test/java/org/zwobble/hobgoblin/compiler/ast/untyped/UntypedTableDefinitionNodeMatcher.java
+++ b/src/test/java/org/zwobble/hobgoblin/compiler/ast/untyped/UntypedTableDefinitionNodeMatcher.java
@@ -42,7 +42,7 @@ public final class UntypedTableDefinitionNodeMatcher implements org.zwobble.prec
         return new org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTableDefinitionNodeMatcher(submatchers);
     }
 
-    public org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTableDefinitionNodeMatcher withPrimaryKey(org.zwobble.precisely.Matcher<? super String> primaryKey) {
+    public org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTableDefinitionNodeMatcher withPrimaryKey(org.zwobble.precisely.Matcher<? super org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTypeLevelReferenceNode> primaryKey) {
         var submatchers = new java.util.ArrayList<org.zwobble.precisely.Matcher<? super org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTableDefinitionNode>>(this.submatchers);
         submatchers.add(org.zwobble.precisely.Matchers.has("primaryKey", org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTableDefinitionNode::primaryKey, (org.zwobble.precisely.Matcher) primaryKey));
         return new org.zwobble.hobgoblin.compiler.ast.untyped.UntypedTableDefinitionNodeMatcher(submatchers);
diff --git a/src/test/java/org/zwobble/hobgoblin/compiler/parser/ParserTableDefinitionTests.java b/src/test/java/org/zwobble/hobgoblin/compiler/parser/ParserTableDefinitionTests.java
index 3c3b259..cb1e1cb 100644
--- a/src/test/java/org/zwobble/hobgoblin/compiler/parser/ParserTableDefinitionTests.java
+++ b/src/test/java/org/zwobble/hobgoblin/compiler/parser/ParserTableDefinitionTests.java
@@ -25,7 +25,7 @@ public class ParserTableDefinitionTests {
         assertThat(node, isUntypedTableDefinitionNode()
             .withName(equalTo("posts"))
             .withRowType(isUntypedTypeLevelReferenceNode().withName(equalTo("Post")))
-            .withPrimaryKey(equalTo("id"))
+            .withPrimaryKey(isUntypedTypeLevelReferenceNode().withName(equalTo("id")))
         );
     }
 }
diff --git a/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerDatabaseSchemaDefinitionTests.java b/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerDatabaseSchemaDefinitionTests.java
index daecb54..63ed810 100644
--- a/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerDatabaseSchemaDefinitionTests.java
+++ b/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerDatabaseSchemaDefinitionTests.java
@@ -27,9 +27,8 @@ public class TypeCheckerDatabaseSchemaDefinitionTests {
                 .withName("authors")
                 .withRowType(UntypedTypeLevelReferenceNode.arbitrary()
                     .withName("Author")
-                    .build()
                 )
-                .withPrimaryKey("id")
+                .withPrimaryKey(UntypedTypeLevelReferenceNode.arbitrary().withName("id"))
                 .build()
             )
             .build();
diff --git a/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerTableDefinitionTests.java b/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerTableDefinitionTests.java
index 0dd911b..9dbee97 100644
--- a/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerTableDefinitionTests.java
+++ b/src/test/java/org/zwobble/hobgoblin/compiler/typechecker/TypeCheckerTableDefinitionTests.java
@@ -26,9 +26,8 @@ public class TypeCheckerTableDefinitionTests {
             .withName("authors")
             .withRowType(UntypedTypeLevelReferenceNode.arbitrary()
                 .withName("Author")
-                .build()
             )
-            .withPrimaryKey("id")
+            .withPrimaryKey(UntypedTypeLevelReferenceNode.arbitrary().withName("id"))
             .build();
         var globalContext = TypeCheckerContextArb.globalContext();
         var namespaceContext = globalContext.enterNamespace(NamespaceName.of("App"));
@@ -85,9 +84,8 @@ public class TypeCheckerTableDefinitionTests {
             .withName("authors")
             .withRowType(UntypedTypeLevelReferenceNode.arbitrary()
                 .withName("Author")
-                .build()
             )
-            .withPrimaryKey("id")
+            .withPrimaryKey(UntypedTypeLevelReferenceNode.arbitrary().withName("id"))
             .build();
         var globalContext = TypeCheckerContextArb.globalContext();
         var namespaceContext = globalContext.enterNamespace(NamespaceName.of("App"));