402 lines
13 KiB
Python
402 lines
13 KiB
Python
"""
|
||
Unit tests for distlib CWE-407 defects.
|
||
|
||
distlib-0001: get_dependent_dists — dep list O(N²) membership check
|
||
File: distlib/database.py ~line 1284
|
||
Fix: parallel dep_set for O(1) succ-not-in-dep guard
|
||
|
||
distlib-0002: get_steps — result.remove O(N) + todo.pop(0) O(N)
|
||
File: distlib/util.py ~line 1127
|
||
Fix: OrderedDict.move_to_end O(1) + deque.popleft O(1)
|
||
|
||
Tests use operation counters to verify algorithmic complexity.
|
||
No distlib import needed — implementations are inlined as stubs.
|
||
"""
|
||
|
||
import sys
|
||
import time
|
||
from collections import deque, OrderedDict
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# distlib-0001: get_dependent_dists
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def get_dependent_dists_before(reverse_list, dist):
|
||
"""
|
||
Defective implementation: O(N²) due to `succ not in dep` list scan.
|
||
Returns (result_list, op_count).
|
||
"""
|
||
ops = 0
|
||
dep = [dist]
|
||
todo = list(reverse_list.get(dist, []))
|
||
while todo:
|
||
d = todo.pop()
|
||
dep.append(d)
|
||
for succ in reverse_list.get(d, []):
|
||
ops += len(dep) # cost of `succ not in dep`
|
||
if succ not in dep:
|
||
todo.append(succ)
|
||
dep.pop(0)
|
||
return dep, ops
|
||
|
||
|
||
def get_dependent_dists_after(reverse_list, dist):
|
||
"""
|
||
Fixed implementation: O(N) with parallel dep_set for O(1) membership.
|
||
Returns (result_list, op_count).
|
||
"""
|
||
ops = 0
|
||
dep = [dist]
|
||
dep_set = {dist}
|
||
todo = list(reverse_list.get(dist, []))
|
||
while todo:
|
||
d = todo.pop()
|
||
dep.append(d)
|
||
dep_set.add(d)
|
||
for succ in reverse_list.get(d, []):
|
||
ops += 1 # O(1) set lookup
|
||
if succ not in dep_set:
|
||
todo.append(succ)
|
||
dep.pop(0)
|
||
return dep, ops
|
||
|
||
|
||
def make_fan_in_graph(n):
|
||
"""
|
||
Build a reverse_list that forces O(N²) membership checks.
|
||
|
||
Structure: a "collector" node c is depended on by p0..p(N-1) (in
|
||
reverse_list terms: reverse_list[c] = [p0..p(N-1)]).
|
||
Each pi is in turn depended on by a unique "leaf" node li
|
||
(reverse_list[pi] = [li]).
|
||
Additionally, every leaf li lists *all* pj (j!=i) as its successors,
|
||
so when we visit li, we check all N pj nodes against dep.
|
||
By the time we visit leaf li, dep already contains c, p0..p(i), and
|
||
the previously-visited leaves, so each `pj not in dep` check runs
|
||
against a growing list.
|
||
|
||
Total membership checks: N leaves × N pj nodes = O(N²).
|
||
All checks correctly return False the first time (pj nodes added from
|
||
c's reverse_list are in dep) or True when pj is already there.
|
||
Result size = 1 (c) + N (pi) + N (li) = 2N+1 dependents.
|
||
"""
|
||
reverse_list = {}
|
||
# c's direct dependents: p0..p(N-1)
|
||
reverse_list["c"] = [f"p{i}" for i in range(n)]
|
||
for i in range(n):
|
||
# pi is depended on by leaf li
|
||
reverse_list[f"p{i}"] = [f"l{i}"]
|
||
# leaf li lists all pj as successors (all already in dep when li visited)
|
||
reverse_list[f"l{i}"] = [f"p{j}" for j in range(n)]
|
||
return reverse_list
|
||
|
||
|
||
def make_simple_chain(n):
|
||
"""
|
||
Simple linear chain: base <- p0 <- p1 <- ... <- p(n-1).
|
||
reverse_list[base] = [p0], reverse_list[pi] = [p(i+1)], last has [].
|
||
Result of get_dependent_dists(base) = [p0, p1, ..., p(n-1)].
|
||
"""
|
||
rl = {"base": ["p0"]}
|
||
for i in range(n - 1):
|
||
rl[f"p{i}"] = [f"p{i+1}"]
|
||
rl[f"p{n-1}"] = []
|
||
return rl
|
||
|
||
|
||
def test_distlib_0001_correctness():
|
||
"""Both implementations produce the same result on a simple chain."""
|
||
N = 20
|
||
rl = make_simple_chain(N)
|
||
result_before, _ = get_dependent_dists_before(rl, "base")
|
||
result_after, _ = get_dependent_dists_after(rl, "base")
|
||
assert sorted(result_before) == sorted(result_after), (
|
||
f"Results differ:\n before={sorted(result_before)}\n after={sorted(result_after)}"
|
||
)
|
||
assert len(result_before) == N, (
|
||
f"Expected {N} dependents, got {len(result_before)}"
|
||
)
|
||
print(f" correctness: both return {len(result_before)} dependents")
|
||
|
||
|
||
def test_distlib_0001_complexity():
|
||
"""
|
||
BEFORE: op count grows as O(N²); AFTER: op count grows as O(N).
|
||
|
||
At N=500, each of N leaves checks N peers against a dep list of size ~N
|
||
→ ~N² = 250,000 membership-check ops in BEFORE; each check is O(1) in
|
||
AFTER → ~N² total ops but each counted as 1 → ratio reflects list vs set.
|
||
Expected ratio: >= 50x.
|
||
"""
|
||
N = 500
|
||
rl = make_fan_in_graph(N)
|
||
|
||
_, ops_before = get_dependent_dists_before(rl, "c")
|
||
_, ops_after = get_dependent_dists_after(rl, "c")
|
||
|
||
ratio = ops_before / ops_after if ops_after > 0 else float("inf")
|
||
print(f" N={N}: ops_before={ops_before:,}, ops_after={ops_after:,}, ratio={ratio:.1f}x")
|
||
|
||
assert ratio >= 50, (
|
||
f"Expected >= 50x op ratio at N={N}, got {ratio:.1f}x "
|
||
f"(before={ops_before:,}, after={ops_after:,})"
|
||
)
|
||
print(f" PASS: {ratio:.1f}x operation-count reduction")
|
||
|
||
|
||
def test_distlib_0001_no_infinite_loop():
|
||
"""
|
||
dep / dep_set prevent revisiting nodes in a graph with shared sub-deps.
|
||
Diamond: base <- {a, b}, a <- c, b <- c (c depends on both a and b).
|
||
"""
|
||
reverse_list = {
|
||
"base": ["a", "b"],
|
||
"a": ["c"],
|
||
"b": ["c"],
|
||
"c": [],
|
||
}
|
||
result_before, _ = get_dependent_dists_before(reverse_list, "base")
|
||
result_after, _ = get_dependent_dists_after(reverse_list, "base")
|
||
|
||
# c should appear exactly once in both (dedup prevents double-add)
|
||
assert result_before.count("c") == 1, f"BEFORE: c appears {result_before.count('c')} times"
|
||
assert result_after.count("c") == 1, f"AFTER: c appears {result_after.count('c')} times"
|
||
print(f" diamond dedup: BEFORE={result_before}, AFTER={result_after}")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# distlib-0002: get_steps
|
||
# ---------------------------------------------------------------------------
|
||
|
||
def get_steps_before(preds, final):
|
||
"""
|
||
Defective implementation:
|
||
- result.remove(step) is O(N) list scan
|
||
- todo.pop(0) is O(N) list shift
|
||
Returns (result_list, op_count).
|
||
"""
|
||
ops = 0
|
||
result = []
|
||
todo = []
|
||
seen = set()
|
||
todo.append(final)
|
||
while todo:
|
||
ops += len(todo) # cost of pop(0): shifts all remaining elements
|
||
step = todo.pop(0)
|
||
if step in seen:
|
||
if step != final:
|
||
ops += len(result) # cost of result.remove(step)
|
||
result.remove(step)
|
||
result.append(step)
|
||
else:
|
||
seen.add(step)
|
||
result.append(step)
|
||
preds_list = preds.get(step, ())
|
||
todo.extend(preds_list)
|
||
return list(reversed(result)), ops
|
||
|
||
|
||
def get_steps_after(preds, final):
|
||
"""
|
||
Fixed implementation:
|
||
- result is OrderedDict; move_to_end is O(1)
|
||
- todo is deque; popleft is O(1)
|
||
Returns (result_list, op_count).
|
||
"""
|
||
ops = 0
|
||
result = OrderedDict()
|
||
todo = deque()
|
||
seen = set()
|
||
todo.append(final)
|
||
while todo:
|
||
ops += 1 # O(1) deque.popleft
|
||
step = todo.popleft()
|
||
if step in seen:
|
||
if step != final:
|
||
ops += 1 # O(1) OrderedDict.move_to_end
|
||
result.move_to_end(step)
|
||
else:
|
||
seen.add(step)
|
||
result[step] = None
|
||
preds_list = preds.get(step, ())
|
||
todo.extend(preds_list)
|
||
return list(reversed(list(result))), ops
|
||
|
||
|
||
def make_diamond_steps():
|
||
"""
|
||
Diamond dependency graph for steps:
|
||
final -> a -> c
|
||
final -> b -> c
|
||
Step c appears as a predecessor of both a and b.
|
||
When c is first seen via a, then revisited via b, result.remove(c)
|
||
is called in the BEFORE version.
|
||
"""
|
||
return {
|
||
"final": ("a", "b"),
|
||
"a": ("c",),
|
||
"b": ("c",),
|
||
"c": (),
|
||
}
|
||
|
||
|
||
def make_chain_steps(n):
|
||
"""
|
||
Linear chain: final -> s0 -> s1 -> ... -> s(n-1).
|
||
Every step is visited exactly once; no remove() calls.
|
||
Used to stress-test pop(0) vs popleft().
|
||
"""
|
||
preds = {"final": (f"s0",)}
|
||
for i in range(n - 1):
|
||
preds[f"s{i}"] = (f"s{i+1}",)
|
||
preds[f"s{n-1}"] = ()
|
||
return preds
|
||
|
||
|
||
def make_wide_diamond_steps(width, depth):
|
||
"""
|
||
Wide diamond: final has `width` direct predecessors, each of which
|
||
has the same `depth` common predecessors. Each common predecessor
|
||
is revisited `width` times — triggering result.remove() in BEFORE.
|
||
"""
|
||
preds = {}
|
||
shared = [f"shared_{d}" for d in range(depth)]
|
||
branches = [f"branch_{w}" for w in range(width)]
|
||
|
||
preds["final"] = tuple(branches)
|
||
for b in branches:
|
||
preds[b] = tuple(shared)
|
||
for s in shared:
|
||
preds[s] = ()
|
||
return preds, shared
|
||
|
||
|
||
def test_distlib_0002_correctness_diamond():
|
||
"""Both implementations produce the same topological order on a diamond."""
|
||
preds = make_diamond_steps()
|
||
result_before, _ = get_steps_before(preds, "final")
|
||
result_after, _ = get_steps_after(preds, "final")
|
||
|
||
assert result_before == result_after, (
|
||
f"Results differ:\n before={result_before}\n after={result_after}"
|
||
)
|
||
# c should be first (deepest common dep), final last
|
||
assert result_before[0] == "c", f"Expected 'c' first, got {result_before[0]}"
|
||
assert result_before[-1] == "final", f"Expected 'final' last, got {result_before[-1]}"
|
||
print(f" diamond order: {result_before}")
|
||
|
||
|
||
def test_distlib_0002_correctness_chain():
|
||
"""Linear chain: order must be [s(n-1), ..., s0, final]."""
|
||
N = 20
|
||
preds = make_chain_steps(N)
|
||
result_before, _ = get_steps_before(preds, "final")
|
||
result_after, _ = get_steps_after(preds, "final")
|
||
|
||
assert result_before == result_after, (
|
||
f"Chain results differ:\n before={result_before}\n after={result_after}"
|
||
)
|
||
expected_first = f"s{N-1}"
|
||
assert result_before[0] == expected_first, (
|
||
f"Expected '{expected_first}' first, got {result_before[0]}"
|
||
)
|
||
print(f" chain N={N}: first={result_before[0]}, last={result_before[-1]}")
|
||
|
||
|
||
def test_distlib_0002_complexity():
|
||
"""
|
||
BEFORE: op count grows as O(N²) for wide-diamond graphs.
|
||
AFTER: op count grows as O(N).
|
||
|
||
Build a graph with width=200 branches all sharing depth=100 common steps.
|
||
BEFORE: each of 100 shared nodes is remove()'d 199 times → ~100×200 = 20,000
|
||
remove ops, each scanning a result list of up to 300 items → ~6,000,000 ops.
|
||
AFTER: each move_to_end is O(1) → ~200×100 = 20,000 ops total.
|
||
Expected ratio: >= 50x.
|
||
"""
|
||
WIDTH = 200
|
||
DEPTH = 100
|
||
preds, _ = make_wide_diamond_steps(WIDTH, DEPTH)
|
||
|
||
_, ops_before = get_steps_before(preds, "final")
|
||
_, ops_after = get_steps_after(preds, "final")
|
||
|
||
ratio = ops_before / ops_after if ops_after > 0 else float("inf")
|
||
print(f" width={WIDTH}, depth={DEPTH}: ops_before={ops_before:,}, "
|
||
f"ops_after={ops_after:,}, ratio={ratio:.1f}x")
|
||
|
||
assert ratio >= 50, (
|
||
f"Expected >= 50x op ratio, got {ratio:.1f}x "
|
||
f"(before={ops_before:,}, after={ops_after:,})"
|
||
)
|
||
print(f" PASS: {ratio:.1f}x operation-count reduction")
|
||
|
||
|
||
def test_distlib_0002_popleft_complexity():
|
||
"""
|
||
todo.pop(0) on a list is O(N); deque.popleft() is O(1).
|
||
|
||
To expose the difference, we need a graph where todo stays large during
|
||
traversal. A wide fan-out graph achieves this: final has N direct
|
||
predecessors, all leaf nodes. When each branch is popped from todo, the
|
||
remaining todo list still holds all the sibling branches — pop(0) must
|
||
shift them all. At N predecessors the total shift-work is O(N²).
|
||
|
||
We count shifts explicitly by accumulating len(todo) before each pop(0).
|
||
"""
|
||
N = 2000
|
||
# final -> p0, p1, ..., p(N-1); each pi is a leaf
|
||
preds = {"final": tuple(f"p{i}" for i in range(N))}
|
||
for i in range(N):
|
||
preds[f"p{i}"] = ()
|
||
|
||
_, ops_before = get_steps_before(preds, "final")
|
||
_, ops_after = get_steps_after(preds, "final")
|
||
|
||
ratio = ops_before / ops_after if ops_after > 0 else float("inf")
|
||
print(f" fan-out N={N}: ops_before={ops_before:,}, ops_after={ops_after:,}, ratio={ratio:.1f}x")
|
||
|
||
assert ratio >= 50, (
|
||
f"Expected >= 50x op ratio at N={N}, got {ratio:.1f}x "
|
||
f"(before={ops_before:,}, after={ops_after:,})"
|
||
)
|
||
print(f" PASS: {ratio:.1f}x operation-count reduction")
|
||
|
||
|
||
# ---------------------------------------------------------------------------
|
||
# Runner
|
||
# ---------------------------------------------------------------------------
|
||
|
||
TESTS = [
|
||
# distlib-0001
|
||
test_distlib_0001_correctness,
|
||
test_distlib_0001_complexity,
|
||
test_distlib_0001_no_infinite_loop,
|
||
# distlib-0002
|
||
test_distlib_0002_correctness_diamond,
|
||
test_distlib_0002_correctness_chain,
|
||
test_distlib_0002_complexity,
|
||
test_distlib_0002_popleft_complexity,
|
||
]
|
||
|
||
if __name__ == "__main__":
|
||
failed = 0
|
||
for t in TESTS:
|
||
print(f"=== {t.__name__} ===")
|
||
try:
|
||
t()
|
||
print(" PASS")
|
||
except AssertionError as e:
|
||
print(f" FAIL: {e}", file=sys.stderr)
|
||
failed += 1
|
||
except Exception as e:
|
||
print(f" ERROR: {e}", file=sys.stderr)
|
||
failed += 1
|
||
|
||
print()
|
||
if failed:
|
||
print(f"FAILED: {failed}/{len(TESTS)} tests failed", file=sys.stderr)
|
||
sys.exit(1)
|
||
else:
|
||
print(f"ALL {len(TESTS)} TESTS PASSED")
|