Use precomputed index to get mutation coordinate

Signed-off-by: Keshav Priyadarshi <git@keshav.space>

Author: keshav-space

Makefile

@@ -19,11 +19,11 @@ dev:
 
 isort:
 	@echo "-> Apply isort changes to ensure proper imports ordering"
-	${VENV}/bin/isort --sl -l 100 src tests setup.py --skip="tests/testfiles/stemming/"
+	${VENV}/bin/isort --sl -l 100 src tests setup.py --skip="tests/testfiles/"
 
 black:
 	@echo "-> Apply black code formatter"
-	${VENV}/bin/black -l 100 src tests setup.py --exclude="tests/testfiles/stemming/"
+	${VENV}/bin/black -l 100 src tests setup.py --exclude="tests/testfiles/"
 
 doc8:
 	@echo "-> Run doc8 validation"
@@ -35,9 +35,9 @@ check:
 	@echo "-> Run pycodestyle (PEP8) validation"
 	@${ACTIVATE} pycodestyle --max-line-length=100 --exclude=.eggs,venv,lib,thirdparty,docs,migrations,settings.py,.cache,tests/testfiles/stemming/ .
 	@echo "-> Run isort imports ordering validation"
-	@${ACTIVATE} isort --sl --check-only -l 100 setup.py src tests . --skip="tests/testfiles/stemming/"
+	@${ACTIVATE} isort --sl --check-only -l 100 setup.py src tests . --skip="tests/testfiles/"
 	@echo "-> Run black validation"
-	@${ACTIVATE} black --check --check -l 100 src tests setup.py --exclude="tests/testfiles/stemming/"
+	@${ACTIVATE} black --check --check -l 100 src tests setup.py --exclude="tests/testfiles/"
 
 clean:
 	@echo "-> Clean the Python env"

src/matchcode_toolkit/stemming.py

@@ -59,7 +59,7 @@ class TreeSitterWheelNotInstalled(Exception):
 
 def get_parser(location):
     """
-    Get the appropriate tree-sitter parser and string identifier for
+    Get the appropriate tree-sitter parser and grammar config for
     file at location.
     """
     file_type = Type(location)
@@ -108,27 +108,34 @@ def traverse(node, language_info, mutation_index):
         traverse(child, language_info, mutation_index)
 
 
-def apply_mutation(text, start_point, end_point, replacement):
+def apply_mutation(text, start_point, end_point, replacement, successive_line_count):
+    """Mutate tokens between start and end points with replacement string."""
+
     start_row, start_col = start_point
     end_row, end_col = end_point
 
-    lines = text.splitlines()
-
-    # Compute the start and end indices, +1 for newline.
-    start_index = sum(len(line) + 1 for line in lines[:start_row]) + start_col
-    end_index = sum(len(line) + 1 for line in lines[:end_row]) + end_col
+    # Compute 1D mutation position from 2D coordinates
+    start_index = successive_line_count[start_row] + start_col
+    end_index = successive_line_count[end_row] + end_col
 
     modified_text = text[:start_index] + replacement + text[end_index:]
+    modified_lines = modified_text.splitlines(keepends=True)
 
-    modified_lines = modified_text.splitlines()
     # Remove empty comment lines.
     if not replacement and modified_lines[start_row].strip() == "":
         del modified_lines[start_row]
 
-    return "\n".join(modified_lines)
+    return "".join(modified_lines)
 
 
 def get_stem_code(location):
+    """
+    Return the stemmed code for the code file at the specified `location`.
+
+    Parse the code using tree-sitter, create a mutation index for tokens that
+    need to be replaced or removed, and apply these mutations bottom-up to
+    generate the stemmed code.
+    """
     parser_result = get_parser(location)
     if not parser_result:
         return
@@ -143,11 +150,17 @@ def get_stem_code(location):
     # Apply mutations bottom-up
     mutations = dict(sorted(mutations.items(), reverse=True))
     text = source.decode()
+    cur_count = 0
+    lines = text.splitlines(keepends=True)
+    successive_line_count = [cur_count := cur_count + len(line) for line in lines]
+    successive_line_count.insert(0, 0)
+
     for value in mutations.values():
         text = apply_mutation(
             text=text,
             end_point=value["end_point"],
             start_point=value["start_point"],
             replacement=("idf" if value["type"] == "identifier" else ""),
+            successive_line_count=successive_line_count,
         )
     return text

tests/testfiles/stemming/c/main-stemmed.c

@@ -68,4 +68,4 @@ int idf idf(HINSTANCE idf, HINSTANCE idf,
 	idf();
 
 	return idf;
-}
+}

tests/testfiles/stemming/cpp/string-stemmed.cpp

@@ -2220,4 +2220,4 @@ wxString* wxCArrayString::idf()
     wxString *idf = idf();
     idf = NULL;
     return idf;
-}
+}

tests/testfiles/stemming/golang/utils-stemmed.go

@@ -124,4 +124,4 @@ func idf(idf []string, idf bool) {
 
 	}
 
-}
+}

tests/testfiles/stemming/java/contenttype-stemmed.java

@@ -113,4 +113,4 @@ boolean idf() {
     return (idf+1 == idf);
   }
 
-}
+}

tests/testfiles/stemming/javascript/utils-stemmed.js

@@ -43,4 +43,4 @@ function idf(idf, idf) {
 idf.exports = {
   addMetaDataFilesToPackage,
   buildPackageArchive,
-};
+};

tests/testfiles/stemming/python/sync_scancode_scans-stemmed.py

@@ -100,4 +100,4 @@ def idf(idf, idf):
         idf, idf = idf.idf.idf(idf=idf(idf), idf=idf.idf)
         idf = idf.idf(idf=idf)
         idf.idf(idf=idf, idf=idf)
-    return idf
+    return idf

tests/testfiles/stemming/rust/metrics-stemmeds.rs

@@ -524,4 +524,4 @@ mod idf {
             "GRPC_ENDPOINT_WHITELIST must be sorted in code to allow binary search"
         );
     }
-}
+}