Collapse nodes where appropriate

4 files changed, 108 insertions, 1 deletions

diff --git a/src/main/java/org/zwobble/mammoth/html/Html.java b/src/main/java/org/zwobble/mammoth/html/Html.java
index 14298b7..c0b234b 100644
--- a/src/main/java/org/zwobble/mammoth/html/Html.java
+++ b/src/main/java/org/zwobble/mammoth/html/Html.java
@@ -1,10 +1,14 @@
 package org.zwobble.mammoth.html;
 
 import org.zwobble.mammoth.util.MammothLists;
+import org.zwobble.mammoth.util.MammothOptionals;
 
+import java.util.ArrayList;
 import java.util.List;
 import java.util.Map;
 
+import static org.zwobble.mammoth.util.Casts.tryCast;
+import static org.zwobble.mammoth.util.MammothIterables.tryGetLast;
 import static org.zwobble.mammoth.util.MammothLists.list;
 import static org.zwobble.mammoth.util.MammothMaps.map;
 
@@ -37,6 +41,10 @@ public class Html {
         return new HtmlElement(tagName, attributes, children, false);
     }
 
+    public static HtmlNode collapsibleElement(String tagName) {
+        return collapsibleElement(tagName, list());
+    }
+
     public static HtmlNode collapsibleElement(String tagName, List<HtmlNode> children) {
         return collapsibleElement(tagName, map(), children);
     }
@@ -95,6 +103,69 @@ public class Html {
     }
 
     public static List<HtmlNode> collapse(List<HtmlNode> nodes) {
-        return nodes;
+        List<HtmlNode> collapsed = new ArrayList<>();
+
+        for (HtmlNode node : nodes) {
+            collapsingAdd(collapsed, node);
+        }
+
+        return collapsed;
+    }
+
+    private static void collapsingAdd(List<HtmlNode> collapsed, HtmlNode node) {
+        HtmlNode collapsedNode = collapse(node);
+        if (!tryCollapse(collapsed, collapsedNode)) {
+            collapsed.add(collapsedNode);
+        }
+    }
+
+    private static HtmlNode collapse(HtmlNode node) {
+        return node.accept(new HtmlNode.Mapper<HtmlNode>() {
+            @Override
+            public HtmlNode visit(HtmlElement element) {
+                return new HtmlElement(
+                    element.getTagName(),
+                    element.getAttributes(),
+                    collapse(element.getChildren()),
+                    element.isCollapsible());
+            }
+
+            @Override
+            public HtmlNode visit(HtmlSelfClosingElement element) {
+                return element;
+            }
+
+            @Override
+            public HtmlNode visit(HtmlTextNode node) {
+                return node;
+            }
+
+            @Override
+            public HtmlNode visit(HtmlForceWrite forceWrite) {
+                return forceWrite;
+            }
+        });
+    }
+
+    private static boolean tryCollapse(List<HtmlNode> collapsed, HtmlNode node) {
+        return MammothOptionals.map(
+            tryGetLast(collapsed).flatMap(last -> tryCast(HtmlElement.class, last)),
+            tryCast(HtmlElement.class, node),
+            (last, next) -> {
+                if (next.isCollapsible() && isMatch(last, next)) {
+                    for (HtmlNode child : next.getChildren()) {
+                        collapsingAdd(last.getChildren(), child);
+                    }
+                    return true;
+                } else {
+                    return false;
+                }
+            }
+        ).orElse(false);
+    }
+
+    private static boolean isMatch(HtmlElement first, HtmlElement second) {
+        return first.getTagName().equals(second.getTagName()) &&
+            first.getAttributes().equals(second.getAttributes());
     }
 }
diff --git a/src/main/java/org/zwobble/mammoth/util/Casts.java b/src/main/java/org/zwobble/mammoth/util/Casts.java
new file mode 100644
index 0000000..0925990
--- /dev/null
+++ b/src/main/java/org/zwobble/mammoth/util/Casts.java
@@ -0,0 +1,13 @@
+package org.zwobble.mammoth.util;
+
+import java.util.Optional;
+
+public class Casts {
+    public static <T> Optional<T> tryCast(Class<T> type, Object value) {
+        try {
+            return Optional.of(type.cast(value));
+        } catch (ClassCastException exception) {
+            return Optional.empty();
+        }
+    }
+}
diff --git a/src/main/java/org/zwobble/mammoth/util/MammothIterables.java b/src/main/java/org/zwobble/mammoth/util/MammothIterables.java
new file mode 100644
index 0000000..89caf61
--- /dev/null
+++ b/src/main/java/org/zwobble/mammoth/util/MammothIterables.java
@@ -0,0 +1,11 @@
+package org.zwobble.mammoth.util;
+
+import com.google.common.collect.Iterables;
+
+import java.util.Optional;
+
+public class MammothIterables {
+    public static <T> Optional<T> tryGetLast(Iterable<T> iterable) {
+        return Optional.ofNullable(Iterables.getLast(iterable, null));
+    }
+}
diff --git a/src/main/java/org/zwobble/mammoth/util/MammothOptionals.java b/src/main/java/org/zwobble/mammoth/util/MammothOptionals.java
index e0e4876..6890f16 100644
--- a/src/main/java/org/zwobble/mammoth/util/MammothOptionals.java
+++ b/src/main/java/org/zwobble/mammoth/util/MammothOptionals.java
@@ -15,4 +15,16 @@ public class MammothOptionals {
             return Optional.empty();
         }
     }
+
+    public static <T1, T2, R> Optional<R> map(
+        Optional<T1> first,
+        Optional<T2> second,
+        BiFunction<T1, T2, R> function)
+    {
+        if (first.isPresent() && second.isPresent()) {
+            return Optional.of(function.apply(first.get(), second.get()));
+        } else {
+            return Optional.empty();
+        }
+    }
 }