diff --git a/android/src/main/java/com/swmansion/enriched/markdown/EnrichedMarkdown.kt b/android/src/main/java/com/swmansion/enriched/markdown/EnrichedMarkdown.kt
index 7351c32c..3008664c 100644
--- a/android/src/main/java/com/swmansion/enriched/markdown/EnrichedMarkdown.kt
+++ b/android/src/main/java/com/swmansion/enriched/markdown/EnrichedMarkdown.kt
@@ -19,6 +19,8 @@ import com.swmansion.enriched.markdown.spans.ImageSpan
 import com.swmansion.enriched.markdown.spoiler.SpoilerMode
 import com.swmansion.enriched.markdown.styles.StyleConfig
 import com.swmansion.enriched.markdown.utils.common.FeatureFlags
+import com.swmansion.enriched.markdown.utils.common.MarkdownSegment
+import com.swmansion.enriched.markdown.utils.common.splitASTIntoSegments
 import com.swmansion.enriched.markdown.utils.text.view.emitLinkLongPressEvent
 import com.swmansion.enriched.markdown.utils.text.view.emitLinkPressEvent
 import com.swmansion.enriched.markdown.views.BlockSegmentView
@@ -199,38 +201,11 @@ class EnrichedMarkdown
             }
 
           val processedSegments =
-            splitASTIntoSegments(ast).map { segmentNode ->
-              when (segmentNode) {
-                is MarkdownASTNode -> {
-                  when (segmentNode.type) {
-                    MarkdownASTNode.NodeType.Table -> {
-                      RenderSegment.Table(segmentNode)
-                    }
-
-                    MarkdownASTNode.NodeType.LatexMathDisplay -> {
-                      val latex =
-                        if (segmentNode.children.isNotEmpty()) {
-                          segmentNode.children.first().content
-                        } else {
-                          segmentNode.content
-                        }
-                      RenderSegment.Math(latex)
-                    }
-
-                    else -> {
-                      renderTextSegment(listOf(segmentNode), style)
-                    }
-                  }
-                }
-
-                is List<*> -> {
-                  @Suppress("UNCHECKED_CAST")
-                  renderTextSegment(segmentNode as List<MarkdownASTNode>, style)
-                }
-
-                else -> {
-                  throw IllegalArgumentException("Unknown segment type")
-                }
+            splitASTIntoSegments(ast).map { segment ->
+              when (segment) {
+                is MarkdownSegment.Text -> renderTextSegment(segment.nodes, style)
+                is MarkdownSegment.Table -> RenderSegment.Table(segment.node)
+                is MarkdownSegment.Math -> RenderSegment.Math(segment.latex)
               }
             }
 
@@ -324,32 +299,6 @@ class EnrichedMarkdown
       }
     }
 
-    private fun splitASTIntoSegments(root: MarkdownASTNode): List<Any> {
-      val segments = mutableListOf<Any>()
-      val currentTextBuffer = mutableListOf<MarkdownASTNode>()
-
-      fun flushTextBuffer() {
-        if (currentTextBuffer.isNotEmpty()) {
-          segments.add(currentTextBuffer.toList())
-          currentTextBuffer.clear()
-        }
-      }
-
-      root.children.forEach { child ->
-        if (child.type == MarkdownASTNode.NodeType.Table) {
-          flushTextBuffer()
-          segments.add(child)
-        } else if (child.type == MarkdownASTNode.NodeType.LatexMathDisplay) {
-          flushTextBuffer()
-          segments.add(child)
-        } else {
-          currentTextBuffer.add(child)
-        }
-      }
-      flushTextBuffer()
-      return segments
-    }
-
     private fun postToMain(
       renderId: Long,
       action: () -> Unit,
diff --git a/android/src/main/java/com/swmansion/enriched/markdown/MeasurementStore.kt b/android/src/main/java/com/swmansion/enriched/markdown/MeasurementStore.kt
index 0f728876..8d381e15 100644
--- a/android/src/main/java/com/swmansion/enriched/markdown/MeasurementStore.kt
+++ b/android/src/main/java/com/swmansion/enriched/markdown/MeasurementStore.kt
@@ -21,9 +21,11 @@ import com.swmansion.enriched.markdown.spans.MathMetrics
 import com.swmansion.enriched.markdown.spans.MathRenderMode
 import com.swmansion.enriched.markdown.styles.StyleConfig
 import com.swmansion.enriched.markdown.utils.common.FeatureFlags
+import com.swmansion.enriched.markdown.utils.common.MarkdownSegment
 import com.swmansion.enriched.markdown.utils.common.getBooleanOrDefault
 import com.swmansion.enriched.markdown.utils.common.getMapOrNull
 import com.swmansion.enriched.markdown.utils.common.getStringOrDefault
+import com.swmansion.enriched.markdown.utils.common.splitASTIntoSegments
 import com.swmansion.enriched.markdown.utils.text.extensions.replaceMathSpansWithPlaceholders
 import com.swmansion.enriched.markdown.views.TableContainerView
 import java.util.concurrent.ConcurrentHashMap
@@ -275,22 +277,6 @@ object MeasurementStore {
     return adjustedSize
   }
 
-/** Sealed interface for type-safe segment handling */
-  private sealed interface MarkdownSegment {
-    data class Text(
-      val nodes: List<MarkdownASTNode>,
-    ) : MarkdownSegment
-
-    data class Table(
-      val node: MarkdownASTNode,
-    ) : MarkdownSegment
-
-    data class Math(
-      val latex: String,
-      val node: MarkdownASTNode,
-    ) : MarkdownSegment
-  }
-
   private fun measureAndCacheSplit(
     context: Context,
     id: Int?,
@@ -427,44 +413,6 @@ object MeasurementStore {
       }.build()
   }
 
-  private fun splitASTIntoSegments(root: MarkdownASTNode): List<MarkdownSegment> {
-    val segments = mutableListOf<MarkdownSegment>()
-    val currentTextNodes = mutableListOf<MarkdownASTNode>()
-
-    fun flushTextNodes() {
-      if (currentTextNodes.isNotEmpty()) {
-        segments.add(MarkdownSegment.Text(currentTextNodes.toList()))
-        currentTextNodes.clear()
-      }
-    }
-
-    for (child in root.children) {
-      when (child.type) {
-        MarkdownASTNode.NodeType.Table -> {
-          flushTextNodes()
-          segments.add(MarkdownSegment.Table(child))
-        }
-
-        MarkdownASTNode.NodeType.LatexMathDisplay -> {
-          flushTextNodes()
-          val latex =
-            if (child.children.isNotEmpty()) {
-              child.children.first().content
-            } else {
-              child.content
-            }
-          segments.add(MarkdownSegment.Math(latex, child))
-        }
-
-        else -> {
-          currentTextNodes.add(child)
-        }
-      }
-    }
-    flushTextNodes()
-    return segments
-  }
-
   private fun tryRenderMarkdown(
     markdown: String,
     styleMap: ReadableMap?,
diff --git a/android/src/main/java/com/swmansion/enriched/markdown/utils/common/MarkdownSegment.kt b/android/src/main/java/com/swmansion/enriched/markdown/utils/common/MarkdownSegment.kt
new file mode 100644
index 00000000..bc2ac1bd
--- /dev/null
+++ b/android/src/main/java/com/swmansion/enriched/markdown/utils/common/MarkdownSegment.kt
@@ -0,0 +1,56 @@
+package com.swmansion.enriched.markdown.utils.common
+
+import com.swmansion.enriched.markdown.parser.MarkdownASTNode
+
+sealed interface MarkdownSegment {
+  data class Text(
+    val nodes: List<MarkdownASTNode>,
+  ) : MarkdownSegment
+
+  data class Table(
+    val node: MarkdownASTNode,
+  ) : MarkdownSegment
+
+  data class Math(
+    val latex: String,
+    val node: MarkdownASTNode,
+  ) : MarkdownSegment
+}
+
+fun splitASTIntoSegments(root: MarkdownASTNode): List<MarkdownSegment> {
+  val segments = mutableListOf<MarkdownSegment>()
+  val currentTextNodes = mutableListOf<MarkdownASTNode>()
+
+  fun flushTextNodes() {
+    if (currentTextNodes.isNotEmpty()) {
+      segments.add(MarkdownSegment.Text(currentTextNodes.toList()))
+      currentTextNodes.clear()
+    }
+  }
+
+  for (child in root.children) {
+    when (child.type) {
+      MarkdownASTNode.NodeType.Table -> {
+        flushTextNodes()
+        segments.add(MarkdownSegment.Table(child))
+      }
+
+      MarkdownASTNode.NodeType.LatexMathDisplay -> {
+        flushTextNodes()
+        val latex =
+          if (child.children.isNotEmpty()) {
+            child.children.first().content
+          } else {
+            child.content
+          }
+        segments.add(MarkdownSegment.Math(latex, child))
+      }
+
+      else -> {
+        currentTextNodes.add(child)
+      }
+    }
+  }
+  flushTextNodes()
+  return segments
+}