fix: Kotlin regex, PySBD line tracking, tree-sitter container headers

Author: nikolay-e

docs/Fragmentation.md …-Selection-for-Git-Diff/Fragmentation.mddocs/Fragmentation.md renamed to docs/Context-Selection-for-Git-Diff/Fragmentation.md docs/Fragmentation.md renamed to docs/Context-Selection-for-Git-Diff/Fragmentation.md

@@ -96,7 +96,7 @@ class ParsedArgs:
     force_stdout: bool
     diff_range: str | None = None
     budget: int | None = None
-    alpha: float = 0.55
+    alpha: float = 0.60
     tau: float = 0.08
     full_diff: bool = False
 

src/treemapper/cli.py

@@ -198,6 +198,7 @@ def build_diff_context(
     full: bool = False,
 ) -> dict[str, Any]:
     _validate_inputs(root_dir, alpha, tau, budget_tokens)
+    root_dir = root_dir.resolve()
 
     hunks = parse_diff(root_dir, diff_range)
     if not hunks:
@@ -265,6 +266,9 @@ def _validate_inputs(root_dir: Path, alpha: float, tau: float, budget_tokens: in
         raise ValueError(f"budget_tokens must be > 0, got {budget_tokens}")
 
 
+_CONTAINER_FRAGMENT_KINDS = frozenset({"class", "interface", "struct"})
+
+
 def _identify_core_fragments(hunks: list[DiffHunk], all_fragments: list[Fragment]) -> set[FragmentId]:
     frags_by_path: dict[Path, list[Fragment]] = defaultdict(list)
     for frag in all_fragments:
@@ -276,9 +280,25 @@ def _identify_core_fragments(hunks: list[DiffHunk], all_fragments: list[Fragment
         if frags:
             h_start, h_end = h.core_selection_range
             core_ids.update(_find_core_for_hunk(frags, h_start, h_end))
+
+    _add_container_headers(core_ids, frags_by_path)
     return core_ids
 
 
+def _add_container_headers(core_ids: set[FragmentId], frags_by_path: dict[Path, list[Fragment]]) -> None:
+    core_paths = {fid.path for fid in core_ids}
+    headers_to_add: list[FragmentId] = []
+    for path in core_paths:
+        for frag in frags_by_path.get(path, []):
+            if frag.kind not in _CONTAINER_FRAGMENT_KINDS or frag.id in core_ids:
+                continue
+            for core_id in core_ids:
+                if core_id.path == path and core_id.start_line > frag.end_line:
+                    headers_to_add.append(frag.id)
+                    break
+    core_ids.update(headers_to_add)
+
+
 def _log_full_mode(selected: list[Fragment]) -> None:
     try:
         used = sum(f.token_count for f in selected)
@@ -409,21 +429,6 @@ def _collect_expansion_files(
     return list(expansion_files)
 
 
-def _expand_universe_by_rare_identifiers(
-    root_dir: Path,
-    concepts: frozenset[str],
-    already_included: list[Path],
-    combined_spec: pathspec.PathSpec,
-) -> list[Path]:
-    if not concepts:
-        return []
-
-    included_set = set(already_included)
-    files = _collect_candidate_files(root_dir, included_set, combined_spec)
-    inverted_index = _build_ident_index(files, concepts)
-    return _collect_expansion_files(inverted_index, concepts, included_set)
-
-
 def _filter_ignored(
     files: list[Path],
     root_dir: Path,
@@ -436,10 +441,25 @@ def _filter_ignored(
             if not should_ignore(rel_path, combined_spec):
                 result.append(file_path)
         except ValueError:
-            result.append(file_path)
+            pass
     return result
 
 
+def _expand_universe_by_rare_identifiers(
+    root_dir: Path,
+    concepts: frozenset[str],
+    already_included: list[Path],
+    combined_spec: pathspec.PathSpec,
+) -> list[Path]:
+    if not concepts:
+        return []
+
+    included_set = set(already_included)
+    files = _collect_candidate_files(root_dir, included_set, combined_spec)
+    inverted_index = _build_ident_index(files, concepts)
+    return _collect_expansion_files(inverted_index, concepts, included_set)
+
+
 def _empty_tree(root_dir: Path) -> dict[str, Any]:
     return {
         "name": root_dir.name,

src/treemapper/diffctx/__init__.py

@@ -16,27 +16,27 @@
 
 KOTLIN_EXTENSIONS = frozenset({".kt", ".kts"})
 
-SCALA_EXTENSIONS = frozenset({".scala"})
+SCALA_EXTENSIONS = frozenset({".scala", ".sc"})
 
 JVM_EXTENSIONS = JAVA_EXTENSIONS | KOTLIN_EXTENSIONS | SCALA_EXTENSIONS
 
 C_EXTENSIONS = frozenset({".c", ".h"})
 
-CPP_EXTENSIONS = frozenset({".cpp", ".hpp", ".cc", ".hh", ".cxx", ".hxx"})
+CPP_EXTENSIONS = frozenset({".cpp", ".hpp", ".cc", ".hh", ".cxx", ".hxx", ".c++", ".h++", ".ipp", ".tpp"})
 
 C_FAMILY_EXTENSIONS = C_EXTENSIONS | CPP_EXTENSIONS | frozenset({".m", ".mm"})
 
 CSHARP_EXTENSIONS = frozenset({".cs"})
 
-FSHARP_EXTENSIONS = frozenset({".fs", ".fsx"})
+FSHARP_EXTENSIONS = frozenset({".fs", ".fsi", ".fsx"})
 
 DOTNET_EXTENSIONS = CSHARP_EXTENSIONS | FSHARP_EXTENSIONS
 
 RUBY_EXTENSIONS = frozenset({".rb", ".rake", ".gemspec"})
 
-PHP_EXTENSIONS = frozenset({".php"})
+PHP_EXTENSIONS = frozenset({".php", ".phtml", ".php3", ".php4", ".php5", ".php7", ".phps"})
 
-SHELL_EXTENSIONS = frozenset({".sh", ".bash", ".zsh", ".ps1"})
+SHELL_EXTENSIONS = frozenset({".sh", ".bash", ".zsh", ".ksh", ".fish", ".ps1", ".psm1", ".psd1"})
 
 SWIFT_EXTENSIONS = frozenset({".swift"})
 

src/treemapper/diffctx/config/extensions.py

@@ -65,7 +65,9 @@ class LangWeights:
 LANG_WEIGHTS: dict[str, LangWeights] = {
     "python": LangWeights(0.55, 0.60, 0.50, 0.20, 0.35),
     "javascript": LangWeights(0.50, 0.55, 0.45, 0.25, 0.35),
+    "jsx": LangWeights(0.50, 0.55, 0.45, 0.25, 0.35),
     "typescript": LangWeights(0.70, 0.75, 0.65, 0.15, 0.25),
+    "tsx": LangWeights(0.70, 0.75, 0.65, 0.15, 0.25),
     "rust": LangWeights(0.90, 0.95, 0.85, 0.10, 0.15),
     "java": LangWeights(0.85, 0.90, 0.80, 0.10, 0.15),
     "kotlin": LangWeights(0.80, 0.85, 0.75, 0.12, 0.18),

src/treemapper/diffctx/config/weights.py

@@ -23,9 +23,12 @@ def path_to_module(path: Path, repo_root: Path | None = None) -> str:
             parts = parts[i + 1 :]
             break
 
-    if parts and parts[-1].endswith(".py"):
-        parts[-1] = parts[-1][:-3]
-        if parts[-1] == "__init__":
+    if parts:
+        for _ext in (".pyw", ".pyi", ".py"):
+            if parts[-1].endswith(_ext):
+                parts[-1] = parts[-1][: -len(_ext)]
+                break
+        if parts and parts[-1] == "__init__":
             parts = parts[:-1]
 
     return ".".join(parts) if parts else ""
@@ -127,6 +130,26 @@ def add_ref_edges(
             edges[(dst, src_id)] = max(edges.get((dst, src_id), 0.0), weight * reverse_factor)
 
 
+def add_semantic_edges(
+    edges: EdgeDict,
+    src_id: FragmentId,
+    info: object,
+    name_to_defs: dict[str, list[FragmentId]],
+    call_weight: float,
+    ref_weight: float,
+    type_weight: float,
+    reverse_factor: float,
+    self_defs: set[str],
+) -> None:
+    add_ref_edges(edges, src_id, set(info.calls), name_to_defs, call_weight, reverse_factor=reverse_factor)  # type: ignore[attr-defined]
+    add_ref_edges(
+        edges, src_id, set(info.references), name_to_defs, ref_weight, reverse_factor=reverse_factor, skip_self_defs=self_defs  # type: ignore[attr-defined]
+    )
+    add_ref_edges(
+        edges, src_id, set(info.type_refs), name_to_defs, type_weight, reverse_factor=reverse_factor, skip_self_defs=self_defs  # type: ignore[attr-defined]
+    )
+
+
 class EdgeBuilder(ABC):
     weight: float = 0.5
     reverse_weight_factor: float = 0.7

src/treemapper/diffctx/edges/base.py

@@ -17,7 +17,8 @@
     r"^\s{0,20}-?\s{0,5}(?:local|project|remote|template):\s{0,5}['\"]?([^'\"#\n]{1,300})", re.MULTILINE
 )
 
-_JENKINS_SH_RE = re.compile(r"sh\s*(?:\(['\"]|['\"])(.+?)['\"]\)?", re.MULTILINE | re.DOTALL)
+_JENKINS_SH_DOUBLE_RE = re.compile(r'sh\s*\(?"([^"]*)"', re.MULTILINE)
+_JENKINS_SH_SINGLE_RE = re.compile(r"sh\s*\(?'([^']*)'", re.MULTILINE)
 _JENKINS_SCRIPT_RE = re.compile(r"script\s*\{([^}]+)\}", re.MULTILINE | re.DOTALL)
 
 _SCRIPT_CALL_TOOLS = frozenset(
@@ -174,9 +175,11 @@ def _extract_gitlab_refs(content: str) -> set[str]:
 def _extract_jenkins_refs(content: str) -> set[str]:
     refs: set[str] = set()
 
-    for match in _JENKINS_SH_RE.finditer(content):
-        sh_content = match.group(1)
-        refs.update(_extract_script_refs(sh_content))
+    for match in _JENKINS_SH_DOUBLE_RE.finditer(content):
+        refs.update(_extract_script_refs(match.group(1)))
+
+    for match in _JENKINS_SH_SINGLE_RE.finditer(content):
+        refs.update(_extract_script_refs(match.group(1)))
 
     for match in _JENKINS_SCRIPT_RE.finditer(content):
         refs.update(_extract_script_refs(match.group(1)))
@@ -328,7 +331,11 @@ def _link_by_name(self, ci_id: FragmentId, ref: str, idx: FragmentIndex, edges:
         ref_name = ref.split("/")[-1].lower()
         ref_base = ref_name.split(".")[0]
         for name, frag_ids in idx.by_name.items():
-            if name == ref_name or name.startswith(ref_base):
+            if name == ref_name or (
+                ref_base
+                and len(ref_base) >= 3
+                and (name == ref_base or name.startswith(ref_base + "_") or name.startswith(ref_base + "."))
+            ):
                 for fid in frag_ids:
                     if fid != ci_id:
                         self.add_edge(edges, ci_id, fid, self.script_weight)

src/treemapper/diffctx/edges/config/cicd.py

@@ -31,10 +31,15 @@ def _is_compose_file(path: Path) -> bool:
     return path.name.lower() in _COMPOSE_NAMES
 
 
+def _strip_dot_slash(s: str) -> str:
+    while s.startswith("./"):
+        s = s[2:]
+    return s
+
+
 def _normalize_path(base_dir: Path, rel_path: str) -> Path:
     rel_path = rel_path.strip().strip("'\"")
-    if rel_path.startswith("./"):
-        rel_path = rel_path[2:]
+    rel_path = _strip_dot_slash(rel_path)
     return base_dir / rel_path
 
 
@@ -57,14 +62,14 @@ def _collect_dockerfile_refs(content: str, refs: set[str]) -> None:
     for match in _DOCKERFILE_COPY_RE.finditer(content):
         src = match.group(1)
         if not src.startswith("--") and not src.startswith("$"):
-            refs.add(src.strip().strip("'\"").lstrip("./"))
+            refs.add(_strip_dot_slash(src.strip().strip("'\"")))
 
 
 def _collect_compose_refs(content: str, refs: set[str]) -> None:
     for match in _COMPOSE_BUILD_RE.finditer(content):
-        refs.add(match.group(1).strip().lstrip("./"))
+        refs.add(_strip_dot_slash(match.group(1).strip()))
     for match in _COMPOSE_VOLUME_RE.finditer(content):
-        refs.add(match.group(1).strip().lstrip("./"))
+        refs.add(_strip_dot_slash(match.group(1).strip()))
 
 
 class DockerEdgeBuilder(EdgeBuilder):
@@ -134,7 +139,9 @@ def _link_copy_source(
 
         if "*" in src_path:
             return
-        suffix = src_path.lstrip("./")
+        suffix = _strip_dot_slash(src_path)
+        if not suffix or suffix == ".":
+            return
         for p, frag_ids in path_to_frags.items():
             if str(p).endswith(suffix):
                 for frag_id in frag_ids:

src/treemapper/diffctx/edges/config/docker.py

@@ -9,6 +9,66 @@
 
 _CONFIG_EXTENSIONS = {".yaml", ".yml", ".json", ".toml", ".ini", ".env"}
 
+_CONFIG_KEY_STOPWORDS = frozenset(
+    {
+        "action",
+        "actions",
+        "assert",
+        "author",
+        "before",
+        "branch",
+        "change",
+        "client",
+        "config",
+        "create",
+        "default",
+        "delete",
+        "deploy",
+        "description",
+        "enable",
+        "engine",
+        "engines",
+        "export",
+        "exports",
+        "format",
+        "health",
+        "ignore",
+        "import",
+        "inputs",
+        "keywords",
+        "module",
+        "modules",
+        "number",
+        "object",
+        "openapi",
+        "option",
+        "options",
+        "output",
+        "outputs",
+        "params",
+        "plugin",
+        "plugins",
+        "private",
+        "public",
+        "remove",
+        "render",
+        "report",
+        "require",
+        "result",
+        "return",
+        "script",
+        "scripts",
+        "server",
+        "source",
+        "status",
+        "string",
+        "target",
+        "update",
+        "verbose",
+        "version",
+    }
+)
+
 _CONFIG_KEY_RE = re.compile(r"^\s*([a-zA-Z_][a-zA-Z0-9_-]*)\s*:", re.MULTILINE)
 _JSON_KEY_RE = re.compile(r'"([a-zA-Z_][a-zA-Z0-9_-]*)"\s*:')
 _TOML_KEY_RE = re.compile(r"^\s*([a-zA-Z_][a-zA-Z0-9_-]*)\s*=", re.MULTILINE)
@@ -86,7 +146,7 @@ def discover_related_files(
     def _build_key_patterns(self, keys: set[str]) -> dict[str, re.Pattern[str]]:
         patterns: dict[str, re.Pattern[str]] = {}
         for key in keys:
-            if len(key) >= 6:
+            if len(key) >= 6 and key not in _CONFIG_KEY_STOPWORDS:
                 patterns[key] = re.compile(rf"\b{re.escape(key)}\b", re.IGNORECASE)
         return patterns
 

src/treemapper/diffctx/edges/config/generic.py

@@ -7,7 +7,7 @@
 from ...types import Fragment, FragmentId
 from ..base import EdgeBuilder, EdgeDict
 
-_HELM_VALUES_RE = re.compile(r"\{\{\s*\.Values\.([a-zA-Z0-9_.]+)\s*\}\}")
+_HELM_VALUES_RE = re.compile(r"\{\{-?\s*\.Values\.([a-zA-Z0-9_.]+)")
 _HELM_INCLUDE_RE = re.compile(r'\{\{\s*(?:include|template)\s+"([^"]+)"')
 _HELM_DEFINE_RE = re.compile(r'\{\{-?\s*define\s+"([^"]+)"')
 _HELM_RELEASE_RE = re.compile(r"\{\{\s*\.Release\.(\w+)\s*\}\}")