Use HtmlElement for void elements

1 files changed, 11 insertions, 13 deletions

diff --git a/src/main/java/org/zwobble/mammoth/internal/html/HtmlWriter.java b/src/main/java/org/zwobble/mammoth/internal/html/HtmlWriter.java
index 59bb90a..910fd81 100644
--- a/src/main/java/org/zwobble/mammoth/internal/html/HtmlWriter.java
+++ b/src/main/java/org/zwobble/mammoth/internal/html/HtmlWriter.java
@@ -2,6 +2,7 @@ package org.zwobble.mammoth.internal.html;
 
 import java.util.Map;
 
+import static org.zwobble.mammoth.internal.html.Html.isVoidElement;
 import static org.zwobble.mammoth.internal.util.Lists.orderedBy;
 
 public class HtmlWriter {
@@ -13,21 +14,18 @@ public class HtmlWriter {
 
                 HtmlWriter.generateAttributes(element.getAttributes(), builder);
 
-                builder.append(">");
+                if (isVoidElement(element)) {
+                    builder.append(" />");
+                } else {
+                    builder.append(">");
 
-                element.getChildren().forEach(child -> write(child, builder));
+                    element.getChildren().forEach(child -> write(child, builder));
 
-                builder
-                    .append("</")
-                    .append(element.getTagName())
-                    .append(">");
-            }
-
-            @Override
-            public void visit(HtmlSelfClosingElement element) {
-                builder.append("<").append(element.getTagName());
-                HtmlWriter.generateAttributes(element.getAttributes(), builder);
-                builder.append(" />");
+                    builder
+                        .append("</")
+                        .append(element.getTagName())
+                        .append(">");
+                }
             }
 
             @Override