Write thead and th for headers in table

18 files changed, 141 insertions, 80 deletions

diff --git a/src/main/java/org/zwobble/mammoth/internal/conversion/DocumentToHtml.java b/src/main/java/org/zwobble/mammoth/internal/conversion/DocumentToHtml.java
index bdb3d01..9de8224 100644
--- a/src/main/java/org/zwobble/mammoth/internal/conversion/DocumentToHtml.java
+++ b/src/main/java/org/zwobble/mammoth/internal/conversion/DocumentToHtml.java
@@ -15,6 +15,8 @@ import java.io.InputStream;
 import java.util.*;
 import java.util.function.Supplier;
 
+import static org.zwobble.mammoth.internal.util.Casts.tryCast;
+import static org.zwobble.mammoth.internal.util.Iterables.findIndex;
 import static org.zwobble.mammoth.internal.util.Lists.*;
 import static org.zwobble.mammoth.internal.util.Maps.lookup;
 import static org.zwobble.mammoth.internal.util.Maps.map;
@@ -23,7 +25,7 @@ public class DocumentToHtml {
     public static InternalResult<List<HtmlNode>> convertToHtml(Document document, DocumentToHtmlOptions options) {
         DocumentToHtml documentConverter = new DocumentToHtml(options, document.getComments());
         return new InternalResult<>(
-            documentConverter.convertToHtml(document),
+            documentConverter.convertToHtml(document, new Context()),
             documentConverter.warnings);
     }
 
@@ -37,7 +39,7 @@ public class DocumentToHtml {
     public static InternalResult<List<HtmlNode>> convertToHtml(DocumentElement element, DocumentToHtmlOptions options) {
         DocumentToHtml documentConverter = new DocumentToHtml(options, list());
         return new InternalResult<>(
-            documentConverter.convertToHtml(element),
+            documentConverter.convertToHtml(element, new Context()),
             documentConverter.warnings);
     }
 
@@ -60,6 +62,22 @@ public class DocumentToHtml {
     private final List<ReferencedComment> referencedComments = new ArrayList<>();
     private final Set<String> warnings = new HashSet<>();
 
+    private static class Context {
+        private final boolean isHeader;
+
+        Context() {
+            this(false);
+        }
+
+        Context(boolean isHeader) {
+            this.isHeader = isHeader;
+        }
+
+        Context isHeader(boolean isHeader) {
+            return new Context(isHeader);
+        }
+    }
+
     private DocumentToHtml(DocumentToHtmlOptions options, List<Comment> comments) {
         this.idPrefix = options.idPrefix();
         this.preserveEmptyParagraphs = options.shouldPreserveEmptyParagraphs();
@@ -68,26 +86,26 @@ public class DocumentToHtml {
         this.comments = Maps.toMapWithKey(comments, Comment::getCommentId);
     }
 
-    private List<HtmlNode> convertToHtml(Document document) {
-        List<HtmlNode> mainBody = convertChildrenToHtml(document);
+    private List<HtmlNode> convertToHtml(Document document, Context context) {
+        List<HtmlNode> mainBody = convertChildrenToHtml(document, context);
         // TODO: can you have note references inside a note?
         List<Note> notes = findNotes(document, noteReferences);
 
         List<HtmlNode> noteNodes = notes.isEmpty()
             ? list()
-            : list(Html.element("ol", eagerMap(notes, this::convertToHtml)));
+            : list(Html.element("ol", eagerMap(notes, note -> convertToHtml(note, context))));
 
         List<HtmlNode> commentNodes = referencedComments.isEmpty()
             ? list()
-            : list(Html.element("dl", eagerFlatMap(referencedComments, this::convertToHtml)));
+            : list(Html.element("dl", eagerFlatMap(referencedComments, comment -> convertToHtml(comment, context))));
 
         return eagerConcat(mainBody, noteNodes, commentNodes);
     }
 
-    private HtmlNode convertToHtml(Note note) {
+    private HtmlNode convertToHtml(Note note, Context context) {
         String id = generateNoteHtmlId(note.getNoteType(), note.getId());
         String referenceId = generateNoteRefHtmlId(note.getNoteType(), note.getId());
-        List<HtmlNode> noteBody = convertToHtml(note.getBody());
+        List<HtmlNode> noteBody = convertToHtml(note.getBody(), context);
         // TODO: we probably want this to collapse more eagerly than other collapsible elements
         // -- for instance, any paragraph will probably do, regardless of attributes. (Possible other elements will do too.)
         HtmlNode backLink = Html.collapsibleElement("p", list(
@@ -96,10 +114,10 @@ public class DocumentToHtml {
         return Html.element("li", map("id", id), eagerConcat(noteBody, list(backLink)));
     }
 
-    private List<HtmlNode> convertToHtml(ReferencedComment referencedComment) {
+    private List<HtmlNode> convertToHtml(ReferencedComment referencedComment, Context context) {
         // TODO: remove duplication with notes
         String commentId = referencedComment.comment.getCommentId();
-        List<HtmlNode> body = convertToHtml(referencedComment.comment.getBody());
+        List<HtmlNode> body = convertToHtml(referencedComment.comment.getBody(), context);
         // TODO: we probably want this to collapse more eagerly than other collapsible elements
         // -- for instance, any paragraph will probably do, regardless of attributes. (Possible other elements will do too.)
         HtmlNode backLink = Html.collapsibleElement("p", list(
@@ -115,22 +133,23 @@ public class DocumentToHtml {
                 eagerConcat(body, list(backLink))));
     }
 
-    private List<HtmlNode> convertToHtml(List<DocumentElement> elements) {
+    private List<HtmlNode> convertToHtml(List<DocumentElement> elements, Context context) {
         return eagerFlatMap(
             elements,
-            this::convertToHtml);
+            element -> convertToHtml(element, context)
+        );
     }
 
-    private List<HtmlNode> convertChildrenToHtml(HasChildren element) {
-        return convertToHtml(element.getChildren());
+    private List<HtmlNode> convertChildrenToHtml(HasChildren element, Context context) {
+        return convertToHtml(element.getChildren(), context);
     }
 
-    private List<HtmlNode> convertToHtml(DocumentElement element) {
-        return element.accept(new DocumentElementVisitor<List<HtmlNode>>() {
+    private List<HtmlNode> convertToHtml(DocumentElement element, Context context) {
+        return element.accept(new DocumentElementVisitor<List<HtmlNode>, Context>() {
             @Override
-            public List<HtmlNode> visit(Paragraph paragraph) {
+            public List<HtmlNode> visit(Paragraph paragraph, Context context) {
                 Supplier<List<HtmlNode>> children = () -> {
-                    List<HtmlNode> content = convertChildrenToHtml(paragraph);
+                    List<HtmlNode> content = convertChildrenToHtml(paragraph, context);
                     return preserveEmptyParagraphs ? cons(Html.FORCE_WRITE, content) : content;
                 };
                 HtmlPath mapping = styleMap.getParagraphHtmlPath(paragraph)
@@ -144,8 +163,8 @@ public class DocumentToHtml {
             }
 
             @Override
-            public List<HtmlNode> visit(Run run) {
-                Supplier<List<HtmlNode>> nodes = () -> convertChildrenToHtml(run);
+            public List<HtmlNode> visit(Run run, Context context) {
+                Supplier<List<HtmlNode>> nodes = () -> convertChildrenToHtml(run, context);
                 if (run.isStrikethrough()) {
                     nodes = styleMap.getStrikethrough().orElse(HtmlPath.collapsibleElement("s")).wrap(nodes);
                 }
@@ -175,7 +194,7 @@ public class DocumentToHtml {
             }
 
             @Override
-            public List<HtmlNode> visit(Text text) {
+            public List<HtmlNode> visit(Text text, Context context) {
                 if (text.getValue().isEmpty()) {
                     return list();
                 } else {
@@ -184,12 +203,12 @@ public class DocumentToHtml {
             }
 
             @Override
-            public List<HtmlNode> visit(Tab tab) {
+            public List<HtmlNode> visit(Tab tab, Context context) {
                 return list(Html.text("\t"));
             }
 
             @Override
-            public List<HtmlNode> visit(Break breakElement) {
+            public List<HtmlNode> visit(Break breakElement, Context context) {
                 HtmlPath mapping = styleMap.getBreakHtmlPath(breakElement)
                     .orElseGet(() -> {
                         if (breakElement.getType() == Break.Type.LINE) {
@@ -202,17 +221,46 @@ public class DocumentToHtml {
             }
 
             @Override
-            public List<HtmlNode> visit(Table table) {
-                return list(Html.element("table", convertChildrenToHtml(table)));
+            public List<HtmlNode> visit(Table table, Context context) {
+                List<HtmlNode> children = generateTableChildren(table, context);
+                return list(Html.element("table", children));
+            }
+
+            private List<HtmlNode> generateTableChildren(Table table, Context context) {
+                int bodyIndex = findIndex(table.getChildren(), child -> !isHeader(child))
+                    .orElse(table.getChildren().size());
+                if (bodyIndex == 0) {
+                    return convertToHtml(table.getChildren(), context.isHeader(false));
+                } else {
+                    List<HtmlNode> headRows = convertToHtml(
+                        table.getChildren().subList(0, bodyIndex),
+                        context.isHeader(true)
+                    );
+                    List<HtmlNode> bodyRows = convertToHtml(
+                        table.getChildren().subList(bodyIndex, table.getChildren().size()),
+                        context.isHeader(false)
+                    );
+                    return list(
+                        Html.element("thead", headRows),
+                        Html.element("tbody", bodyRows)
+                    );
+                }
+            }
+
+            private boolean isHeader(DocumentElement child) {
+                return tryCast(TableRow.class, child)
+                    .map(TableRow::isHeader)
+                    .orElse(false);
             }
 
             @Override
-            public List<HtmlNode> visit(TableRow tableRow) {
-                return list(Html.element("tr", convertChildrenToHtml(tableRow)));
+            public List<HtmlNode> visit(TableRow tableRow, Context context) {
+                return list(Html.element("tr", convertChildrenToHtml(tableRow, context)));
             }
 
             @Override
-            public List<HtmlNode> visit(TableCell tableCell) {
+            public List<HtmlNode> visit(TableCell tableCell, Context context) {
+                String tagName = context.isHeader ? "th" : "td";
                 Map<String, String> attributes = new HashMap<>();
                 if (tableCell.getColspan() != 1) {
                     attributes.put("colspan", Integer.toString(tableCell.getColspan()));
@@ -220,14 +268,14 @@ public class DocumentToHtml {
                 if (tableCell.getRowspan() != 1) {
                     attributes.put("rowspan", Integer.toString(tableCell.getRowspan()));
                 }
-                return list(Html.element("td", attributes,
-                    Lists.cons(Html.FORCE_WRITE, convertChildrenToHtml(tableCell))));
+                return list(Html.element(tagName, attributes,
+                    Lists.cons(Html.FORCE_WRITE, convertChildrenToHtml(tableCell, context))));
             }
 
             @Override
-            public List<HtmlNode> visit(Hyperlink hyperlink) {
+            public List<HtmlNode> visit(Hyperlink hyperlink, Context context) {
                 Map<String, String> attributes = map("href", generateHref(hyperlink));
-                return list(Html.collapsibleElement("a", attributes, convertChildrenToHtml(hyperlink)));
+                return list(Html.collapsibleElement("a", attributes, convertChildrenToHtml(hyperlink, context)));
             }
 
             private String generateHref(Hyperlink hyperlink) {
@@ -241,12 +289,12 @@ public class DocumentToHtml {
             }
 
             @Override
-            public List<HtmlNode> visit(Bookmark bookmark) {
+            public List<HtmlNode> visit(Bookmark bookmark, Context context) {
                 return list(Html.element("a", map("id", generateId(bookmark.getName())), list(Html.FORCE_WRITE)));
             }
 
             @Override
-            public List<HtmlNode> visit(NoteReference noteReference) {
+            public List<HtmlNode> visit(NoteReference noteReference, Context context) {
                 noteReferences.add(noteReference);
                 String noteAnchor = generateNoteHtmlId(noteReference.getNoteType(), noteReference.getNoteId());
                 String noteReferenceAnchor = generateNoteRefHtmlId(noteReference.getNoteType(), noteReference.getNoteId());
@@ -256,7 +304,7 @@ public class DocumentToHtml {
             }
 
             @Override
-            public List<HtmlNode> visit(CommentReference commentReference) {
+            public List<HtmlNode> visit(CommentReference commentReference, Context context) {
                 return styleMap.getCommentReference().orElse(HtmlPath.IGNORE).wrap(() -> {
                     String commentId = commentReference.getCommentId();
                     Comment comment = lookup(comments, commentId)
@@ -275,7 +323,7 @@ public class DocumentToHtml {
             }
 
             @Override
-            public List<HtmlNode> visit(Image image) {
+            public List<HtmlNode> visit(Image image, Context context) {
                 // TODO: custom image handlers
                 // TODO: handle empty content type
                 return image.getContentType()
@@ -306,7 +354,7 @@ public class DocumentToHtml {
                     })
                     .orElse(list());
             }
-        });
+        }, context);
     }
 
     private String generateNoteHtmlId(NoteType noteType, String noteId) {
diff --git a/src/main/java/org/zwobble/mammoth/internal/documents/Bookmark.java b/src/main/java/org/zwobble/mammoth/internal/documents/Bookmark.java
index cca6317..a2890c4 100644
--- a/src/main/java/org/zwobble/mammoth/internal/documents/Bookmark.java
+++ b/src/main/java/org/zwobble/mammoth/internal/documents/Bookmark.java
@@ -12,7 +12,7 @@ public class Bookmark implements DocumentElement {
     }
 
     @Override
-    public <T> T accept(DocumentElementVisitor<T> visitor) {
-        return visitor.visit(this);
+    public <T, U> T accept(DocumentElementVisitor<T, U> visitor, U context) {
+        return visitor.visit(this, context);
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/documents/Break.java b/src/main/java/org/zwobble/mammoth/internal/documents/Break.java
index 0764221..71e129d 100644
--- a/src/main/java/org/zwobble/mammoth/internal/documents/Break.java
+++ b/src/main/java/org/zwobble/mammoth/internal/documents/Break.java
@@ -22,7 +22,7 @@ public class Break implements DocumentElement {
     }
 
     @Override
-    public <T> T accept(DocumentElementVisitor<T> visitor) {
-        return visitor.visit(this);
+    public <T, U> T accept(DocumentElementVisitor<T, U> visitor, U context) {
+        return visitor.visit(this, context);
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/documents/CommentReference.java b/src/main/java/org/zwobble/mammoth/internal/documents/CommentReference.java
index cea65f4..7db9343 100644
--- a/src/main/java/org/zwobble/mammoth/internal/documents/CommentReference.java
+++ b/src/main/java/org/zwobble/mammoth/internal/documents/CommentReference.java
@@ -12,7 +12,7 @@ public class CommentReference implements DocumentElement {
     }
 
     @Override
-    public <T> T accept(DocumentElementVisitor<T> visitor) {
-        return visitor.visit(this);
+    public <T, U> T accept(DocumentElementVisitor<T, U> visitor, U context) {
+        return visitor.visit(this, context);
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/documents/DocumentElement.java b/src/main/java/org/zwobble/mammoth/internal/documents/DocumentElement.java
index b3360ed..7b440cf 100644
--- a/src/main/java/org/zwobble/mammoth/internal/documents/DocumentElement.java
+++ b/src/main/java/org/zwobble/mammoth/internal/documents/DocumentElement.java
@@ -1,5 +1,5 @@
 package org.zwobble.mammoth.internal.documents;
 
 public interface DocumentElement {
-    <T> T accept(DocumentElementVisitor<T> visitor);
+    <T, U> T accept(DocumentElementVisitor<T, U> visitor, U context);
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/documents/DocumentElementVisitor.java b/src/main/java/org/zwobble/mammoth/internal/documents/DocumentElementVisitor.java
index c83e3c5..908a9b7 100644
--- a/src/main/java/org/zwobble/mammoth/internal/documents/DocumentElementVisitor.java
+++ b/src/main/java/org/zwobble/mammoth/internal/documents/DocumentElementVisitor.java
@@ -1,22 +1,22 @@
 package org.zwobble.mammoth.internal.documents;
 
-public interface DocumentElementVisitor<T> {
-    T visit(Paragraph paragraph);
-    T visit(Run run);
-    T visit(Text text);
+public interface DocumentElementVisitor<T, U> {
+    T visit(Paragraph paragraph, U context);
+    T visit(Run run, U context);
+    T visit(Text text, U context);
 
-    T visit(Tab tab);
-    T visit(Break lineBreak);
+    T visit(Tab tab, U context);
+    T visit(Break lineBreak, U context);
 
-    T visit(Table table);
-    T visit(TableRow tableRow);
-    T visit(TableCell tableCell);
+    T visit(Table table, U context);
+    T visit(TableRow tableRow, U context);
+    T visit(TableCell tableCell, U context);
 
-    T visit(Hyperlink hyperlink);
-    T visit(Bookmark bookmark);
-    T visit(NoteReference noteReference);
-    T visit(CommentReference commentReference);
+    T visit(Hyperlink hyperlink, U context);
+    T visit(Bookmark bookmark, U context);
+    T visit(NoteReference noteReference, U context);
+    T visit(CommentReference commentReference, U context);
 
-    T visit(Image image);
+    T visit(Image image, U context);
 
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/documents/Hyperlink.java b/src/main/java/org/zwobble/mammoth/internal/documents/Hyperlink.java
index ceb558f..f78b7fe 100644
--- a/src/main/java/org/zwobble/mammoth/internal/documents/Hyperlink.java
+++ b/src/main/java/org/zwobble/mammoth/internal/documents/Hyperlink.java
@@ -36,7 +36,7 @@ public class Hyperlink implements DocumentElement, HasChildren {
     }
 
     @Override
-    public <T> T accept(DocumentElementVisitor<T> visitor) {
-        return visitor.visit(this);
+    public <T, U> T accept(DocumentElementVisitor<T, U> visitor, U context) {
+        return visitor.visit(this, context);
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/documents/Image.java b/src/main/java/org/zwobble/mammoth/internal/documents/Image.java
index 8ceee7c..f112eed 100644
--- a/src/main/java/org/zwobble/mammoth/internal/documents/Image.java
+++ b/src/main/java/org/zwobble/mammoth/internal/documents/Image.java
@@ -31,7 +31,7 @@ public class Image implements DocumentElement {
     }
 
     @Override
-    public <T> T accept(DocumentElementVisitor<T> visitor) {
-        return visitor.visit(this);
+    public <T, U> T accept(DocumentElementVisitor<T, U> visitor, U context) {
+        return visitor.visit(this, context);
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/documents/NoteReference.java b/src/main/java/org/zwobble/mammoth/internal/documents/NoteReference.java
index 7d0beb2..3a5047e 100644
--- a/src/main/java/org/zwobble/mammoth/internal/documents/NoteReference.java
+++ b/src/main/java/org/zwobble/mammoth/internal/documents/NoteReference.java
@@ -26,7 +26,7 @@ public class NoteReference implements DocumentElement {
     }
 
     @Override
-    public <T> T accept(DocumentElementVisitor<T> visitor) {
-        return visitor.visit(this);
+    public <T, U> T accept(DocumentElementVisitor<T, U> visitor, U context) {
+        return visitor.visit(this, context);
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/documents/Paragraph.java b/src/main/java/org/zwobble/mammoth/internal/documents/Paragraph.java
index 66ea009..a404ae6 100644
--- a/src/main/java/org/zwobble/mammoth/internal/documents/Paragraph.java
+++ b/src/main/java/org/zwobble/mammoth/internal/documents/Paragraph.java
@@ -27,7 +27,7 @@ public class Paragraph implements DocumentElement, HasChildren {
     }
 
     @Override
-    public <T> T accept(DocumentElementVisitor<T> visitor) {
-        return visitor.visit(this);
+    public <T, U> T accept(DocumentElementVisitor<T, U> visitor, U context) {
+        return visitor.visit(this, context);
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/documents/Run.java b/src/main/java/org/zwobble/mammoth/internal/documents/Run.java
index 61f865f..29415c7 100644
--- a/src/main/java/org/zwobble/mammoth/internal/documents/Run.java
+++ b/src/main/java/org/zwobble/mammoth/internal/documents/Run.java
@@ -59,7 +59,7 @@ public class Run implements DocumentElement, HasChildren {
     }
 
     @Override
-    public <T> T accept(DocumentElementVisitor<T> visitor) {
-        return visitor.visit(this);
+    public <T, U> T accept(DocumentElementVisitor<T, U> visitor, U context) {
+        return visitor.visit(this, context);
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/documents/Tab.java b/src/main/java/org/zwobble/mammoth/internal/documents/Tab.java
index b977d12..a9a59c8 100644
--- a/src/main/java/org/zwobble/mammoth/internal/documents/Tab.java
+++ b/src/main/java/org/zwobble/mammoth/internal/documents/Tab.java
@@ -7,7 +7,7 @@ public class Tab implements DocumentElement {
     }
 
     @Override
-    public <T> T accept(DocumentElementVisitor<T> visitor) {
-        return visitor.visit(this);
+    public <T, U> T accept(DocumentElementVisitor<T, U> visitor, U context) {
+        return visitor.visit(this, context);
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/documents/Table.java b/src/main/java/org/zwobble/mammoth/internal/documents/Table.java
index 7381c7d..031b382 100644
--- a/src/main/java/org/zwobble/mammoth/internal/documents/Table.java
+++ b/src/main/java/org/zwobble/mammoth/internal/documents/Table.java
@@ -15,7 +15,7 @@ public class Table implements DocumentElement, HasChildren {
     }
 
     @Override
-    public <T> T accept(DocumentElementVisitor<T> visitor) {
-        return visitor.visit(this);
+    public <T, U> T accept(DocumentElementVisitor<T, U> visitor, U context) {
+        return visitor.visit(this, context);
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/documents/TableCell.java b/src/main/java/org/zwobble/mammoth/internal/documents/TableCell.java
index 0874109..2837725 100644
--- a/src/main/java/org/zwobble/mammoth/internal/documents/TableCell.java
+++ b/src/main/java/org/zwobble/mammoth/internal/documents/TableCell.java
@@ -27,7 +27,7 @@ public class TableCell implements DocumentElement, HasChildren {
     }
 
     @Override
-    public <T> T accept(DocumentElementVisitor<T> visitor) {
-        return visitor.visit(this);
+    public <T, U> T accept(DocumentElementVisitor<T, U> visitor, U context) {
+        return visitor.visit(this, context);
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/documents/TableRow.java b/src/main/java/org/zwobble/mammoth/internal/documents/TableRow.java
index f00ec58..0b5721d 100644
--- a/src/main/java/org/zwobble/mammoth/internal/documents/TableRow.java
+++ b/src/main/java/org/zwobble/mammoth/internal/documents/TableRow.java
@@ -17,8 +17,8 @@ public class TableRow implements DocumentElement, HasChildren {
     }
 
     @Override
-    public <T> T accept(DocumentElementVisitor<T> visitor) {
-        return visitor.visit(this);
+    public <T, U> T accept(DocumentElementVisitor<T, U> visitor, U context) {
+        return visitor.visit(this, context);
     }
 
     public boolean isHeader() {
diff --git a/src/main/java/org/zwobble/mammoth/internal/documents/Text.java b/src/main/java/org/zwobble/mammoth/internal/documents/Text.java
index 64123e9..cf1ab6c 100644
--- a/src/main/java/org/zwobble/mammoth/internal/documents/Text.java
+++ b/src/main/java/org/zwobble/mammoth/internal/documents/Text.java
@@ -12,7 +12,7 @@ public class Text implements DocumentElement {
     }
 
     @Override
-    public <T> T accept(DocumentElementVisitor<T> visitor) {
-        return visitor.visit(this);
+    public <T, U> T accept(DocumentElementVisitor<T, U> visitor, U context) {
+        return visitor.visit(this, context);
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/internal/docx/StatefulBodyXmlReader.java b/src/main/java/org/zwobble/mammoth/internal/docx/StatefulBodyXmlReader.java
index 76ea8ea..f4eda9e 100644
--- a/src/main/java/org/zwobble/mammoth/internal/docx/StatefulBodyXmlReader.java
+++ b/src/main/java/org/zwobble/mammoth/internal/docx/StatefulBodyXmlReader.java
@@ -434,8 +434,8 @@ class StatefulBodyXmlReader {
         }
 
         @Override
-        public <T> T accept(DocumentElementVisitor<T> visitor) {
-            return visitor.visit(new TableCell(1, colspan, children));
+        public <T, U> T accept(DocumentElementVisitor<T, U> visitor, U context) {
+            return visitor.visit(new TableCell(1, colspan, children), context);
         }
     }
 
diff --git a/src/main/java/org/zwobble/mammoth/internal/util/Optionals.java b/src/main/java/org/zwobble/mammoth/internal/util/Optionals.java
index d98b579..d0ca67e 100644
--- a/src/main/java/org/zwobble/mammoth/internal/util/Optionals.java
+++ b/src/main/java/org/zwobble/mammoth/internal/util/Optionals.java
@@ -1,7 +1,9 @@
 package org.zwobble.mammoth.internal.util;
 
 import java.util.Optional;
+import java.util.OptionalInt;
 import java.util.function.BiFunction;
+import java.util.function.IntFunction;
 
 public class Optionals {
     public static <T> Optional<T> first(Optional<T> first, Optional<T> second) {
@@ -35,4 +37,15 @@ public class Optionals {
             return Optional.empty();
         }
     }
+
+    public static <R> Optional<R> map(
+        OptionalInt first,
+        IntFunction<R> function)
+    {
+        if (first.isPresent()) {
+            return Optional.of(function.apply(first.getAsInt()));
+        } else {
+            return Optional.empty();
+        }
+    }
 }