fix: deep audit — correctness, dead code, perf, and edge builder fixes

Author: nikolay-e

.pre-commit-config.yaml

@@ -79,13 +79,8 @@ repos:
         name: Lint GitHub Actions
 
   # ============================================================================
-  # CODE FORMATTING (Order: isort → black to avoid conflicts)
+  # CODE FORMATTING
   # ============================================================================
-  - repo: https://github.com/pycqa/isort
-    rev: 7.0.0
-    hooks:
-      - id: isort
-
   - repo: https://github.com/psf/black
     rev: 26.1.0
     hooks:
@@ -97,7 +92,7 @@ repos:
     hooks:
       - id: pyupgrade
         name: pyupgrade (auto-upgrade Python syntax)
-        args: ['--py312-plus']
+        args: ['--py310-plus']
         files: ^src/.*\.py$
 
   # ============================================================================

.secrets.baseline

@@ -139,9 +139,9 @@
         "filename": ".pre-commit-config.yaml",
         "hashed_secret": "861d9d607b1513c2096a6c1495043e2cc31a8f41",
         "is_verified": false,
-        "line_number": 161
+        "line_number": 156
       }
     ]
   },
-  "generated_at": "2026-03-15T15:14:51Z"
+  "generated_at": "2026-03-21T11:42:35Z"
 }

pyproject.toml

@@ -52,45 +52,44 @@ classifiers = [
 ]
 dynamic = [ "version" ]  # Version is still managed in version.py
 dependencies = [
-  "lxml>=5.0,<7.0",
-  "mistune>=3.0,<4.0",
   "pathspec>=0.11,<2.0",
-  "pyyaml>=6.0.2,<8.0",
-  "ruamel-yaml>=0.18,<1.0",
   "tiktoken>=0.7,<1.0",
 ]
 optional-dependencies.dev = [
-  "autoflake>=2.0,<3.0",
   "black>=23.0.0,<27.0",
   # Build and release
   "build>=0.10,<2.0",
   "coverage>=7.0,<8.0",
   "hypothesis>=6.0,<7.0",
   "import-linter>=2.0,<3.0",
   "isort>=5.12,<9.0",
+  "lxml>=5.0,<7.0",
+  "mistune>=3.0,<4.0",
   "mutmut>=3.5,<4.0",
   "mypy>=1.0,<2.0",
   "pre-commit>=3.0,<5.0",
   "pyinstaller>=5.0,<7.0",
   "pytest>=7.0,<10.0",
   "pytest-cov>=3.0,<8.0",
   "pytest-xdist>=3.0,<4.0",
+  "pyyaml>=6.0.2,<8.0",
   "radon>=6.0,<7.0",
+  "ruamel-yaml>=0.18,<1.0",
   # Quality checks
   "ruff>=0.4,<1.0",
   "treemapper[tree-sitter]",
-  "twine>=4.0,<7.0",
   # Type stubs
   "types-pyyaml>=6.0,<7.0",
-  # Security fixes
-  "urllib3>=2.6.0",
 ]
 optional-dependencies.embeddings = [
   # Semantic code embeddings for diffctx
   "sentence-transformers>=3.0,<6.0",
 ]
 optional-dependencies.full = [
+  "lxml>=5.0,<7.0",
+  "mistune>=3.0,<4.0",
   "pysbd>=0.3,<1.0",
+  "ruamel-yaml>=0.18,<1.0",
   "treemapper[tree-sitter]",
 ]
 optional-dependencies.nlp = [

src/treemapper/diffctx/__init__.py

@@ -205,7 +205,7 @@ def _truncate_generated_fragments(file_frags: list[Fragment]) -> list[Fragment]:
         truncated_content = "\n".join(lines) + f"\n# ... [{remaining} more lines]"
         truncated.append(
             Fragment(
-                id=FragmentId(frag.path, frag.start_line, frag.start_line + len(lines)),
+                id=FragmentId(frag.path, frag.start_line, frag.start_line + len(lines) - 1),
                 kind=frag.kind,
                 content=truncated_content,
                 identifiers=extract_identifiers(truncated_content),
@@ -871,16 +871,19 @@ def _count_brackets_outside_strings(line: str) -> tuple[int, int, int, int]:
     open_braces = 0
     close_braces = 0
     in_string: str | None = None
-    prev = ""
+    escaped = False
     for ch in line:
         if in_string is not None:
-            if ch == in_string and prev != "\\":
+            if escaped:
+                escaped = False
+            elif ch == "\\":
+                escaped = True
+            elif ch == in_string:
                 in_string = None
-            prev = ch
             continue
         if ch in ("'", '"', "`"):
             in_string = ch
-            prev = ch
+            escaped = False
             continue
         if ch == "(":
             open_parens += 1
@@ -890,7 +893,6 @@ def _count_brackets_outside_strings(line: str) -> tuple[int, int, int, int]:
             open_braces += 1
         elif ch == "}":
             close_braces += 1
-        prev = ch
     return open_parens, close_parens, open_braces, close_braces
 
 
@@ -965,16 +967,12 @@ def _add_container_headers(core_ids: set[FragmentId], frags_by_path: dict[Path,
 
 
 def _log_full_mode(selected: list[Fragment]) -> None:
-    try:
-        used = sum(f.token_count for f in selected)
-        logger.info(
-            "diffctx: full mode selected=%d from changed files used=%d tokens",
-            len(selected),
-            used,
-        )
-    except (TypeError, AttributeError) as e:
-        # nosemgrep: python-logger-credential-disclosure
-        logger.debug("diffctx: failed to compute token count: %s", e)
+    used = sum(f.token_count for f in selected)
+    logger.info(
+        "diffctx: full mode selected=%d from changed files used=%d tokens",
+        len(selected),
+        used,
+    )
 
 
 def _log_ppr_mode(
@@ -985,23 +983,19 @@ def _log_ppr_mode(
     alpha: float,
     tau: float,
 ) -> None:
-    try:
-        used = sum(f.token_count for f in selected)
-        budget_str = str(budget_tokens) if budget_tokens is not None else "unlimited"
-        logger.info(
-            "diffctx: selected=%d core=%d used=%d/%s reason=%s utility=%.4f alpha=%.3f tau=%.3f",
-            len(selected),
-            len(core_ids),
-            used,
-            budget_str,
-            result.reason,
-            result.utility,
-            alpha,
-            tau,
-        )
-    except (TypeError, AttributeError) as e:
-        # nosemgrep: python-logger-credential-disclosure
-        logger.debug("diffctx: failed to compute token count: %s", e)
+    used = sum(f.token_count for f in selected)
+    budget_str = str(budget_tokens) if budget_tokens is not None else "unlimited"
+    logger.info(
+        "diffctx: selected=%d core=%d used=%d/%s reason=%s utility=%.4f alpha=%.3f tau=%.3f",
+        len(selected),
+        len(core_ids),
+        used,
+        budget_str,
+        result.reason,
+        result.utility,
+        alpha,
+        tau,
+    )
 
 
 _MAX_FILE_SIZE = LIMITS.max_file_size

src/treemapper/diffctx/edges/base.py

@@ -10,14 +10,7 @@
 EdgeDict = dict[tuple[FragmentId, FragmentId], float]
 
 
-_STRIP_EXTENSIONS = CODE_EXTENSIONS | frozenset(
-    {
-        ".sc",
-        ".fs",
-        ".fsi",
-        ".fsx",
-    }
-)
+_STRIP_EXTENSIONS = CODE_EXTENSIONS
 
 _INDEX_FILE_STEMS = frozenset({"__init__", "index", "mod"})
 

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

@@ -9,21 +9,11 @@
 _MAKEFILE_NAMES = {"makefile", "gnumakefile"}
 _MAKEFILE_EXTS = {".mk", ".mak", ".make"}
 
-_MAKE_TARGET_RE = re.compile(r"^([a-zA-Z_][a-zA-Z0-9_.-]{0,100})\s{0,20}:(?!=)", re.MULTILINE)
 _MAKE_INCLUDE_RE = re.compile(r"^(?:-)?include\s+([^\n]{1,500})$", re.MULTILINE)
-_MAKE_VAR_RE = re.compile(r"^\s{0,20}([A-Z_][A-Z0-9_]{0,100})\s{0,20}[:?]?=", re.MULTILINE)
 _MAKE_RECIPE_RE = re.compile(r"^\t([^\n]{1,1000})$", re.MULTILINE)
 
-_CMAKE_ADD_EXE_RE = re.compile(r"add_executable\s{0,10}\(\s{0,10}(\w{1,100})", re.IGNORECASE)
-_CMAKE_ADD_LIB_RE = re.compile(r"add_library\s{0,10}\(\s{0,10}(\w{1,100})", re.IGNORECASE)
-_CMAKE_TARGET_LINK_RE = re.compile(
-    r"target_link_libraries\s{0,10}\(\s{0,10}(\w{1,100})\s{1,20}(?:PUBLIC|PRIVATE|INTERFACE)?\s{0,10}([^)]{1,500})\)",
-    re.IGNORECASE,
-)
 _CMAKE_INCLUDE_RE = re.compile(r"include\s{0,10}\(\s{0,10}([^)]{1,300})\)", re.IGNORECASE)
 _CMAKE_ADD_SUBDIR_RE = re.compile(r"add_subdirectory\s{0,10}\(\s{0,10}([^\)\s]{1,200})", re.IGNORECASE)
-_CMAKE_FIND_PKG_RE = re.compile(r"find_package\s{0,10}\(\s{0,10}(\w{1,100})", re.IGNORECASE)
-_CMAKE_SET_RE = re.compile(r"set\s{0,10}\(\s{0,10}([A-Z_][A-Z0-9_]{0,100})", re.IGNORECASE)
 
 _SCRIPT_CALL_RE = re.compile(r"(?:bash|sh|python|python3|\.\/scripts\/|\.\/bin\/)([a-zA-Z0-9_.-]+)")
 _SOURCE_FILE_RE = re.compile(r"\b([a-zA-Z_]\w*\.(?:c|cpp|cc|cxx|h|hpp|hxx|py|sh|go|rs|java))\b")
@@ -39,13 +29,9 @@ def _is_cmake(path: Path) -> bool:
     return name == "cmakelists.txt" or path.suffix.lower() == ".cmake"
 
 
-def _extract_make_refs(content: str) -> tuple[set[str], set[str]]:
-    targets: set[str] = set()
+def _extract_make_refs(content: str) -> set[str]:
     file_refs: set[str] = set()
 
-    for match in _MAKE_TARGET_RE.finditer(content):
-        targets.add(match.group(1))
-
     for match in _MAKE_INCLUDE_RE.finditer(content):
         includes = match.group(1).split()
         for inc in includes:
@@ -60,22 +46,12 @@ def _extract_make_refs(content: str) -> tuple[set[str], set[str]]:
 
     file_refs.update(_SOURCE_FILE_RE.findall(content))
 
-    return targets, file_refs
+    return file_refs
 
 
-def _extract_cmake_refs(content: str) -> tuple[set[str], set[str]]:
-    targets: set[str] = set()
+def _extract_cmake_refs(content: str) -> set[str]:
     file_refs: set[str] = set()
 
-    for pattern in [_CMAKE_ADD_EXE_RE, _CMAKE_ADD_LIB_RE]:
-        for match in pattern.finditer(content):
-            targets.add(match.group(1))
-
-    for match in _CMAKE_TARGET_LINK_RE.finditer(content):
-        targets.add(match.group(1))
-        deps = match.group(2).split()
-        targets.update(d for d in deps if d and not d.startswith("$"))
-
     for match in _CMAKE_INCLUDE_RE.finditer(content):
         file_refs.add(match.group(1).strip())
 
@@ -86,7 +62,7 @@ def _extract_cmake_refs(content: str) -> tuple[set[str], set[str]]:
 
     file_refs.update(_SOURCE_FILE_RE.findall(content))
 
-    return targets, file_refs
+    return file_refs
 
 
 def _collect_build_refs(make_files: list[Path], cmake_files: list[Path]) -> set[str]:
@@ -95,16 +71,14 @@ def _collect_build_refs(make_files: list[Path], cmake_files: list[Path]) -> set[
     for mf in make_files:
         try:
             content = mf.read_text(encoding="utf-8")
-            _, file_refs = _extract_make_refs(content)
-            refs.update(file_refs)
+            refs.update(_extract_make_refs(content))
         except (OSError, UnicodeDecodeError):
             continue
 
     for cf in cmake_files:
         try:
             content = cf.read_text(encoding="utf-8")
-            _, file_refs = _extract_cmake_refs(content)
-            refs.update(file_refs)
+            refs.update(_extract_cmake_refs(content))
         except (OSError, UnicodeDecodeError):
             continue
 
@@ -150,7 +124,7 @@ def build(self, fragments: list[Fragment], repo_root: Path | None = None) -> Edg
         return edges
 
     def _add_makefile_edges(self, mf: Fragment, cmake_frags: list[Fragment], idx: FragmentIndex, edges: EdgeDict) -> None:
-        _, file_refs = _extract_make_refs(mf.content)
+        file_refs = _extract_make_refs(mf.content)
 
         for ref in file_refs:
             self._link_ref(mf.id, ref, idx, edges)
@@ -160,7 +134,7 @@ def _add_makefile_edges(self, mf: Fragment, cmake_frags: list[Fragment], idx: Fr
                 self.add_edge(edges, mf.id, cf.id, self.weight * 0.7)
 
     def _add_cmake_edges(self, cf: Fragment, fragments: list[Fragment], idx: FragmentIndex, edges: EdgeDict) -> None:
-        _, file_refs = _extract_cmake_refs(cf.content)
+        file_refs = _extract_cmake_refs(cf.content)
 
         for ref in file_refs:
             self._link_ref(cf.id, ref, idx, edges)

src/treemapper/diffctx/edges/semantic/go.py

@@ -14,7 +14,6 @@
 
 _GO_FUNC_RE = re.compile(r"^func\s+(?:\([^)]+\)\s+)?(\w+)\s*\(", re.MULTILINE)
 _GO_TYPE_RE = re.compile(r"^type\s+(\w+)\s+", re.MULTILINE)
-_GO_CONST_VAR_RE = re.compile(r"^(?:const|var)\s+(\w+)\s+", re.MULTILINE)
 
 _GO_FUNC_CALL_RE = re.compile(r"\b([a-zA-Z_]\w*)\s*\(")
 _GO_KEYWORDS = frozenset(
@@ -49,6 +48,7 @@
 _GO_TYPE_REF_RE = re.compile(r"\*?([A-Z]\w*)\b")
 _GO_PKG_CALL_RE = re.compile(r"\b(\w+)\.([A-Z]\w*)")
 _GO_EMBED_RE = re.compile(r"//go:embed\s+(\S+)", re.MULTILINE)
+_GO_PKG_DECL_RE = re.compile(r"^package\s+(\w+)", re.MULTILINE)
 
 
 def _extract_imports(content: str) -> set[str]:
@@ -65,11 +65,10 @@ def _extract_imports(content: str) -> set[str]:
     return imports
 
 
-def _extract_definitions(content: str) -> tuple[set[str], set[str], set[str]]:
+def _extract_definitions(content: str) -> tuple[set[str], set[str]]:
     funcs = {m.group(1) for m in _GO_FUNC_RE.finditer(content)}
     types = {m.group(1) for m in _GO_TYPE_RE.finditer(content)}
-    consts_vars = {m.group(1) for m in _GO_CONST_VAR_RE.finditer(content)}
-    return funcs, types, consts_vars
+    return funcs, types
 
 
 def _extract_references(content: str) -> tuple[set[str], set[str], set[tuple[str, str]]]:
@@ -83,7 +82,10 @@ def _is_go_file(path: Path) -> bool:
     return path.suffix.lower() == ".go"
 
 
-def _get_package_name(path: Path) -> str:
+def _get_package_name_from_content(content: str, path: Path) -> str:
+    match = _GO_PKG_DECL_RE.search(content)
+    if match:
+        return match.group(1)
     return path.parent.name
 
 
@@ -207,7 +209,7 @@ def _build_indices(
         func_defs: dict[str, list[FragmentId]] = defaultdict(list)
 
         for f in go_frags:
-            pkg = _get_package_name(f.path).lower()
+            pkg = _get_package_name_from_content(f.content, f.path).lower()
             pkg_to_frags[pkg].append(f.id)
 
             if repo_root:
@@ -217,7 +219,7 @@ def _build_indices(
                 except ValueError:
                     pass
 
-            funcs, types, _ = _extract_definitions(f.content)
+            funcs, types = _extract_definitions(f.content)
             for t in types:
                 type_defs[t.lower()].append(f.id)
             for fn in funcs:
@@ -275,7 +277,7 @@ def _link_import_by_path(
         edges: EdgeDict,
     ) -> None:
         for path_str, frag_ids in path_to_frags.items():
-            if path_str in imp or imp.endswith(path_str):
+            if f"/{path_str}" in imp or imp == path_str or imp.endswith(f"/{path_str}"):
                 self.add_edges_from_ids(gf_id, frag_ids, self.import_weight, edges)
 
     def _link_refs(
@@ -309,7 +311,7 @@ def _link_same_package(
         pkg_to_frags: dict[str, list[FragmentId]],
         edges: EdgeDict,
     ) -> None:
-        current_pkg = _get_package_name(gf.path).lower()
+        current_pkg = _get_package_name_from_content(gf.content, gf.path).lower()
         for fid in pkg_to_frags.get(current_pkg, []):
             if fid != gf.id:
                 self.add_edge(edges, gf.id, fid, self.same_package_weight)