diff --git a/trace-diff/src/main/java/io/github/chains_project/tracediff/statediff/ui/StateDiffUIManipulator.java b/trace-diff/src/main/java/io/github/chains_project/tracediff/statediff/ui/StateDiffUIManipulator.java
index 2c66bf01..eb145dc7 100755
--- a/trace-diff/src/main/java/io/github/chains_project/tracediff/statediff/ui/StateDiffUIManipulator.java
+++ b/trace-diff/src/main/java/io/github/chains_project/tracediff/statediff/ui/StateDiffUIManipulator.java
@@ -52,16 +52,17 @@ public void addStateDiffToExecDiffUI(
         long processStartTime = new Date().getTime();
         boolean isHitDataIncluded = ghFullDiff != null;
 
-        Map<Integer, Integer> returnSrcToDstMappings = new HashMap<>(), returnDstToSrcMappings = new HashMap<>();
-
-        extractReturnMappings(srcFile, dstFile, returnSrcToDstMappings, returnDstToSrcMappings);
-
         SourceInfo srcInfo = new SourceInfo(srcFile), dstInfo = new SourceInfo(dstFile);
 
         if (ghFullDiff == null) {
             ghFullDiff = GHHelper.getGHDiff(slug, commit, srcInfo, dstInfo);
         }
 
+        Map<Integer, Integer> returnSrcToDstMappings = new HashMap<>(), returnDstToSrcMappings = new HashMap<>();
+
+        extractReturnMappings(
+                srcFile, dstFile, returnSrcToDstMappings, returnDstToSrcMappings, isHitDataIncluded, ghFullDiff);
+
         Pair<Map<Integer, Integer>, Map<Integer, Integer>> lineMappings =
                 ExecDiffHelper.getMappingFromExecDiff(ghFullDiff, isHitDataIncluded);
 
@@ -110,8 +111,12 @@ private void extractReturnMappings(
             File srcFile,
             File dstFile,
             Map<Integer, Integer> returnSrcToDstMappings,
-            Map<Integer, Integer> returnDstToSrcMappings)
+            Map<Integer, Integer> returnDstToSrcMappings,
+            boolean isHitDataIncluded,
+            File ghFullDiff)
             throws Exception {
+        Pair<Set<Integer>, Set<Integer>> validLines = ExecDiffHelper.getValidLines(ghFullDiff, isHitDataIncluded);
+
         Diff diff = new AstComparator().compare(srcFile, dstFile);
         Iterator<Mapping> mappings = diff.getMappingsComp().iterator();
         while (mappings.hasNext()) {
@@ -123,6 +128,11 @@ private void extractReturnMappings(
                 int srcLine = srcElem.getPosition().getLine(),
                         dstLine = dstElem.getPosition().getLine();
 
+                if (!validLines.getLeft().contains(srcLine)
+                        || !validLines.getRight().contains(dstLine)) {
+                    continue;
+                }
+
                 returnSrcToDstMappings.put(srcLine, dstLine);
                 returnDstToSrcMappings.put(dstLine, srcLine);
             }
diff --git a/trace-diff/src/main/java/io/github/chains_project/tracediff/statediff/utils/ExecDiffHelper.java b/trace-diff/src/main/java/io/github/chains_project/tracediff/statediff/utils/ExecDiffHelper.java
index df6a3b12..de6f22a3 100755
--- a/trace-diff/src/main/java/io/github/chains_project/tracediff/statediff/utils/ExecDiffHelper.java
+++ b/trace-diff/src/main/java/io/github/chains_project/tracediff/statediff/utils/ExecDiffHelper.java
@@ -3,7 +3,9 @@
 import java.io.File;
 import java.io.IOException;
 import java.util.HashMap;
+import java.util.HashSet;
 import java.util.Map;
+import java.util.Set;
 import org.apache.commons.io.FileUtils;
 import org.apache.commons.lang3.tuple.Pair;
 import org.jsoup.Jsoup;
@@ -34,6 +36,31 @@ public static Pair<Map<Integer, Integer>, Map<Integer, Integer>> getMappingFromE
         return Pair.of(leftRightMapping, rightLeftMapping);
     }
 
+    public static Pair<Set<Integer>, Set<Integer>> getValidLines(File execDiffReport, boolean isHitDataIncluded)
+            throws IOException {
+        Set<Integer> leftValidLines = new HashSet<>(), rightValidLines = new HashSet<>();
+        Document doc = Jsoup.parse(execDiffReport, "UTF-8");
+        Elements srcRows = doc.selectFirst("tbody").children();
+
+        srcRows.forEach(tr -> {
+            Elements cols = tr.children();
+            int leftLineNumCol = isHitDataIncluded ? 1 : 0;
+            try {
+                int srcLine = Integer.parseInt(cols.get(leftLineNumCol).attr("data-line-number"));
+                leftValidLines.add(srcLine);
+            } catch (NumberFormatException e) {
+            }
+
+            try {
+                int dstLine = Integer.parseInt(cols.get(leftLineNumCol + 1).attr("data-line-number"));
+                rightValidLines.add(dstLine);
+            } catch (NumberFormatException e) {
+            }
+        });
+
+        return Pair.of(leftValidLines, rightValidLines);
+    }
+
     public static void addLineInfoAfter(
             Integer line, String infoHtml, File ghDiff, boolean isOriginalLine, boolean isHitDataIncluded)
             throws Exception {