From c80874b8d7734c7f50e150c408666e1725956b71 Mon Sep 17 00:00:00 2001
From: Knut Wannheden <knut@moderne.io>
Date: Wed, 15 Jan 2025 22:49:39 +0100
Subject: [PATCH] TOML model cleanups (#4907)

* TOML model cleanups

Added a few missing `@Nullable` annotations and added accessors for `Toml.Table#name`.

* Add two more overrides to `TomlIsoVisitor`
---
 .../org/openrewrite/toml/TomlIsoVisitor.java  | 10 ++++++++
 .../toml/internal/TomlPrinter.java            |  4 +--
 .../java/org/openrewrite/toml/tree/Toml.java  | 25 +++++++++++++++++--
 3 files changed, 35 insertions(+), 4 deletions(-)

diff --git a/rewrite-toml/src/main/java/org/openrewrite/toml/TomlIsoVisitor.java b/rewrite-toml/src/main/java/org/openrewrite/toml/TomlIsoVisitor.java
index fff819fd1af..850c76832a1 100644
--- a/rewrite-toml/src/main/java/org/openrewrite/toml/TomlIsoVisitor.java
+++ b/rewrite-toml/src/main/java/org/openrewrite/toml/TomlIsoVisitor.java
@@ -29,6 +29,11 @@ public Toml.Document visitDocument(Toml.Document document, P p) {
         return (Toml.Document) super.visitDocument(document, p);
     }
 
+    @Override
+    public Toml.Empty visitEmpty(Toml.Empty empty, P p) {
+        return (Toml.Empty) super.visitEmpty(empty, p);
+    }
+
     @Override
     public Toml.Identifier visitIdentifier(Toml.Identifier identifier, P p) {
         return (Toml.Identifier) super.visitIdentifier(identifier, p);
@@ -43,4 +48,9 @@ public Toml.KeyValue visitKeyValue(Toml.KeyValue keyValue, P p) {
     public Toml.Literal visitLiteral(Toml.Literal literal, P p) {
         return (Toml.Literal) super.visitLiteral(literal, p);
     }
+
+    @Override
+    public Toml.Table visitTable(Toml.Table table, P p) {
+        return (Toml.Table) super.visitTable(table, p);
+    }
 }
diff --git a/rewrite-toml/src/main/java/org/openrewrite/toml/internal/TomlPrinter.java b/rewrite-toml/src/main/java/org/openrewrite/toml/internal/TomlPrinter.java
index af61c294a17..b0bdeceed5f 100644
--- a/rewrite-toml/src/main/java/org/openrewrite/toml/internal/TomlPrinter.java
+++ b/rewrite-toml/src/main/java/org/openrewrite/toml/internal/TomlPrinter.java
@@ -103,12 +103,12 @@ public Toml visitTable(Toml.Table table, PrintOutputCapture<P> p) {
             p.append("}");
         } else if (table.getMarkers().findFirst(ArrayTable.class).isPresent()) {
             p.append("[[");
-            visitRightPadded(table.getName(), p);
+            visitRightPadded(table.getPadding().getName(), p);
             p.append("]]");
             visitRightPadded(table.getPadding().getValues(), "", p);
         } else {
             p.append("[");
-            visitRightPadded(table.getName(), p);
+            visitRightPadded(table.getPadding().getName(), p);
             p.append("]");
             visitRightPadded(table.getPadding().getValues(), "", p);
         }
diff --git a/rewrite-toml/src/main/java/org/openrewrite/toml/tree/Toml.java b/rewrite-toml/src/main/java/org/openrewrite/toml/tree/Toml.java
index 8d6037316fe..f429cd99f56 100644
--- a/rewrite-toml/src/main/java/org/openrewrite/toml/tree/Toml.java
+++ b/rewrite-toml/src/main/java/org/openrewrite/toml/tree/Toml.java
@@ -128,7 +128,11 @@ class Document implements Toml, SourceFile {
         String charsetName;
 
         boolean charsetBomMarked;
+
+        @Nullable
         Checksum checksum;
+
+        @Nullable
         FileAttributes fileAttributes;
 
         @Override
@@ -137,7 +141,8 @@ public Charset getCharset() {
         }
 
         @Override
-        public SourceFile withCharset(Charset charset) {
+        @SuppressWarnings("unchecked")
+        public Document withCharset(Charset charset) {
             return withCharsetName(charset.name());
         }
 
@@ -309,9 +314,17 @@ class Table implements TomlValue {
         @With
         Markers markers;
 
-        @With
+        @Nullable
         TomlRightPadded<Toml.Identifier> name;
 
+        public Toml.@Nullable Identifier getName() {
+            return name != null ? name.getElement() : null;
+        }
+
+        public Table withName(Toml.@Nullable Identifier name) {
+            return getPadding().withName(TomlRightPadded.withElement(this.name, name));
+        }
+
         List<TomlRightPadded<Toml>> values;
 
         public List<Toml> getValues() {
@@ -353,6 +366,14 @@ public List<TomlRightPadded<Toml>> getValues() {
             public Table withValues(List<TomlRightPadded<Toml>> values) {
                 return t.values == values ? t : new Table(t.id, t.prefix, t.markers, t.name, values);
             }
+
+            public @Nullable TomlRightPadded<Toml.Identifier> getName() {
+                return t.name;
+            }
+
+            public Table withName(@Nullable TomlRightPadded<Identifier> name) {
+                return t.name == name ? t : new Table(t.id, t.prefix, t.markers, name, t.values);
+            }
         }
     }
 }