Switch to iterative version of WKT format parser

Signed-off-by: Heemin Kim <heemin@amazon.com>
heemin32 · Jun 12, 2024 · 89c92b7 · 89c92b7
1 parent 10c0b77
commit 89c92b7
Show file tree

Hide file tree

Showing 3 changed files with 65 additions and 6 deletions.
diff --git a/CHANGELOG.md b/CHANGELOG.md
@@ -17,6 +17,7 @@ The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.0.0/),
 ### Removed
 
 ### Fixed
+- Switch to iterative version of WKT format parser ([#14086](https://github.com/opensearch-project/OpenSearch/pull/14086))
 
 ### Security
 

diff --git a/libs/geo/src/main/java/org/opensearch/geometry/utils/WellKnownText.java b/libs/geo/src/main/java/org/opensearch/geometry/utils/WellKnownText.java
@@ -53,6 +53,7 @@
 import java.util.Collections;
 import java.util.List;
 import java.util.Locale;
+import java.util.Stack;
 
 /**
  * Utility class for converting to and from WKT
@@ -278,6 +279,16 @@ public Geometry fromWKT(String wkt) throws IOException, ParseException {
      */
     private Geometry parseGeometry(StreamTokenizer stream) throws IOException, ParseException {
         final String type = nextWord(stream).toLowerCase(Locale.ROOT);
+        switch (type) {
+            case "geometrycollection":
+                return parseGeometryCollection(stream);
+            default:
+                return parseSimpleGeometry(stream, type);
+        }
+    }
+
+    private Geometry parseSimpleGeometry(StreamTokenizer stream, String type) throws IOException, ParseException {
+        assert "geometrycollection".equals(type) == false;
         switch (type) {
             case "point":
                 return parsePoint(stream);
@@ -294,23 +305,67 @@ private Geometry parseGeometry(StreamTokenizer stream) throws IOException, Parse
             case "bbox":
                 return parseBBox(stream);
             case "geometrycollection":
-                return parseGeometryCollection(stream);
+                throw new IllegalStateException("Unexpected type: geometrycollection");
             case "circle": // Not part of the standard, but we need it for internal serialization
                 return parseCircle(stream);
         }
         throw new IllegalArgumentException("Unknown geometry type: " + type);
     }
 
+    /**
+     * Parse geometry collection iteratively
+     *
+     * Parsing geometry collection recursively can lead to StackOverflowError when there is a deeply nested structure of GeometryCollection
+     */
     private GeometryCollection<Geometry> parseGeometryCollection(StreamTokenizer stream) throws IOException, ParseException {
         if (nextEmptyOrOpen(stream).equals(EMPTY)) {
             return GeometryCollection.EMPTY;
         }
-        List<Geometry> shapes = new ArrayList<>();
-        shapes.add(parseGeometry(stream));
-        while (nextCloserOrComma(stream).equals(COMMA)) {
-            shapes.add(parseGeometry(stream));
+
+        List<Geometry> topLevelShapes = new ArrayList<>();
+        Stack<List<Geometry>> stack = new Stack<>();
+        stack.push(topLevelShapes);
+        boolean isFirstIteration = true;
+        List<Geometry> currentLevelShapes = null;
+        while (!stack.isEmpty()) {
+            List<Geometry> previousShapes = stack.pop();
+            if (currentLevelShapes != null) {
+                previousShapes.add(new GeometryCollection<>(currentLevelShapes));
+            }
+            currentLevelShapes = previousShapes;
+
+            if (isFirstIteration == true) {
+                isFirstIteration = false;
+            } else {
+                if (!nextCloserOrComma(stream).equals(COMMA)) {
+                    // Done with current level, continue with parent level
+                    continue;
+                }
+            }
+            while (true) {
+                final String type = nextWord(stream).toLowerCase(Locale.ROOT);
+                switch (type) {
+                    case "geometrycollection":
+                        if (nextEmptyOrOpen(stream).equals(EMPTY)) {
+                            currentLevelShapes.add(GeometryCollection.EMPTY);
+                            break;
+                        } else {
+                            stack.push(currentLevelShapes);
+                            currentLevelShapes = new ArrayList<>();
+                            continue;
+                        }
+                    default:
+                        currentLevelShapes.add(parseSimpleGeometry(stream, type));
+                        break;
+                }
+
+                if (!nextCloserOrComma(stream).equals(COMMA)) {
+                    break;
+                }
+            }
         }
-        return new GeometryCollection<>(shapes);
+
+        return new GeometryCollection<>(topLevelShapes);
     }
 
     private Point parsePoint(StreamTokenizer stream) throws IOException, ParseException {

diff --git a/libs/geo/src/test/java/org/opensearch/geometry/GeometryCollectionTests.java b/libs/geo/src/test/java/org/opensearch/geometry/GeometryCollectionTests.java
@@ -62,6 +62,9 @@ public void testBasicSerialization() throws IOException, ParseException {
 
         assertEquals("GEOMETRYCOLLECTION EMPTY", wkt.toWKT(GeometryCollection.EMPTY));
         assertEquals(GeometryCollection.EMPTY, wkt.fromWKT("GEOMETRYCOLLECTION EMPTY)"));
+
+        assertEquals(new GeometryCollection<Geometry>(Arrays.asList(GeometryCollection.EMPTY)),
+            wkt.fromWKT("GEOMETRYCOLLECTION (GEOMETRYCOLLECTION EMPTY)"));
     }
 
     @SuppressWarnings("ConstantConditions")