diff --git a/pkl-core/src/main/java/org/pkl/core/ast/builder/AstBuilder.java b/pkl-core/src/main/java/org/pkl/core/ast/builder/AstBuilder.java
index cd871708..e7a795e9 100644
--- a/pkl-core/src/main/java/org/pkl/core/ast/builder/AstBuilder.java
+++ b/pkl-core/src/main/java/org/pkl/core/ast/builder/AstBuilder.java
@@ -361,7 +361,7 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
     var identifier = type.getName();
     var args = type.getArgs();
 
-    if (args.isEmpty()) {
+    if (args == null) {
       if (identifier.getIdentifiers().size() == 1) {
         var text = identifier.getIdentifiers().get(0).getValue();
         var typeParameter = symbolTable.findTypeParameter(text);
@@ -374,9 +374,10 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
           createSourceSection(type), doVisitTypeName(identifier));
     }
 
-    var argTypes = new UnresolvedTypeNode[args.size()];
-    for (var i = 0; i < args.size(); i++) {
-      argTypes[i] = visitType(args.get(i));
+    var targs = args.getTypes();
+    var argTypes = new UnresolvedTypeNode[targs.size()];
+    for (var i = 0; i < targs.size(); i++) {
+      argTypes[i] = visitType(targs.get(i));
     }
 
     return new UnresolvedTypeNode.Parameterized(
@@ -839,7 +840,7 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
                 createSourceSection(newExpr),
                 parentType,
                 symbolTable.getCurrentScope().getQualifiedName()));
-    if (type instanceof DeclaredType declaredType && !declaredType.getArgs().isEmpty()) {
+    if (type instanceof DeclaredType declaredType && declaredType.getArgs() != null) {
       return new TypeCastNode(parentType.getSourceSection(), expr, parentType);
     }
     return expr;
@@ -1259,14 +1260,14 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
   @Override
   public GeneratorMemberNode visitWhenGenerator(WhenGenerator member) {
     var sourceSection = createSourceSection(member);
-    var thenNodes = doVisitForWhenBody(member.getBody());
+    var thenNodes = doVisitForWhenBody(member.getThenClause());
     var elseNodes =
         member.getElseClause() == null
             ? new GeneratorMemberNode[0]
             : doVisitForWhenBody(member.getElseClause());
 
     return new GeneratorWhenNode(
-        sourceSection, visitExpr(member.getThenClause()), thenNodes, elseNodes);
+        sourceSection, visitExpr(member.getPredicate()), thenNodes, elseNodes);
   }
 
   private GeneratorMemberNode[] doVisitForWhenBody(ObjectBody body) {
@@ -1454,7 +1455,7 @@ public class AstBuilder extends AbstractAstBuilder<Object> {
 
     var moduleNode =
         AmendModuleNodeGen.create(
-            moduleInfo.getSourceSection(),
+            moduleInfo.getHeaderSection(),
             language,
             annotationNodes,
             moduleProperties,
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/BaseParserVisitor.java b/pkl-core/src/main/java/org/pkl/core/parser/BaseParserVisitor.java
index 931b1799..1c9cb32a 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/BaseParserVisitor.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/BaseParserVisitor.java
@@ -55,6 +55,7 @@ import org.pkl.core.parser.ast.Expr.UnqualifiedAccessExpr;
 import org.pkl.core.parser.ast.ExtendsOrAmendsClause;
 import org.pkl.core.parser.ast.Identifier;
 import org.pkl.core.parser.ast.ImportClause;
+import org.pkl.core.parser.ast.Keyword;
 import org.pkl.core.parser.ast.Modifier;
 import org.pkl.core.parser.ast.ModuleDecl;
 import org.pkl.core.parser.ast.Node;
@@ -86,6 +87,7 @@ import org.pkl.core.parser.ast.Type.UnionType;
 import org.pkl.core.parser.ast.Type.UnknownType;
 import org.pkl.core.parser.ast.TypeAlias;
 import org.pkl.core.parser.ast.TypeAnnotation;
+import org.pkl.core.parser.ast.TypeArgumentList;
 import org.pkl.core.parser.ast.TypeParameter;
 import org.pkl.core.parser.ast.TypeParameterList;
 
@@ -456,6 +458,16 @@ public abstract class BaseParserVisitor<T> implements ParserVisitor<T> {
     return visitChildren(replInput);
   }
 
+  @Override
+  public T visitKeyword(Keyword keyword) {
+    return defaultValue();
+  }
+
+  @Override
+  public T visitTypeArgumentList(TypeArgumentList typeArgumentList) {
+    return visitChildren(typeArgumentList);
+  }
+
   private T visitChildren(Node node) {
     T result = defaultValue();
     var children = node.children();
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/Parser.java b/pkl-core/src/main/java/org/pkl/core/parser/Parser.java
index 1b50cff4..bf9b5a9f 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/Parser.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/Parser.java
@@ -60,6 +60,7 @@ import org.pkl.core.parser.ast.Expr.UnqualifiedAccessExpr;
 import org.pkl.core.parser.ast.ExtendsOrAmendsClause;
 import org.pkl.core.parser.ast.Identifier;
 import org.pkl.core.parser.ast.ImportClause;
+import org.pkl.core.parser.ast.Keyword;
 import org.pkl.core.parser.ast.Modifier;
 import org.pkl.core.parser.ast.Module;
 import org.pkl.core.parser.ast.ModuleDecl;
@@ -85,6 +86,7 @@ import org.pkl.core.parser.ast.Type.ParenthesizedType;
 import org.pkl.core.parser.ast.Type.StringConstantType;
 import org.pkl.core.parser.ast.TypeAlias;
 import org.pkl.core.parser.ast.TypeAnnotation;
+import org.pkl.core.parser.ast.TypeArgumentList;
 import org.pkl.core.parser.ast.TypeParameter;
 import org.pkl.core.parser.ast.TypeParameterList;
 import org.pkl.core.util.ErrorMessages;
@@ -147,8 +149,7 @@ public class Parser {
         header = parseMemberHeader();
         end = parseModuleMember(header, nodes);
       }
-      assert end != null;
-      return new Module(nodes, start.endWith(end));
+      return new Module(nodes, start.endWith(spanLookahead));
     } catch (ParserError pe) {
       var spanEnd = end != null ? end : start;
       pe.setPartialParseResult(new Module(nodes, start.endWith(spanEnd)));
@@ -208,10 +209,15 @@ public class Parser {
 
   private @Nullable ModuleDecl parseModuleDecl(MemberHeader header) {
     QualifiedIdentifier moduleName = null;
-    Span start = null;
+    Keyword moduleKeyword = null;
+    var start = header.span();
     Span end = null;
     if (lookahead == Token.MODULE) {
-      start = expect(Token.MODULE, "unexpectedToken", "module").span;
+      var module = expect(Token.MODULE, "unexpectedToken", "module");
+      moduleKeyword = new Keyword(module.span);
+      if (start == null) {
+        start = module.span;
+      }
       moduleName = parseQualifiedIdentifier();
       end = moduleName.span();
     }
@@ -229,6 +235,7 @@ public class Parser {
       var modifiersOffset = children.size();
       children.addAll(header.modifiers);
       var nameOffset = children.size();
+      children.add(moduleKeyword);
       children.add(moduleName);
       children.add(extendsOrAmendsDecl);
       return new ModuleDecl(children, modifiersOffset, nameOffset, start.endWith(end));
@@ -353,8 +360,8 @@ public class Parser {
   }
 
   private TypeAlias parseTypeAlias(MemberHeader header) {
-    var start = next().span;
-    var startSpan = header.span(start);
+    var typeAlias = next().span;
+    var startSpan = header.span(typeAlias);
 
     var identifier = parseIdentifier();
     TypeParameterList typePars = null;
@@ -363,12 +370,13 @@ public class Parser {
     }
     expect(Token.ASSIGN, "unexpectedToken", "=");
     var type = parseType();
-    var children = new ArrayList<Node>(header.annotations.size() + header.modifiers.size() + 4);
+    var children = new ArrayList<Node>(header.annotations.size() + header.modifiers.size() + 5);
     children.add(header.docComment);
     children.addAll(header.annotations);
     var modifiersOffset = header.annotations.size() + 1;
     children.addAll(header.modifiers);
     var nameOffset = modifiersOffset + header.modifiers.size();
+    children.add(new Keyword(typeAlias));
     children.add(identifier);
     children.add(typePars);
     children.add(type);
@@ -376,14 +384,15 @@ public class Parser {
   }
 
   private Class parseClass(MemberHeader header) {
-    var start = next().span;
-    var startSpan = header.span(start);
+    var classKeyword = next();
+    var startSpan = header.span(classKeyword.span);
     var children = new ArrayList<Node>();
     children.add(header.docComment);
     children.addAll(header.annotations);
     var modifiersOffset = header.annotations.size() + 1;
     children.addAll(header.modifiers);
     var nameOffset = modifiersOffset + header.modifiers.size();
+    children.add(new Keyword(classKeyword.span));
     var name = parseIdentifier();
     children.add(name);
     TypeParameterList typePars = null;
@@ -550,13 +559,17 @@ public class Parser {
         if (lookahead == Token.COMMA) {
           // it's a parameter
           next();
-          nodes.add(new TypedIdentifier(identifier, typeAnnotation, identifier.span()));
+          nodes.add(
+              new TypedIdentifier(
+                  identifier, typeAnnotation, identifier.span().endWith(type.span())));
           nodes.addAll(parseListOfParameter(Token.COMMA));
           expect(Token.ARROW, "unexpectedToken2", ",", "->");
         } else if (lookahead == Token.ARROW) {
           // it's a parameter
           next();
-          nodes.add(new TypedIdentifier(identifier, typeAnnotation, identifier.span()));
+          nodes.add(
+              new TypedIdentifier(
+                  identifier, typeAnnotation, identifier.span().endWith(type.span())));
         } else {
           // it's a member
           expect(Token.ASSIGN, "unexpectedToken", "=");
@@ -637,6 +650,9 @@ public class Parser {
 
   private ObjectMember parseObjectProperty(@Nullable List<Modifier> modifiers) {
     var start = spanLookahead;
+    if (modifiers != null && !modifiers.isEmpty()) {
+      start = modifiers.get(0).span();
+    }
     var allModifiers = modifiers;
     if (allModifiers == null) {
       allModifiers = parseModifierList();
@@ -670,7 +686,10 @@ public class Parser {
 
   private ObjectMember.ObjectMethod parseObjectMethod(List<Modifier> modifiers) {
     var start = spanLookahead;
-    expect(Token.FUNCTION, "unexpectedToken", "function");
+    if (!modifiers.isEmpty()) {
+      start = modifiers.get(0).span();
+    }
+    var function = expect(Token.FUNCTION, "unexpectedToken", "function").span;
     var identifier = parseIdentifier();
     TypeParameterList params = null;
     if (lookahead == Token.LT) {
@@ -683,8 +702,9 @@ public class Parser {
     }
     expect(Token.ASSIGN, "unexpectedToken", "=");
     var expr = parseExpr("}");
-    var nodes = new ArrayList<Node>(modifiers.size() + 5);
+    var nodes = new ArrayList<Node>(modifiers.size() + 6);
     nodes.addAll(modifiers);
+    nodes.add(new Keyword(function));
     nodes.add(identifier);
     nodes.add(params);
     nodes.add(args);
@@ -1147,7 +1167,9 @@ public class Parser {
       case COLON -> {
         var typeAnnotation = parseTypeAnnotation();
         var params = new ArrayList<Parameter>();
-        params.add(new TypedIdentifier(identifier, typeAnnotation, identifier.span()));
+        params.add(
+            new TypedIdentifier(
+                identifier, typeAnnotation, identifier.span().endWith(typeAnnotation.span())));
         if (lookahead == Token.COMMA) {
           next();
           params.addAll(parseListOfParameter(Token.COMMA));
@@ -1257,23 +1279,21 @@ public class Parser {
           expect(Token.ARROW, "unexpectedToken", "->");
           var ret = parseType(expectation);
           children.add(ret);
-          typ = new Type.FunctionType(children, tk.span.endWith(end));
+          typ = new Type.FunctionType(children, tk.span.endWith(ret.span()));
         } else {
           typ = new ParenthesizedType((Type) children.get(0), tk.span.endWith(end));
         }
       }
       case IDENTIFIER -> {
         var start = spanLookahead;
-        var nodes = new ArrayList<Node>(1);
         var name = parseQualifiedIdentifier();
         var end = name.span();
-        nodes.add(name);
+        TypeArgumentList typeArgumentList = null;
         if (lookahead == Token.LT) {
-          next();
-          nodes.addAll(parseListOf(Token.COMMA, () -> parseType(">")));
-          end = expect(Token.GT, "unexpectedToken2", ",", ">").span;
+          typeArgumentList = parseTypeArgumentList();
+          end = typeArgumentList.span();
         }
-        typ = new DeclaredType(nodes, start.endWith(end));
+        typ = new DeclaredType(name, typeArgumentList, start.endWith(end));
       }
       case STRING_START -> {
         var str = parseStringConstant();
@@ -1389,6 +1409,13 @@ public class Parser {
     return new TypeParameterList(pars, start.endWith(end));
   }
 
+  private TypeArgumentList parseTypeArgumentList() {
+    var start = expect(Token.LT, "unexpectedToken", "<").span;
+    var pars = parseListOf(Token.COMMA, this::parseType);
+    var end = expect(Token.GT, "unexpectedToken2", ",", ">").span;
+    return new TypeArgumentList(pars, start.endWith(end));
+  }
+
   private ArgumentList parseArgumentList() {
     var start = expect(Token.LPAREN, "unexpectedToken", "(").span;
     if (lookahead == Token.RPAREN) {
@@ -1524,20 +1551,21 @@ public class Parser {
       return !(docComment == null && annotations.isEmpty() && modifiers.isEmpty());
     }
 
+    @SuppressWarnings("DataFlowIssue")
+    @Nullable
+    Span span() {
+      return span(null);
+    }
+
     Span span(Span or) {
-      Span start = null;
-      Span end = null;
+      if (docComment != null) {
+        return docComment.span();
+      }
       if (!annotations().isEmpty()) {
-        start = annotations.get(0).span();
-        end = annotations.get(annotations.size() - 1).span();
+        return annotations.get(0).span();
       }
       if (!modifiers().isEmpty()) {
-        if (start == null) start = modifiers.get(0).span();
-        end = modifiers.get(modifiers.size() - 1).span();
-        return start.endWith(end);
-      }
-      if (end != null) {
-        return start.endWith(end);
+        return modifiers.get(0).span();
       }
       return or;
     }
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ParserVisitor.java b/pkl-core/src/main/java/org/pkl/core/parser/ParserVisitor.java
index e9b4a61d..fb6385cf 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ParserVisitor.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ParserVisitor.java
@@ -55,6 +55,7 @@ import org.pkl.core.parser.ast.Expr.UnqualifiedAccessExpr;
 import org.pkl.core.parser.ast.ExtendsOrAmendsClause;
 import org.pkl.core.parser.ast.Identifier;
 import org.pkl.core.parser.ast.ImportClause;
+import org.pkl.core.parser.ast.Keyword;
 import org.pkl.core.parser.ast.Modifier;
 import org.pkl.core.parser.ast.Module;
 import org.pkl.core.parser.ast.ModuleDecl;
@@ -70,6 +71,7 @@ import org.pkl.core.parser.ast.StringPart;
 import org.pkl.core.parser.ast.Type;
 import org.pkl.core.parser.ast.TypeAlias;
 import org.pkl.core.parser.ast.TypeAnnotation;
+import org.pkl.core.parser.ast.TypeArgumentList;
 import org.pkl.core.parser.ast.TypeParameter;
 import org.pkl.core.parser.ast.TypeParameterList;
 
@@ -220,4 +222,8 @@ public interface ParserVisitor<Result> {
   Result visitObjectBody(ObjectBody objectBody);
 
   Result visitReplInput(ReplInput replInput);
+
+  Result visitKeyword(Keyword keyword);
+
+  Result visitTypeArgumentList(TypeArgumentList typeArgumentList);
 }
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/ArgumentList.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/ArgumentList.java
index a00948b5..c57542b6 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/ArgumentList.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/ArgumentList.java
@@ -20,7 +20,6 @@ import org.pkl.core.parser.ParserVisitor;
 import org.pkl.core.parser.Span;
 import org.pkl.core.util.Nullable;
 
-@SuppressWarnings("unchecked")
 public class ArgumentList extends AbstractNode {
 
   public ArgumentList(List<Expr> arguments, Span span) {
@@ -32,6 +31,7 @@ public class ArgumentList extends AbstractNode {
     return visitor.visitArgumentList(this);
   }
 
+  @SuppressWarnings("unchecked")
   public List<Expr> getArguments() {
     assert children != null;
     return (List<Expr>) children;
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/Class.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/Class.java
index 009b8ac4..ea5c524d 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/Class.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/Class.java
@@ -20,51 +20,74 @@ import org.pkl.core.parser.ParserVisitor;
 import org.pkl.core.parser.Span;
 import org.pkl.core.util.Nullable;
 
-@SuppressWarnings({"unchecked", "DataFlowIssue"})
 public final class Class extends AbstractNode {
   private final int modifiersOffset;
-  private final int nameOffset;
+  private final int keywordOffset;
 
-  public Class(List<Node> nodes, int modifiersOffset, int nameOffset, Span span) {
+  public Class(List<Node> nodes, int modifiersOffset, int keywordOffset, Span span) {
     super(span, nodes);
     this.modifiersOffset = modifiersOffset;
-    this.nameOffset = nameOffset;
+    this.keywordOffset = keywordOffset;
   }
 
   @Override
-  public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+  public <T> T accept(ParserVisitor<? extends T> visitor) {
     return visitor.visitClass(this);
   }
 
   public @Nullable DocComment getDocComment() {
+    assert children != null;
     return (DocComment) children.get(0);
   }
 
+  @SuppressWarnings("unchecked")
   public List<Annotation> getAnnotations() {
+    assert children != null;
     return (List<Annotation>) children.subList(1, modifiersOffset);
   }
 
+  @SuppressWarnings("unchecked")
   public List<Modifier> getModifiers() {
-    return (List<Modifier>) children.subList(modifiersOffset, nameOffset);
+    assert children != null;
+    return (List<Modifier>) children.subList(modifiersOffset, keywordOffset);
+  }
+
+  public Keyword getClassKeyword() {
+    assert children != null;
+    return (Keyword) children.get(keywordOffset);
   }
 
   public Identifier getName() {
-    return (Identifier) children.get(nameOffset);
+    assert children != null;
+    return (Identifier) children.get(keywordOffset + 1);
   }
 
   public @Nullable TypeParameterList getTypeParameterList() {
-    return (TypeParameterList) children.get(nameOffset + 1);
+    assert children != null;
+    return (TypeParameterList) children.get(keywordOffset + 2);
   }
 
   public @Nullable Type getSuperClass() {
-    return (Type) children.get(nameOffset + 2);
+    assert children != null;
+    return (Type) children.get(keywordOffset + 3);
   }
 
   public @Nullable ClassBody getBody() {
-    return (ClassBody) children.get(nameOffset + 3);
+    assert children != null;
+    return (ClassBody) children.get(keywordOffset + 4);
   }
 
+  @SuppressWarnings("DuplicatedCode")
   public Span getHeaderSpan() {
+    Span start = null;
+    assert children != null;
+    for (var i = modifiersOffset; i < children.size(); i++) {
+      var child = children.get(i);
+      if (child != null) {
+        start = child.span();
+        break;
+      }
+    }
     Span end;
     if (getSuperClass() != null) {
       end = getSuperClass().span();
@@ -73,6 +96,7 @@ public final class Class extends AbstractNode {
     } else {
       end = getName().span();
     }
-    return span.endWith(end);
+    assert start != null;
+    return start.endWith(end);
   }
 }
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/ClassMethod.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/ClassMethod.java
index d86070e8..68d78209 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/ClassMethod.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/ClassMethod.java
@@ -20,7 +20,6 @@ import org.pkl.core.parser.ParserVisitor;
 import org.pkl.core.parser.Span;
 import org.pkl.core.util.Nullable;
 
-@SuppressWarnings("unchecked")
 public class ClassMethod extends AbstractNode {
   private final int modifiersOffset;
   private final int nameOffset;
@@ -44,11 +43,13 @@ public class ClassMethod extends AbstractNode {
     return (DocComment) children.get(0);
   }
 
+  @SuppressWarnings("unchecked")
   public List<Annotation> getAnnotations() {
     assert children != null;
     return (List<Annotation>) children.subList(1, modifiersOffset);
   }
 
+  @SuppressWarnings("unchecked")
   public List<Modifier> getModifiers() {
     assert children != null;
     return (List<Modifier>) children.subList(modifiersOffset, nameOffset);
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/ClassProperty.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/ClassProperty.java
index 1934d53a..7bff4b66 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/ClassProperty.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/ClassProperty.java
@@ -20,7 +20,6 @@ import org.pkl.core.parser.ParserVisitor;
 import org.pkl.core.parser.Span;
 import org.pkl.core.util.Nullable;
 
-@SuppressWarnings({"DuplicatedCode", "unchecked"})
 public final class ClassProperty extends AbstractNode {
   private final int modifiersOffset;
   private final int nameOffset;
@@ -32,7 +31,7 @@ public final class ClassProperty extends AbstractNode {
   }
 
   @Override
-  public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+  public <T> T accept(ParserVisitor<? extends T> visitor) {
     return visitor.visitClassProperty(this);
   }
 
@@ -41,11 +40,13 @@ public final class ClassProperty extends AbstractNode {
     return (DocComment) children.get(0);
   }
 
+  @SuppressWarnings("unchecked")
   public List<Annotation> getAnnotations() {
     assert children != null;
     return (List<Annotation>) children.subList(1, modifiersOffset);
   }
 
+  @SuppressWarnings("unchecked")
   public List<Modifier> getModifiers() {
     assert children != null;
     return (List<Modifier>) children.subList(modifiersOffset, nameOffset);
@@ -66,6 +67,7 @@ public final class ClassProperty extends AbstractNode {
     return (Expr) children.get(nameOffset + 2);
   }
 
+  @SuppressWarnings("unchecked")
   public List<ObjectBody> getBodyList() {
     assert children != null;
     return (List<ObjectBody>) children.subList(nameOffset + 3, children.size());
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/Expr.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/Expr.java
index 616d8573..a3784c1a 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/Expr.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/Expr.java
@@ -23,7 +23,6 @@ import org.pkl.core.parser.ParserVisitor;
 import org.pkl.core.parser.Span;
 import org.pkl.core.util.Nullable;
 
-@SuppressWarnings("ALL")
 public abstract sealed class Expr extends AbstractNode {
 
   public Expr(Span span, @Nullable List<? extends @Nullable Node> children) {
@@ -36,7 +35,7 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitThisExpr(this);
     }
   }
@@ -47,7 +46,7 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitOuterExpr(this);
     }
   }
@@ -58,7 +57,7 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitModuleExpr(this);
     }
   }
@@ -69,7 +68,7 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitNullLiteralExpr(this);
     }
   }
@@ -83,7 +82,7 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitBoolLiteralExpr(this);
     }
 
@@ -101,7 +100,7 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitIntLiteralExpr(this);
     }
 
@@ -119,7 +118,7 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitFloatLiteralExpr(this);
     }
 
@@ -140,10 +139,11 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitSingleLineStringLiteralExpr(this);
     }
 
+    @SuppressWarnings("unchecked")
     public List<StringPart> getParts() {
       assert children != null;
       return (List<StringPart>) children;
@@ -170,11 +170,13 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitMultiLineStringLiteralExpr(this);
     }
 
+    @SuppressWarnings("unchecked")
     public List<StringPart> getParts() {
+      assert children != null;
       return (List<StringPart>) children;
     }
 
@@ -193,11 +195,12 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitThrowExpr(this);
     }
 
     public Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(0);
     }
   }
@@ -208,11 +211,12 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitTraceExpr(this);
     }
 
     public Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(0);
     }
   }
@@ -226,11 +230,12 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitImportExpr(this);
     }
 
     public StringConstant getImportStr() {
+      assert children != null;
       return (StringConstant) children.get(0);
     }
 
@@ -248,11 +253,12 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitReadExpr(this);
     }
 
     public Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(0);
     }
 
@@ -274,15 +280,17 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitUnqualifiedAccessExpr(this);
     }
 
     public Identifier getIdentifier() {
+      assert children != null;
       return (Identifier) children.get(0);
     }
 
     public @Nullable ArgumentList getArgumentList() {
+      assert children != null;
       return (ArgumentList) children.get(1);
     }
   }
@@ -301,15 +309,17 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitQualifiedAccessExpr(this);
     }
 
     public Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(0);
     }
 
     public Identifier getIdentifier() {
+      assert children != null;
       return (Identifier) children.get(1);
     }
 
@@ -318,6 +328,7 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     public @Nullable ArgumentList getArgumentList() {
+      assert children != null;
       return (ArgumentList) children.get(2);
     }
   }
@@ -328,15 +339,17 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitSuperAccessExpr(this);
     }
 
     public Identifier getIdentifier() {
+      assert children != null;
       return (Identifier) children.get(0);
     }
 
     public @Nullable ArgumentList getArgumentList() {
+      assert children != null;
       return (ArgumentList) children.get(1);
     }
   }
@@ -347,11 +360,12 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitSuperSubscriptExpr(this);
     }
 
     public Expr getArg() {
+      assert children != null;
       return (Expr) children.get(0);
     }
   }
@@ -362,15 +376,17 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitSubscriptExpr(this);
     }
 
     public Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(0);
     }
 
     public Expr getArg() {
+      assert children != null;
       return (Expr) children.get(1);
     }
   }
@@ -381,19 +397,22 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitIfExpr(this);
     }
 
     public Expr getCond() {
+      assert children != null;
       return (Expr) children.get(0);
     }
 
     public Expr getThen() {
+      assert children != null;
       return (Expr) children.get(1);
     }
 
     public Expr getEls() {
+      assert children != null;
       return (Expr) children.get(2);
     }
   }
@@ -404,19 +423,22 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitLetExpr(this);
     }
 
     public Parameter getParameter() {
+      assert children != null;
       return (Parameter) children.get(0);
     }
 
     public Expr getBindingExpr() {
+      assert children != null;
       return (Expr) children.get(1);
     }
 
     public Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(2);
     }
   }
@@ -427,15 +449,17 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitFunctionLiteralExpr(this);
     }
 
     public ParameterList getParameterList() {
+      assert children != null;
       return (ParameterList) children.get(0);
     }
 
     public Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(1);
     }
   }
@@ -446,11 +470,12 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitParenthesizedExpr(this);
     }
 
     public Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(0);
     }
   }
@@ -461,15 +486,17 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitNewExpr(this);
     }
 
     public @Nullable Type getType() {
+      assert children != null;
       return (Type) children.get(0);
     }
 
     public ObjectBody getBody() {
+      assert children != null;
       return (ObjectBody) children.get(1);
     }
 
@@ -484,15 +511,17 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitAmendsExpr(this);
     }
 
     public Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(0);
     }
 
     public ObjectBody getBody() {
+      assert children != null;
       return (ObjectBody) children.get(1);
     }
   }
@@ -503,11 +532,12 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitNonNullExpr(this);
     }
 
     public Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(0);
     }
   }
@@ -518,11 +548,12 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitUnaryMinusExpr(this);
     }
 
     public Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(0);
     }
   }
@@ -533,11 +564,12 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitLogicalNotExpr(this);
     }
 
     public Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(0);
     }
   }
@@ -551,15 +583,17 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitBinaryOperatorExpr(this);
     }
 
     public Expr getLeft() {
+      assert children != null;
       return (Expr) children.get(0);
     }
 
     public Expr getRight() {
+      assert children != null;
       return (Expr) children.get(1);
     }
 
@@ -569,9 +603,10 @@ public abstract sealed class Expr extends AbstractNode {
 
     @Override
     public String toString() {
-      return "BinaryOp{" + "children=" + children + ", op=" + op + ", span=" + span + '}';
+      return "BinaryOp{children=" + children + ", op=" + op + ", span=" + span + '}';
     }
 
+    @SuppressWarnings("ConstantValue")
     @Override
     public boolean equals(Object o) {
       if (this == o) {
@@ -598,15 +633,17 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitTypeCheckExpr(this);
     }
 
     public Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(0);
     }
 
     public Type getType() {
+      assert children != null;
       return (Type) children.get(1);
     }
   }
@@ -617,15 +654,17 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitTypeCastExpr(this);
     }
 
     public Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(0);
     }
 
     public Type getType() {
+      assert children != null;
       return (Type) children.get(1);
     }
   }
@@ -640,7 +679,7 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       // should never be called
       throw PklBugException.unreachableCode();
     }
@@ -654,6 +693,7 @@ public abstract sealed class Expr extends AbstractNode {
       return "OperatorExpr{op=" + op + ", span=" + span + '}';
     }
 
+    @SuppressWarnings("ConstantValue")
     @Override
     public boolean equals(Object o) {
       if (this == o) {
@@ -679,12 +719,13 @@ public abstract sealed class Expr extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       // should never be called
       throw PklBugException.unreachableCode();
     }
 
     public Type getType() {
+      assert children != null;
       return (Type) children.get(0);
     }
   }
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/ImportClause.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/ImportClause.java
index 635d10ec..d8a57d69 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/ImportClause.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/ImportClause.java
@@ -21,7 +21,6 @@ import org.pkl.core.parser.ParserVisitor;
 import org.pkl.core.parser.Span;
 import org.pkl.core.util.Nullable;
 
-@SuppressWarnings("DataFlowIssue")
 public final class ImportClause extends AbstractNode {
   private final boolean isGlob;
 
@@ -32,11 +31,12 @@ public final class ImportClause extends AbstractNode {
   }
 
   @Override
-  public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+  public <T> T accept(ParserVisitor<? extends T> visitor) {
     return visitor.visitImportClause(this);
   }
 
   public StringConstant getImportStr() {
+    assert children != null;
     return (StringConstant) children.get(0);
   }
 
@@ -45,6 +45,7 @@ public final class ImportClause extends AbstractNode {
   }
 
   public @Nullable Identifier getAlias() {
+    assert children != null;
     return (Identifier) children.get(1);
   }
 
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/Keyword.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/Keyword.java
new file mode 100644
index 00000000..cadda0a3
--- /dev/null
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/Keyword.java
@@ -0,0 +1,31 @@
+/*
+ * Copyright © 2025 Apple Inc. and the Pkl project authors. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.pkl.core.parser.ast;
+
+import org.pkl.core.parser.ParserVisitor;
+import org.pkl.core.parser.Span;
+
+public class Keyword extends AbstractNode {
+
+  public Keyword(Span span) {
+    super(span, null);
+  }
+
+  @Override
+  public <T> T accept(ParserVisitor<? extends T> visitor) {
+    return visitor.visitKeyword(this);
+  }
+}
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/Modifier.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/Modifier.java
index 998529bd..3f1b92a4 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/Modifier.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/Modifier.java
@@ -18,7 +18,6 @@ package org.pkl.core.parser.ast;
 import java.util.Objects;
 import org.pkl.core.parser.ParserVisitor;
 import org.pkl.core.parser.Span;
-import org.pkl.core.util.Nullable;
 
 public final class Modifier extends AbstractNode {
   private final ModifierValue value;
@@ -29,7 +28,7 @@ public final class Modifier extends AbstractNode {
   }
 
   @Override
-  public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+  public <T> T accept(ParserVisitor<? extends T> visitor) {
     return visitor.visitModifier(this);
   }
 
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/Module.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/Module.java
index cda1e24e..9445d493 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/Module.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/Module.java
@@ -21,22 +21,23 @@ import org.pkl.core.parser.ParserVisitor;
 import org.pkl.core.parser.Span;
 import org.pkl.core.util.Nullable;
 
-@SuppressWarnings("DataFlowIssue")
 public final class Module extends AbstractNode {
   public Module(List<Node> nodes, Span span) {
     super(span, nodes);
   }
 
   @Override
-  public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+  public <T> T accept(ParserVisitor<? extends T> visitor) {
     return visitor.visitModule(this);
   }
 
   public @Nullable ModuleDecl getDecl() {
+    assert children != null;
     return (ModuleDecl) children.get(0);
   }
 
   public List<ImportClause> getImports() {
+    assert children != null;
     if (children.size() < 2) return List.of();
     var res = new ArrayList<ImportClause>();
     for (int i = 1; i < children.size(); i++) {
@@ -53,6 +54,7 @@ public final class Module extends AbstractNode {
 
   public List<Class> getClasses() {
     var res = new ArrayList<Class>();
+    assert children != null;
     for (var child : children) {
       if (child instanceof Class clazz) {
         res.add(clazz);
@@ -63,6 +65,7 @@ public final class Module extends AbstractNode {
 
   public List<TypeAlias> getTypeAliases() {
     var res = new ArrayList<TypeAlias>();
+    assert children != null;
     for (var child : children) {
       if (child instanceof TypeAlias typeAlias) {
         res.add(typeAlias);
@@ -73,6 +76,7 @@ public final class Module extends AbstractNode {
 
   public List<ClassProperty> getProperties() {
     var res = new ArrayList<ClassProperty>();
+    assert children != null;
     for (var child : children) {
       if (child instanceof ClassProperty classProperty) {
         res.add(classProperty);
@@ -83,6 +87,7 @@ public final class Module extends AbstractNode {
 
   public List<ClassMethod> getMethods() {
     var res = new ArrayList<ClassMethod>();
+    assert children != null;
     for (var child : children) {
       if (child instanceof ClassMethod classMethod) {
         res.add(classMethod);
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/ModuleDecl.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/ModuleDecl.java
index bef1e004..c1f1035a 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/ModuleDecl.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/ModuleDecl.java
@@ -20,7 +20,6 @@ import org.pkl.core.parser.ParserVisitor;
 import org.pkl.core.parser.Span;
 import org.pkl.core.util.Nullable;
 
-@SuppressWarnings({"DataFlowIssue", "unchecked"})
 public final class ModuleDecl extends AbstractNode {
   private final int modifiersOffset;
   private final int nameOffset;
@@ -32,42 +31,58 @@ public final class ModuleDecl extends AbstractNode {
   }
 
   @Override
-  public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+  public <T> T accept(ParserVisitor<? extends T> visitor) {
     return visitor.visitModuleDecl(this);
   }
 
   public @Nullable DocComment getDocComment() {
+    assert children != null;
     return (DocComment) children.get(0);
   }
 
+  @SuppressWarnings("unchecked")
   public List<Annotation> getAnnotations() {
+    assert children != null;
     return (List<Annotation>) children.subList(1, modifiersOffset);
   }
 
+  @SuppressWarnings("unchecked")
   public List<Modifier> getModifiers() {
+    assert children != null;
     return (List<Modifier>) children.subList(modifiersOffset, nameOffset);
   }
 
+  public @Nullable Keyword getModuleKeyword() {
+    assert children != null;
+    return (Keyword) children.get(nameOffset);
+  }
+
   public @Nullable QualifiedIdentifier getName() {
-    return (QualifiedIdentifier) children.get(nameOffset);
+    assert children != null;
+    return (QualifiedIdentifier) children.get(nameOffset + 1);
   }
 
   public @Nullable ExtendsOrAmendsClause getExtendsOrAmendsDecl() {
-    return (ExtendsOrAmendsClause) children.get(nameOffset + 1);
+    assert children != null;
+    return (ExtendsOrAmendsClause) children.get(nameOffset + 2);
   }
 
+  @SuppressWarnings("DuplicatedCode")
   public Span headerSpan() {
     Span start = null;
-    Span end = null;
+    assert children != null;
     for (var i = modifiersOffset; i < children.size(); i++) {
       var child = children.get(i);
       if (child != null) {
-        if (start == null) {
-          start = child.span();
-        }
-        end = child.span();
+        start = child.span();
+        break;
       }
     }
-    return start.endWith(end);
+    var extendsOrAmends = children.get(nameOffset + 2);
+    assert start != null;
+    if (extendsOrAmends != null) {
+      return start.endWith(extendsOrAmends.span());
+    }
+    return start.endWith(children.get(nameOffset + 1).span());
   }
 }
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/ObjectBody.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/ObjectBody.java
index 24677c61..83041c3c 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/ObjectBody.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/ObjectBody.java
@@ -18,9 +18,7 @@ package org.pkl.core.parser.ast;
 import java.util.List;
 import org.pkl.core.parser.ParserVisitor;
 import org.pkl.core.parser.Span;
-import org.pkl.core.util.Nullable;
 
-@SuppressWarnings({"unchecked", "DataFlowIssue"})
 public final class ObjectBody extends AbstractNode {
   private final int membersOffset;
 
@@ -30,15 +28,19 @@ public final class ObjectBody extends AbstractNode {
   }
 
   @Override
-  public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+  public <T> T accept(ParserVisitor<? extends T> visitor) {
     return visitor.visitObjectBody(this);
   }
 
+  @SuppressWarnings("unchecked")
   public List<Parameter> getParameters() {
+    assert children != null;
     return (List<Parameter>) children.subList(0, membersOffset);
   }
 
+  @SuppressWarnings("unchecked")
   public List<ObjectMember> getMembers() {
+    assert children != null;
     return (List<ObjectMember>) children.subList(membersOffset, children.size());
   }
 }
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/ObjectMember.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/ObjectMember.java
index 123e5d8c..04c829e7 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/ObjectMember.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/ObjectMember.java
@@ -22,7 +22,6 @@ import org.pkl.core.parser.ParserVisitor;
 import org.pkl.core.parser.Span;
 import org.pkl.core.util.Nullable;
 
-@SuppressWarnings("ALL")
 public abstract sealed class ObjectMember extends AbstractNode {
 
   public ObjectMember(Span span, @Nullable List<? extends @Nullable Node> children) {
@@ -35,11 +34,12 @@ public abstract sealed class ObjectMember extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitObjectElement(this);
     }
 
     public Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(0);
     }
   }
@@ -53,27 +53,34 @@ public abstract sealed class ObjectMember extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitObjectProperty(this);
     }
 
+    @SuppressWarnings("unchecked")
     public List<Modifier> getModifiers() {
+      assert children != null;
       return (List<Modifier>) children.subList(0, identifierOffset);
     }
 
     public Identifier getIdentifier() {
+      assert children != null;
       return (Identifier) children.get(identifierOffset);
     }
 
     public @Nullable TypeAnnotation getTypeAnnotation() {
+      assert children != null;
       return (TypeAnnotation) children.get(identifierOffset + 1);
     }
 
     public @Nullable Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(identifierOffset + 2);
     }
 
-    public @Nullable List<ObjectBody> getBodyList() {
+    @SuppressWarnings("unchecked")
+    public List<ObjectBody> getBodyList() {
+      assert children != null;
       return (List<ObjectBody>) children.subList(identifierOffset + 3, children.size());
     }
   }
@@ -87,43 +94,65 @@ public abstract sealed class ObjectMember extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitObjectMethod(this);
     }
 
+    @SuppressWarnings("unchecked")
     public List<Modifier> getModifiers() {
+      assert children != null;
       return (List<Modifier>) children.subList(0, identifierOffset);
     }
 
+    public Keyword getFunctionKeyword() {
+      assert children != null;
+      return (Keyword) children.get(identifierOffset);
+    }
+
     public Identifier getIdentifier() {
-      return (Identifier) children.get(identifierOffset);
+      assert children != null;
+      return (Identifier) children.get(identifierOffset + 1);
     }
 
     public @Nullable TypeParameterList getTypeParameterList() {
-      return (TypeParameterList) children.get(identifierOffset + 1);
+      assert children != null;
+      return (TypeParameterList) children.get(identifierOffset + 2);
     }
 
     public ParameterList getParamList() {
-      return (ParameterList) children.get(identifierOffset + 2);
+      assert children != null;
+      return (ParameterList) children.get(identifierOffset + 3);
     }
 
     public @Nullable TypeAnnotation getTypeAnnotation() {
-      return (TypeAnnotation) children.get(identifierOffset + 3);
+      assert children != null;
+      return (TypeAnnotation) children.get(identifierOffset + 4);
     }
 
     public Expr getExpr() {
-      return (Expr) children.get(identifierOffset + 4);
+      assert children != null;
+      return (Expr) children.get(identifierOffset + 5);
     }
 
+    @SuppressWarnings("DuplicatedCode")
     public Span headerSpan() {
+      Span start = null;
+      assert children != null;
+      for (var child : children) {
+        if (child != null) {
+          start = child.span();
+          break;
+        }
+      }
       Span end;
-      var typeAnnotation = children.get(identifierOffset + 3);
+      var typeAnnotation = children.get(identifierOffset + 4);
       if (typeAnnotation == null) {
-        end = children.get(identifierOffset + 2).span();
+        end = children.get(identifierOffset + 3).span();
       } else {
         end = typeAnnotation.span();
       }
-      return span.endWith(end);
+      assert start != null;
+      return start.endWith(end);
     }
   }
 
@@ -133,19 +162,23 @@ public abstract sealed class ObjectMember extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitMemberPredicate(this);
     }
 
     public Expr getPred() {
+      assert children != null;
       return (Expr) children.get(0);
     }
 
     public @Nullable Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(1);
     }
 
+    @SuppressWarnings("unchecked")
     public List<ObjectBody> getBodyList() {
+      assert children != null;
       return (List<ObjectBody>) children.subList(2, children.size());
     }
   }
@@ -156,19 +189,23 @@ public abstract sealed class ObjectMember extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitObjectEntry(this);
     }
 
     public Expr getKey() {
+      assert children != null;
       return (Expr) children.get(0);
     }
 
     public @Nullable Expr getValue() {
+      assert children != null;
       return (Expr) children.get(1);
     }
 
+    @SuppressWarnings("unchecked")
     public List<ObjectBody> getBodyList() {
+      assert children != null;
       return (List<ObjectBody>) children.subList(2, children.size());
     }
   }
@@ -182,11 +219,12 @@ public abstract sealed class ObjectMember extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitObjectSpread(this);
     }
 
     public Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(0);
     }
 
@@ -206,6 +244,7 @@ public abstract sealed class ObjectMember extends AbstractNode {
           + '}';
     }
 
+    @SuppressWarnings("ConstantValue")
     @Override
     public boolean equals(Object o) {
       if (this == o) {
@@ -229,24 +268,27 @@ public abstract sealed class ObjectMember extends AbstractNode {
 
   public static final class WhenGenerator extends ObjectMember {
     public WhenGenerator(
-        Expr thenClause, ObjectBody body, @Nullable ObjectBody elseClause, Span span) {
-      super(span, Arrays.asList(thenClause, body, elseClause));
+        Expr predicate, ObjectBody thenClause, @Nullable ObjectBody elseClause, Span span) {
+      super(span, Arrays.asList(predicate, thenClause, elseClause));
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitWhenGenerator(this);
     }
 
-    public Expr getThenClause() {
+    public Expr getPredicate() {
+      assert children != null;
       return (Expr) children.get(0);
     }
 
-    public ObjectBody getBody() {
+    public ObjectBody getThenClause() {
+      assert children != null;
       return (ObjectBody) children.get(1);
     }
 
     public @Nullable ObjectBody getElseClause() {
+      assert children != null;
       return (ObjectBody) children.get(2);
     }
   }
@@ -258,23 +300,27 @@ public abstract sealed class ObjectMember extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitForGenerator(this);
     }
 
     public Parameter getP1() {
+      assert children != null;
       return (Parameter) children.get(0);
     }
 
     public @Nullable Parameter getP2() {
+      assert children != null;
       return (Parameter) children.get(1);
     }
 
     public Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(2);
     }
 
     public ObjectBody getBody() {
+      assert children != null;
       return (ObjectBody) children.get(3);
     }
 
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/Parameter.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/Parameter.java
index c90aebbf..151ac89d 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/Parameter.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/Parameter.java
@@ -21,7 +21,6 @@ import org.pkl.core.parser.ParserVisitor;
 import org.pkl.core.parser.Span;
 import org.pkl.core.util.Nullable;
 
-@SuppressWarnings("ALL")
 public abstract sealed class Parameter extends AbstractNode {
 
   public Parameter(Span span, @Nullable List<? extends @Nullable Node> children) {
@@ -46,10 +45,12 @@ public abstract sealed class Parameter extends AbstractNode {
     }
 
     public Identifier getIdentifier() {
+      assert children != null;
       return (Identifier) children.get(0);
     }
 
     public @Nullable TypeAnnotation getTypeAnnotation() {
+      assert children != null;
       return (TypeAnnotation) children.get(1);
     }
   }
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/ParameterList.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/ParameterList.java
index f5cc0d07..b1a5a45e 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/ParameterList.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/ParameterList.java
@@ -20,7 +20,6 @@ import org.pkl.core.parser.ParserVisitor;
 import org.pkl.core.parser.Span;
 import org.pkl.core.util.Nullable;
 
-@SuppressWarnings("unchecked")
 public class ParameterList extends AbstractNode {
   public ParameterList(List<Parameter> parameters, Span span) {
     super(span, parameters);
@@ -31,6 +30,7 @@ public class ParameterList extends AbstractNode {
     return visitor.visitParameterList(this);
   }
 
+  @SuppressWarnings("unchecked")
   public List<Parameter> getParameters() {
     assert children != null;
     return (List<Parameter>) children;
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/QualifiedIdentifier.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/QualifiedIdentifier.java
index ae71f866..43d7d5dc 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/QualifiedIdentifier.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/QualifiedIdentifier.java
@@ -18,7 +18,6 @@ package org.pkl.core.parser.ast;
 import java.util.List;
 import java.util.stream.Collectors;
 import org.pkl.core.parser.ParserVisitor;
-import org.pkl.core.util.Nullable;
 
 public final class QualifiedIdentifier extends AbstractNode {
   public QualifiedIdentifier(List<Identifier> identifiers) {
@@ -27,12 +26,13 @@ public final class QualifiedIdentifier extends AbstractNode {
   }
 
   @Override
-  public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+  public <T> T accept(ParserVisitor<? extends T> visitor) {
     return visitor.visitQualifiedIdentifier(this);
   }
 
-  @SuppressWarnings({"unchecked", "DataFlowIssue"})
+  @SuppressWarnings("unchecked")
   public List<Identifier> getIdentifiers() {
+    assert children != null;
     return (List<Identifier>) children;
   }
 
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/StringConstantPart.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/StringConstantPart.java
index 20a7d903..a55b0aab 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/StringConstantPart.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/StringConstantPart.java
@@ -21,7 +21,6 @@ import org.pkl.core.parser.ParserVisitor;
 import org.pkl.core.parser.Span;
 import org.pkl.core.util.Nullable;
 
-@SuppressWarnings("ALL")
 public abstract sealed class StringConstantPart extends AbstractNode {
 
   public StringConstantPart(Span span, @Nullable List<? extends @Nullable Node> children) {
@@ -56,6 +55,7 @@ public abstract sealed class StringConstantPart extends AbstractNode {
       return "ConstantPart{str='" + str + '\'' + ", span=" + span + '}';
     }
 
+    @SuppressWarnings("ConstantValue")
     @Override
     public boolean equals(Object o) {
       if (this == o) {
@@ -91,6 +91,7 @@ public abstract sealed class StringConstantPart extends AbstractNode {
       return "StringUnicodeEscape{escape='" + escape + '\'' + ", span=" + span + '}';
     }
 
+    @SuppressWarnings("ConstantValue")
     @Override
     public boolean equals(Object o) {
       if (this == o) {
@@ -126,6 +127,7 @@ public abstract sealed class StringConstantPart extends AbstractNode {
       return "StringEscape{type=" + type + ", span=" + span + '}';
     }
 
+    @SuppressWarnings("ConstantValue")
     @Override
     public boolean equals(Object o) {
       if (this == o) {
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/StringPart.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/StringPart.java
index 5847ad60..4e6d8e65 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/StringPart.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/StringPart.java
@@ -20,7 +20,6 @@ import org.pkl.core.parser.ParserVisitor;
 import org.pkl.core.parser.Span;
 import org.pkl.core.util.Nullable;
 
-@SuppressWarnings("ALL")
 public abstract sealed class StringPart extends AbstractNode {
 
   public StringPart(Span span, @Nullable List<? extends @Nullable Node> children) {
@@ -37,7 +36,9 @@ public abstract sealed class StringPart extends AbstractNode {
       super(span, parts);
     }
 
+    @SuppressWarnings("unchecked")
     public List<StringConstantPart> getParts() {
+      assert children != null;
       return (List<StringConstantPart>) children;
     }
   }
@@ -48,6 +49,7 @@ public abstract sealed class StringPart extends AbstractNode {
     }
 
     public Expr getExpr() {
+      assert children != null;
       return (Expr) children.get(0);
     }
   }
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/Type.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/Type.java
index aefc3916..e76255af 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/Type.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/Type.java
@@ -15,13 +15,13 @@
  */
 package org.pkl.core.parser.ast;
 
+import java.util.Arrays;
 import java.util.List;
 import java.util.Objects;
 import org.pkl.core.parser.ParserVisitor;
 import org.pkl.core.parser.Span;
 import org.pkl.core.util.Nullable;
 
-@SuppressWarnings("ALL")
 public abstract sealed class Type extends AbstractNode {
 
   public Type(Span span, @Nullable List<? extends @Nullable Node> children) {
@@ -34,7 +34,7 @@ public abstract sealed class Type extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitUnknownType(this);
     }
   }
@@ -45,7 +45,7 @@ public abstract sealed class Type extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitNothingType(this);
     }
   }
@@ -56,7 +56,7 @@ public abstract sealed class Type extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitModuleType(this);
     }
   }
@@ -67,31 +67,34 @@ public abstract sealed class Type extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitStringConstantType(this);
     }
 
     public StringConstant getStr() {
+      assert children != null;
       return (StringConstant) children.get(0);
     }
   }
 
   public static final class DeclaredType extends Type {
-    public DeclaredType(List<Node> nodes, Span span) {
-      super(span, nodes);
+    public DeclaredType(QualifiedIdentifier name, @Nullable TypeArgumentList typeArgs, Span span) {
+      super(span, Arrays.asList(name, typeArgs));
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitDeclaredType(this);
     }
 
     public QualifiedIdentifier getName() {
+      assert children != null;
       return (QualifiedIdentifier) children.get(0);
     }
 
-    public List<Type> getArgs() {
-      return (List<Type>) children.subList(1, children.size());
+    public @Nullable TypeArgumentList getArgs() {
+      assert children != null;
+      return (TypeArgumentList) children.get(1);
     }
   }
 
@@ -101,11 +104,12 @@ public abstract sealed class Type extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitParenthesizedType(this);
     }
 
     public Type getType() {
+      assert children != null;
       return (Type) children.get(0);
     }
   }
@@ -116,11 +120,12 @@ public abstract sealed class Type extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitNullableType(this);
     }
 
     public Type getType() {
+      assert children != null;
       return (Type) children.get(0);
     }
   }
@@ -131,15 +136,18 @@ public abstract sealed class Type extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitConstrainedType(this);
     }
 
     public Type getType() {
+      assert children != null;
       return (Type) children.get(0);
     }
 
+    @SuppressWarnings("unchecked")
     public List<Expr> getExprs() {
+      assert children != null;
       return (List<Expr>) children.subList(1, children.size());
     }
   }
@@ -153,11 +161,13 @@ public abstract sealed class Type extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitUnionType(this);
     }
 
+    @SuppressWarnings("unchecked")
     public List<Type> getTypes() {
+      assert children != null;
       return (List<Type>) children;
     }
 
@@ -177,6 +187,7 @@ public abstract sealed class Type extends AbstractNode {
           + '}';
     }
 
+    @SuppressWarnings("ConstantValue")
     @Override
     public boolean equals(Object o) {
       if (this == o) {
@@ -204,15 +215,18 @@ public abstract sealed class Type extends AbstractNode {
     }
 
     @Override
-    public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+    public <T> T accept(ParserVisitor<? extends T> visitor) {
       return visitor.visitFunctionType(this);
     }
 
+    @SuppressWarnings("unchecked")
     public List<Type> getArgs() {
+      assert children != null;
       return (List<Type>) children.subList(0, children.size() - 1);
     }
 
     public Type getRet() {
+      assert children != null;
       return (Type) children.get(children.size() - 1);
     }
   }
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/TypeAlias.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/TypeAlias.java
index be0f8246..14f90c13 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/TypeAlias.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/TypeAlias.java
@@ -20,7 +20,6 @@ import org.pkl.core.parser.ParserVisitor;
 import org.pkl.core.parser.Span;
 import org.pkl.core.util.Nullable;
 
-@SuppressWarnings("ALL")
 public final class TypeAlias extends AbstractNode {
   private final int modifiersOffset;
   private final int nameOffset;
@@ -32,40 +31,64 @@ public final class TypeAlias extends AbstractNode {
   }
 
   @Override
-  public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+  public <T> T accept(ParserVisitor<? extends T> visitor) {
     return visitor.visitTypeAlias(this);
   }
 
   public @Nullable DocComment getDocComment() {
+    assert children != null;
     return (DocComment) children.get(0);
   }
 
+  @SuppressWarnings("unchecked")
   public List<Annotation> getAnnotations() {
+    assert children != null;
     return (List<Annotation>) children.subList(1, modifiersOffset);
   }
 
+  @SuppressWarnings("unchecked")
   public List<Modifier> getModifiers() {
+    assert children != null;
     return (List<Modifier>) children.subList(modifiersOffset, nameOffset);
   }
 
+  public Keyword getTypealiasKeyword() {
+    assert children != null;
+    return (Keyword) children.get(nameOffset);
+  }
+
   public Identifier getName() {
-    return (Identifier) children.get(nameOffset);
+    assert children != null;
+    return (Identifier) children.get(nameOffset + 1);
   }
 
   public @Nullable TypeParameterList getTypeParameterList() {
-    return (TypeParameterList) children.get(nameOffset + 1);
+    assert children != null;
+    return (TypeParameterList) children.get(nameOffset + 2);
   }
 
   public Type getType() {
-    return (Type) children.get(nameOffset + 2);
+    assert children != null;
+    return (Type) children.get(nameOffset + 3);
   }
 
+  @SuppressWarnings("DuplicatedCode")
   public Span getHeaderSpan() {
-    var end = children.get(nameOffset).span();
-    var tparList = children.get(nameOffset + 1);
+    Span start = null;
+    assert children != null;
+    for (var i = modifiersOffset; i < children.size(); i++) {
+      var child = children.get(i);
+      if (child != null) {
+        start = child.span();
+        break;
+      }
+    }
+    var end = children.get(nameOffset + 1).span();
+    var tparList = children.get(nameOffset + 2);
     if (tparList != null) {
       end = tparList.span();
     }
-    return span.endWith(end);
+    assert start != null;
+    return start.endWith(end);
   }
 }
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/TypeArgumentList.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/TypeArgumentList.java
new file mode 100644
index 00000000..563b54fa
--- /dev/null
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/TypeArgumentList.java
@@ -0,0 +1,38 @@
+/*
+ * Copyright © 2025 Apple Inc. and the Pkl project authors. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.pkl.core.parser.ast;
+
+import java.util.List;
+import org.pkl.core.parser.ParserVisitor;
+import org.pkl.core.parser.Span;
+
+public class TypeArgumentList extends AbstractNode {
+
+  public TypeArgumentList(List<Type> children, Span span) {
+    super(span, children);
+  }
+
+  @SuppressWarnings("unchecked")
+  public List<Type> getTypes() {
+    assert children != null;
+    return (List<Type>) children;
+  }
+
+  @Override
+  public <T> T accept(ParserVisitor<? extends T> visitor) {
+    return visitor.visitTypeArgumentList(this);
+  }
+}
diff --git a/pkl-core/src/main/java/org/pkl/core/parser/ast/TypeParameter.java b/pkl-core/src/main/java/org/pkl/core/parser/ast/TypeParameter.java
index 098111ce..312d06a4 100644
--- a/pkl-core/src/main/java/org/pkl/core/parser/ast/TypeParameter.java
+++ b/pkl-core/src/main/java/org/pkl/core/parser/ast/TypeParameter.java
@@ -30,7 +30,7 @@ public final class TypeParameter extends AbstractNode {
   }
 
   @Override
-  public <T> @Nullable T accept(ParserVisitor<? extends T> visitor) {
+  public <T> T accept(ParserVisitor<? extends T> visitor) {
     return visitor.visitTypeParameter(this);
   }
 
@@ -55,6 +55,7 @@ public final class TypeParameter extends AbstractNode {
         + '}';
   }
 
+  @SuppressWarnings("ConstantValue")
   @Override
   public boolean equals(Object o) {
     if (this == o) {
diff --git a/pkl-core/src/test/files/LanguageSnippetTests/output/projects/badPklProject1/bug.err b/pkl-core/src/test/files/LanguageSnippetTests/output/projects/badPklProject1/bug.err
index f3e580d3..8dc07f5c 100644
--- a/pkl-core/src/test/files/LanguageSnippetTests/output/projects/badPklProject1/bug.err
+++ b/pkl-core/src/test/files/LanguageSnippetTests/output/projects/badPklProject1/bug.err
@@ -2,7 +2,7 @@
 Expected `output.value` of module `file:///$snippetsDir/input/projects/badPklProject1/PklProject` to be of type `pkl.Project`, but got type `invalid.project.Module`.
 
 x | module invalid.project.Module
-           ^^^^^^^^^^^^^^^^^^^^^^
+    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 at invalid.project.Module (file:///$snippetsDir/input/projects/badPklProject1/PklProject)
 
 Try adding `amends "pkl:Project"` to the module header.
diff --git a/pkl-core/src/test/kotlin/org/pkl/core/parser/ANTLRSexpRenderer.kt b/pkl-core/src/test/kotlin/org/pkl/core/parser/ANTLRSexpRenderer.kt
index 0f6493bd..a330359a 100644
--- a/pkl-core/src/test/kotlin/org/pkl/core/parser/ANTLRSexpRenderer.kt
+++ b/pkl-core/src/test/kotlin/org/pkl/core/parser/ANTLRSexpRenderer.kt
@@ -356,10 +356,20 @@ class ANTLRSexpRenderer {
     renderQualifiedIdent(type.qualifiedIdentifier())
     val args = type.typeArgumentList()
     if (args != null) {
-      for (arg in args.type()) {
-        buf.append('\n')
-        renderType(arg)
-      }
+      buf.append('\n')
+      renderTypeArgumentList(args)
+    }
+    buf.append(')')
+    tab = oldTab
+  }
+
+  fun renderTypeArgumentList(ctx: TypeArgumentListContext) {
+    buf.append(tab)
+    buf.append("(typeArgumentList")
+    val oldTab = increaseTab()
+    for (arg in ctx.type()) {
+      buf.append('\n')
+      renderType(arg)
     }
     buf.append(')')
     tab = oldTab
@@ -422,21 +432,6 @@ class ANTLRSexpRenderer {
     tab = oldTab
   }
 
-  private fun flattenUnion(type: UnionTypeContext): List<TypeContext> {
-    val types = mutableListOf<TypeContext>()
-    val toCheck = mutableListOf(type.l, type.r)
-    while (toCheck.isNotEmpty()) {
-      val typ = toCheck.removeAt(0)
-      if (typ is UnionTypeContext) {
-        toCheck.add(0, typ.r)
-        toCheck.add(0, typ.l)
-      } else {
-        types += typ
-      }
-    }
-    return types
-  }
-
   fun renderFunctionType(type: FunctionTypeContext) {
     buf.append(tab)
     buf.append("(functionType")
@@ -1054,5 +1049,20 @@ class ANTLRSexpRenderer {
       res.sortWith(compareBy { it.sourceInterval.a })
       return res
     }
+
+    fun flattenUnion(type: UnionTypeContext): List<TypeContext> {
+      val types = mutableListOf<TypeContext>()
+      val toCheck = mutableListOf(type.l, type.r)
+      while (toCheck.isNotEmpty()) {
+        val typ = toCheck.removeAt(0)
+        if (typ is UnionTypeContext) {
+          toCheck.add(0, typ.r)
+          toCheck.add(0, typ.l)
+        } else {
+          types += typ
+        }
+      }
+      return types
+    }
   }
 }
diff --git a/pkl-core/src/test/kotlin/org/pkl/core/parser/ParserComparisonTestInterface.kt b/pkl-core/src/test/kotlin/org/pkl/core/parser/ParserComparisonTestInterface.kt
index 0b2b132f..b10f2d0f 100644
--- a/pkl-core/src/test/kotlin/org/pkl/core/parser/ParserComparisonTestInterface.kt
+++ b/pkl-core/src/test/kotlin/org/pkl/core/parser/ParserComparisonTestInterface.kt
@@ -47,6 +47,23 @@ interface ParserComparisonTestInterface {
     }
   }
 
+  @Test
+  @Execution(ExecutionMode.CONCURRENT)
+  fun compareSnippetTestsSpans() {
+    SoftAssertions.assertSoftly { softly ->
+      getSnippets()
+        .parallelStream()
+        .map { Pair(it.pathString, it.readText()) }
+        .forEach { (path, snippet) ->
+          try {
+            compareSpans(snippet, path, softly)
+          } catch (e: ParserError) {
+            softly.fail("path: $path. Message: ${e.message}", e)
+          }
+        }
+    }
+  }
+
   fun getSnippets(): List<Path>
 
   fun compare(code: String, path: String? = null, softly: SoftAssertions? = null) {
@@ -58,6 +75,22 @@ interface ParserComparisonTestInterface {
     }
   }
 
+  fun compareSpans(code: String, path: String, softly: SoftAssertions) {
+    // Our ANTLR grammar always start doc comment spans in the beginning of the line,
+    // even though they may have leading spaces.
+    // This is a regression, but it's a bugfix
+    if (path.endsWith("annotation1.pkl")) return
+
+    val parser = Parser()
+    val mod = parser.parseModule(code)
+    val lexer = PklLexer(ANTLRInputStream(code))
+    val antlr = PklParser(CommonTokenStream(lexer))
+    val antlrMod = antlr.module()
+
+    val comparer = SpanComparison(path, softly)
+    comparer.compare(mod, antlrMod)
+  }
+
   fun renderBoth(code: String): Pair<String, String> = Pair(renderCode(code), renderANTLRCode(code))
 
   companion object {
diff --git a/pkl-core/src/test/kotlin/org/pkl/core/parser/SexpRenderer.kt b/pkl-core/src/test/kotlin/org/pkl/core/parser/SexpRenderer.kt
index 9250c192..a2b24175 100644
--- a/pkl-core/src/test/kotlin/org/pkl/core/parser/SexpRenderer.kt
+++ b/pkl-core/src/test/kotlin/org/pkl/core/parser/SexpRenderer.kt
@@ -763,7 +763,19 @@ class SexpRenderer {
     val oldTab = increaseTab()
     buf.append('\n')
     renderQualifiedIdent(type.name)
-    for (arg in type.args) {
+    if (type.args !== null) {
+      buf.append('\n')
+      renderTypeArgumentList(type.args!!)
+    }
+    buf.append(')')
+    tab = oldTab
+  }
+
+  fun renderTypeArgumentList(typeArgumentList: TypeArgumentList) {
+    buf.append(tab)
+    buf.append("(typeArgumentList")
+    val oldTab = increaseTab()
+    for (arg in typeArgumentList.types) {
       buf.append('\n')
       renderType(arg)
     }
@@ -976,9 +988,9 @@ class SexpRenderer {
     buf.append("(whenGenerator")
     val oldTab = increaseTab()
     buf.append('\n')
-    renderExpr(generator.thenClause)
+    renderExpr(generator.predicate)
     buf.append('\n')
-    renderObjectBody(generator.body)
+    renderObjectBody(generator.thenClause)
     if (generator.elseClause !== null) {
       buf.append('\n')
       renderObjectBody(generator.elseClause!!)
diff --git a/pkl-core/src/test/kotlin/org/pkl/core/parser/SpanComparison.kt b/pkl-core/src/test/kotlin/org/pkl/core/parser/SpanComparison.kt
new file mode 100644
index 00000000..269beae7
--- /dev/null
+++ b/pkl-core/src/test/kotlin/org/pkl/core/parser/SpanComparison.kt
@@ -0,0 +1,463 @@
+/*
+ * Copyright © 2025 Apple Inc. and the Pkl project authors. All rights reserved.
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ *     https://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+package org.pkl.core.parser
+
+import org.antlr.v4.runtime.ParserRuleContext
+import org.antlr.v4.runtime.tree.TerminalNode
+import org.assertj.core.api.SoftAssertions
+import org.pkl.core.parser.antlr.PklParser.*
+import org.pkl.core.parser.ast.*
+import org.pkl.core.parser.ast.Annotation
+import org.pkl.core.parser.ast.Expr.AmendsExpr
+import org.pkl.core.parser.ast.Expr.BinaryOperatorExpr
+import org.pkl.core.parser.ast.Expr.FunctionLiteralExpr
+import org.pkl.core.parser.ast.Expr.IfExpr
+import org.pkl.core.parser.ast.Expr.ImportExpr
+import org.pkl.core.parser.ast.Expr.LetExpr
+import org.pkl.core.parser.ast.Expr.LogicalNotExpr
+import org.pkl.core.parser.ast.Expr.MultiLineStringLiteralExpr
+import org.pkl.core.parser.ast.Expr.NewExpr
+import org.pkl.core.parser.ast.Expr.NonNullExpr
+import org.pkl.core.parser.ast.Expr.ParenthesizedExpr
+import org.pkl.core.parser.ast.Expr.QualifiedAccessExpr
+import org.pkl.core.parser.ast.Expr.ReadExpr
+import org.pkl.core.parser.ast.Expr.SingleLineStringLiteralExpr
+import org.pkl.core.parser.ast.Expr.SubscriptExpr
+import org.pkl.core.parser.ast.Expr.SuperAccessExpr
+import org.pkl.core.parser.ast.Expr.SuperSubscriptExpr
+import org.pkl.core.parser.ast.Expr.ThrowExpr
+import org.pkl.core.parser.ast.Expr.TraceExpr
+import org.pkl.core.parser.ast.Expr.TypeCastExpr
+import org.pkl.core.parser.ast.Expr.TypeCheckExpr
+import org.pkl.core.parser.ast.Expr.UnaryMinusExpr
+import org.pkl.core.parser.ast.Expr.UnqualifiedAccessExpr
+import org.pkl.core.parser.ast.ObjectMember.ForGenerator
+import org.pkl.core.parser.ast.ObjectMember.MemberPredicate
+import org.pkl.core.parser.ast.ObjectMember.ObjectElement
+import org.pkl.core.parser.ast.ObjectMember.ObjectEntry
+import org.pkl.core.parser.ast.ObjectMember.ObjectMethod
+import org.pkl.core.parser.ast.ObjectMember.ObjectProperty
+import org.pkl.core.parser.ast.ObjectMember.ObjectSpread
+import org.pkl.core.parser.ast.ObjectMember.WhenGenerator
+import org.pkl.core.parser.ast.Parameter.TypedIdentifier
+import org.pkl.core.parser.ast.Type.ConstrainedType
+import org.pkl.core.parser.ast.Type.DeclaredType
+import org.pkl.core.parser.ast.Type.FunctionType
+import org.pkl.core.parser.ast.Type.NullableType
+import org.pkl.core.parser.ast.Type.ParenthesizedType
+import org.pkl.core.parser.ast.Type.StringConstantType
+import org.pkl.core.parser.ast.Type.UnionType
+
+class SpanComparison(val path: String, private val softly: SoftAssertions) {
+
+  fun compare(module: Module, modCtx: ModuleContext) {
+    compareSpan(module, modCtx)
+    if (module.decl !== null) {
+      compareModuleDecl(module.decl!!, modCtx.moduleDecl())
+    }
+    module.imports.zip(modCtx.importClause()).forEach { (i1, i2) -> compareImport(i1, i2) }
+    module.classes.zip(modCtx.clazz()).forEach { (class1, class2) -> compareClass(class1, class2) }
+    module.typeAliases.zip(modCtx.typeAlias()).forEach { (ta1, ta2) -> compareTypealias(ta1, ta2) }
+    module.properties.zip(modCtx.classProperty()).forEach { (prop1, prop2) ->
+      compareProperty(prop1, prop2)
+    }
+    module.methods.zip(modCtx.classMethod()).forEach { (m1, m2) -> compareMethod(m1, m2) }
+  }
+
+  private fun compareModuleDecl(node: ModuleDecl, ctx: ModuleDeclContext) {
+    compareSpan(node, ctx)
+    compareDocComment(node.docComment, ctx.DocComment())
+    node.annotations.zip(ctx.annotation()).forEach { (a1, a2) -> compareAnnotation(a1, a2) }
+    val header = ctx.moduleHeader()
+    node.modifiers.zip(header.modifier()).forEach { (m1, m2) -> compareSpan(m1, m2) }
+    compareQualifiedIdentifier(node.name, header.qualifiedIdentifier())
+    compareExtendsOrAmendsClause(node.extendsOrAmendsDecl, header.moduleExtendsOrAmendsClause())
+  }
+
+  private fun compareImport(node: ImportClause, ctx: ImportClauseContext) {
+    compareSpan(node, ctx)
+    compareSpan(node.importStr, ctx.stringConstant())
+    compareSpan(node.alias, ctx.Identifier())
+  }
+
+  private fun compareClass(node: Class, ctx: ClazzContext) {
+    compareSpan(node, ctx)
+    compareDocComment(node.docComment, ctx.DocComment())
+    node.annotations.zip(ctx.annotation()).forEach { (a1, a2) -> compareAnnotation(a1, a2) }
+    val header = ctx.classHeader()
+    node.modifiers.zip(header.modifier()).forEach { (m1, m2) -> compareSpan(m1, m2) }
+    compareSpan(node.classKeyword, header.CLASS())
+    compareSpan(node.name, header.Identifier())
+    compareTypeParameterList(node.typeParameterList, header.typeParameterList())
+    compareType(node.superClass, header.type())
+    compareClassBody(node.body, ctx.classBody())
+    compareSpan(node.headerSpan, header)
+  }
+
+  private fun compareTypealias(node: TypeAlias, ctx: TypeAliasContext) {
+    compareSpan(node, ctx)
+    compareDocComment(node.docComment, ctx.DocComment())
+    node.annotations.zip(ctx.annotation()).forEach { (a1, a2) -> compareAnnotation(a1, a2) }
+    val header = ctx.typeAliasHeader()
+    node.modifiers.zip(header.modifier()).forEach { (m1, m2) -> compareSpan(m1, m2) }
+    compareSpan(node.typealiasKeyword, header.TYPE_ALIAS())
+    compareSpan(node.name, header.Identifier())
+    compareTypeParameterList(node.typeParameterList, header.typeParameterList())
+    compareType(node.type, ctx.type())
+    compareSpan(node.headerSpan, header)
+  }
+
+  private fun compareProperty(node: ClassProperty, ctx: ClassPropertyContext) {
+    if (node.docComment === null) {
+      compareSpan(node, ctx)
+    }
+    compareDocComment(node.docComment, ctx.DocComment())
+    node.annotations.zip(ctx.annotation()).forEach { (a1, a2) -> compareAnnotation(a1, a2) }
+    node.modifiers.zip(ctx.modifier()).forEach { (m1, m2) -> compareSpan(m1, m2) }
+    compareSpan(node.name, ctx.Identifier())
+    compareTypeAnnotation(node.typeAnnotation, ctx.typeAnnotation())
+    compareExpr(node.expr, ctx.expr())
+    node.bodyList.zip(ctx.objectBody()).forEach { (b1, b2) -> compareObjectBody(b1, b2) }
+  }
+
+  private fun compareMethod(node: ClassMethod, ctx: ClassMethodContext) {
+    compareSpan(node, ctx)
+    compareDocComment(node.docComment, ctx.DocComment())
+    node.annotations.zip(ctx.annotation()).forEach { (a1, a2) -> compareAnnotation(a1, a2) }
+    val header = ctx.methodHeader()
+    compareSpan(node.headerSpan, header)
+    node.modifiers.zip(header.modifier()).forEach { (m1, m2) -> compareSpan(m1, m2) }
+    compareSpan(node.name, header.Identifier())
+    compareTypeParameterList(node.typeParameterList, header.typeParameterList())
+    compareParameterList(node.parameterList, header.parameterList())
+    compareTypeAnnotation(node.typeAnnotation, header.typeAnnotation())
+    compareExpr(node.expr, ctx.expr())
+  }
+
+  private fun compareAnnotation(node: Annotation, ctx: AnnotationContext) {
+    compareSpan(node, ctx)
+    compareType(node.type, ctx.type())
+    compareObjectBody(node.body, ctx.objectBody())
+  }
+
+  private fun compareParameterList(node: ParameterList?, ctx: ParameterListContext?) {
+    if (node === null) return
+    compareSpan(node, ctx!!)
+    node.parameters.zip(ctx.parameter()).forEach { (p1, p2) -> compareParameter(p1, p2) }
+  }
+
+  private fun compareParameter(node: Parameter?, ctx: ParameterContext?) {
+    if (node === null) return
+    compareSpan(node, ctx!!)
+    if (node is TypedIdentifier) {
+      val tident = ctx.typedIdentifier()
+      compareSpan(node.identifier, tident.Identifier())
+      compareTypeAnnotation(node.typeAnnotation, tident.typeAnnotation())
+    }
+  }
+
+  private fun compareTypeParameterList(node: TypeParameterList?, ctx: TypeParameterListContext?) {
+    if (node === null) return
+    compareSpan(node, ctx!!)
+    node.parameters.zip(ctx.typeParameter()).forEach { (p1, p2) -> compareTypeParameter(p1, p2) }
+  }
+
+  private fun compareTypeParameter(node: TypeParameter, ctx: TypeParameterContext) {
+    compareSpan(node, ctx)
+    compareSpan(node.identifier, ctx.Identifier())
+  }
+
+  private fun compareClassBody(node: ClassBody?, ctx: ClassBodyContext?) {
+    if (node === null) return
+    compareSpan(node, ctx!!)
+    node.properties.zip(ctx.classProperty()).forEach { (p1, p2) -> compareProperty(p1, p2) }
+    node.methods.zip(ctx.classMethod()).forEach { (m1, m2) -> compareMethod(m1, m2) }
+  }
+
+  private fun compareTypeAnnotation(node: TypeAnnotation?, ctx: TypeAnnotationContext?) {
+    if (node === null) return
+    compareSpan(node, ctx!!)
+    compareSpan(node.type, ctx.type())
+  }
+
+  private fun compareObjectBody(node: ObjectBody?, ctx: ObjectBodyContext?) {
+    if (node === null) return
+    compareSpan(node, ctx!!)
+    node.parameters.zip(ctx.parameter()).forEach { (p1, p2) -> compareParameter(p1, p2) }
+    node.members.zip(ctx.objectMember()).forEach { (m1, m2) -> compareObjectMember(m1, m2) }
+  }
+
+  private fun compareType(node: Type?, actx: TypeContext?) {
+    if (node === null) return
+    val ctx = if (actx is DefaultUnionTypeContext) actx.type() else actx
+    compareSpan(node, ctx!!)
+    when (node) {
+      is StringConstantType ->
+        compareSpan(node.str, (ctx as StringLiteralTypeContext).stringConstant())
+      is DeclaredType -> {
+        val decl = ctx as DeclaredTypeContext
+        compareQualifiedIdentifier(node.name, decl.qualifiedIdentifier())
+        compareTypeArgumentList(node.args, ctx.typeArgumentList())
+      }
+      is ParenthesizedType -> compareType(node.type, (ctx as ParenthesizedTypeContext).type())
+      is NullableType -> compareType(node.type, (ctx as NullableTypeContext).type())
+      is ConstrainedType -> {
+        val cons = ctx as ConstrainedTypeContext
+        compareType(node.type, cons.type())
+        node.exprs.zip(cons.expr()).forEach { (e1, e2) -> compareExpr(e1, e2) }
+      }
+      is UnionType -> {
+        val flattened = ANTLRSexpRenderer.flattenUnion(ctx as UnionTypeContext)
+        node.types.zip(flattened).forEach { (t1, t2) -> compareType(t1, t2) }
+      }
+      is FunctionType -> {
+        val func = ctx as FunctionTypeContext
+        node.args.zip(func.ps).forEach { (t1, t2) -> compareType(t1, t2) }
+        compareType(node.ret, func.r)
+      }
+      else -> {}
+    }
+  }
+
+  private fun compareObjectMember(node: ObjectMember, ctx: ObjectMemberContext) {
+    compareSpan(node, ctx)
+    when (node) {
+      is ObjectElement -> compareExpr(node.expr, (ctx as ObjectElementContext).expr())
+      is ObjectProperty -> {
+        ctx as ObjectPropertyContext
+        node.modifiers.zip(ctx.modifier()).forEach { (m1, m2) -> compareSpan(m1, m2) }
+        compareSpan(node.identifier, ctx.Identifier())
+        compareTypeAnnotation(node.typeAnnotation, ctx.typeAnnotation())
+        compareExpr(node.expr, ctx.expr())
+        if (node.bodyList.isNotEmpty()) {
+          node.bodyList.zip(ctx.objectBody()).forEach { (b1, b2) -> compareObjectBody(b1, b2) }
+        }
+      }
+      is ObjectMethod -> {
+        ctx as ObjectMethodContext
+        val header = ctx.methodHeader()
+        node.modifiers.zip(header.modifier()).forEach { (m1, m2) -> compareSpan(m1, m2) }
+        compareSpan(node.functionKeyword, header.FUNCTION())
+        compareSpan(node.identifier, header.Identifier())
+        compareTypeParameterList(node.typeParameterList, header.typeParameterList())
+        compareParameterList(node.paramList, header.parameterList())
+        compareTypeAnnotation(node.typeAnnotation, header.typeAnnotation())
+        compareExpr(node.expr, ctx.expr())
+        compareSpan(node.headerSpan(), header)
+      }
+      is MemberPredicate -> {
+        ctx as MemberPredicateContext
+        compareExpr(node.pred, ctx.k)
+        compareExpr(node.expr, ctx.v)
+        if (node.bodyList.isNotEmpty()) {
+          node.bodyList.zip(ctx.objectBody()).forEach { (b1, b2) -> compareObjectBody(b1, b2) }
+        }
+      }
+      is ObjectEntry -> {
+        ctx as ObjectEntryContext
+        compareExpr(node.key, ctx.k)
+        compareExpr(node.value, ctx.v)
+        if (node.bodyList.isNotEmpty()) {
+          node.bodyList.zip(ctx.objectBody()).forEach { (b1, b2) -> compareObjectBody(b1, b2) }
+        }
+      }
+      is ObjectSpread -> compareExpr(node.expr, (ctx as ObjectSpreadContext).expr())
+      is WhenGenerator -> {
+        ctx as WhenGeneratorContext
+        compareExpr(node.predicate, ctx.expr())
+        compareObjectBody(node.thenClause, ctx.b1)
+        compareObjectBody(node.elseClause, ctx.b2)
+      }
+      is ForGenerator -> {
+        ctx as ForGeneratorContext
+        compareParameter(node.p1, ctx.t1)
+        compareParameter(node.p2, ctx.t2)
+        compareExpr(node.expr, ctx.expr())
+        compareObjectBody(node.body, ctx.objectBody())
+      }
+    }
+  }
+
+  private fun compareExpr(node: Expr?, ctx: ExprContext?) {
+    if (node === null) return
+    compareSpan(node, ctx!!)
+    when (node) {
+      is SingleLineStringLiteralExpr -> {
+        node.parts.zip((ctx as SingleLineStringLiteralContext).singleLineStringPart()).forEach {
+          (s1, s2) ->
+          compareSpan(s1, s2)
+        }
+      }
+      is MultiLineStringLiteralExpr -> {
+        node.parts.zip((ctx as MultiLineStringLiteralContext).multiLineStringPart()).forEach {
+          (s1, s2) ->
+          compareSpan(s1, s2)
+        }
+      }
+      is ThrowExpr -> compareExpr(node.expr, (ctx as ThrowExprContext).expr())
+      is TraceExpr -> compareExpr(node.expr, (ctx as TraceExprContext).expr())
+      is ImportExpr -> compareSpan(node.importStr, (ctx as ImportExprContext).stringConstant())
+      is ReadExpr -> compareExpr(node.expr, (ctx as ReadExprContext).expr())
+      is UnqualifiedAccessExpr -> {
+        ctx as UnqualifiedAccessExprContext
+        compareSpan(node.identifier, ctx.Identifier())
+        compareArgumentList(node.argumentList, ctx.argumentList())
+      }
+      is QualifiedAccessExpr -> {
+        ctx as QualifiedAccessExprContext
+        compareExpr(node.expr, ctx.expr())
+        compareSpan(node.identifier, ctx.Identifier())
+        compareArgumentList(node.argumentList, ctx.argumentList())
+      }
+      is SuperAccessExpr -> {
+        ctx as SuperAccessExprContext
+        compareSpan(node.identifier, ctx.Identifier())
+        compareArgumentList(node.argumentList, ctx.argumentList())
+      }
+      is SuperSubscriptExpr -> compareExpr(node.arg, (ctx as SuperSubscriptExprContext).expr())
+      is SubscriptExpr -> {
+        ctx as SubscriptExprContext
+        compareExpr(node.expr, ctx.l)
+        compareExpr(node.arg, ctx.r)
+      }
+      is IfExpr -> {
+        ctx as IfExprContext
+        compareExpr(node.cond, ctx.c)
+        compareExpr(node.then, ctx.l)
+        compareExpr(node.els, ctx.r)
+      }
+      is LetExpr -> {
+        ctx as LetExprContext
+        compareParameter(node.parameter, ctx.parameter())
+        compareExpr(node.bindingExpr, ctx.l)
+        compareExpr(node.expr, ctx.r)
+      }
+      is FunctionLiteralExpr -> {
+        ctx as FunctionLiteralContext
+        compareParameterList(node.parameterList, ctx.parameterList())
+        compareExpr(node.expr, ctx.expr())
+      }
+      is ParenthesizedExpr -> compareExpr(node.expr, (ctx as ParenthesizedExprContext).expr())
+      is NewExpr -> {
+        ctx as NewExprContext
+        compareType(node.type, ctx.type())
+        compareObjectBody(node.body, ctx.objectBody())
+      }
+      is AmendsExpr -> {
+        ctx as AmendExprContext
+        compareExpr(node.expr, ctx.expr())
+        compareObjectBody(node.body, ctx.objectBody())
+      }
+      is NonNullExpr -> compareExpr(node.expr, (ctx as NonNullExprContext).expr())
+      is UnaryMinusExpr -> compareExpr(node.expr, (ctx as UnaryMinusExprContext).expr())
+      is LogicalNotExpr -> compareExpr(node.expr, (ctx as LogicalNotExprContext).expr())
+      is BinaryOperatorExpr -> {
+        val (l, r) =
+          when (ctx) {
+            is ExponentiationExprContext -> ctx.l to ctx.r
+            is MultiplicativeExprContext -> ctx.l to ctx.r
+            is AdditiveExprContext -> ctx.l to ctx.r
+            is ComparisonExprContext -> ctx.l to ctx.r
+            is EqualityExprContext -> ctx.l to ctx.r
+            is LogicalAndExprContext -> ctx.l to ctx.r
+            is LogicalOrExprContext -> ctx.l to ctx.r
+            is PipeExprContext -> ctx.l to ctx.r
+            is NullCoalesceExprContext -> ctx.l to ctx.r
+            else -> throw RuntimeException("unreacheable code")
+          }
+        compareExpr(node.left, l)
+        compareExpr(node.right, r)
+      }
+      is TypeCheckExpr -> {
+        ctx as TypeTestExprContext
+        compareExpr(node.expr, ctx.expr())
+        compareType(node.type, ctx.type())
+      }
+      is TypeCastExpr -> {
+        ctx as TypeTestExprContext
+        compareExpr(node.expr, ctx.expr())
+        compareType(node.type, ctx.type())
+      }
+      else -> {}
+    }
+  }
+
+  private fun compareArgumentList(node: ArgumentList?, ctx: ArgumentListContext?) {
+    if (node === null) return
+    compareSpan(node, ctx!!)
+    node.arguments.zip(ctx.expr()).forEach { (e1, e2) -> compareExpr(e1, e2) }
+  }
+
+  private fun compareTypeArgumentList(node: TypeArgumentList?, ctx: TypeArgumentListContext?) {
+    if (node === null) return
+    compareSpan(node, ctx!!)
+    node.types.zip(ctx.type()).forEach { (t1, t2) -> compareType(t1, t2) }
+  }
+
+  private fun compareQualifiedIdentifier(
+    node: QualifiedIdentifier?,
+    ctx: QualifiedIdentifierContext?,
+  ) {
+    if (node === null) return
+    compareSpan(node, ctx!!)
+    node.identifiers.zip(ctx.Identifier()).forEach { (id1, id2) -> compareSpan(id1, id2) }
+  }
+
+  private fun compareExtendsOrAmendsClause(
+    node: ExtendsOrAmendsClause?,
+    ctx: ModuleExtendsOrAmendsClauseContext?,
+  ) {
+    if (node === null) return
+    compareSpan(node, ctx!!)
+    compareSpan(node.url, ctx.stringConstant())
+  }
+
+  private fun compareSpan(node: Node?, ctx: ParserRuleContext) {
+    if (node != null) {
+      compareSpan(node, ctx.start.startIndex, ctx.stop.stopIndex)
+    }
+  }
+
+  private fun compareSpan(node: Node?, ctx: TerminalNode?) {
+    if (node != null) {
+      compareSpan(node, ctx!!.symbol.startIndex, ctx.symbol.stopIndex)
+    }
+  }
+
+  private fun compareSpan(node: Node, charIndex: Int, tokenStop: Int) {
+    compareSpan(node.span(), charIndex, tokenStop)
+  }
+
+  private fun compareSpan(span: Span, ctx: ParserRuleContext) {
+    compareSpan(span, ctx.start.startIndex, ctx.stop.stopIndex)
+  }
+
+  private fun compareSpan(span: Span, charIndex: Int, tokenStop: Int) {
+    val length = tokenStop - charIndex + 1
+    softly.assertThat(span.charIndex).`as`("$span, index for path: $path").isEqualTo(charIndex)
+    softly.assertThat(span.length).`as`("$span, length for path: $path").isEqualTo(length)
+  }
+
+  private fun compareDocComment(node: Node?, ctx: TerminalNode?) {
+    if (node == null) return
+    val charIndex = ctx!!.symbol.startIndex
+    // for some reason antlr's doc coments are off by one
+    val length = ctx.symbol.stopIndex - charIndex
+    val span = node.span()
+    softly.assertThat(span.charIndex).`as`("$span, index for path: $path").isEqualTo(charIndex)
+    softly.assertThat(span.length).`as`("$span, length for path: $path").isEqualTo(length)
+  }
+}